TabFM: Tabular Foundation Models

TabFM (Tabular Foundation Model) is a scikit-learn compatible tabular foundation model. It allows you to perform zero-shot classification and regression on tabular datasets with mixed column types out-of-the-box.

At inference time, TabFM does not require training parameters on your dataset; instead, it leverages in-context learning by reading your training data as "context" to make instant predictions on new test samples.

This is not an officially supported Google product.

---

Installation

To install TabFM, clone the repository and install it locally with the backend of your choice:

JAX (CPU):

git clone https://github.com/google-research/tabfm.git
cd tabfm
pip install -e .[jax]

JAX (GPU):

git clone https://github.com/google-research/tabfm.git
cd tabfm
pip install -e .[jax,cuda]

PyTorch (CPU/GPU):

git clone https://github.com/google-research/tabfm.git
cd tabfm
pip install -e .[pytorch]

Note: For PyTorch with GPU support, ensure you have the appropriate PyTorch version installed for your CUDA version before installing TabFM.

Requirements

For a complete list of pinned dependencies and versions, please see requirements.txt. The core requirements depend on the backend you choose:

• Python >= 3.11
• Hugging Face Hub (for downloading pre-trained weights)
• JAX Backend:
  • JAX (specifically jax==0.10.1)
  • Flax (specifically flax==0.12.7, using the modern flax.nnx API)
• PyTorch Backend:
  • PyTorch (specifically torch==2.12.1+cpu or a GPU version)

---

License notice for pretrained weights

Important: The TabFM source code in this repository is licensed under Apache-2.0. However, the default Quick Start calls tabfm_v1_0_0.load(), which automatically downloads pretrained weights from Hugging Face. Those pretrained weights are distributed under the separate tabfm-non-commercial-v1.0 license and are restricted to non-commercial, non-production use. Commercial or production use of the default pretrained weights is not permitted.

---

Quick Start (TabFM v1.0.0)

We provide pre-trained weights for the TabFM v1.0.0 release. The library handles downloading and loading these weights automatically from Hugging Face. These default weights are governed by the separate tabfm-non-commercial-v1.0 license described above.

1. Classification Example

import numpy as np
import pandas as pd
from tabfm import TabFMClassifier

# Choose your backend:

# OPTION A: JAX Backend
from tabfm import tabfm_v1_0_0_jax as tabfm_v1_0_0
model = tabfm_v1_0_0.load()

# OPTION B: PyTorch Backend
# from tabfm import tabfm_v1_0_0_pytorch as tabfm_v1_0_0
# model = tabfm_v1_0_0.load()

# Initialize scikit-learn compatible classifier (works with either backend model)
clf = TabFMClassifier(model=model)

# Prepare your dataset (supports mixed numerical and categorical features)
X_train = pd.DataFrame({
    "age": [25.0, 45.0, 35.0, 50.0],
    "job": ["engineer", "manager", "engineer", "manager"],
    "income": [80000, 120000, 90000, 130000]
})
y_train = np.array(["low_risk", "high_risk", "low_risk", "high_risk"])

X_test = pd.DataFrame({
    "age": [30.0, 48.0],
    "job": ["engineer", "manager"],
    "income": [85000, 125000]
})

# Fit classifier (prepares ordinal encoders and numerical scalers)
clf.fit(X_train, y_train)

# Predict classes and probabilities
predictions = clf.predict(X_test)
probabilities = clf.predict_proba(X_test)

print("Predictions:", predictions)
print("Class Probabilities:\n", probabilities)

2. Regression Example

import numpy as np
import pandas as pd
from tabfm import TabFMRegressor

# Choose your backend:

# OPTION A: JAX Backend
from tabfm import tabfm_v1_0_0_jax as tabfm_v1_0_0
model = tabfm_v1_0_0.load(model_type="regression")

# OPTION B: PyTorch Backend
# from tabfm import tabfm_v1_0_0_pytorch as tabfm_v1_0_0
# model = tabfm_v1_0_0.load(model_type="regression")

# Initialize scikit-learn compatible regressor (works with either backend model)
reg = TabFMRegressor(model=model)

# Prepare your dataset
X_train = pd.DataFrame({
    "sqft": [1200, 2500, 1500, 3000],
    "neighborhood": ["A", "B", "A", "C"]
})
y_train = np.array([250000, 550000, 310000, 620000])

X_test = pd.DataFrame({
    "sqft": [1800, 2800],
    "neighborhood": ["A", "B"]
})

# Fit and Predict
reg.fit(X_train, y_train)
predictions = reg.predict(X_test)

print("Predicted Prices:", predictions)

---

Multi-GPU inference (PyTorch)

TabFM reads the whole training fold as one in-context sequence, so very large contexts (roughly >450k rows on an 80GB GPU) exceed a single device's memory. tabfm.src.pytorch.seqpar shards the rows of the sequence across the ranks of a torch.distributed process group with mathematically exact attention (no approximation; results match the single-device path up to bf16 summation order). Launch one process per GPU, e.g. with torchrun:

import torch.distributed as dist
from tabfm import TabFMRegressor, tabfm_v1_0_0_pytorch
from tabfm.src.pytorch import seqpar

dist.init_process_group("nccl")
model = tabfm_v1_0_0_pytorch.load(model_type="regression", device=f"cuda:{rank}")
reg = TabFMRegressor(model=model, n_estimators=1)
reg.fit(X_train, y_train)          # cheap: no GPU forward
preds = seqpar.predict(reg, X_test)  # collective call; every rank returns preds

seqpar.predict_proba is the classification equivalent. A runnable script is provided in examples/seqpar_regression_example.py.

The JAX backend has an equivalent, tabfm.src.jax.seqpar: a single process shards the rows across all local devices with jax.shard_map, with the same seqpar.predict(reg, X_test) / seqpar.predict_proba(clf, X_test) API (no torchrun needed). A 1M-row context fits in ~35GB/GPU on 4 devices (a single 80GB device cannot run it at all), and at single-device-feasible sizes the sharded path is ~5x faster on 4 GPUs.

Examples Directory

You can find runnable scripts for both classification and regression under the examples/ folder:

• classification_example.py
• regression_example.py

To run them, simply execute:

python examples/classification_example.py

(You can edit these files to switch between JAX and PyTorch backends as shown in the comments inside them).

---

Evaluation Results

Our model evaluation results can be found in results/.

---

FAQ

Is there a maximum table size for TabFM inputs?

TabFM uses in-context learning over a bounded context window, so very large tables should be sampled or split before inference. The scikit-learn estimators expose the main practical limits through max_num_features and max_num_rows (defaults are 500 features and 100 context rows), plus n_estimators for ensembling over multiple sampled contexts and inference_batch_size for memory control. If your dataset is larger than these limits, TabFM will work with the sampled/context rows rather than consuming the full table at once.

Is there a TabFM technical report or paper?

A technical report is not included in this repository at this time. If a report or paper describing the architecture, training pipeline, datasets, and evaluation methodology is released, this README will be updated with a link.

---

Running Tests

You can run the unit tests directly using Python's unittest module:

# Run all tests (requires both JAX and PyTorch installed)
PYTHONPATH=. python3 -m unittest discover -s tabfm/src/ -p "*_test.py"

# Or run specific test files:
PYTHONPATH=. python3 -m unittest tabfm/src/pytorch/model_test.py
PYTHONPATH=. python3 -m unittest tabfm/src/classifier_and_regressor_pytorch_test.py

Alternatively, if you have Bazel installed, you can run tests with:

bazel test //...