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
1pip install tabfm[jax]
2# GPU/TPU: pip install tabfm[jax,cuda]
Classification:
1from tabfm import TabFMClassifier, tabfm_v1_0_0_jax as tabfm_v1_0_0
2
3model = tabfm_v1_0_0.load(model_type="classification")
4clf = TabFMClassifier(model=model)
5clf.fit(X_train, y_train)
6probs = clf.predict_proba(X_test)
Regression:
1from tabfm import TabFMRegressor, tabfm_v1_0_0_jax as tabfm_v1_0_0
2
3model = tabfm_v1_0_0.load(model_type="regression")
4reg = TabFMRegressor(model=model)
5reg.fit(X_train, y_train)
6preds = reg.predict(X_test)
You can also load directly using the HuggingFace Hub API:
1from tabfm.src.jax.tabfm_v1_0_0 import TabFM_HF
2
3clf_model = TabFM_HF.from_pretrained(
4 "google/tabfm-1.0.0-jax", model_type="classification"
5)
6reg_model = TabFM_HF.from_pretrained(
7 "google/tabfm-1.0.0-jax", model_type="regression"
8)
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:
- Column attention (Set Transformer): embeds each cell using Fourier features and
a per-group linear projection, then aggregates across rows via induced self-attention
- Row compression: CLS tokens summarise each row into a dense vector via row-level
attention with Rotary Position Embedding (RoPE)
- 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):
1model = tabfm_v1_0_0.load(
2 col_attention_impl="flash", # default: flash (memory-efficient for wide tables)
3 row_attention_impl="jax", # default: jax (row attn is over ~8 CLS tokens)
4 icl_attention_impl="flash", # default: flash (ICL attention is memory-critical)
5)
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
1@article{tabfm2026,
2 title = {TabFM: A Zero-Shot Foundation Model for Tabular Data},
3 author = {Google Research},
4 year = {2026},
5 url = {https://research.google/blog/introducing-tabfm-a-zero-shot-foundation-model-for-tabular-data/}
6}