google-research / google-research/tabfm

Expose the existing prefill/decode KV cache through the sklearn API

Open
#60 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
2.6k
Forks
270
Avg merge
1d 7h
Merged PRs (30d)
1

Description

## Summary

The JAX model already implements context caching: `TabFM.prefill()` and
`TabFM.decode()` in `tabfm/src/jax/model.py`, with `ICLearningCache` and the
`cache_icl_input_only` option, covered by
`model_test.py::test_prefill_decode_consistency`. The public `TabFMClassifier`
and `TabFMRegressor` never call them. Every `predict_proba` call re-encodes
the full training context, once per ensemble member.

## What this costs

Measured on an RTX 5090 (PyTorch backend, bf16, `n_estimators=4`, 40k-row
context, creditcard dataset from OpenML), steady state:

| queries per predict call | wall time |
|---|---|
| 1 | 25.7 s |
| 100 | 26.1 s |
| 1,000 | 27.1 s |
| 30,000 | 37.5 s |

Predict cost is a fixed cost of about 25 s to re-encode the context, plus
about 1.4 ms per query row. A single prediction costs the same as a thousand.
`prefill`/`decode` exists to move that fixed cost to fit time.

## Proposal

Add an option to the sklearn wrappers, using TabPFN's naming so users can
find it:

```python
clf = TabFMClassifier(model=model, fit_mode="fit_with_cache")
clf.fit(X_train, y_train) # runs prefill once per ensemble member view
clf.predict_proba(X_test) # runs decode against the cached context
```

The default `fit_mode="fit"` keeps current behavior.

We have a working branch with tests: 60 passing, parity against the default
path at about 1e-7 in fp32 including the `ensemble()` preset, and a mean
probability difference of about 1e-4 in bf16 on the released checkpoint.
Wrapper-level measurements on the JAX backend at a 10k-row context show 1.3x
to 2.4x faster predict calls after a one-time 28.9 s prefill. The gain grows
with context length: the default path's fixed cost grows from about 2.5 s at
10k rows to about 35 s at 40k rows per member, while the cached path does not
depend on context length. We are happy to send the PR if this direction is
welcome.

## Two findings from the implementation

1. Wrapping `decode` in `nnx.jit` does not work at real scale. The jit copies
the multi-GB caches into the executable arena and fuses a transpose across
all 24 layers that runs out of memory at a 20k context and fails XLA
autotuning at 40k. Our branch calls `decode` eagerly, which matches how
`model_test.py` uses it.
2. `decode` builds an internal per-head key/value tensor that spans all 24
layers at once (shape `[heads, layers, B, T, head_dim]`, about 1 GB per 10k
context rows per member). On a 32 GB GPU this caps cached prediction at
roughly a 10k context with one member. Chunking that tensor per layer would
unlock the larger contexts where caching helps most. We can file this
separately.

A PyTorch-backend equivalent would be a natural follow-up. `prefill`/`decode`
currently exist only in the JAX model.

Contributor guide

Open the contributing guide

Research direction

Start with TabFM.prefill() and TabFM.decode() in tabfm/src/jax/model.py, then run model_test.py::test_prefill_decode_consistency to understand the existing cache contract. Trace how TabFMClassifier and TabFMRegressor implement fit and predict_proba, and verify that a fit_mode="fit_with_cache" path preserves prediction parity while avoiding repeated context encoding.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
api, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.