PaperlayerView exact revision on arXiv

Source-linked AI summary

FlatQuant: Flatness Matters for LLM Quantization

Yuxuan Sun, Ruikang Liu, Haoli Bai, Han Bao, Kang Zhao, Yuening Li, Jiaxin Hu, Xianzhi Yu, Lu Hou, Chun Yuan, Xin Jiang, Wulong Liu, Jun Yao

arXiv:2410.09426v4cs.CLcs.LG

TL;DR

Outliers can leave weights and activations poorly suited to quantization, and existing transformations may not adequately flatten their distributions. FLATQUANT learns affine transformations for each linear layer and fuses transformation with quantization. It reports ≤1% accuracy drop for W4A4 quantization on LLaMA-3-70B, alongside inference speedups over FP16.

  • Problem

    Existing pre-quantization transformations can leave steep, dispersed weight and activation distributions and may add inference overhead.

    Source: §1 · p. 1

  • Method

    FLATQUANT learns affine transformations for each linear layer to flatten weights and activations, using Kronecker products and kernel fusion to limit runtime overhead.

    Source: §1 · p. 2

  • Results

    ≤1% accuracy drop is reported for round-to-nearest W4A4 quantization on LLaMA-3-70B, and FLATQUANT establishes state-of-the-art quantization results.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Fusing affine transformation and quantization yields 2.3x prefill speedup and 1.7x decoding speedup under W4A4 versus the FP16 baseline.

    Source: §1 · p. 2

    The optimal choice of calibration set remains open, and the study does not examine integration with newer data types such as MXFP4.

    Source: E. Limitations · p. 24

Abstract

from arXiv · show

Recently, quantization has been widely used for the compression and acceleration of large language models (LLMs). Due to the outliers in LLMs, it is crucial to flatten weights and activations to minimize quantization error with equally spaced quantization points. Prior research explores various pre-quantization transformations to suppress outliers, such as per-channel scaling and Hadamard transformation. However, we observe that these transformed weights and activations can still exhibit steep and dispersed distributions. In this paper, we propose FlatQuant (Fast and Learnable Affine Transformation), a new post-training quantization approach that enhances the flatness of weights and activations. Our approach identifies optimal affine transformations for each linear layer, calibrated in hours via a lightweight objective. To reduce runtime overhead of affine transformation, we apply Kronecker product with two lightweight matrices, and fuse all operations in FlatQuant into a single kernel. Extensive experiments demonstrate that FlatQuant establishes a new state-of-the-art benchmark for quantization. For example, it achieves less than 1\% accuracy drop for W4A4 quantization on the LLaMA-3-70B model, surpassing SpinQuant by 7.5\%. Additionally, it provides up to 2.3x prefill speedup and 1.7x decoding speedup compared to the FP16 model. Code is available at: https://github.com/ruikangliu/FlatQuant.

1. Introduction

FLATQUANT targets quantization error from outliers by learning affine transformations that flatten weights and activations for each linear layer. The paper reports state-of-the-art quantization results, including low-bit accuracy and inference-speed gains.

1. Introduction

  • Existing pre-quantization transformations remain sub-optimal: scaling does not spread outliers across non-outlier channels, while Hadamard transformation ignores individual layer characteristics.These transformations also add inference overhead.Source: §1 · p. 1
  • FLATQUANT learns a fast affine transformation for each linear layer to flatten weights and activations, using lightweight block-wise training on calibration data.Kronecker products with two lightweight matrices and kernel fusion are used to reduce transformation overhead.Source: §1 · p. 2
  • Flat weights and activations facilitate quantization and reduce error propagation across Transformer layers.The paper presents this as a central motivation for optimizing flatness.Source: §1 · pp. 1–2
  • ≤1% accuracy drop is achieved with round-to-nearest W4A4 quantization on LLaMA-3-70B.The authors describe this as the first such result to their knowledge.Source: §1 · p. 2
  • 2.3x prefill speedup and 1.7x decoding speedup are reported under W4A4 quantization versus the FP16 baseline.The speedups use a kernel that fuses affine transformation and quantization.Source: §1 · p. 2

2. Motivation

The section introduces LLM quantization and compares transformations intended to mitigate outliers. It argues that FLATQUANT produces flatter weight and activation distributions and controls quantization-error propagation more effectively than the alternatives discussed.

2.1. Preliminaries on LLM Quantization

  • Quantization converts full-precision weights or activations, and optionally the KV cache, into low-bit representations.Weight quantization is represented using a step size, a projection function, and a set of b-bit integer points.Source: §2.1 · p. 2
  • Per-channel scaling smooths activations by jointly scaling weights and activations, whereas Hadamard transformation rotates channels to redistribute outliers across them.The scaled weights can be preprocessed offline; Hadamard’s orthogonality preserves the linear-layer multiplication.Source: §2.1 · p. 2

2.2. The Flatness for Quantization

  • FLATQUANT consistently flattens both weights and activations, while per-channel scaling can make weight envelopes steeper and Hadamard transformation can still produce unsatisfactory distributions.Flat distributions are described as preferable for quantization.Source: §2.2 · p. 3, Figure 1 · p. 3
  • FLATQUANT controls layerwise quantization-error propagation best, followed by Hadamard transformation and then per-channel scaling.The comparison also reports lower error at pivot tokens for FLATQUANT than for the other two methods.Source: §2.2 · p. 3, Figure 2 · p. 3

3. Method

FLATQUANT seeks layer-specific learnable affine transformations that flatten weights and activations while keeping transformation overhead low. It integrates those transformations into Transformer layers and fuses transformation and quantization operations in a single kernel.

3.1. Fast and Learnable Affine Transformation

  • FLATQUANT replaces each large affine matrix P with P1 ⊗ P2, two smaller invertible matrices, to reduce transformation memory and computation.The stated memory and computation savings are largest when n1 = n2 = √n.Source: §3.1 · p. 4
  • Layer-specific scaling balances weight and activation outliers, while learned clipping thresholds further limit outliers after transformation.The scaling vector can be merged pair-wise into preceding layer normalization or linear layers without additional inference overhead.Source: §3.1 · p. 4
  • It sequentially minimizes block-level quantization MSE on calibration data, jointly learning transformation matrices, scaling, and clipping thresholds; training yields flat weights and activations.The calibration data may include 128 randomly sampled sentences per Transformer block.Source: §3.1 · pp. 4–5

3.2. Integration with the Transformer Architecture

  • In self-attention, Pa flattens query, key, and value inputs, Po smooths output-projection inputs, and Ph and Pv transform the key and value caches head by head.FLATQUANT leaves Ph and Pv undecomposed because per-head quantization permits cheaper transformations, and fuses Po with Pv without additional accuracy loss in the reported experiments.Source: §3.2 · p. 5
  • For the FFN, Pug and Pd flatten inputs to the up and down projections, while Pd's per-channel scaling is merged into the up-projection weights.FLATQUANT preserves LayerNorm and applies separate learnable affine transformations after it, rather than requiring all Transformer blocks to share one transformation.Source: §3.2 · p. 5

3.3. Efficient Kernel Design

  • FLATQUANT fuses affine transformations and quantization into one kernel, avoiding intermediate global-memory accesses and reducing kernel-launch overhead.The kernel keeps intermediate results in SRAM before writing them back to global memory.Source: §3.3 · p. 5
  • The fused kernel loads both small matrices into SRAM, transforms tiled activations, and quantizes the results before writing them to global memory.The low computational intensity of the Kronecker transformation and the memory-bound quantization motivate the fusion.Source: §3.3 · p. 5

4. Experiments

FlatQuant is evaluated on language-generation and zero-shot QA benchmarks, with inference-latency and component analyses examining its accuracy, runtime, and flatness. The reported results include improved benchmark performance, speedups from kernel fusion, and lower perplexity from learnable transformations.

4.1. Settings

  • Experiments use LLaMA-2 and LLaMA-3 models, with WikiText-2 and C4 perplexity and six zero-shot commonsense tasks; calibration uses 128 WikiText-2 sentences.The default LLaMA-3-8B calibration takes about 0.9 hours on one GPU; weights and activations use per-channel and per-token symmetric quantization, respectively.Source: §4.1 · pp. 5–6

4.2. Main Results

  • On LLaMA-3-8B, FLATQUANT reduces PPL from 7.39 with SpinQuant to 6.98, and on LLaMA-2-70B achieves PPL just 0.23 above the FP16 baseline.With RTN, FLATQUANT performs comparably to GPTQ while taking significantly less calibration time.Source: §4.2 · p. 6
  • 0.94% accuracy loss on LLaMA-3-70B and 2.00% on LLaMA-3-8B show FlatQuant narrows the gap to FP16 on zero-shot QA tasks.The authors also report that RTN is comparable to GPTQ for FlatQuant.Source: §4.2 · p. 6, Table 2 · p. 7

4.3. Inference Latency

  • With kernel fusion, FlatQuant reaches 2.30x prefill and 1.76x decoding speedup over QuaRot at batch size 64.The evaluation uses a 2048-token prefill and 256-token decoding sequence.Source: §4.3 · p. 7, Figure 4 · p. 8
  • Equal-sized 64 × 64 Kronecker factors maximize prefill speedup, while factor sizes have limited impact on WikiText2 perplexity.Speedup decreases quickly when n2 exceeds 64, attributed to irregular activation memory access.Source: §4.3 · p. 8, Figure 5 · p. 8
  • Five online transformations add a 0.07x end-to-end slowdown, compared with 0.26x for QuaRot’s three Hadamard transformations.Pd contributes 0.04x slowdown and Po 0.01x; Pa and Pug each contribute less than 0.01x.Source: §4.3 · p. 8, Figure 6 · p. 8

4.4. Discussions

  • On LLaMA-3-8B, learnable transformations reduce WikiText-2 PPL from 1266.60 to 8.50; per-channel scaling and learnable clipping further improve PPL by 0.55 and 0.84.The ablation compares learnable transformations, per-channel scaling, and learnable clipping thresholds against an RTN baseline.Source: Table 3 · p. 8, §4.4 · p. 9
  • As training loss decreases, channel distributions become flatter, and the authors report this contributes to smaller quantization error.Their flatness metric is the Euclidean distance between the observed channel-magnitude distribution and an ideal flat distribution.Source: §4.4 · p. 9, Figure 7 · p. 9

5. Conclusions

FlatQuant uses learnable affine transformations to promote flatness in weights and activations, with Kronecker products and kernel fusion to reduce inference overhead. The authors report less than 1% accuracy drop for LLaMA-3-70B W4A4 quantization and speedups over FP16 inference.

5. Conclusions

  • FlatQuant reports less than 1% accuracy drop for W4A4 quantization on LLaMA-3-70B and up to 2.3x prefill and 1.7x decoding speedup over FP16.Kernel fusion integrates affine transformation and quantization.Source: §5 · p. 9

Impact Statement

The impact statement presents quantization as a potential way to lower LLM operational costs and energy consumption, while noting that it does not address biases in training data. The related-work discussion covers scaling, Hadamard and orthogonal transformations, and affine transformations for quantization.

Impact Statement

  • Quantization improvements may lower operational costs and energy consumption, but do not address biases in training data.The impact statement calls for careful consideration and responsible deployment to mitigate ethical risks and support fair use.Source: Impact Statement · p. 9

A. Related Work

  • Per-channel scaling methods may impair weight quantization and struggle in low-bit settings such as W4A4.The discussion includes SmoothQuant, channel-wise shifting, and OmniQuant’s learned scaling and shifting.Source: A. Related Work · p. 13
  • Hadamard transformations redistribute outliers across channels, while SpinQuant uses learnable orthogonal matrices with model-level loss.QuaRot applies Hadamard transformation in the LLM W4A4 post-training quantization setting.Source: A. Related Work · p. 13
  • AffineQuant learns full-size diagonally dominant matrices but applies transformations to only a small fraction of linear layers because of matrix-multiplication overhead.FlatQuant is described as using fast, learnable affine transformations without those limitations.Source: A. Related Work · p. 13
  • Other quantization settings use activation-aware scaling in AWQ, rotations or Hadamard transforms in QUIP and QUIP#, and block-diagonal Hadamard transforms in fully quantized training.Source: A. Related Work · p. 13

B.1. Matrix Inversion and Training Cost

FlatQuant’s training requires accurate computation of the inverse affine transformation matrix, but direct FP16 inversion is imprecise. The paper uses SVD with AMP as its default recipe, reducing calibration costs while retaining nearly lossless accuracy in most cases.

B.1. Matrix Inversion and Training Cost

  • At FP16 precision, direct inversion leaves off-diagonal elements of PP−1 around 1 × 10−3, negatively affecting early training.Training in FP32 avoids the precision problem but increases training time and GPU memory use.Source: §B.1 · p. 13
  • AMP with SVD reduces training time and memory usage by 50% while maintaining nearly lossless accuracy in most cases.SVD reduces the off-diagonal elements of PP−1 to 1 × 10−6 at FP16 precision, enabling AMP during calibration.Source: §B.1 · p. 14
  • FP32 training takes more than twice as long as AMP training and uses 1.28x more GPU memory, with relatively close performance.The authors select SVD with AMP as the default, while noting that full precision is necessary for some models or extremely low-bit scenarios because AMP may have numerical errors.Source: §B.1 · p. 14, Table 4 · p. 14
  • FLATQUANT requires significantly less calibration time and computational resources than SpinQuant and QAT methods while delivering superior performance.Weight-only calibration takes less time than weight-activation calibration, and FLATQUANT does not require GPTQ to achieve optimal performance.Source: §B.1 · p. 14, Table 5 · p. 14

B.2. Overhead Analysis of Affine Transformations

The online affine transformations are designed to add modest inference costs, while kernel fusion combines transformation and quantization operations. The implementation also includes alternate designs for cases where thread-block SRAM cannot hold intermediate results.

B.2. Overhead Analysis of Affine Transformations

  • For LLaMA-2-7B at sequence length 2048, online transformations account for about 2.61% of FP16 model FLOPs and add about 3.41MB of inference memory.The FLOPs estimate assumes the hidden and intermediate dimensions decompose into square factors.Source: §B.2 · pp. 14–15

B.3. Detailed Design of Kernel Fusion

  • The default fused kernel combines affine transformations and quantization, with revised designs for cases where shared memory cannot hold intermediate tensors.Figure 8 depicts the default design and two SRAM-constrained corner cases; the corner-case designs tile computations to manage the memory constraint.Source: Figure 8 · p. 15, §B.3 · p. 15
  • Kernel fusion achieves 1.5x–3x prefill speedup and 1.2x–4x decoding speedup across tested hidden dimensions and batch sizes.The default design fits the tested settings on an NVIDIA RTX 3090; tested corner cases with hidden size 28762 each achieved 2.3x speedup.Source: §B.3 · p. 16, Table 6 · p. 16

C.1. Results on Other LLM Architectures

FlatQuant is evaluated on instruction-tuned and MoE models, alternative quantization settings, ablations, and inference memory. The reported results show strong performance across these settings, with stable calibration-set results and memory use comparable to INT4 quantization.

C.1. Results on Other LLM Architectures

  • FlatQuant outperforms QuaRot on LLaMA-3.1-8B-Instruct, while Qwen-2.5-Instruct results remain close to FP16 and DeepSeek models retain strong 4-bit performance.Qwen averages are 68.62 for 7B and 74.89 for 32B, with the latter reported as a 0.21% drop; the DeepSeek evaluation covers two 671B-parameter MoE models.Source: §C.1 · p. 17, Table 7 · p. 17, Table 8 · p. 17

C.2. Results on MT-Bench

  • On MT-Bench, FlatQuant consistently outperforms QuaRot with GPTQ across categories and matches FP16 on math, exceeding QuaRot by 1.9 points.The model is LLaMA-3.1-8B-Instruct, and GPT-4o is used as evaluator.Source: §C.2 · p. 17, Table 10 · p. 18

C.3. Extension to More Quantization Settings

  • FlatQuant performs strongly in weight-only, KV-cache, and 3-bit settings, and transformations learned for W4A4KV4 transfer to other settings.With keys and values both quantized to 2 bits, it beats QuaRot by 2.57 in perplexity on LLaMA-2-7B; the authors still describe 4-bit quantization as the current better balance of efficiency and degradation.Source: §C.3 · pp. 18–19, Table 13 · p. 19

C.4. Detailed Ablation Study

  • Learnable transformations play a central role in the ablations, while other components further improve performance when added on top.The ablation study is on LLaMA-3-8B.Source: §C.4 · p. 19, Table 16 · p. 20

C.5. Additional Discussions

  • Calibration results are stable across WikiText-2, C4, and Pile, while post-transformation clipping and selective higher-bit layers improve practical quantization performance.WikiText-2 PPL ranges from 6.98 to 7.04, C4 PPL from 11.05 to 11.13, and QA accuracy from 71.04% to 71.23%; W8A8 is applied to the top five important layers and all down-projection layers.Source: §C.5 · pp. 20–21, Table 17 · p. 20, Table 18 · p. 20

C.6. Experiment Details of Figure 5

  • The Figure 5 setup decomposes hidden dimension 4096 into n1 × n2 and fixes the intermediate-dimension decomposition at 64 × 172 for its evaluations.PPL is evaluated by quantizing only the last Transformer block; speedup uses naive PyTorch matrix multiplication rather than the online transformation kernel.Source: §C.6 · p. 21

C.7. Analysis of Inference Memory

  • FlatQuant matches INT4 decoding memory across KV-cache lengths and reduces memory use by over 3.3× versus FP16.The measurements are peak memory during single-token decoding on one LLaMA-2-7B layer at batch size 1; the authors report no additional memory cost from online processing.Source: §C.7 · p. 21, Table 20 · p. 21

C.8. Additional Analysis of Inference Latency

FlatQuant’s latency results vary with batch size and sequence length, and the paper compares its transformations with AffineQuant and earlier distribution-smoothing methods. The reported visualizations show flatter weight and activation distributions after FlatQuant’s affine transformations.

C.8. Additional Analysis of Inference Latency

  • At batch size 64, FlatQuant reaches 2.30x prefill and 1.76x decoding speedup, with 0.07x speedup loss versus naive INT4 for both stages.For batch sizes below 16, quantization overhead outweighs KV-cache memory benefits during decoding, resulting in less than 1x speedup for both INT4 and FlatQuant.Source: §C.8 · p. 22, Figure 9 · p. 22, Figure 10 · p. 22
  • Across LLaMA-3-8B sequence lengths, FlatQuant achieves 2.12× prefill speedup at length 2048 and 1.80× at length 16384, and decoding closely approaches INT4 while surpassing QuaRot.Prefill figures are for batch size 1; decoding results are for batch size 64 and consistently surpass QuaRot across KV-cache lengths.Source: §C.8 · p. 22, Table 21 · p. 23, Table 22 · p. 23

C.9. Comparison with AffineQuant

  • FlatQuant’s Kronecker-structured transformations remain comparable in accuracy to full-size affine transformations and can be applied efficiently across all linear layers.AffineQuant’s full-size matrices are described as feasible only for output projections, with simpler per-channel scaling used for other layers.Source: §C.9 · p. 23

D.1. More Visualizations of Weight and Activation Distributions

  • Visualizations show that FlatQuant produces flatter weight and activation distributions than per-channel scaling or Hadamard transformation, which leave outlier channels.The remaining outliers are especially evident in activations.Source: §D.1 · p. 24

D.2. More Visualizations of Quantization Error Landscapes

The visualizations compare quantization-error landscapes across LLaMA models, showing flatter landscapes with FlatQuant than with per-channel scaling or Hadamard transformation.

D.2. More Visualizations of Quantization Error Landscapes

  • The per-channel scaling baseline follows SmoothQuant and is applied only to inputs of self-attention and feed-forward modules.The Hadamard comparison replaces FlatQuant’s affine transformation with a fixed Hadamard transformation; quantization settings match Section 4.1.Source: §D.2 · p. 24
  • FlatQuant produces a flatter quantization-error landscape, suppressing errors at pivot tokens and easing error propagation compared with per-channel scaling and Hadamard transformation.The visualizations cover LLaMA models and use 128 randomly sampled C4 examples for average mean squared error.Source: §D.2 · p. 24

E. Limitations

The paper identifies unresolved calibration-set selection and limits its examination primarily to INT4, without testing integration with newer data types such as MXFP4.

E. Limitations

  • The optimal selection of calibration sets remains an open question, despite reported robustness across various data sources.The study follows previous work in building its calibration set.Source: E. Limitations · p. 24
  • The study focuses primarily on INT4 and does not examine FlatQuant’s integration with newer data types such as MXFP4.The paper notes that MXFP4 may offer advantages over INT4.Source: E. Limitations · p. 24
Loading 2410.09426v4…