1. Bibliographic Information

1.1. Title

The central topic of the paper is FlatQuant: Flatness Matters for LLM Quantization. It proposes a novel approach to improve Large Language Model (LLM) quantization by enhancing the flatness of weights and activations.

1.2. Authors

The authors are Yuxuan Sun, Ruikang Liu, Haoli Bai, Han Bao, Kang Zhao, Yuening Li, Jiaxin Hu, Xianzhi Yu, Lu Hou, Chun Yuan, Xin Jiang, Wulong Liu, and Jun Yao. They are affiliated with various institutions, with a significant presence from what appears to be a research institution or university (indicated by '1', '2', '3' affiliations, typically denoting different departments or organizations). For example, Yuxuan Sun and Haoli Bai are from affiliation '1', Ruikang Liu from '2', and Yuening Li from '3'. The specific organizations are not explicitly named in the provided text but are typical for academic research collaborations.

1.3. Journal/Conference

The paper is published as a preprint on arXiv (arXiv preprint arXiv:2410.09426v4). While arXiv is not a peer-reviewed journal or conference, it is a highly influential open-access repository for scholarly articles, particularly in physics, mathematics, computer science, quantitative biology, quantitative finance, statistics, electrical engineering and systems science, and economics. Papers on arXiv are widely circulated and serve as a crucial platform for disseminating cutting-edge research before or during the peer-review process, allowing for rapid sharing of findings and feedback from the scientific community.

1.4. Publication Year

The paper was published at 2024-10-12T08:10:28.000Z (UTC), which corresponds to October 12, 2024.

1.5. Abstract

This paper introduces FlatQuant (Fast and Learnable Affine Transformation), a new post-training quantization (PTQ) method for compressing and accelerating Large Language Models (LLMs). The core idea is that flatness (i.e., less steep and dispersed distributions) of weights and activations is crucial for minimizing quantization error, especially given the presence of outliers in LLMs. While prior methods use pre-quantization transformations like per-channel scaling and Hadamard transformations, the authors observe these can still result in sub-optimal distributions. FlatQuant proposes to learn optimal affine transformations for each linear layer, calibrated efficiently within hours using a lightweight objective function. To mitigate runtime overhead, it employs a Kronecker product with two lightweight matrices and fuses all operations into a single kernel. Extensive experiments demonstrate that FlatQuant achieves state-of-the-art results, such as less than 1% accuracy drop for W4A4 quantization on the LLaMA-3-70B model, outperforming SpinQuant by 7.5%. Furthermore, it provides up to 2.3x prefill speedup and 1.7x decoding speedup compared to the FP16 model.

1.6. Original Source Link

The official source link is: https://arxiv.org/abs/2410.09426v4 The PDF link is: https://arxiv.org/pdf/2410.09426v4.pdf The paper is available as a preprint on arXiv.

2. Executive Summary

2.1. Background & Motivation

The rapid growth in the size of Large Language Models (LLMs) (e.g., billions of parameters) has led to significant computational and memory overhead during inference. This makes deploying LLMs on resource-constrained devices or in high-throughput scenarios challenging and expensive. Quantization is a highly effective technique to address this by reducing the precision of model parameters (weights) and intermediate computations (activations) from full-precision (e.g., FP16 or FP32) to lower-precision formats (e.g., INT4 or INT8).

The core problem in LLM quantization is minimizing quantization error. This error arises because low-bit quantization uses a limited number of discrete points to represent a continuous range of values. A critical challenge for LLMs, specifically, is the presence of outliers—extreme values in both weights and activations. These outliers can cause large quantization errors when mapped to the same low-bit integer points, leading to significant performance degradation.

Prior research has explored various pre-quantization transformations to suppress outliers and flatten distributions, such as per-channel scaling (e.g., SmoothQuant) and Hadamard transformation (e.g., QuaRot, SpinQuant). While these methods offer improvements, the authors observe that they often remain sub-optimal, with transformed weights and activations still exhibiting steep and dispersed distributions. This sub-optimality limits the achievable accuracy and inference speedup.

The paper's entry point and innovative idea is to explicitly focus on enhancing the flatness of weight and activation distributions through learnable affine transformations. By making these transformations fast and adaptive for each linear layer, FlatQuant aims to systematically mitigate outliers and create distributions that are inherently easier to quantize with minimal error. This adaptive, learnable approach, combined with efficient implementation strategies, seeks to overcome the limitations of fixed or less flexible pre-quantization methods.

2.2. Main Contributions / Findings

The paper makes several significant contributions:

• Highlighting Flatness Significance: It rigorously demonstrates that achieving flatter distributions of weights and activations is crucial for effective LLM quantization, as it reduces quantization error and controls its propagation across Transformer layers.
• Introducing FlatQuant: The paper proposes FlatQuant, a novel post-training quantization method. This approach uses fast and learnable affine transformations optimized specifically for each linear layer. Empirically, FlatQuant is shown to significantly enhance the flatness of both weights and activations in LLMs.
• New State-of-the-Art Accuracy: FlatQuant establishes new state-of-the-art results for LLM quantization accuracy. Notably, it is presented as the first method to achieve less than 1% accuracy drop with simple round-to-nearest (RTN) W4A4 quantization on the LLaMA-3-70B model, a highly challenging benchmark. This significantly surpasses previous methods like SpinQuant by a large margin (e.g., 7.5%).
• Efficient Kernel Design and Speedup: The authors design an efficient kernel that fuses the affine transformations and quantization operations. This design minimizes global memory access and kernel launch overhead, leading to substantial inference speedups. Specifically, FlatQuant achieves up to 2.3x prefill speedup and 1.7x decoding speedup under W4A4 quantization when compared to an FP16 baseline.
• Robustness and Versatility: FlatQuant demonstrates robustness across different LLM families (LLaMA, Qwen, DeepSeek), various tasks (language modeling, QA, MT-bench), and is compatible with different quantization settings (weight-only, KV cache quantization). Its calibration process is also efficient, taking only hours on a single GPU.

3. Prerequisite Knowledge & Related Work

3.1. Foundational Concepts

To understand FlatQuant, a novice reader should grasp the following foundational concepts:

• Large Language Models (LLMs): These are advanced artificial intelligence models, typically based on the Transformer architecture, trained on vast amounts of text data. They are designed to understand, generate, and process human language, excelling at tasks like text generation, question answering, summarization, and translation. Their parameter count (weights and biases) can range from billions to trillions, making them computationally intensive and memory hungry.
• Quantization: This is a model compression technique that reduces the precision of numbers used to represent a neural network's weights and activations. For instance, FP16 (16-bit floating-point) or FP32 (32-bit floating-point) values might be converted to INT4 (4-bit integer) or INT8 (8-bit integer). The primary goals are to:
  • Reduce Memory Footprint: Lower-precision numbers require less storage, allowing larger models to fit into available memory or run on devices with limited resources.
  • Accelerate Inference: Low-bit arithmetic operations are generally faster and consume less power than high-precision floating-point operations, leading to quicker model responses.
  • Quantization Error: The process of mapping continuous high-precision values to discrete low-precision values inevitably introduces quantization error. The goal of quantization research is to minimize this error so that the low-precision model's performance (e.g., accuracy, perplexity) remains close to that of its full-precision counterpart.
• Weights and Activations:
  • Weights ($\mathbf{W}$): These are the learnable parameters within a neural network. They determine the strength of connections between neurons and are adjusted during the training process.
  • Activations ($\mathbf{X}$): These are the outputs of neurons or layers at each step of the network's computation. They represent the data flowing through the network. In an LLM, activations are typically the output of a linear layer or the attention mechanism.
• Outliers: In the context of LLMs, outliers refer to a small number of extremely large or small values present in the weight matrices or activation tensors. These extreme values are problematic for quantization because a fixed-range, equally-spaced quantization scheme will allocate many of its limited discrete points to represent these few outliers, leaving fewer points for the majority of "normal" values. This leads to high quantization error for the non-outlier values, significantly degrading model performance.
• Flatness of Distributions: A flat distribution (also referred to as having low kurtosis) means that the values within a tensor (weights or activations) are relatively uniformly distributed across their range, without excessively sharp peaks or dispersed extreme values. Conversely, a steep or dispersed distribution implies the presence of many outliers or values concentrated in a narrow range with a long tail. For quantization, a flatter distribution is highly desirable because:
  • It allows equally spaced quantization points to cover the value range more effectively.
  • It minimizes the impact of outliers, as values are more evenly spread out.
  • This directly leads to lower quantization error.
• Transformer Architecture: The foundational architecture for most modern LLMs. Key components include:
  • Linear Layers: Also known as fully connected layers, these perform matrix multiplication (e.g., $\mathbf{Y} = \mathbf{X}\mathbf{W}^{\top} + \mathbf{b}$). They are a primary target for quantization.
  • Self-Attention: A mechanism that allows the model to weigh the importance of different parts of the input sequence when processing each token. It involves Query (Q), Key (K), and Value (V) projections, which are typically implemented using linear layers.
  • Feed-Forward Networks (FFN): A stack of linear layers (often two or three) with non-linear activation functions, applied independently to each position in the sequence.
  • Layer Normalization (LayerNorm): A technique used to stabilize the training of deep neural networks by normalizing the inputs to each layer across the features.
  • Residual Connections: Add the input of a sub-layer to its output, helping to mitigate the vanishing gradient problem in deep networks.
  • KV Cache: During LLM inference, especially in the decoding stage (when generating one token at a time), the Key and Value tensors computed for previous tokens are stored in a cache (KV Cache) to avoid recomputing them. Quantizing the KV Cache is crucial for reducing memory usage, especially for long sequences.
• Perplexity (PPL): A common intrinsic evaluation metric for language models. It quantifies how well a probability model predicts a sample. A lower PPL indicates that the model is better at predicting the next word in a sequence, implying higher quality language generation.
• Accuracy: An extrinsic evaluation metric, typically used for classification tasks (like zero-shot QA). It measures the proportion of correctly predicted instances out of the total instances. Higher accuracy means better performance.
• Kronecker Product ($\otimes$): A mathematical operation that combines two matrices of arbitrary size to form a larger matrix. If $\mathbf{A}$ is an $m \times n$ matrix and $\mathbf{B}$ is a $p \times q$ matrix, their Kronecker product $\mathbf{A} \otimes \mathbf{B}$ is an $mp \times nq$ block matrix. FlatQuant uses this to construct a large affine transformation matrix from two smaller, lightweight matrices, significantly reducing computational and memory costs.
• Singular Value Decomposition (SVD): A powerful matrix factorization technique. Any real matrix $\mathbf{P}$ can be decomposed into $\mathbf{U} \mathbf{\Sigma} \mathbf{V}^{\top}$, where $\mathbf{U}$ and $\mathbf{V}$ are orthogonal matrices and $\mathbf{\Sigma}$ is a diagonal matrix of singular values. SVD is numerically stable and can be used to accurately compute the inverse of a matrix ($\mathbf{P}^{-1} = \mathbf{V} \mathbf{\Sigma}^{-1} \mathbf{U}^{\top}$), which is essential for FlatQuant's learnable transformations.
• Automatic Mixed Precision (AMP): A training technique that uses both FP16 and FP32 precision during model training. It typically performs most operations in FP16 for speed and memory efficiency, while keeping certain critical operations (e.g., loss calculation, weight updates) in FP32 to maintain numerical stability and prevent gradient underflow/overflow. This allows for faster training with less memory consumption.
• CUDA, Triton, CUTLASS, FlashInfer: These are technologies related to GPU programming and high-performance computing for deep learning:
  • CUDA: NVIDIA's parallel computing platform and programming model for GPUs.
  • Triton (OpenAI Triton): A domain-specific language and compiler for writing highly optimized GPU kernels. It simplifies the process of achieving high performance for deep learning operations, often outperforming manually optimized CUDA kernels. FlatQuant uses it for kernel fusion.
  • CUTLASS: A CUDA template library for highly optimized GEMM (General Matrix Multiply) operations on NVIDIA GPUs. It provides high-performance routines for various data types, including low-precision integers (INT4). FlatQuant adopts CUTLASS for INT4 matrix multiplication.
  • FlashInfer: A kernel library specifically designed for efficient LLM serving, including optimized KV cache operations and quantization. FlatQuant uses FlashInfer for KV cache quantization.

3.2. Previous Works

The paper discusses several categories of prior work related to LLM quantization:

• General LLM Quantization Methods:

  • GPTQ (Frantar et al., 2022): A widely used post-training quantization method for Generative Pre-trained Transformers. It performs weight-only quantization by sequentially quantizing weights layer-by-layer while minimizing the mean squared error (MSE) in the output of each layer. It uses a Hessian-based approach for efficient weight updates. GPTQ is often combined with other techniques for better results.
  • AWQ (Activation-aware Weight Quantization) (Lin et al., 2023): Focuses on the observation that only a small percentage of weights are salient (critical for performance) and sensitive to quantization. AWQ proposes to protect these salient weights from quantization error by scaling them, while quantizing the rest aggressively.
  • QUIK-4B (Ashkboos et al., 2023): An end-to-end 4-bit inference method for generative LLMs, aiming for high accuracy with very low precision.
  • QuIP (Chee et al., 2024) / QuIP# (Tseng et al., 2024): These methods propose 2-bit quantization with guarantees, utilizing techniques like Hadamard incoherence and lattice codebooks to achieve high performance at extremely low bit-widths.

• Pre-Quantization Transformations for Outlier Suppression:

  • Per-channel Scaling (e.g., SmoothQuant by Xiao et al., 2023; OmniQuant by Shao et al., 2023; Outlier Suppression+ by Wei et al., 2023): This is a popular technique to mitigate the impact of outliers in activations. The core idea is to move the quantization difficulty from activations to weights by applying a channel-wise scaling factor.
    • The linear layer operation is $\mathbf{Y} = \mathbf{X} \mathbf{W}^{\top}$.
    • With per-channel scaling, this is transformed to $\mathbf{Y} = (\mathbf{X} \mathrm{diag}(\mathbf{c})^{-1}) \cdot (\mathrm{diag}(\mathbf{c}) \mathbf{W}^{\top})$, where $\mathbf{c} \in \mathbb{R}^n$ is the channel-wise scaling factor.
    • The inverse scaling factor $\mathrm{diag}(\mathbf{c})^{-1}$ is applied to activations, and the scaling factor $\mathrm{diag}(\mathbf{c})$ is applied to weights.
    • The scaled weights $\mathrm{diag}(\mathbf{c}) \mathbf{W}^{\top}$ can be pre-computed and stored offline, so only $\mathbf{X} \mathrm{diag}(\mathbf{c})^{-1}$ incurs runtime overhead.
    • SmoothQuant specifically aims to balance the magnitudes of activations and weights to make both easier to quantize. The scaling factor $\mathbf{c}_j$ for channel $j$ is often determined by $\mathbf{c}_j = \mathrm{max}(|\mathbf{X}_j|^\alpha) / \mathrm{max}(|\mathbf{W}_j|^{1-\alpha})$, where $\alpha$ is a hyperparameter.
    • OmniQuant (Shao et al., 2023) further treats both scaling and shifting factors as learnable parameters.
  • Hadamard Transformation (e.g., QuaRot by Ashkboos et al., 2024; SpinQuant by Liu et al., 2024c; Training Transformers with 4-bit Integers by Xi et al., 2023): These methods use Hadamard matrices to rotate the channels of both activations and weights.
    • A Hadamard matrix $\mathbf{H} \in \{+1, -1\}^{n \times n}$ is an orthogonal matrix (i.e., $\mathbf{H}^{\top}\mathbf{H} = \mathbf{I}$).
    • The core idea is that $\mathbf{Y} = \mathbf{X} \mathbf{W}^{\top} = (\mathbf{X} \mathbf{H}) (\mathbf{H}^{\top} \mathbf{W}^{\top})$.
    • By multiplying activations with $\mathbf{H}$ and weights with $\mathbf{H}^{\top}$, outliers are redistributed more evenly across channels, leading to flatter distributions.
    • Similar to per-channel scaling, the transformed weight $\mathbf{W}_{\text{transformed}} = \mathbf{H}^{\top} \mathbf{W}^{\top}$ can be pre-computed offline. The online operation is $\mathbf{X} \mathbf{H}$.
    • QuaRot and SpinQuant build on this, often combining Hadamard rotations with other techniques like learnable rotations and modifications to LayerNorm for efficiency. SpinQuant specifically uses learned rotations to alleviate outliers.
  • Affine Transformation Quantization (AffineQuant by Ma et al., 2024): This work also explores affine transformations to improve quantization. It aims to find optimal invertible matrices $P$ for each layer to minimize quantization error. FlatQuant can be seen as building upon this idea but with a strong focus on efficiency (Kronecker product) and flatness.

3.3. Technological Evolution

The field of LLM quantization has rapidly evolved to address the growing size and computational demands of LLMs.

1. Early Methods (e.g., RTN): Started with simple Round-To-Nearest quantization, which often led to significant accuracy drops, especially for lower bit-widths.

2. Outlier-Aware Techniques: The discovery of outliers in LLM activations became a key focus. Methods like SmoothQuant introduced per-channel scaling to smooth distributions by shifting quantization difficulty from activations to weights.

3. Rotation-Based Transformations: Recognizing the limitations of simple scaling, Hadamard transformations and orthogonal transformations (e.g., QuaRot, SpinQuant) emerged. These methods rotate the data in the channel dimension to redistribute outliers and achieve better flatness, often combined with learnable components.

4. Hardware-Aware Optimization: Alongside algorithmic improvements, significant effort has gone into optimizing quantized inference on hardware. This includes designing custom kernels (e.g., using CUTLASS, Triton) and specialized libraries (FlashInfer) to make low-bit operations truly fast by reducing memory access overheads and maximizing GPU utilization.

5. Learnable & Adaptive Transformations: The latest trend, exemplified by AffineQuant and FlatQuant, involves learning the optimal transformations rather than relying on fixed or heuristically derived ones. This allows for greater adaptability to the unique characteristics of each LLM layer.

   FlatQuant fits into this evolution by pushing the boundaries of learnable transformations. It explicitly targets flatness as a primary objective for quantization error reduction and combines this with highly efficient implementation strategies (Kronecker product, kernel fusion) to ensure practical speedups without sacrificing accuracy. It represents a significant step towards deploying highly accurate, ultra-low-bit LLMs in real-world applications.

3.4. Differentiation Analysis

Compared to the main methods in related work, FlatQuant introduces several core differences and innovations:

• Explicit Focus on Flatness as Objective: While previous methods (e.g., SmoothQuant, Hadamard transformation methods) implicitly aim for flatter distributions by mitigating outliers, FlatQuant explicitly identifies flatness as a crucial factor for minimizing quantization error and directly optimizes for it. This objective-driven approach allows for more targeted and effective transformations.

• Learnable Affine Transformations per Layer:

  • Unlike per-channel scaling (e.g., SmoothQuant), which applies a scalar diagonal transformation, FlatQuant uses full affine transformation matrices ($\mathbf{P}$) that are learnable and layer-specific. This offers significantly more expressiveness and adaptability to the unique distribution characteristics of each individual linear layer within the LLM. SmoothQuant typically uses a fixed formula or a limited set of learnable scalars.
  • Compared to Hadamard transformations (e.g., QuaRot, SpinQuant), which use fixed orthogonal matrices or learned rotations, FlatQuant's affine transformations are more general. While Hadamard matrices are orthogonal, FlatQuant learns a broader class of invertible matrices, providing greater flexibility to shape distributions. Moreover, existing Hadamard-based methods often apply a single transformation across many layers or share transformations, whereas FlatQuant optimizes them per linear layer.
  • AffineQuant also uses affine transformations, but FlatQuant differentiates itself through its efficiency design and kernel fusion.

• Efficiency through Kronecker Product: A major innovation of FlatQuant is the use of Kronecker product ($\mathbf{P} = \mathbf{P}_1 \otimes \mathbf{P}_2$) to construct the large affine transformation matrix from two much smaller matrices. This dramatically reduces the memory footprint and computational overhead of the transformation compared to a full-size matrix $\mathbf{P}$. This addresses a key practical limitation of general affine transformations, making them viable for large LLMs.

• Optimized Kernel Fusion: FlatQuant integrates the affine transformation and quantization operations into a single fused kernel using OpenAI Triton. This kernel fusion is crucial for minimizing global memory access (a common bottleneck for memory-bound operations) and kernel launch overhead, which translates directly into superior end-to-end inference speedup. This hardware-aware optimization sets it apart from methods that might perform transformations and quantization as separate, sequential operations.

• Compatibility and Robustness: FlatQuant is shown to be highly compatible with various quantization techniques (e.g., learnable clipping, RTN, GPTQ) and can be applied to different settings (e.g., weight-only, KV cache quantization). Its ability to achieve high accuracy with simple RTN quantization reduces calibration complexity and time compared to methods that heavily rely on GPTQ for good performance.

  In essence, FlatQuant moves beyond fixed or simple transformations by learning more powerful affine transformations for each layer, while simultaneously ensuring these transformations are computationally and memory-efficient through mathematical decomposition (Kronecker product) and hardware-optimized kernel design (kernel fusion). This combined approach enables it to achieve superior accuracy-speed trade-offs for LLM quantization.

4. Methodology

4.1. Principles

The core principle behind FlatQuant is that flatness in the distributions of weights and activations is paramount for minimizing quantization error in Large Language Models (LLMs). When values are tightly clustered or contain extreme outliers, mapping them to a limited set of equally spaced quantization points (as is common in many quantization schemes) leads to significant information loss and performance degradation. FlatQuant's key insight is to systematically apply fast and learnable affine transformations to each linear layer within the LLM to explicitly enhance this flatness.

The theoretical basis is rooted in the understanding that transformations can reshape data distributions. By learning optimal affine matrices, the method can rotate, scale, and shift the data in a way that spreads out outliers and makes the overall distribution more uniform across the available quantization range. This reduces the mean squared error (MSE) between the original and quantized values. To ensure these powerful transformations do not introduce prohibitive runtime overhead, FlatQuant employs two crucial efficiency mechanisms:

1. Kronecker Product Decomposition: Instead of learning a single large transformation matrix for each high-dimensional layer, it decomposes this matrix into a Kronecker product of two much smaller, lightweight matrices. This significantly reduces the parameter count and computational cost of applying the transformation.

2. Kernel Fusion: The transformation operations, which are memory-bound, are fused together with the quantization steps into a single custom kernel. This minimizes global memory accesses and kernel launch overhead, ensuring that the algorithmic benefits translate into practical inference speedups.

   In essence, FlatQuant seeks to intelligently preprocess the numerical data within an LLM to be maximally amenable to low-bit quantization, balancing expressive power with computational efficiency.

4.2. Core Methodology In-depth (Layer by Layer)

4.2.1. Preliminaries on LLM Quantization

LLM inference generally consists of two stages:

1. Prefill Stage: Processing the input sequence to build a Key-Value cache (KV Cache) layer by layer.

2. Decoding Stage: Autoregressively generating new tokens based on the accumulated KV Cache.

   Quantization is applied to reduce the precision of weights $\mathbf{W} \in \mathbb{R}^{m \times n}$ and activations $\mathbf{X} \in \mathbb{R}^{k \times n}$ in linear layers (which typically perform the operation $\mathbf{Y} = \mathbf{X} \mathbf{W}^{\top}$), and optionally to the KV Cache.

For $b$-bit weight quantization, the process can be represented as: $$\hat{\mathbf{W}} = \mathcal{Q}_b(\mathbf{W}) = s \cdot \Pi_{\Omega(b)}(\mathbf{W}/s)$$ Here:

• $\hat{\mathbf{W}}$ represents the quantized weight matrix.
• $\mathcal{Q}_b(\cdot)$ is the general $b$-bit quantization function.
• $s$ is the quantization step size (or scale), which determines the interval between quantization points. It effectively sets the range of values that can be represented.
• $\Pi(\cdot)$ is the projection function that maps a real-valued number to the nearest integer quantization point. For round-to-nearest (RTN) quantization, it simply rounds to the closest integer.
• $\Omega(b) = \{0, 1, ..., 2^b - 1\}$ is the set of $b$-bit integer points. For example, for $b=4$, it's $\{0, 1, ..., 15\}$. For simplicity, $\mathcal{Q}(\cdot)$ is used to denote the general quantization function.

As LLMs are known to have persistent outliers in activations, two common pre-quantization transformations are discussed:

• Per-channel Scaling: This method addresses outliers in input activations $\mathbf{X}$ by applying channel-wise scaling. The original matrix multiplication $\mathbf{Y} = \mathbf{X} \mathbf{W}^{\top}$ is reformulated as: $ \mathbf{Y} = (\mathbf{X} \mathrm{diag}(\mathbf{c})^{-1}) \cdot (\mathrm{diag}(\mathbf{c}) \mathbf{W}^{\top}) $ Here:

  • $\mathbf{c} \in \mathbb{R}^n$ is the channel-wise scaling factor vector.
  • $\mathrm{diag}(\mathbf{c})$ creates a diagonal matrix with elements of $\mathbf{c}$ on its diagonal.
  • The term $(\mathbf{X} \mathrm{diag}(\mathbf{c})^{-1})$ scales the activations by $1/\mathbf{c}_j$ for each channel $j$.
  • The term $(\mathrm{diag}(\mathbf{c}) \mathbf{W}^{\top})$ scales the weights by $\mathbf{c}_j$ for each channel $j$.
  • The scaling vector $\mathbf{c}$ is designed to smooth the activations by considering the magnitudes of both activations and weights, e.g., $\mathbf{c}_j = \mathrm{max}(|\mathbf{X}_j|^\alpha) / \mathrm{max}(|\mathbf{W}_j|^{1-\alpha})$. The scaled weights $\mathrm{diag}(\mathbf{c}) \mathbf{W}^{\top}$ can be merged and pre-processed offline, eliminating runtime computation overhead for the weight scaling.
  • Variants like Outlier Suppression+ (Wei et al., 2023) introduce channel-wise shifting, and OmniQuant (Shao et al., 2023) makes $\mathrm{diag}(\mathbf{c})$ and the shift $z$ learnable parameters.

• Hadamard Transformation: This method uses Hadamard matrices $\mathbf{H} \in \{+1, -1\}^{n \times n}$ to rotate the channels of both activations and weights, thereby redistributing outliers. Due to the orthogonality of Hadamard matrices ($\mathbf{H}^{\top}\mathbf{H} = \mathbf{I}$), the following equivalence holds: $ \mathbf{Y} = \mathbf{X} \mathbf{W}^{\top} = (\mathbf{X} \mathbf{H}) (\mathbf{H}^{\top} \mathbf{W}^{\top}) $

  • Activations $\mathbf{X}$ are transformed by multiplying with $\mathbf{H}$ (i.e., $\mathbf{X} \mathbf{H}$).
  • Weights $\mathbf{W}^{\top}$ are transformed by multiplying with $\mathbf{H}^{\top}$ (i.e., $\mathbf{H}^{\top} \mathbf{W}^{\top}$).
  • The transformed weights $\mathbf{H}^{\top} \mathbf{W}^{\top}$ can also be pre-processed offline. This rotation aims to more effectively eliminate outliers by spreading them across all channels.

4.2.2. The Flatness for Quantization

The paper motivates its approach by emphasizing the importance of flatness for quantization. Flat tensors (those with low kurtosis) are intuitively easier to quantize after outlier removal. Figure 1 visually demonstrates this:

As shown in Figure 1, FlatQuant (green lines) consistently achieves flatter distributions for both weights and activations compared to other methods. The original distributions (gray) are very steep and dispersed, indicating severe outliers. Per-channel scaling (orange) flattens activations but can lead to steeper weight envelopes. Hadamard transformation (blue) generally performs better than per-channel scaling but can still be unsatisfactory.

The following figure (Figure 1 from the original paper) shows distributions of weights and inputs from LLaMA-3-8B and LLaMA-3-70B:

该图像是一个图表，展示了 LLaMA-3-8B 和 LLaMA-3-70B 模型中权重和输入的分布。左侧显示的是第 10 层 Transformer 的 $W_o$（a）和 $X_o$（b）；右侧则是第 30 层 Transformer 的 $W_g$（c）和 $X_g$（d）。数据按照信号幅值（Frobenius 范数）降序排列，图中比较了原始数据、每通道缩放、Hadamard 变换和 FlatQuant 方法的效果。 Figure 1: Distributions of weights and inputs from LLaMA-3-8B and LLaMA-3-70B, sorted by the channel magnitude (i.e., the Frobenius norm) in descending order. In a Transformer layer, $\mathbf{W}_o$ and $\mathbf{X}_o$ denote the weight matrix and input of the output projection layer in the self-attention layer, respectively. $\mathbf{W}_g$ and $\mathbf{X}_g$ denote the weight and input of the gated linear layer of the feed-forward network, respectively. More visualizations can be found in Appendix D.

Beyond individual layer flatness, FlatQuant also addresses the flatness of the quantization error landscape. Quantization error accumulates and propagates across Transformer layers. Figure 2 illustrates the mean squared error (MSE) landscape. Initially, massive quantization errors occur at pivot tokens (early tokens in a sequence), which are rich in outliers. Neither per-channel scaling nor Hadamard transformation effectively mitigates these initial errors (Figure 2a-2b). In contrast, FlatQuant (Figure 2c) shows significantly lower error at pivot tokens. Furthermore, while quantization error generally increases layer-wise, FlatQuant is shown to be most effective at controlling this error propagation (Figure 2d), followed by Hadamard transformation and then per-channel scaling.

The following figure (Figure 2 from the original paper) shows the mean squared quantization error landscape across Transformer layers and input sequence tokens for LLaMA-3-8B:

Figure 2: The mean squared quantization error (MSE) landscapes across Transformer layers and input sequence tokens for LLaMA-3-8B. The MSE is normalized with respect to the FP16 baseline. FlatQuant shows much lower error at these pivot tokens from Figure 2c. According to Figure 2d, FlatQuant is the best in controlling the error propagation.

4.2.3. Fast and Learnable Affine Transformation

FlatQuant's goal is to find the optimal affine transformation for each linear layer. For a linear layer computing $\mathbf{Y} = \mathbf{X} \mathbf{W}^{\top}$, one ideally seeks an optimal invertible matrix $\mathbf{P}^* \in \mathbb{R}^{n \times n}$ by minimizing the Frobenius norm of the difference between the full-precision output and the quantized transformed output: $$\mathbf{P}^* = \arg \min_{\mathbf{P}} \| \mathbf{Y} - \mathcal{Q}(\mathbf{X} \mathbf{P}) \mathcal{Q}(\mathbf{P}^{-1} \mathbf{W}^{\top}) \|_F^2$$ Here:

• $\mathbf{P}^*$ is the optimal invertible transformation matrix.
• $\mathbf{Y}$ is the original full-precision output of the linear layer.
• $\mathcal{Q}(\mathbf{X} \mathbf{P})$ is the quantized version of the activations transformed by $\mathbf{P}$.
• $\mathcal{Q}(\mathbf{P}^{-1} \mathbf{W}^{\top})$ is the quantized version of the weights transformed by the inverse of $\mathbf{P}$.
• $\|\cdot\|_F^2$ is the squared Frobenius norm, which measures the mean squared error (MSE). The transformed weights $\mathbf{P}^{-1} \mathbf{W}^{\top}$ can be pre-computed offline. However, maintaining and applying individual full-size $\mathbf{P}$ matrices for each layer would be computationally and memory-wise expensive.

Kronecker Product

FlatQuant addresses the cost of full-size affine transformations by using a Kronecker product of two lightweight matrices as an efficient substitute. Specifically, the transformation matrix $\mathbf{P}$ is constructed as $\mathbf{P} = \mathbf{P}_1 \otimes \mathbf{P}_2$, where $\mathbf{P}_1 \in \mathbb{R}^{n_1 \times n_1}$ and $\mathbf{P}_2 \in \mathbb{R}^{n_2 \times n_2}$ are invertible matrices of smaller sizes, such that the original dimension $n = n_1 n_2$.

Leveraging the vectorization trick of the Kronecker product (i.e., $\mathrm{vec}(\mathbf{V}) (\mathbf{P}_1 \otimes \mathbf{P}_2) = \mathrm{vec}(\mathbf{P}_1^{\top} \mathbf{V} \mathbf{P}_2)$ for some $\mathbf{V} \in \mathbb{R}^{n_1 \times n_2}$), the matrix multiplication in Equation 2 can be rewritten as: $$\mathcal{Q}(\mathbf{X} \mathbf{P}) \mathcal{Q}(\mathbf{P}^{-1} \mathbf{W}^{\top}) = \mathcal{Q}(\mathbf{P}_{1}^{\top} \times_{1} \tilde{\mathbf{X}} \times_{2} \mathbf{P}_{2}) \times \mathcal{Q}(\mathbf{P}_{1}^{-1} \times_{1} \tilde{\mathbf{W}} \times_{2} (\mathbf{P}_{2}^{-1})^{\top})^{\top}$$ Here:

• $\tilde{\mathbf{X}} \in \mathbb{R}^{k \times n_1 \times n_2}$ and $\tilde{\mathbf{W}} \in \mathbb{R}^{m \times n_1 \times n_2}$ are the activations $\mathbf{X}$ and weights $\mathbf{W}$ reshaped accordingly from their original matrix forms to 3D tensors compatible with the Kronecker product transformation.
• $\times_i$ denotes the reduction (matrix multiplication) over the $i$-th axis (dimension). This implies the operations are applied across the reshaped dimensions of $\tilde{\mathbf{X}}$ and $\tilde{\mathbf{W}}$.
• The terms $\mathbf{P}_{1}^{\top} \times_{1} \tilde{\mathbf{X}} \times_{2} \mathbf{P}_{2}$ and $\mathbf{P}_{1}^{-1} \times_{1} \tilde{\mathbf{W}} \times_{2} (\mathbf{P}_{2}^{-1})^{\top}$ represent the transformations applied to the reshaped activations and weights, respectively.
• The final $(\cdot)^{\top}$ indicates a transpose operation, restoring the matrix multiplication structure. This design significantly saves memory (up to $n/2$ times) and computation (up to $\sqrt{n}/2$ times) when $n_1 = n_2 = \sqrt{n}$. For instance, for $n=8192$, the optimal configuration is $(n_1^*, n_2^*) = (64, 128)$. These affine transformations are very lightweight, adding minimal FLOPs and memory.

Per-channel Scaling

To further enhance outlier balancing, FlatQuant explicitly incorporates a learnable scaling vector $\mathrm{diag}(\mathbf{c}) \in \mathbb{R}^n$ before the pre-quantization transformation. This vector can be merged with preceding layer normalization or linear layers, incurring no additional inference overhead, similar to SmoothQuant.

Learnable Clipping Thresholds

FlatQuant also includes learnable clipping thresholds $\alpha_w, \alpha_a \in (0, 1)$ after sigmoid functions for both weight and activation quantization in each linear layer, as well as for the KV cache. While grid search is common for finding these thresholds, FlatQuant learns them, leading to better results. These parameters are layer-specific and optimized jointly with $\mathbf{P}$ and $\mathrm{diag}(\mathbf{c})$.

The Training Objective

FlatQuant uses a post-training quantization approach, minimizing mean squared error (MSE) sequentially for each Transformer block over a small calibration dataset (e.g., 128 sentences). The objective for the $l$-th Transformer block is: $$\min_{\Theta} \Big\| \mathcal{F}_l \big( \mathbf{X} \big) - \hat{\mathcal{F}}_l \big( \mathbf{X} ; \Theta \big) \Big\|_F^2$$ Here:

• $\mathcal{F}_l(\cdot)$ denotes the original full-precision output of the $l$-th Transformer block.
• $\hat{\mathcal{F}}_l(\mathbf{X} ; \Theta)$ denotes the output of the $l$-th quantized Transformer block with learnable parameters $\Theta$.
• $\Theta = \{\mathbf{P}, \mathbf{c}, \alpha_a, \alpha_w\}$ represents all learnable parameters within that block: the affine transformation matrices $\mathbf{P}$ (implicitly $\mathbf{P}_1, \mathbf{P}_2$), the scaling vector $\mathbf{c}$, and the activation ($\alpha_a$) and weight ($\alpha_w$) clipping thresholds. The optimization is performed using singular value decomposition (SVD) for accurate and efficient matrix inversion and automatic mixed precision (AMP) for training stability and efficiency (details in Appendix B.1).

The overall framework of FlatQuant is depicted in Figure 3. It shows the integration of learnable affine transformations, per-channel scaling, and learnable clipping thresholds within the Transformer architecture.

The following figure (Figure 3 from the original paper) illustrates the overall framework of FlatQuant:

该图像是图表，展示了LLama-3-70B模型中多个Transformer层的参数分布。图中包括四个子图：(a)和(b)分别展示了第10个Transformer层中$W_o$和$X_o$的幅度随通道变化的情况；(c)和(d)展示了第30个Transformer层中$W_g$和$X_g$的幅度变化情况。各子图的横轴为通道数，纵轴为幅度，显示了不同参数的分布特点。 Figure 3: The overall framework of FlatQuant. (a) Necessary notations for FlatQuant; (b) The integration of FlatQuant with a conventional LLaMA layer, where merged parameters are grouped in red; (c) Online transformation and down-projection layer, where the scaling vector $\mathrm{diag}(\mathbf{c})$ over $\tilde{\mathbf{X}}$ is merged to $\mathbf{W}_u$ in practice.

4.2.4. Integration with the Transformer Architecture

FlatQuant is integrated into a Transformer block (e.g., LLaMA-like architecture) as follows:

• Quantized Layers: All linear layers utilize low-bit matrix multiplications.
• FP16 Layers: Layer normalization layers, pre-quantization transformations, RoPE embeddings, and attention scores remain in FP16 for precision.

Self-Attention

The self-attention module incorporates four transformation matrices:

• $\mathbf{P}_a$: Applied to flatten the input activation for the query, key, and value projections.
• $\mathbf{P}_o$: Smooths the input activation for the output projection.
• $\mathbf{P}_h$: Transforms the key cache head by head.
• $\mathbf{P}_v$: Transforms the value cache head by head. Notably, only $\mathbf{P}_a$ and $\mathbf{P}_o$ are decomposed using the Kronecker product because the head size in per-head quantization for KV cache is already small, making full transformations for $\mathbf{P}_h$ and $\mathbf{P}_v$ computationally cheap. $\mathbf{P}_o$ is further fused with $\mathbf{P}_v$ to reduce overhead, inspired by QuaRot.

Feed-forward Network (FFN)

The FFN utilizes two transformation matrices:

• $\mathbf{P}_{ug}$: Applied to flatten the input of the FFN after layer normalization.
• $\mathbf{P}_d$: Flattens the input for the down-projection layer. Both $\mathbf{P}_{ug}$ and $\mathbf{P}_d$ are decomposed to minimize inference overhead. The per-channel scaling of $\mathbf{P}_d$ is merged into the weight of the up-projection layer to avoid additional computational cost.

Layer Normalization

Unlike some methods (e.g., QuaRot, SpinQuant) that modify LayerNorm to RMSNorm and merge transformations into preceding layers, FlatQuant preserves the original LayerNorm. This allows for distinct fast and learnable affine transformations to be applied after LayerNorm for different layers, increasing the model's expressiveness and adaptive capacity.

4.2.5. Efficient Kernel Design

FlatQuant implements an efficient kernel that fuses the affine transformations and quantization into a single operation using OpenAI Triton. This design is motivated by:

1. Memory-Bound Operations: The Kronecker product transformations ($\mathbf{P}_{1}^{\top} \times_{1} \tilde{\mathbf{X}} \times_{2} \mathbf{P}_{2}$) are low computational intensity operations, making them memory-bound. Similarly, quantization itself is also memory-bound.

2. Minimizing Overhead: Fusing these operations into a single kernel eliminates redundant global memory accesses for intermediate results and reduces kernel launch overhead, which is critical for achieving substantial speedups.

   The fusion process involves:

• Loading the entire P1 and P2 matrices into SRAM (on-chip memory, much faster than global memory).

• Each thread block processes a small tiling block $\bar{\mathbf{X}}$ from the reshaped activations $\tilde{\mathbf{X}}$.

• It performs the matrix multiplication $\mathbf{P}_{1} \bar{\mathbf{X}} \mathbf{P}_{2}$ and quantizes the results on the fly.

• All intermediate results are kept in SRAM before the final quantized output is written back to global memory.

  For INT4 matrix multiplication, FlatQuant follows QuaRot by adopting the CUTLASS kernel. For KV cache quantization, it uses FlashInfer. Details on handling corner cases where SRAM might not be large enough for very large tensors are provided in Appendix B.3, involving strategies like tiling non-reduction dimensions.

5. Experimental Setup

5.1. Datasets

The experiments primarily evaluate FlatQuant on various Large Language Models (LLMs) from the LLaMA-2 (Touvron et al., 2023) and LLaMA-3 (Dubey et al., 2024) series. Additional results on Qwen (Yang et al., 2024) and DeepSeek families are provided in Appendix C.1.

The datasets used for evaluation are:

• Language Generation Tasks:
  • WikiText-2 (Merity et al., 2016): A dataset of over 100 million tokens extracted from "Good" and "Featured" articles on Wikipedia. It is commonly used for evaluating language models, especially for perplexity (PPL) due to its coherent and diverse text.
  • C4 (Colossal Clean Crawled Corpus) (Raffel et al., 2020): A massive, cleaned corpus of text derived from web pages, comprising hundreds of gigabytes of text. It's widely used for training and evaluating large language models, providing a broad and diverse range of text styles and topics.
• Commonsense Reasoning (Zero-shot QA) Tasks: Six popular benchmarks are used:

  • ARC-Challenge and ARC-Easy (Clark et al., 2018): Datasets designed to test a machine's ability to answer questions requiring commonsense reasoning. ARC-Challenge contains questions that are harder for algorithms to answer.

  • HellaSwag (Zellers et al., 2019): A dataset for commonsense natural language inference, where models must choose the most plausible ending to a given sentence. It's designed to be difficult for models that rely purely on statistical cues.

  • LAMBADA (Paperno et al., 2016): A dataset for word prediction that specifically requires broad discourse context. Models must predict the last word of a sentence, where the necessary context often spans multiple sentences.

  • PIQA (Bisk et al., 2020): Physical Interaction QA, a commonsense reasoning benchmark focused on questions about physical commonsense, requiring models to choose the correct plan to achieve a goal.

  • WinoGrande (Sakaguchi et al., 2021): An adversarial Winograd Schema Challenge at scale, designed to test commonsense reasoning by resolving ambiguous pronouns in sentences.

    These datasets are chosen because they are standard benchmarks in LLM research, providing a comprehensive assessment of both language generation quality and reasoning abilities, which are critical for validating the effectiveness of quantization methods.

5.2. Evaluation Metrics

For every evaluation metric mentioned in the paper, here is a complete explanation:

5.2.1. Perplexity (PPL)

• Conceptual Definition: Perplexity (PPL) is a measure of how well a probability model predicts a sample. In the context of language models, it quantifies how "surprised" the model is by a given text sequence. A lower PPL indicates that the model assigns a higher probability to the observed sequence, meaning it predicts the text more accurately and is a better language model. It is the inverse probability of the test set, normalized by the number of words.
• Mathematical Formula: $ PPL(W) = \exp\left(-\frac{1}{N} \sum_{i=1}^{N} \log P(w_i | w_{1...i-1})\right) $
• Symbol Explanation:
  • $W = (w_1, w_2, ..., w_N)$: The sequence of words (or tokens) in the test set.
  • $N$: The total number of words (or tokens) in the sequence $W$.
  • $P(w_i | w_{1...i-1})$: The probability of the $i$-th word $w_i$ given all preceding words $w_1, ..., w_{i-1}$, as predicted by the language model.
  • $\log$: The natural logarithm.
  • $\exp(\cdot)$: The exponential function (inverse of natural logarithm).

5.2.2. Accuracy

• Conceptual Definition: Accuracy is a common metric for classification tasks, such as zero-shot question answering (QA). It represents the proportion of correct predictions made by the model out of the total number of predictions. A higher accuracy percentage indicates better performance.
• Mathematical Formula: $ Accuracy = \frac{\text{Number of Correct Predictions}}{\text{Total Number of Predictions}} $
• Symbol Explanation:
  • Number of Correct Predictions: The count of instances where the model's output (e.g., the chosen answer in a QA task) perfectly matches the true, ground-truth label.
  • Total Number of Predictions: The total number of instances (e.g., questions) in the evaluation dataset.

5.3. Baselines

FlatQuant is compared against several popular INT4 post-training quantization (PTQ) methods, which are representative of the state-of-the-art in the field:

• SmoothQuant (Xiao et al., 2023): A widely recognized method that uses per-channel scaling to smooth activations by transferring quantization difficulty to weights.

• OmniQuant (Shao et al., 2023): An extension that makes scaling and shifting factors learnable parameters for omnidirectional calibration.

• AffineQuant (Ma et al., 2024): A recent method that also employs affine transformations for quantization, serving as a direct comparison point for FlatQuant's transformation approach.

• QUIK-4B (Ashkboos et al., 2023): A method focused on end-to-end 4-bit inference for generative LLMs.

• QuaRot (Ashkboos et al., 2024): A very recent state-of-the-art method that utilizes Hadamard transformations (rotations) to achieve outlier-free 4-bit inference.

• SpinQuant (Liu et al., 2024c): Another recent state-of-the-art method that uses learned rotations for LLM quantization.

  These baselines are chosen because they represent the current leading approaches in LLM quantization, covering different strategies like scaling, fixed rotations, learned rotations, and general affine transformations, providing a robust comparison for FlatQuant's performance.

5.4. Implementation Details

The implementation specifics are as follows:

• Frameworks: Huggingface Transformers (Wolf, 2019) and PyTorch (Paszke et al., 2019).
• Optimizer: AdamW optimizer.
• Learning Rates: Initial learning rate of 5e-3 for the main learnable parameters, and 5e-2 for clipping thresholds.
• Learning Rate Schedule: Cosine annealing learning rate decay.
• Calibration:
  • 15 epochs of training.
  • Calibration set: 128 sentences sampled from WikiText-2.
  • Sequence length: Each sentence is processed with 2048 tokens.
  • Batch size: 4.
• Computational Resources: Calibration for LLaMA-3-8B requires approximately 26GB of GPU memory and takes about 0.9 hours on a single GPU.
• Initialization: FlatQuant is robust to initialization, using random affine transformation matrices as a starting point.
• Matrix Inversion and Training: Singular Value Decomposition (SVD) is used for efficient and accurate matrix inversion (for $\mathbf{P}^{-1}$), combined with Automatic Mixed Precision (AMP) to reduce training time and memory usage while maintaining accuracy. (Appendix B.1).

5.5. Quantization Scheme

• Weights and Activations: Per-channel symmetric quantization for weights and per-token symmetric quantization for activations.
• Weight Quantizer: For fair comparison, FlatQuant is evaluated with both round-to-nearest (RTN) and GPTQ as the weight quantizer. When GPTQ is used, it shares the same calibration data as FlatQuant for both its closed-form weight updates and training. The paper notes that FlatQuant with RTN is often sufficient and competitive.
• KV Cache Quantization: Group-wise asymmetric quantization is applied to the KV cache, with a group size of 128. This choice matches the head dimension of LLaMA models and leverages the memory-bound characteristics of self-attention.

6. Results & Analysis

6.1. Main Results

The experiments demonstrate FlatQuant's superior performance in both accuracy and inference latency across various LLM families and tasks.

6.1.1. Results on Language Generation Tasks

The following are the results from Table 1 of the original paper:

Table 1: WikiText-2 and C4 perplexity of 4-bit weight & activation quantized LLaMA models.

Table 1 presents the Perplexity (PPL) results for 4-bit weight & activation quantized LLaMA models on WikiText-2 and C4 datasets. Lower PPL indicates better performance.

• Superiority of FlatQuant (RTN): FlatQuant using round-to-nearest (RTN) as the weight quantizer consistently outperforms all previous state-of-the-art methods across all LLaMA models and both datasets. For instance, on LLaMA-2-70B, FlatQuant-RTN achieves a PPL of 3.55 on WikiText-2 and 5.91 on C4, which is very close to the FP16 baseline (3.32 and 5.71 respectively). This implies a minimal performance degradation.

• LLaMA-3 Models: FlatQuant shows strong performance on the newer LLaMA-3 models. For LLaMA-3-8B, FlatQuant-RTN achieves 6.98 PPL on WikiText-2, significantly better than SpinQuant-RTN's 7.96 and QuaRot-RTN's 10.60.

• GPTQ vs. RTN: Remarkably, FlatQuant-RTN's performance is often comparable to, and sometimes even surpasses, GPTQ-based methods from other baselines. For example, FlatQuant-RTN on LLaMA-3-70B on WikiText-2 (PPL 3.78) is better than QuaRot-GPTQ (PPL 6.60) and SpinQuant-GPTQ (PPL 6.21). This highlights that FlatQuant's learnable transformations are highly effective even with a simpler quantizer, reducing calibration time. FlatQuant-GPTQ provides a slight additional boost, but the difference is small, reinforcing the strength of its core method.

  These results strongly validate that FlatQuant's approach of enhancing flatness through learnable transformations effectively mitigates quantization error, setting a new benchmark for low-bit LLM quantization.

6.1.2. Results on Zero-shot QA Tasks

The following are the results from Table 2 of the original paper:

Table 2: Zero-shot QA task results of 4-bit weight & activation quantized LLaMA models.

Table 2 shows the zero-shot QA task results (accuracy) for 4-bit weight & activation quantized LLaMA models across six benchmarks. Higher accuracy indicates better performance.

• Significant Accuracy Gains: FlatQuant consistently achieves higher average accuracy compared to other quantization methods. For the challenging LLaMA-3 models, FlatQuant significantly narrows the performance gap to the FP16 baseline. For LLaMA-3-8B, it achieves an average accuracy of 71.23% (RTN) with an accuracy loss of only 2.00% compared to FP16 (73.23%). For LLaMA-3-70B, FlatQuant-RTN achieves 79.01% (note: there is a discrepancy in the table for LLaMA-3-70B FP16 avg where 71.33 and 79.95 are both present; FlatQuant-RTN is 79.01 which is a 0.94% drop from 79.95).

• Robustness of FlatQuant-RTN: Similar to the PPL results, FlatQuant-RTN often performs comparably to, or even better than, baselines using GPTQ. For LLaMA-3-8B, FlatQuant-RTN (71.23%) outperforms QuaRot-GPTQ (65.79%) and SpinQuant-GPTQ (68.70%). This further underscores FlatQuant's ability to achieve high performance without the additional overhead of GPTQ's weight updates.

• Challenging LLaMA-3 Quantization: The paper notes that LLaMA-3 models are particularly challenging for quantization, which is evident from the larger performance drops for other methods, especially QuaRot-RTN on LLaMA-3-70B (average 35.36% vs FP16 79.95%). FlatQuant's ability to maintain high accuracy here is a strong indicator of its effectiveness.

  The consistent outperformance on both language generation (PPL) and zero-shot QA (accuracy) tasks, particularly for the latest LLaMA-3 models and with simple RTN quantization, establishes FlatQuant as a new state-of-the-art.

6.1.3. Additional Experiments (Appendix C)

Results on Other LLM Architectures (Appendix C.1)

• LLaMA-3.1-8B-Instruct (Table 7): FlatQuant also outperforms QuaRot on LLaMA-3.1-8B-Instruct for both WikiText-2 PPL and C4 PPL, and significantly higher average accuracy on zero-shot QA tasks. For example, FlatQuant achieves an average QA accuracy of 72.03% compared to QuaRot's 66.84% (FP16 is 73.69%).
• Qwen-2.5-Instruct (Table 8): FlatQuant demonstrates competitive performance on Qwen-2.5-Instruct models. For the Qwen-2.5-32B model, FlatQuant-RTN achieves an average QA score of 74.89%, which is only a 0.21% drop from the FP16 baseline (75.10%) and higher than QuaRot-GPTQ (72.25%).
• DeepSeek V3-Base and DeepSeek R1 (Table 9): FlatQuant-W4A4 shows strong results on large-scale Mixture-of-Experts (MoE) models like DeepSeek V3-Base and DeepSeek R1, demonstrating its applicability beyond standard dense LLMs. For DeepSeek V3-Base, FlatQuant-W4A4 achieves 89.59 C-Eval and 86.32 MMLU, very close to FP8 (90.10 and 87.10 respectively).

Results on MT-Bench (Appendix C.2)

• LLaMA-3.1-8B-Instruct (Table 10): On MT-Bench, FlatQuant (average 6.94) significantly outperforms QuaRot (5.99), narrowing the gap to FP16 (7.60). In some categories like Math, FlatQuant (7.20) even surpasses FP16 (7.00).

Extension to More Quantization Settings (Appendix C.3)

• Weight-Only Quantization (Table 11): FlatQuant-RTN shows competitive performance in weight-only settings (W4A16, W3A16). For W4A16 on WikiText-2, FlatQuant-RTN achieves 6.54 PPL, comparable to GPTQ-g128 (6.50) and QuIP (6.50), and significantly better than plain GPTQ (7.00).

• KV Cache Quantization (Table 12, 13): FlatQuant effectively quantizes the KV cache to very low bit-widths. For LLaMA-3-8B with K4V4 (4-bit Key, 4-bit Value) quantization, FlatQuant achieves a WikiText-2 PPL of 6.20 and average QA accuracy of 73.12%, very close to FP16. Even at K2V2 (2-bit Key, 2-bit Value), FlatQuant maintains significantly better performance than QuaRot (Table 13). For LLaMA-2-7B K2V2, FlatQuant has a PPL of 6.66 compared to QuaRot's 9.23.

• Extreme Low Bit Quantization (Table 14): For W3A3KV3 quantization on LLaMA-3-8B, FlatQuant (10.82 WikiText-2 PPL, 58.45% Avg QA) vastly outperforms QuaRot (686.54 WikiText-2 PPL, 30.33% Avg QA), demonstrating its robustness in highly aggressive quantization scenarios.

• Flexible Quantization Settings (Table 15): The learnable transformations in FlatQuant are flexible. A single set of transformation matrices can be used for different quantization settings (e.g., W4, A4, KV4 independently), maintaining high accuracy.

  Overall, the results consistently demonstrate FlatQuant's ability to achieve state-of-the-art accuracy, even with simpler RTN quantization, across a wide range of LLMs, tasks, and quantization settings.

6.2. Inference Latency

All inference latency experiments were conducted on an RTX3090 GPU.

6.2.1. End-to-end Speedup

The following figure (Figure 4 from the original paper) shows the prefill and decoding speedup of FlatQuant across different batch sizes:

Figure 4: Prefill and decoding speedup of FlatQuant for LLaMA-2-7B on an RTX3090 GPU. We evaluate prefill on a sequence length of 2048 and decoding on 256 tokens. FlatQuant with kernel fusion (green line) achieves up to 2.3x prefill speedup and 1.76x decoding speedup compared to FP16.

Figure 4 illustrates the prefill and decoding speedup of FlatQuant compared to FP16, INT4, and QuaRot across various batch sizes (up to 64).

• Kernel Fusion Impact: Even without kernel fusion, FlatQuant achieves comparable speedup to QuaRot, validating the efficiency of the Kronecker product approach.
• Superior Speedup with Kernel Fusion: With kernel fusion, FlatQuant significantly outperforms QuaRot. It achieves up to 2.30x prefill speedup and 1.76x decoding speedup under a batch size of 64 (for prefill and decoding respectively), which is faster than QuaRot. While there's a minor gap compared to vanilla INT4 quantization (which lacks the complex transformations), FlatQuant's speedup is substantial given its superior accuracy. This makes it highly practical for deploying INT4 LLMs.

6.2.2. Kronecker Product: Sizes and Perplexities

The following figure (Figure 5 from the original paper) examines the impact of different decomposed matrix sizes in Equation 3 on model performance and speedup:

该图像是图表，展示了在 LLaMA-2-7B 模型上，不同降维矩阵大小对预填充速度和 WikiText2 PPL 结果的影响。x 轴为降维大小，y 轴分别表示速度加速比和 PPL 值，速度加速比用红色星形标记，PPL 值用蓝色圆点标记。 Figure 5: Prefill speedup and WikiText2 PPL results of different decomposed matrix sizes on LLaMA-2-7B model. We decompose the hidden dimension 4096 into $n_1 \times n_2$ and range $n_1$ from 1 to 2048, where $n_1 = 1$ amounts to maintaining a full-size transformation matrix. More details can be found in Appendix C.6.

Figure 5 shows the impact of different decomposed matrix sizes ($n_1 \times n_2$, where $n=4096$) for the Kronecker product on prefill speedup and WikiText-2 PPL for the LLaMA-2-7B model.

• Optimal Speedup: The speedup peaks when $\mathbf{P}_1$ and $\mathbf{P}_2$ are of approximately equal size (i.e., $n_1 = n_2 = \sqrt{n} = 64$). This aligns with the theoretical analysis in Section 3.1, which predicts optimal computation saving under this condition.

• Impact on PPL: Importantly, variations in decomposed matrix sizes have limited impact on perplexity, indicating the robustness of FlatQuant's transformation capability regardless of the exact decomposition, as long as the total transformation is effective.

• Performance Drop: When $n_2$ (the second dimension of the decomposition) exceeds 64, the speedup decreases. This is attributed to irregular memory access patterns for activations when the dimensions become highly unbalanced.

  These results validate FlatQuant's effectiveness in minimizing inference overhead through the Kronecker product while preserving quantization accuracy.

6.2.3. Overhead of Each Online Transformation

The following figure (Figure 6 from the original paper) investigates the impact of the five online transformations in FlatQuant on the overall speedup:

Figure 6: Prefill speedup of LLaMA-2-7B on a sequence length of 2048 under a batch size of 64 after applying different online transformations. We incorporate different online transformations sequentially to gauge their impact on the final speedup. Each point on the $\mathbf{X}$ axis indicates adding a new online transformation.

Figure 6 analyzes the individual contributions of the five online transformations ($$\mathbf{P}_h`,`\mathbf{P}_o`,`\mathbf{P}_d`,`\mathbf{P}_a`,`\mathbf{P}_{ug}$$) to the overall speedup in FlatQuant.

• Minimal Overall Slowdown: Even with all five per-layer transformations, FlatQuant introduces a minimal 0.07x end-to-end slowdown (when compared to naive INT4), significantly outperforming QuaRot's 0.26x slowdown with only three Hadamard transformations.

• Specific Transformation Impacts:

  • $\mathbf{P}_d$ (down-projection layer) causes the largest individual slowdown for FlatQuant (0.04x), due to the large FFN intermediate sizes. This is still much smaller than QuaRot's corresponding slowdown (0.17x).
  • $\mathbf{P}_o$ (output projection) results in a 0.01x slowdown, again less than QuaRot's 0.1x.
  • The remaining transformations ($\mathbf{P}_a$ for query/key/value projections and $\mathbf{P}_{ug}$ for FFN input) have an insignificant impact (less than 0.01x).

• Efficiency without Kernel Fusion: Even without kernel fusion, the additional transformations in FlatQuant maintain competitive performance relative to QuaRot, primarily due to the efficiency gained from the Kronecker product decomposition.

  This detailed analysis confirms that FlatQuant's transformations are lightweight and efficiently integrated, contributing minimal overhead while delivering significant accuracy benefits.

6.3. Ablation Studies / Parameter Analysis

6.3.1. Ablation Study (Main Components)

The following are the results from Table 3 of the original paper:

Table 3: Ablation study of FLATQUANT's main components on LLaMA-3-8B.

Table 3 presents an ablation study of FlatQuant's main components on the LLaMA-3-8B model, starting from a round-to-nearest (RTN) baseline (no LT, PS, LCT).

• Learnable Transformation (LT): Enabling Learnable Transformation (LT) alone drastically improves performance, reducing WikiText-2 PPL from 1266.60 to 8.50 and C4 PPL from 936.41 to 13.51. This component brings the average QA accuracy from 30.99% to 66.82%. This demonstrates that LT is the most critical component, capable of adaptively flattening distributions and significantly enhancing model accuracy.

• Per-channel Scaling (PS): Adding Per-channel Scaling (PS) on top of LT further improves performance, lowering WikiText-2 PPL to 7.95 and C4 PPL to 12.74, with a slight increase in average QA accuracy to 67.08%. This indicates PS plays a complementary role in balancing outliers.

• Learnable Clipping Thresholds (LCT): Including Learnable Clipping Thresholds (LCT) (either with just LT or with both LT and PS) yields substantial improvements. With $LT + LCT$, WikiText-2 PPL drops to 7.11 and C4 PPL to 11.47, with average QA accuracy reaching 70.72%.

• Full FlatQuant: The combination of all three components ($LT + PS + LCT$) achieves the best results, reaching 6.98 WikiText-2 PPL, 11.13 C4 PPL, and 71.23% average QA accuracy.

  This ablation study clearly demonstrates the necessity and effectiveness of each component (Learnable Transformation, Per-channel Scaling, and Learnable Clipping Thresholds) in contributing to FlatQuant's overall superior performance by collectively enhancing flatness and mitigating quantization error.

6.3.2. FlatQuant Leads to Flatness

The paper quantitatively evaluates flatness by analyzing channel-wise magnitude distributions. Each distribution is represented as a 1D vector $\mathbf{d}$. Flatness is measured by the Euclidean distance between the observed distribution $\mathbf{d}$ and an idealized perfectly flat distribution $\mathbf{d}'$. The idealized flat distribution $\mathbf{d}'$ is defined such that all channels have equal magnitudes and the same $\ell_2$ norm as $\mathbf{d}$, i.e., $\mathbf{d}' = \frac{\|\mathbf{d}\|_2}{\sqrt{N}} \cdot \mathbf{1}_N$, where $N$ is the number of channels and $\mathbf{1}_N$ is an $N$-dimensional vector of ones. A smaller Euclidean distance ${ \|\mathbf{d} - \mathbf{d}'\|_2 }$ indicates greater flatness.

The following figure (Figure 7 from the original paper) visualizes the evolution of flatness and the training objective across different Transformer blocks of LLaMA-3-8B during training:

该图像是一个示意图，展示了不同量化方法的效果，包括(a) 按通道缩放，(b) 哈达玛变换，(c) FlatQuant 方法，以及(d) 叠加视图。每个子图通过三维展示描述了在不同层和 token 数量下的分布情况。 Figure 7: Flatness and mean squared quantization error (MSE) for different Transformer blocks of LLaMA-3-8B during FlatQuant's training process. The metric of flatness is calculated as the sum of Euclidean distances ${ \|\mathbf{d} - \mathbf{d}'\|_2 }$ for all weights and activations within a Transformer block.

Figure 7 shows that as the training progresses and the training loss (Equation 4) decreases, the channel distributions become increasingly flat. This direct correlation indicates that FlatQuant successfully learns transformations that yield flatter distributions, which in turn contributes to smaller quantization error. This empirically validates the core hypothesis that flatness matters for LLM quantization.

6.3.3. Additional Ablation Study and Discussions (Appendix C.4 and C.5)

• Detailed Ablation Study (Table 16): A more comprehensive ablation confirms the individual and combined effectiveness of LT, PS, and LCT on all individual QA tasks, reinforcing their necessity.
• Impact of Calibration Data (Table 17): FlatQuant shows robustness to the choice of calibration dataset. When calibrated on WikiText2, C4, or Pile, FlatQuant maintains stable performance, with only minor variations in PPL and QA accuracy, indicating good generalization.
• Effect of Learnable Clipping Thresholds (Table 18): Learnable Clipping Thresholds (LCT) are most effective when applied after the affine transformation. Applying LCT before Transformation yields worse results (7.37 WikiText-2 PPL) than LCT after Transformation (6.98 PPL), and even QuaRot-style Fixed Threshold (7.25 PPL) is not as effective. This emphasizes the optimal placement and learning of clipping.
• Mixed-Precision Quantization (Table 19): FlatQuant can be combined with mixed-precision schemes. Selectively using 8-bit quantization for certain sensitive layers (e.g., down_proj layers, or Top5 most sensitive layers) can further improve accuracy with minimal impact on speed, indicating flexibility for deployment scenarios where slightly higher precision is acceptable for critical components.

6.4. Inference Memory and Latency (Appendix C.7 and C.8)

• Inference Memory Consumption (Table 20): FlatQuant maintains the memory efficiency of INT4 quantization. For LLaMA-2-7B, it achieves a consistent memory reduction of over 3.3x compared to FP16 for various sequence lengths (batch size 1), with negligible additional memory overhead from its transformations.
• Additional Latency Analysis (Table 21, 22):

  • Prefill Speedup (Table 21, Figure 9): FlatQuant achieves consistent prefill speedups (e.g., 2.12x at 2048 length, 1.80x at 16384 length for batch size 1) comparable to INT4 and outperforming QuaRot.

  • Decoding Speedup (Table 22, Figure 10): For decoding (batch size 64), FlatQuant consistently surpasses QuaRot across all KV cache lengths and closely approaches the efficiency of INT4 quantization.

    The following figure (Figure 9 from the original paper) shows the prefill speedup of LLaMA-2-7B on a sequence length of 2048:

    Figure 9: Prefill speedup of LLaMA-2-7B on a sequence length of 2048. It shows the speedup of INT4, QuaRot, FlatQuant without kernel fusion, and FlatQuant with kernel fusion. FlatQuant with kernel fusion achieves the highest speedup, close to INT4.

The following figure (Figure 10 from the original paper) shows the decoding speedup on LLaMA-2-7B model:

Figure 10: Decoding speedup on LLaMA-2-7B model. We decode 256 tokens after the prefill sequence length 2048. FlatQuant with kernel fusion (green line) shows excellent performance, surpassing QuaRot and approaching INT4.

These results confirm that FlatQuant provides substantial speedup gains across various generation scenarios, including both short and long contexts, without sacrificing memory efficiency.

7. Conclusion & Reflections

7.1. Conclusion Summary

This paper effectively revisits and reinforces the critical role of flat weights and activations in achieving effective Large Language Model (LLM) quantization. It highlights that existing pre-quantization transformations often fall short, still yielding steep and outspread distributions that hinder performance. In response, the authors introduce FlatQuant, a novel post-training quantization method built upon fast and learnable affine transformations. These transformations are optimized for each linear layer to actively promote the flatness of weights and activations.

Through extensive experiments, FlatQuant demonstrates clear superiority, establishing a new state-of-the-art benchmark for LLM quantization. It achieves a remarkably low accuracy drop of less than 1% for W4A4 quantization on the highly challenging LLaMA-3-70B model, significantly outperforming competitors like SpinQuant by 7.5%. Furthermore, its efficient kernel fusion strategy, which integrates affine transformation and quantization operations, delivers substantial inference speedups—up to 2.3x for prefill and 1.7x for decoding compared to the FP16 baseline. This work marks a significant advancement toward the practical application of high-accuracy, low-bit quantization for LLMs.

7.2. Limitations & Future Work

The authors acknowledge a limitation regarding the current scope of FlatQuant primarily focusing on INT4 quantization. They point out that FlatQuant has not yet been rigorously applied to hypothetical new data types, such as MXFP (Mixed-Precision Floating Point). Exploring FlatQuant's compatibility and potential advantages with these emerging data types represents a promising avenue for future research.

7.3. Personal Insights & Critique

FlatQuant offers several compelling insights and makes a strong case for the explicit optimization of distribution flatness in quantization. The primary innovation lies in its learnable affine transformations, which offer greater flexibility than fixed scaling or rotation methods, coupled with a highly engineered Kronecker product decomposition and kernel fusion for efficiency. This combination addresses the common trade-off between expressive power and computational cost in quantization.

Inspirations and Transferability:

• Explicit Flatness Objective: The idea of explicitly optimizing for flatness (quantified by Euclidean distance to an ideal flat distribution) could be transferred to other compression domains beyond LLMs, such as vision transformers or other large neural networks where outliers pose a challenge.
• Learnable Transformations with Efficiency Constraints: The methodology of learning optimal transformations under strict computational and memory constraints (via Kronecker product and kernel fusion) is a powerful paradigm. This could inspire similar approaches in other areas of model compression or hardware-aware neural network design, where custom, efficient operations are critical.
• Robustness to Quantizer and Models: FlatQuant's ability to perform well with simple RTN quantization, and across diverse LLM families and tasks, suggests a strong underlying principle that makes the models inherently "quantization-friendly." This indicates a fundamental improvement in how data is represented for low-bit precision.

Potential Issues, Unverified Assumptions, or Areas for Improvement:

• Generalizability of Kronecker Product: While effective for the tested dimensions, the optimal choice of $n_1, n_2$ and the generalizability of the Kronecker product approach for arbitrary layer dimensions or different Transformer variants could be explored further. The paper states optimal $n_1, n_2$ are sought, but whether this decomposition always maintains expressiveness for drastically different architectures or very small dimensions (where $\sqrt{n}$ might be too small) is worth investigating.

• Complexity of Learnable Transformations: While FlatQuant boasts efficiency, learning these affine transformations adds a calibration step not present in zero-shot PTQ methods. The "hours" of calibration, while significantly less than Quantization-Aware Training (QAT), might still be a factor for extremely rapid deployment or scenarios with very limited compute. Further reducing calibration time or exploring even more lightweight learning strategies could be beneficial.

• Impact of Calibration Data Quality: Although the paper shows robustness to different calibration datasets (WikiText2, C4, Pile), the quality and domain relevance of the calibration data can still be crucial. An LLM fine-tuned for a highly specialized domain might require domain-specific calibration data for optimal FlatQuant performance.

• Extreme Low-Bit Sensitivity: While FlatQuant performs exceptionally well at W4A4KV4 and even W3A3KV3, the performance gap to FP16 inevitably widens at W2A2KV2. The "sweet spot" for practical deployment in terms of bit-width versus accuracy-speed trade-off remains a nuanced decision. Future work could focus on pushing the boundaries of even lower bit-widths while maintaining the same level of relative accuracy.

• Hardware Dependency of Kernel Fusion: The kernel fusion benefits are highly dependent on the target GPU architecture (e.g., NVIDIA with CUDA/Triton). While Triton aims for portability, optimality might still vary across different hardware platforms (e.g., AMD GPUs, custom AI accelerators). Ensuring broad hardware compatibility and performance for the custom kernels would be a continuous effort.

  Overall, FlatQuant presents a robust and highly effective solution to a critical problem in LLM deployment. Its principled approach to flatness and meticulously optimized implementation make it a standout contribution in the field of quantization.