DOCSEARCHbyAlgolia
Submit Docs

TabM

repository·main·Indexed 22 days ago

https://github.com/yandex-research/tabm

A PyTorch-based library for Tabular Deep Learning that implements a parameter-efficient ensembling method. TabM allows for training multiple MLP-like models in parallel using weight sharing to provide ensemble benefits with improved efficiency. The library supports basic model creation via TabM.make(), integration with feature embeddings from rtdl_num_embeddings, and tools for hyperparameter tuning, evaluation, and ensembling.

Tokens
11.8K
Snippets
28
Records
46
Agent score
77%

What's inside tabm

  1. Configure embeddings with `num_embeddings`

    main

    TabM supports different embedding types via num_embeddings:

    • Piecewise-linear embeddings: Historically the more popular choice among users.
    • Periodic embeddings: Can be a better choice for certain specific tasks.

    For detailed hyperparameter tuning recommendations for embeddings, refer to the rtdl_num_embeddings package documentation.

  2. Correct layer ordering for TabM-like models

    main

    When constructing custom TabM-like architectures, the order of layers is critical for effective ensembling.

    The Golden Rule: The $k$ different object representations must be created before the tabular features are mixed with linear layers.

    • GOOD: Use EnsembleView followed by ElementwiseAffine or LinearBatchEnsemble to diversify the $k$ representations before applying standard nn.Linear layers.
    • BAD: Applying nn.Linear to the input before EnsembleView or before the diversification step. This results in $k$ identical representations that are not effectively ensembled.
    • BAD: Using EnsembleView followed by a standard nn.Linear without a diversification step (like ElementwiseAffine), as the $k$ representations remain mathematically identical.
    # GOOD: Diversify representations before linear mixing
model = nn.Sequential(
    tabm.EnsembleView(k=k),
    tabm.ElementwiseAffine((k, d_in), bias=False, scaling_init='normal'),
    nn.Linear(d_in, d),
    ...
)

# BAD: Mixing features before ensemble starts
model = nn.Sequential(
    nn.Linear(d_in, d),
    tabm.EnsembleView(k=k),
    ...
)

  3. Choose an `arch_type` for TabM

    main

    The arch_type parameter determines the model architecture. Choose based on your performance vs. speed requirements:

    • 'tabm': The default value. Expected to provide the best performance in most cases.
    • 'tabm-mini': May result in faster training and/or inference with minimal performance loss. It can occasionally perform better if higher regularization is beneficial, but requires careful tuning of d_block and n_blocks for a given k.
    • 'tabm-packed': Implemented for completeness; typically results in a slower, heavier, and weaker model.

  4. How to build custom TabM-based architectures

    main

    To use TabM as a backbone in a custom model, you should use EnsembleView, make_tabm_backbone, and LinearEnsemble.

    Input Shapes:

    • Standard approach (Recommended): Input x has shape (B, D) during both training and inference. This is simpler and more efficient.
    • Advanced approach: Input x has shape (B, k, D) during training and (B, D) during inference.

    In a custom forward pass, you must call self.ensemble_view(x) to expand the input to the ensemble dimension before passing it to the backbone.

    from tabm import EnsembleView, make_tabm_backbone, LinearEnsemble
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, ...):
        # Create any custom modules.
        ...

        # Create the ensemble input module.
        self.ensemble_view = EnsembleView(...) 
        # Create the backbone.
        self.backbone = make_tabm_backbone(...) 
        # Create the prediction head.
        self.output = LinearEnsemble(...) 

    def forward(self, arg1, arg2, ...):
        # Transform the input as needed to one tensor (e.g., feature embeddings).
        # x shape: (B, D) or (B, k, D)
        x = handle_input(arg1, arg2, ...) 

        # Expand input to ensemble dimension (B, k, D)
        x = self.ensemble_view(x) 
        x = self.backbone(x)
        x = self.output(x)
        return x  # -> (B, k, d_out)

  5. Understand the experiment repository structure

    main

    The exp directory contains the results of experiments, including hyperparameters and metrics. It is organized by model and then by dataset:

    exp/
  <model>/
    <dataset>/       # Or why/<dataset> or tabred/<dataset>
      0-tuning.toml  # The hyperparameter tuning config
      0-tuning/      # The result of the hyperparameter tuning
      0-evaluation/  # The evaluation under multiple random seeds

    Note that some datasets may be prefixed with why/ or tabred/.

  6. Configure the ensemble size `k`

    main

    The parameter k controls the ensemble size. When tuning k, consider these guidelines:

    • Tuning Strategy: Run independent hyperparameter tuning runs with fixed values of k. Changing k can shift the optimal values of other hyperparameters.
    • Performance Threshold: Increasing k improves performance up to a certain threshold, after which it may plateau or decline (especially for arch_type='tabm-mini').
    • Scaling: If you increase k, consider increasing d_block or n_blocks (or both) to provide the base architecture enough capacity to accommodate the larger ensemble via weight sharing.
    • Exploration: Lower values (e.g., 16 or 24) can still yield competitive results.
    • Avoidance: For standard dataset sizes, avoid n_blocks=1 unless you have a high budget for hyperparameter tuning.

  7. How TabM modules and ensembles work

    main

    In conventional MLP models, a module typically operates on a tensor of shape (B, D), where B is batch size and D is latent size. In TabM, modules represent an ensemble of k layers applied in parallel to k inputs.

    Modules in this package operate on tensors of shape (B, k, D), where k represents the ensemble size.

    Key distinction in weight sharing:

    • torch.nn.Linear: Fully shares weights across all ensemble members.
    • tabm.LinearBatchEnsemble: Shares most weights (parameter-efficient).
    • tabm.LinearEnsemble: No weight sharing (independent layers).
    # Example of a module operating on (B, k, D)
# LinearEnsemble is an ensemble of k independent linear layers.
model = tabm.LinearEnsemble(512, d_out, k=k)

  8. Configure Editor settings for VS Code

    main

    To improve the responsiveness of the Python language server, exclude the exp directory (where experiment outputs are stored) from analysis. In VS Code, add this to your .vscode/settings.json:

    {
    "python.analysis.exclude": [ "**/exp" ],
}
    {
    "python.analysis.exclude": [ "**/exp" ],
}

  9. Run training, tuning, evaluation, and ensembling

    main

    The TabM research pipeline consists of several specialized scripts that take TOML configuration files as input.

    Training a single model

    Use bin/model.py with a TOML config. The output (including report.json containing metrics and hyperparameters) will be placed in a directory named after the config file next to the config itself.

    python bin/model.py exp/reproduce/train-once/0.toml

    Hyperparameter tuning

    Use bin/tune.py to find optimal hyperparameters. Use the --continue flag to resume interrupted runs.

    python bin/tune.py exp/reproduce/tabm/california/0-tuning.toml --continue

    Evaluation (Multiple Seeds)

    Use bin/evaluate.py to train a model under multiple random seeds. To evaluate a specific configuration, place it in a directory named <path>-evaluation with the file named 0.toml.

    python bin/evaluate.py exp/reproduce/tabm/california/0-tuning

    Ensembling

    Use bin/ensemble.py to compute metrics for an ensemble of models that have already been trained.

    python bin/ensemble.py exp/reproduce/tabm/california/0-evaluation

    Automated Pipeline

    Use bin/go.py to execute tuning, evaluation, and ensembling in a single automated workflow.

    python bin/go.py exp/reproduce/tabm-go/california/0-tuning --continue

  11. Download and prepare datasets

    main

    Datasets are stored in a data/ directory. Follow these steps to prepare the environment:

    Standard Datasets

    Download the compressed data from Hugging Face and extract it into a data/ folder:

    mkdir local
wget https://huggingface.co/datasets/rototoHF/tabm-data/resolve/main/data.tar -O local/tabm-data.tar.gz
mkdir data
tar -xvf local/tabm-data.tar.gz -C data

    TabReD Benchmark

    1. Download the TabReD benchmark to local/tabred (requires a Kaggle account).
    2. Run the preparation script:
    python tools/prepare_tabred.py local/tabred data
    mkdir local
wget https://huggingface.co/datasets/rototoHF/tabm-data/resolve/main/data.tar -O local/tabm-data.tar.gz
mkdir data
tar -xvf local/tabm-data.tar.gz -C data