[NeurIPS 2026] TabFORGE: A Tabular Foundation Model for Generative Modelling

Related tools: TabCamel for tabular data preparation · TabEval for synthetic-data evaluation · TabStruct for tabular generator benchmarking.

🌐 Overview

TabFORGE learns contextual representations of mixed-type tables with its Structure-aware Feature Encoder, models them with the Score-based Diffusion Transformer, and reconstructs them with the Denoising-aligned Latent Decoder. Its estimator API supports synthetic data generation, classification, regression, embedding, imputation, clustering, and anomaly detection.

Model architecture

✨ Key features

• One estimator interface for pandas DataFrames and NumPy arrays, with schema-aware numerical, categorical, and target handling.
• Pretrained Structure-aware Feature Encoder and leakage-aware held-out-fold embeddings for supervised training rows.
• Latent diffusion with EDM preconditioning, Karras-style schedules, and Euler/Heun sampling.
• Canonical checkpoints for complete estimator restoration or data-free backbone transfer.
• CPU-friendly mock embeddings, single-device execution, and torchrun distributed execution.

🚀 Installation

TabFORGE requires Python 3.10 or later. Install the PyPI distribution tabforge-ml:

pip install tabforge-ml

The Python import name and command-line entry point are both tabforge. For notebooks or optional Weights & Biases logging:

pip install 'tabforge-ml[tutorials]'
pip install 'tabforge-ml[wandb]'

For development, clone the source repository and install it in editable mode:

git clone https://github.com/SilenceX12138/TabFORGE.git
cd TabFORGE
pip install -e .

You can also download the source ZIP.

The default feature encoder and checkpoint="pretrained" use pretrained weights downloaded on first use. For offline examples, configure TabFORGEEmbeddingConfig(model_path="mock") and use checkpoint=None to train all TabFORGE components from scratch.

📊 Logging with W&B

Install tabforge-ml[wandb], authenticate with wandb login, and pass an online logging configuration to an estimator:

from tabforge import TabFORGERuntimeConfig

runtime = TabFORGERuntimeConfig(
    log_wandb=True,
    wandb_project="tabforge",
    wandb_dir="./logs/wandb",
)

Use the estimator's runtime_config=runtime argument, or pass an existing run through logger to keep metrics in your surrounding experiment.

⚡ Quick sanity check

This small CPU example uses deterministic mock embeddings and initializes the TabFORGE model from scratch:

import pandas as pd

from tabforge import (
    TabFORGEArchitectureConfig,
    TabFORGEEmbeddingConfig,
    TabFORGEGenerator,
    TabFORGERuntimeConfig,
    TabFORGETrainingConfig,
)

X = pd.DataFrame({
    "age": [21.0, 35.0, 48.0, 62.0],
    "city": ["Cambridge", "London", "Cambridge", "Oxford"],
})

model = TabFORGEGenerator(
    task="unsupervision",
    embedding_config=TabFORGEEmbeddingConfig(model_path="mock", n_folds=0),
    architecture_config=TabFORGEArchitectureConfig(
        embedding_dimension=8,
        decoder_layers=1,
        decoder_heads=2,
        decoder_ffn_factor=2,
        denoiser_layers=1,
        denoiser_heads=2,
        denoiser_ffn_factor=2,
    ),
    training_config=TabFORGETrainingConfig(max_steps=2, batch_size=4),
    runtime_config=TabFORGERuntimeConfig(device="cpu"),
    random_state=7,
).fit(X, checkpoint=None)

synthetic = model.generate(5, random_state=7)
print(synthetic)

🧩 Example workflows

Prepare train.csv with feature columns in their original order. Categorical DataFrame columns are inferred from their dtypes; pass categorical_features explicitly for integer-coded categories.

import pandas as pd

from tabforge import TabFORGEGenerator

X_train = pd.read_csv("train.csv")
generator = TabFORGEGenerator(task="unsupervision", random_state=7)
generator.fit(X_train)  # Uses the pretrained TabFORGE backbone by default.
synthetic = generator.generate(1_000, random_state=7)
generator.save_checkpoint("checkpoints/generator")

For supervised tasks, fit a classifier or regressor with X_train, y_train, then call predict (and predict_proba for classification, or predict_std and predict_interval for regression). TabFORGEEmbedder.transform extracts embeddings, TabFORGEImputer.transform fills missing values, and the clusterer and anomaly detector expose predict and score_samples respectively. See the linked tutorials for complete examples.

📓 Tutorials

Open a notebook with its Open in Colab badge and follow the tutorial setup guide to install tabforge-ml[tutorials] in the runtime. Select a GPU for practical pretrained workloads; the sanity check above is designed for CPU.

Workflow	Notebook
Unconditional generation	Unconditional generation ·
Pretrained and scratch training	Pretrained and scratch training ·
Devices and distributed execution	Device and distributed execution ·
Classification	Classification ·
Regression	Regression ·
Feature embeddings	Feature embeddings ·
Missing-value imputation	Missing-value imputation ·
Clustering	Clustering ·
Anomaly detection	Anomaly detection ·

💾 Checkpoints and data privacy

The official pretrained backbone is available from the TabFORGE Hugging Face repository. Estimator fitting uses checkpoint="pretrained" by default; use checkpoint=None to train from scratch or pass a canonical checkpoint path to restore model components.

save_checkpoint(path) stores a complete fitted estimator, including preprocessing state and the fitted context required for inference. save_backbone_checkpoint(path) stores reusable model weights without preprocessing state or training rows. Complete fitted checkpoints can contain provider training context; handle them as dataset data. Use a backbone checkpoint for data-free distribution.

📚 Documentation

• Getting started · deployed guide
• Tutorials · deployed gallery
• User guide · deployed guide
• Paper · deployed page
• Estimator API · Configuration · Training · Checkpoint lifecycle

📖 Citation

@inproceedings{jiang2026tabforge,
  title     = {TabFORGE: A Tabular Foundation Model for Generative Modelling},
  author    = {Jiang, Xiangjian and Liu, Mingxuan and Simidjievski, Nikola
               and Klein, Tassilo and Jamnik, Mateja},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2026}
}

@inproceedings{jiang2026generative,
  title     = {A Generative Foundation Model for Heterogeneous Tabular Data},
  author    = {Jiang, Xiangjian and Liu, Mingxuan and Simidjievski, Nikola
               and Klein, Tassilo and Jamnik, Mateja},
  booktitle = {2nd ICML Workshop on Foundation Models for Structured Data},
  year      = {2026},
  url       = {https://openreview.net/forum?id=RcsaxrdpfE}
}