💬 NLP

LATMiX: Learnable Affine Transformations for Microscaling Quantization of LLMs

This paper presents LATMiX, a method that employs learnable invertible affine transformations to optimize activation distributions for micro-scale quantization (MX), thereby significantly improving the accuracy of low-bit Large Language Models compared to existing rotation- or Hadamard-based approaches.

Original authors: Ofir Gordon, Lior Dikstein, Arnon Netzer, Idan Achituve, Hai Victor Habi

Published 2026-04-27

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ofir Gordon, Lior Dikstein, Arnon Netzer, Idan Achituve, Hai Victor Habi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a massive, incredibly detailed library (a Large Language Model, or LLM) capable of answering any question, writing code, or telling stories. The problem is that this library is so vast and heavy that it is expensive to store and slow to read. To make it faster and cheaper, engineers attempt to shrink the books by summarizing them into smaller, simpler formats. This process is called quantization.

However, there is a catch. In these enormous libraries, some pages contain "outliers"—extremely long, complex sentences that do not fit well into the small, simplified format. If you try to force these long sentences into a tiny box, they become squashed and distorted, ruining the story.

The Problem: The "Block" Bottleneck

Recently, a new way of organizing these books called Microscaling (MX) was introduced. Instead of shrinking the entire library with a single massive summary, MX breaks the books into small blocks and gives each block its own tiny ruler (scaling factor) to measure the words. This is great because it handles different parts of the text more flexibly.

But here is the catch: Previous methods tried to solve the "outlier" problem by rotating the entire library like a spinning top. When they attempted to combine this rotation with the new "block" rulers, chaos ensued. The rotation disrupted the block measurements, rendering the library unreadable.

To fix this, earlier researchers tried a workaround: they rotated only the pages within each small block and ignored the rest of the library. Although this stopped the chaos, it was like trying to fix a traffic jam by moving only the cars within a single parking lot while ignoring the cars stuck in the next lane. It did not solve the problem on a large scale.

The Solution: LATMiX (The Intelligent Reorganizer)

The authors of this paper, LATMiX, propose a smarter way to rearrange the books before shrinking them.

Imagine the data in the model as a crowd of people in a room. Some people stand in a huge, dense cluster (the outliers), while others are scattered.

Old methods tried to rotate the entire room or mix people only within small groups.
LATMiX acts like a master choreographer. It learns a custom, flexible dance move (an affine transformation) that distributes the dense cluster evenly across the entire room.

This choreography has two special features:

It is learnable: Instead of using a preset dance (like a simple rotation), LATMiX uses standard AI tools to learn the perfect dance moves specifically for the data it is viewing. It figures out exactly how the energy must be distributed so that no single point becomes too crowded.
It respects the blocks: It knows the "block" rulers (MX format). It does not just mix people randomly; it mixes them in a way that works perfectly with the block system, ensuring the tiny rulers can measure everyone accurately.

How It Works (The Magic Trick)

The paper explains that LATMiX does not actually slow down the library when it is read later. It is like a magician who rearranges the cards before the trick begins and then embeds this rearrangement directly into the deck itself. Once the trick is set up, the magician (the computer) does no extra work to shuffle; the cards are already in the perfect order. This means that speed and memory usage remain just as fast as before, but the quality of the "shrunk" books is much higher.

The Results

The authors tested this on various models (such as Llama and Qwen) and found that LATMiX consistently outperformed other methods.

Better Accuracy: The "shrunk" models made fewer mistakes and understood questions better than models using older shrinking techniques.
Robustness: It worked well across different model sizes, from small to very large models.
Efficiency: It achieved these improvements without adding extra time or cost for running the model.

In short, LATMiX is an intelligent, learnable tool that rearranges data just right before compressing it, ensuring that even the most extreme "outlier" information survives the shrinking process without being lost, while harmonizing seamlessly with modern hardware formats.

1. Problem Statement

Large Language Models (LLMs) are increasingly deployed with Post-Training Quantization (PTQ) to reduce memory and computational costs. A specific data format gaining importance is Microscaling (MX), which divides tensors into blocks and assigns a unique scaling factor to each block to handle dynamic ranges more efficiently than global scaling.

However, combining MX quantization with existing outlier mitigation techniques (which are crucial for low-bitwidth quantization) presents significant challenges:

Activation Outliers: LLM activations contain extreme values (outliers) that dominate the quantization error.
Limitations of Current Methods: Existing approaches use invertible transformations (typically rotations or Hadamard transformations) to redistribute activation energy and suppress outliers.
- Global Transformations: Applying a global rotation to activations prior to MX quantization leads to severe performance degradation, as the global transformation conflicts with the block-wise scaling structure of MX.
- Block-Diagonal Transformations: Recent works propose applying transformations independently within each MX block (block-diagonal matrices). While this avoids inconsistency, it restricts the transformation to small, isolated subspaces and prevents effective redistribution of activation mass across the entire tensor. Consequently, dominant outliers cannot be sufficiently diffused, leading to suboptimal accuracy in low-bitwidth regimes (e.g., 4-bit).

2. Methodology: LATMiX

The authors propose LATMiX (Learnable Affine Transformations for Microscaling), a method that generalizes outlier reduction to learnable, full-rank invertible affine transformations specifically optimized for the MX data format.

A. Theoretical Analysis

The authors derive a theoretical upper bound for the quantization error under MX quantization. The analysis shows that the error depends on two competing factors:

The spectral norm of the inverse transformation ( $||A^{-1}||_\sigma$ ): A measure of how strongly the transformation amplifies errors.
The expected maximum activation value per block: The average of the maximum absolute values within each MX block.

The paper argues that block-diagonal transformations (as used in prior work) are a special case that cannot globally balance these factors. Instead, a full (non-block-diagonal) affine transformation is required to explicitly account for both the MX block structure and the global feature distribution.

B. LATMiX Architecture

LATMiX learns a set of invertible affine transformations ($T(x) = Ax + v$) that can be integrated into the linear layers of the model, resulting in no inference overhead.

Transformation Locations:
- $T_1$ : Applied globally to the residual path (input embeddings).
- $T_2$ : Applied locally to vector values in Multi-Head Attention (MHA) and within Feed-Forward Networks (FFN).
- $T_3$ : An online block-Hadamard transformation (retained for consistency with previous baselines).
Parameterization: To ensure invertibility and enable standard deep learning optimization, the authors propose two parameterizations for the matrix $A$ $A$ :
1. LU Decomposition: $A = P(L + \text{diag}(s))$ , where $L$ is a lower unitriangular matrix.
2. QR Decomposition: $A = Q(R + \text{diag}(s))$ , where $Q$ is orthogonal (via matrix exponential of a skew-symmetric matrix) and $R$ is an upper triangular matrix.
- A volume-preserving regularization is applied to ensure the determinant remains close to 1, preventing numerical instability while allowing simultaneous expansion or contraction in individual directions.
Optimization Objective:
- Distillation Loss: Instead of enforcing strict computational invariance (which requires orthogonal matrices), LATMiX uses a Kullback-Leibler (KL) divergence loss between the full-precision teacher model and the transformed quantized student model. This allows the transformation to deviate from orthogonality to better adapt to the data distribution.
- Regularization: A term to keep the volume of the transformation close to 1.

C. Folding Mechanism

A critical contribution is the ability to fold the learned affine transformations ( $A$ and $v$ ) directly into the weight matrices and bias terms of the LLM.

The inverse transformation is integrated into the input projections (K, Q, V) of the attention mechanism.
The forward transformation is integrated into the output projections of the attention and FFN layers.
This ensures that the model behaves exactly like a standardized quantized model during inference, without additional computational costs.

3. Main Contributions

Theoretical Insight: Providing a theoretical bound showing that quantization error in MX formats depends on both the global feature distribution and the block structure, and proving that block-diagonal transformations are insufficient for optimal error reduction.
LATMiX Algorithm: Introducing a method for learning general invertible affine transformations (not limited to rotations) using LU/QR parameterizations and distillation loss.
Zero-Overhead Deployment: Demonstrating that these complex learned transformations can be mathematically integrated into existing linear layers, preserving the inference speed and memory usage of standardized quantized models.
Comprehensive Evaluation: Validating the method across various model sizes (1B to 32B) and architectures (Llama, Qwen) on diverse benchmarks.

4. Experimental Results

The authors evaluated LATMiX on seven Zero-Shot Reasoning Benchmarks (ARC, HellaSwag, WinoGrande, PIQA, BoolQ, OpenBookQA) and the WikiText2 perplexity task.

Performance Improvement: LATMiX consistently outperformed strong baselines, including:
- Standard MX quantization (RTN, GPTQ).
- Rotation-based methods (QuaRot, SpinQuant).
- Block-diagonal rotation methods (MR-GPTQ, OSTQuant).
- Other transformation-based methods (FlatQuant).
Accuracy Recovery: At MXFP4 quantization (4-bit), LATMiX achieved up to a 4% improvement in average accuracy recovery compared to the best baselines. For example, LATMiX-LU on Llama-3.2-1B achieved a recovery of 93.88% versus 93.07% for MR-GPTQ.
Robustness: The method showed consistent improvements across various model sizes (1B to 32B) and quantization granularities (MXFP4 and MXINT4).
Ablation Studies:
- Full vs. Block: Full (tensor-wise) affine transformations performed significantly better than block-diagonal ones.
- Affine vs. Orthogonal: Learning general affine transformations (with bias) was superior to learning only orthogonal (rotation) matrices.
- Loss Function: The KL distillation loss proved more effective for zero-shot generalization than MSE or Cross-Entropy losses.

5. Significance

Bridging the Gap: LATMiX successfully bridges the gap between the need for aggressive outlier suppression (via global transformations) and the structural constraints of modern hardware-supported MX formats.
Hardware Efficiency: By integrating the transformations, LATMiX enables the use of powerful, learned mathematical operations without sacrificing the inference efficiency gains provided by MX quantization.
Future Direction: The work suggests that restricting transformations to rotations or block-diagonal matrices is suboptimal. It paves the way for learning more expressive, data-dependent transformations that respect the specific data formats of next-generation AI hardware.

In summary, LATMiX demonstrates that by theoretically analyzing the interaction between activation distributions and MX block structures, and by learning flexible affine transformations, near-lossless quantization of LLMs at extremely low bitwidths (4-bit) can be achieved without causing inference overhead.