Attention quantization speeds tabular foundation model inference without accuracy loss
Attention Quantization for Tabular Foundation Models
Machine LearningArtificial Intelligence
Summary
Running big computer models on tables of data can be slow. The authors found that focusing on speeding up a specific part called attention, by using a method called quantization, makes these models run faster. They changed how parts of the model convert data to a simpler form without losing accuracy. Their approach made the models run up to 1.7 times faster while keeping their performance. This helps use these models more efficiently in real-world tasks.
What this means in practice
- •For machine learning engineers: Accelerate inference of tabular transformer models in production by integrating FP8 quantized attention kernels.
- •For cloud service providers: Improve efficiency and reduce latency of tabular model hosting by using optimized attention quantization techniques.$Commercial implications: This enables faster and cheaper cloud AI services specializing in tabular data tasks by reducing computational costs.
Authors
Jonas M. Kübler, Benjamin Jäger, Klemens Flöge, Noah Hollmann, Frank Hutter
Abstract
With the recent rise and adoption of tabular foundation models, optimizing their inference performance becomes an emerging field for efficiency research. While the models are architecturally similar to transformer-based large language models (LLMs), the size and serving patterns differ significantly. We show that the focus should be on the attention calculation and less on weight or KV cache quantization, which are more popular in LLMs. We develop a quantization strategy for queries, keys, and values to FP8 and use explicit FP8 matrix multiplication instructions to speed up the attention calculation. We find that it is crucial to align the quantization error in the test rows with the quantization error in the training rows, as otherwise the accuracy drops drastically. Our Triton kernel achieves a speedup up to 1.7x over regular 16-bit kernels, and we show that on TabPFN-v3 and TabICLv2 there is no relevant accuracy loss across TabArena and BeyondArena.