Reproduce “Attention Quantization for Tabular Foundation Models” on OpenResearch

About this paper

Jonas M. Kübler, Benjamin Jäger, Klemens Flöge, Noah Hollmann, Frank Hutter

The paper presents a gated FP8 attention method for tabular foundation models that coordinates rounding across TabPFN-v3’s training–training and test–training attention calls. It activates above 8,192 training rows, using fused low-precision arithmetic while retaining FP32 accumulation. On one 512,000-row problem, prediction was 1.67× faster, while benchmark Elo stayed close, without establishing exact equivalence on an NVIDIA RTX PRO 6000 Blackwell.

Key insights

• Accuracy depends less on the size of individual FP8 errors than on whether the two related attention calls use consistent perturbations. Quantizing only one call increases relative error by roughly 2.4–2.8%, whereas quantizing both changes it by about 0.01–0.02%.
• The shared-scale design for training and test queries avoids predictions changing when test examples are regrouped into different batches.
• FP8 benefits grow with dataset size: the approach is near parity at 8,192 rows, but increasingly faster for datasets containing hundreds of thousands of training examples.

Summary by alphaXiv.