google/tabfm-1.0.0-jax

TabFM 1.0.0 (JAX/Flax)

TabFM is a zero-shot tabular foundation model from Google Research. It supports classification and regression on structured/tabular data with mixed numerical and categorical columns, requiring no fine-tuning or hyperparameter search - training examples are passed as context and predictions are made in a single forward pass.

This repository contains the JAX/Flax weights stored as Orbax checkpoints. For the PyTorch weights see google/tabfm-1.0.0-pytorch.

Getting Started

pip install tabfm[jax]
# GPU/TPU: pip install tabfm[jax,cuda]

Classification:

from tabfm import TabFMClassifier, tabfm_v1_0_0_jax as tabfm_v1_0_0

model = tabfm_v1_0_0.load(model_type="classification")
clf = TabFMClassifier(model=model)
clf.fit(X_train, y_train)
probs = clf.predict_proba(X_test)

Regression:

from tabfm import TabFMRegressor, tabfm_v1_0_0_jax as tabfm_v1_0_0

model = tabfm_v1_0_0.load(model_type="regression")
reg = TabFMRegressor(model=model)
reg.fit(X_train, y_train)
preds = reg.predict(X_test)

You can also load directly using the HuggingFace Hub API:

from tabfm.src.jax.tabfm_v1_0_0 import TabFM_HF

clf_model = TabFM_HF.from_pretrained(
    "google/tabfm-1.0.0-jax", model_type="classification"
)
reg_model = TabFM_HF.from_pretrained(
    "google/tabfm-1.0.0-jax", model_type="regression"
)

Available Checkpoints

Subfolder	Task	Format
classification/	Classification (up to 10 classes)	Orbax checkpoint
regression/	Regression	Orbax checkpoint

Developers and Affiliations

Developed by the Google Research team.

Intended Use

• Tabular data with numerical and/or categorical columns
• Binary and multiclass classification (up to 10 classes)
• Regression on continuous targets
• Zero-shot inference: no dataset-specific training or hyperparameter tuning
• Works with DataFrames (pandas) or numpy arrays
• GPU/TPU acceleration via JAX (use bfloat16 dtype for efficiency)

Not Intended For

• Images, audio, video, or raw text
• More than 10 output classes (hard model limit)
• Tasks requiring task-specific fine-tuning
• Non-tabular structured data (graphs, sequences)
• Commercial use (see License below)

Model Architecture

TabFM uses alternating row and column attention to capture both feature interactions and row-level patterns:

1. Column attention (Set Transformer): embeds each cell using Fourier features and a per-group linear projection, then aggregates across rows via induced self-attention
2. Row compression: CLS tokens summarise each row into a dense vector via row-level attention with Rotary Position Embedding (RoPE)
3. ICL Transformer: a 24-block causal transformer operates over the compressed row vectors, treating training rows as context and outputting predictions for test rows

Key hyperparameters:

Parameter	Value
Embedding dim	256
Column attention blocks	3 (4 heads, 256 induced points)
Row attention blocks	3 (8 heads, 8 CLS tokens, RoPE base 100k)
ICL transformer blocks	24 (8 heads)
Feed-forward factor	4
Max classes	10
Activation	SwiGLU
Fourier features	32 frequencies
Compute dtype	bfloat16 (default)

Attention implementations can be configured per layer family (flash, jax):

model = tabfm_v1_0_0.load(
    col_attention_impl="flash",   # default: flash (memory-efficient for wide tables)
    row_attention_impl="jax",     # default: jax (row attn is over ~8 CLS tokens)
    icl_attention_impl="flash",   # default: flash (ICL attention is memory-critical)
)

Training Data and Priors

TabFM was trained on hundreds of millions of synthetic datasets generated dynamically using structural causal models (SCMs). Synthetic data was chosen due to the scarcity of diverse, high-quality open-source tabular datasets and to avoid privacy/licensing concerns with real-world industrial data. The SCM prior encodes inductive biases about causal structure and feature relationships typical in tabular tasks.

Performance

TabFM was evaluated on TabArena across 51 datasets (38 classification, 13 regression). In zero-shot mode - a single forward pass with no hyperparameter search - TabFM outperforms heavily-tuned supervised baselines including gradient-boosted trees. The TabFMClassifier.ensemble() preset (feature crosses, SVD features, NNLS blending) yields further improvements.

See the Google Research blog post for full benchmark details.

Ethical Considerations

TabFM was trained entirely on synthetic data. Performance on specific real-world domains, minority groups, or edge distributions is not fully characterised. Users should evaluate the model on held-out data representative of their use case before deploying in high-stakes settings.

Limitations

• Max 10 classes for classification (hard architectural limit)
• Memory usage scales with the number of training rows (all rows are passed as context)
• Optimised for tables up to 500 features; behaviour on very wide tables may degrade
• Performance is not guaranteed to match task-specific, fine-tuned models on all datasets
• Not an officially supported Google product

License

The model weights in this repository are released under the TabFM Non-Commercial License v1.0 - see LICENSE. The source code is Apache 2.0 licensed via google-research/tabfm.

Version

1.0.0

Citation

@article{tabfm2026,
  title  = {TabFM: A Zero-Shot Foundation Model for Tabular Data},
  author = {Google Research},
  year   = {2026},
  url    = {https://research.google/blog/introducing-tabfm-a-zero-shot-foundation-model-for-tabular-data/}
}