README.md
6.2 KB · 189 lines · markdown Raw
1 ---
2 license: other
3 license_name: tabfm-non-commercial-v1.0
4 license_link: https://huggingface.co/google/tabfm-1.0.0-jax/blob/main/LICENSE
5 library_name: tabfm
6 pipeline_tag: tabular-classification
7 tags:
8 - tabular
9 - tabular-regression
10 - zero-shot
11 - in-context-learning
12 - jax
13 - flax
14 - foundation-model
15 ---
16
17 # TabFM 1.0.0 (JAX/Flax)
18
19 TabFM is a zero-shot tabular foundation model from Google Research. It supports
20 classification and regression on structured/tabular data with mixed numerical and
21 categorical columns, requiring no fine-tuning or hyperparameter search - training
22 examples are passed as context and predictions are made in a single forward pass.
23
24 This repository contains the **JAX/Flax** weights stored as Orbax checkpoints. For the
25 PyTorch weights see
26 [google/tabfm-1.0.0-pytorch](https://huggingface.co/google/tabfm-1.0.0-pytorch).
27
28 ## Getting Started
29
30 ```bash
31 pip install tabfm[jax]
32 # GPU/TPU: pip install tabfm[jax,cuda]
33 ```
34
35 **Classification:**
36
37 ```python
38 from tabfm import TabFMClassifier, tabfm_v1_0_0_jax as tabfm_v1_0_0
39
40 model = tabfm_v1_0_0.load(model_type="classification")
41 clf = TabFMClassifier(model=model)
42 clf.fit(X_train, y_train)
43 probs = clf.predict_proba(X_test)
44 ```
45
46 **Regression:**
47
48 ```python
49 from tabfm import TabFMRegressor, tabfm_v1_0_0_jax as tabfm_v1_0_0
50
51 model = tabfm_v1_0_0.load(model_type="regression")
52 reg = TabFMRegressor(model=model)
53 reg.fit(X_train, y_train)
54 preds = reg.predict(X_test)
55 ```
56
57 You can also load directly using the HuggingFace Hub API:
58
59 ```python
60 from tabfm.src.jax.tabfm_v1_0_0 import TabFM_HF
61
62 clf_model = TabFM_HF.from_pretrained(
63 "google/tabfm-1.0.0-jax", model_type="classification"
64 )
65 reg_model = TabFM_HF.from_pretrained(
66 "google/tabfm-1.0.0-jax", model_type="regression"
67 )
68 ```
69
70 ### Available Checkpoints
71
72 | Subfolder | Task | Format |
73 |-----------|------|--------|
74 | `classification/` | Classification (up to 10 classes) | Orbax checkpoint |
75 | `regression/` | Regression | Orbax checkpoint |
76
77 ## Developers and Affiliations
78
79 Developed by the [Google Research](https://research.google) team.
80
81 ## Intended Use
82
83 - Tabular data with numerical and/or categorical columns
84 - Binary and multiclass classification (up to 10 classes)
85 - Regression on continuous targets
86 - Zero-shot inference: no dataset-specific training or hyperparameter tuning
87 - Works with DataFrames (pandas) or numpy arrays
88 - GPU/TPU acceleration via JAX (use `bfloat16` dtype for efficiency)
89
90 ## Not Intended For
91
92 - Images, audio, video, or raw text
93 - More than 10 output classes (hard model limit)
94 - Tasks requiring task-specific fine-tuning
95 - Non-tabular structured data (graphs, sequences)
96 - Commercial use (see License below)
97
98 ## Model Architecture
99
100 TabFM uses alternating row and column attention to capture both feature interactions
101 and row-level patterns:
102
103 1. **Column attention** (Set Transformer): embeds each cell using Fourier features and
104 a per-group linear projection, then aggregates across rows via induced self-attention
105 2. **Row compression**: CLS tokens summarise each row into a dense vector via row-level
106 attention with Rotary Position Embedding (RoPE)
107 3. **ICL Transformer**: a 24-block causal transformer operates over the compressed row
108 vectors, treating training rows as context and outputting predictions for test rows
109
110 Key hyperparameters:
111
112 | Parameter | Value |
113 |-----------|-------|
114 | Embedding dim | 256 |
115 | Column attention blocks | 3 (4 heads, 256 induced points) |
116 | Row attention blocks | 3 (8 heads, 8 CLS tokens, RoPE base 100k) |
117 | ICL transformer blocks | 24 (8 heads) |
118 | Feed-forward factor | 4 |
119 | Max classes | 10 |
120 | Activation | SwiGLU |
121 | Fourier features | 32 frequencies |
122 | Compute dtype | bfloat16 (default) |
123
124 Attention implementations can be configured per layer family (`flash`, `jax`):
125
126 ```python
127 model = tabfm_v1_0_0.load(
128 col_attention_impl="flash", # default: flash (memory-efficient for wide tables)
129 row_attention_impl="jax", # default: jax (row attn is over ~8 CLS tokens)
130 icl_attention_impl="flash", # default: flash (ICL attention is memory-critical)
131 )
132 ```
133
134 ## Training Data and Priors
135
136 TabFM was trained on hundreds of millions of **synthetic** datasets generated
137 dynamically using structural causal models (SCMs). Synthetic data was chosen due to
138 the scarcity of diverse, high-quality open-source tabular datasets and to avoid
139 privacy/licensing concerns with real-world industrial data. The SCM prior encodes
140 inductive biases about causal structure and feature relationships typical in tabular
141 tasks.
142
143 ## Performance
144
145 TabFM was evaluated on [TabArena](https://tabarena.ai) across 51 datasets
146 (38 classification, 13 regression). In zero-shot mode - a single forward pass with no
147 hyperparameter search - TabFM outperforms heavily-tuned supervised baselines including
148 gradient-boosted trees. The `TabFMClassifier.ensemble()` preset (feature crosses,
149 SVD features, NNLS blending) yields further improvements.
150
151 See the [Google Research blog post](https://research.google/blog/introducing-tabfm-a-zero-shot-foundation-model-for-tabular-data/) for full benchmark details.
152
153 ## Ethical Considerations
154
155 TabFM was trained entirely on synthetic data. Performance on specific real-world
156 domains, minority groups, or edge distributions is not fully characterised. Users
157 should evaluate the model on held-out data representative of their use case before
158 deploying in high-stakes settings.
159
160 ## Limitations
161
162 - **Max 10 classes** for classification (hard architectural limit)
163 - Memory usage scales with the number of training rows (all rows are passed as context)
164 - Optimised for tables up to 500 features; behaviour on very wide tables may degrade
165 - Performance is not guaranteed to match task-specific, fine-tuned models on all datasets
166 - Not an officially supported Google product
167
168 ## License
169
170 The model weights in this repository are released under the
171 **TabFM Non-Commercial License v1.0** - see [LICENSE](https://huggingface.co/google/tabfm-1.0.0-jax/blob/main/LICENSE). The source code is
172 Apache 2.0 licensed via [google-research/tabfm](https://github.com/google-research/tabfm).
173
174 ## Version
175
176 1.0.0
177
178 ## Citation
179
180 ```bibtex
181 @article{tabfm2026,
182 title = {TabFM: A Zero-Shot Foundation Model for Tabular Data},
183 author = {Google Research},
184 year = {2026},
185 url = {https://research.google/blog/introducing-tabfm-a-zero-shot-foundation-model-for-tabular-data/}
186 }
187 ```
188
189