KumoTabular#

class sdm.models.KumoTabular(task: Task | str | Iterable[Task | str] | None = None, size: Literal['small', 'large'] = 'large', pretrained: bool = True, device: device | str | None = None)#

Bases: ICLModel

The tabular foundation model from “NVIDIA Kumo Tabular Sets a New Accuracy-Efficiency Frontier for Tabular Prediction”.

../../_images/kumo_tabular.svg

Architecturally, KumoTabular combines the compression-then-ICL structure of TabICLv2 with interleaved row/column attention from TabPFN and Fourier cell embeddings as introduced by TabFM. Numerical and categorical values use separate learned Fourier frequencies and projections. Each cell embedding represents a repeated group of features, with missing values handled via learned missingness projections.

The cell representations are processed by fully interleaved attention stages, scaled up to six stages with 256 hidden cell dimension:

  • Column-wise: Each feature group is processed across rows using induced set attention. Both context and query cells attend only to context-row keys and values.

  • Row-wise: Each row’s feature groups and learnable readout tokens attend to one another, combining feature interactions into a fixed-size row representation.

A final dataset-wise transformer processes the compressed row representations and predicts each query target from the labeled context rows.

The transformer blocks leverage torch.nn.RMSNorm, with normalization applied on query/key/value inputs and before GELU feed-forward networks, and non-affine per-head normalization applied on projected queries and keys.

Additionally, KumoTabular applies learned logarithmic context-length scaling via sdm.nn.LogScale during column-wise and dataset-wise attention, and query-gated logarithmic scaling via sdm.nn.GatedLogScale during row-wise attention.

For regression tasks, KumoTabular predicts 999 quantiles named "q001" through "q999", similar to TabICLv2.

Parameters:
  • task (TaskLike | Iterable[TaskLike] | None) – The tasks to initialize. If None, all tasks supported by this model are initialized.

  • size (Literal['small', 'large']) – The size of the model.

  • pretrained (bool) – Whether to load pretrained checkpoints.

  • device (torch.device | str | None) – The device.

Capabilities#

Supported Input Semantic Types

numerical

Supported Target Semantic Types

categorical, numerical

Supported Prediction Tasks

classification, regression

Multi-Target Support

❌

Related Table Support

❌

Default Recipe#

Recipe(
  features=Sequential(
    StypeDispatch(
      numerical=Sequential(
        Cast(torch.float64),
        DropConstantColumns(),
        Standardize(eps=1e-06),
        Clip(-100.0, 100.0),
        Choice(
          Identity(),
          PowerTransform(),
          Sequential(
            RobustScale(quantile_range=(25.0, 75.0)),
            ClipSoft(3.0),
          ),
          method='round_robin',
        ),
        ClipSigma(threshold=4.0),
        FlipSign(probability=0.5),
      ),
      categorical=Sequential(
        AlignCategories(sort_by='value'),
        AddCategoryCounts(min_cardinality=50),
        ToNumerical(),
        Cast(torch.float64),
        DropConstantColumns(),
        Standardize(eps=1e-06),
        Clip(-100.0, 100.0),
        Choice(
          Identity(),
          PowerTransform(),
          Sequential(
            RobustScale(quantile_range=(25.0, 75.0)),
            ClipSoft(3.0),
          ),
          method='round_robin',
        ),
        ClipSigma(threshold=4.0),
      ),
    ),
    ShuffleColumns(method='latin'),
    SelectColumns(),
    Cast(torch.float32),
  ),
  target=Sequential(
    StypeDispatch(
      numerical=Sequential(
        Standardize(eps=0.0),
        FlipSign(probability=0.5),
      ),
      categorical=Sequential(
        AlignCategories(),
        ShuffleCategories(method='shift'),
      ),
    ),
  ),
  output=TaskDispatch(
    classification=Sequential(
      AverageEstimators(),
      Softmax(),
    ),
    regression=Sequential(
      SortQuantiles(),
      AverageEstimators(trim_fraction=0.2),
    ),
  ),
)