KumoTabular#

class sdm.models.KumoTabular(task: Task | str | Iterable[Task | str] | None = None, size: Literal['small', 'medium', '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

KumoTabular processes a table in two stages. First, an interleaved row/column encoder transforms raw table cells into fixed-size row representations. Numerical and categorical cells use separate learned Fourier frequencies and projections, together with a learned missingness projection. Each encoder stage consists of:

  • 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.

Second, a dataset-wise in-context learning transformer processes the row representations and predicts each query target from the labeled context rows. The "large" model widens this transformer to 1024 channels and 16 attention heads. Its four 256-channel readout tokens concatenate directly to that width without a projection layer.

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

Parameters:

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

  • size (Literal['small', 'medium', 'large']) – The model size, one of "small", "medium", or "large". Defaults to "large".

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

  • device (torch.device | str | None) – The device for model parameters. If None, uses PyTorch’s default device.

Capabilities#

Supported Input Semantic Types

numerical

Supported Target Semantic Types

categorical, numerical

Supported Prediction Tasks

classification, regression

Multi-Target Support

❌

Context-only Column 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),
          ),
          RankGaussian(max_knots=8192),
          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),
          ),
          RankGaussian(max_knots=8192),
          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(sort_by='value'),
        ShuffleCategories(method='shift'),
      ),
    ),
  ),
  output=TaskDispatch(
    classification=Sequential(
      AverageEstimators(),
      Softmax(),
    ),
    regression=Sequential(
      SortQuantiles(),
      AverageEstimators(trim_fraction=0.2),
    ),
  ),
)