Tag: matrix orthogonalization

  • QuaRot: Enabling End-to-End 4-Bit LLM Inference via Randomized Rotations

    QuaRot: Enabling End-to-End 4-Bit LLM Inference via Randomized Rotations

    TL;DR

    • QuaRot enables full 4-bit quantization across weights, activations, and the KV cache without retaining high-precision channels.
    • The framework suppresses disruptive activation outliers by using Randomized Hadamard Rotations via computational invariance to disperse peak magnitudes across dimensions.
    • Offline orthogonal rotations fuse directly into weights and normalization layers with zero inference overhead, while lightweight online Hadamard transforms handle internal activations and attention states.

    The Challenge of Outliers in Low-Bit LLM Quantization

    While weight-only 4-bit quantization (W4A16) is widely adopted, true end-to-end low-precision Post-Training Quantization (PTQ)—quantizing weights, activations, and the KV Cache to 4 bits (W4A4 Quantization with 4-bit KV Cache)—has historically resulted in severe accuracy degradation.

    The primary difficulty stems from Activation Outliers: persistent, extreme-value activations in specific channels with magnitudes orders of magnitude larger than standard features. Conventional uniform quantizers cannot allocate sufficient dynamic range to represent these channels without crushing non-outlier features to zero.

    Prior Post-Training Quantization approaches handle outliers through differing mechanisms that carry distinct trade-offs:

    • Outlier Channel Migration vs. Mixed-Precision Retention: SmoothQuant uses per-channel scaling factors to migrate quantization difficulty from activations to weights, targeting 8-bit quantization without keeping mixed-precision channels. In contrast, frameworks like LLM.int8() (and later works like QUIK and Atom) detect outliers dynamically or offline and retain those sensitive channels in higher precision (such as FP16).
    • Kernel and Hardware Overhead: For methods retaining mixed-precision channels, specialized execution kernels are required, introducing memory layout overhead, non-contiguous indexing, and hardware underutilization during matrix multiplication.

    Sources: GitHub – spcl/QuaRot: Code for Neurips24 paper: QuaRot, an end-to-end 4-bit inference of large language models, microsoft.com/en-us/research/publication/quarot-outlier-free-4-bit-in…, QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs, openreview.net/forum, Vercel Security Checkpoint

    What Is QuaRot and How Randomized Rotations Eliminate Outliers

    To eliminate the need for mixed-precision channels, QuaRot introduces an end-to-end Post-Training Quantization (PTQ) framework that achieves full W4A4 Quantization and 4-bit KV Cache compression. Rather than keeping outlier channels in FP16 during inference, QuaRot suppresses them across the model using Randomized Hadamard Rotations applied via Computational Invariance.

    A randomized Hadamard transformation multiplies a standard Hadamard matrix H by a diagonal matrix diag(s), where s is a vector of random signs drawn from {+1, -1}. A weight matrix W is defined as mu-incoherent if:

    max(W) <= mu * ||W||_F / sqrt(m * n)

    When incoherence is high, large individual elements emerge as outliers relative to the average magnitude. Multiplying by orthogonal randomized Hadamard matrices spreads high-energy spikes across all dimensions, systematically lowering incoherence and shaping activation distributions into uniform forms amenable to low-bit uniform quantizers.

    Under the principle of Computational Invariance, transformers remain invariant to orthogonal transformations across linear and normalization layers. Specifically, multiplying an input projection weight matrix on the left by Q^T and post-multiplying the preceding output projection matrix on the right by Q yields mathematically equivalent outputs. This invariance holds across RMSNorm operations because normalization commutes with orthogonal rotations when rescaling parameters (diag(alpha)) are fused into the subsequent weights:

    RMSNorm(X) = RMSNorm(X * Q^T) * Q

    QuaRot structures its modifications into two stages, illustrated below via mathematical pseudocode:

    # Stage 1: Transformations and Invariance (Mathematical Pseudocode)
# Applied to all input projections following RMSNorm (W_q, W_k, W_v, W_gate, W_up):
W_in      <- Q^T * diag(alpha) * W_in

# Output projections post-multiplied by Q; FFN down-proj absorbs H:
W_down    <- H * W_down * Q

# Attention value and output projections:
W_v       <- W_v * (I (x) H_d_h)
W_out     <- (H_n_h (x) I) * (I (x) H_d_h) * W_out = H * W_out

# Online post-RoPE transformations for queries and keys:
Q_rot     <- Pos(X * W_q) * (I (x) H_d_h)
K_rot     <- Pos(X * W_k) * (I (x) H_d_h)

# Stage 2: Quantization
Scale     <- max(|row|) / 7
Quantized <- round(row / Scale)

    In the first stage, RMSNorm scale parameters (diag(alpha)) are absorbed into adjacent input projection weights (including W_q, W_k, W_v, W_gate, and W_up), and randomized Hadamard rotation Q matches the hidden dimension to rotate hidden states (X <- X * Q). Because Q is merged directly into static weights offline, these inter-block rotations incur zero runtime inference overhead. Within the Feed-Forward Network (FFN), an online Hadamard transform is applied before the down-projection layer and cancelled out by fusing H into W_down. In the Multi-Head Attention module, head-wise Hadamard transforms H_d_h are applied to W_v and W_out, and an online Hadamard heads block applies (H_n_h (x) I) to attention activations Z. For keys and queries, online head-wise Hadamard rotations are applied post-RoPE, preserving attention scores while eliminating outliers from the KV Cache.

    In the second stage, weights are quantized (via GPTQ or round-to-nearest), while activations undergo online symmetric per-token quantization scaled by dividing the absolute maximum by 7. The KV Cache is quantized using asymmetric 4-bit quantization with a group size of 128.

    Empirically, on LLaMA-2 70B, QuaRot achieves full W4A4KV4 inference with a WikiText-2 perplexity loss of at most 0.47 (degrading from 3.32 to 3.79), while retaining over 99% of original zero-shot accuracy. In hardware benchmarks, it achieves prefill speedups of up to 2.16x on LLaMA-2 7B and 3.33x on LLaMA-2 70B (batch size 64, sequence length 2048), alongside up to 3.89x memory savings during decoding due to 4-bit KV caching.

    Sources: QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs

    Closing thoughts

    By leveraging computational invariance, QuaRot treats activation outliers not as an unavoidable feature of large language models, but as a coordinate-alignment artifact that can be systematically rotated away. This contrasts with mixed-precision retention strategies such as LLM.int8(), QUIK, or Atom, which keep sensitive channels in FP16 or INT8 and require specialized, non-contiguous matrix multiplication kernels, as well as SmoothQuant, which migrates difficulty through channel scaling.

    A key distinction in QuaRot is that offline rotations are fused directly into linear weights and normalization layers at zero runtime cost, while online Hadamard operations are restricted to intermediate points like the FFN down-projection and post-RoPE attention heads. Structurally neutralizing incoherence prior to uniform quantization allows standard integer GEMM kernels to execute across all weights, activations, and KV cache states with minimal accuracy degradation.

    Frequently Asked Questions

    What is QuaRot?

    QuaRot is an end-to-end Post-Training Quantization framework that enables 4-bit weight, activation, and KV Cache quantization without retaining high-precision channels. It eliminates activation outliers by applying Randomized Hadamard Rotations through computational invariance.

    Why do activation outliers cause issues in low-bit LLM quantization?

    Activation outliers are extreme values in specific channels that are orders of magnitude larger than standard features. Conventional uniform quantizers cannot allocate sufficient dynamic range for these channels without crushing the remaining non-outlier features to zero.

    How does QuaRot differ from approaches like SmoothQuant and LLM.int8()?

    SmoothQuant uses per-channel scaling to migrate quantization difficulty from activations to weights for 8-bit quantization without keeping mixed-precision channels. LLM.int8() isolates outlier channels dynamically to compute them in FP16. QuaRot takes a different approach: it applies orthogonal Hadamard rotations to disperse peak magnitudes across dimensions, allowing uniform 4-bit quantization across all matrix operations without mixed-precision channels.

    How does QuaRot quantize the KV Cache?

    QuaRot applies online head-wise Randomized Hadamard Rotations post-RoPE to keys and queries to remove outliers directly from the KV Cache. The cache is then compressed using asymmetric 4-bit quantization with a group size of 128.


    Keep reading: LLM Inference & Serving

    All guides in this series: LLM Inference & Serving Guides: vLLM, Ollama & Hugging Face

  • Muon Optimizer: Scaling Matrix Orthogonalization for Efficient Deep Learning Training

    Muon Optimizer: Scaling Matrix Orthogonalization for Efficient Deep Learning Training

    TL;DR

    • The Muon optimizer trains 2D hidden layer weight matrices via steepest descent under the spectral norm using efficient iterative Newton-Schulz orthogonalization.
    • Muon operates as a hybrid optimizer: it manages 2D hidden-layer weight matrices, while AdamW handles everything else — genuinely 1D parameters (biases, normalization scales) plus the input embedding and output layers, which are excluded from Muon for role-based reasons even though they are typically 2D matrices themselves.
    • Compared to AdamW, Muon cuts training compute — about 1.35x faster wall-clock in Keller Jordan’s NanoGPT speedrun and up to 2x fewer FLOPs under Moonlight’s compute-optimal scaling law — while also cutting VRAM overhead by eliminating the second-moment buffer and reducing the high condition number (near low-rank concentration) of raw weight updates.

    Introduction: Understanding the Muon Optimizer and Orthogonal Updates

    The Muon Optimizer—short for MomentUm Orthogonalized by Newton-Schulz—is an optimization method designed for training the Hidden Layers of neural networks. Unlike conventional first-order optimizers such as SGD, Adam, and AdamW, which treat weight parameters as flattened 1D Parameters and apply coordinate-wise scaling, Muon treats parameters as 2D matrices. It performs Steepest Descent with respect to the Spectral Norm (the matrix operator norm), optimizing both the directional geometry and the magnitude of Weight Updates.

    Muon was introduced in late 2024 by Keller Jordan and collaborators, building on foundational theoretical work by Jeremy Bernstein and Laker Newhouse (‘Old Optimizer, New Norm: An Anthology’), which frames this Steepest Descent step under an RMS-to-RMS operator norm — a normalized formulation that is effectively equivalent to the classical Spectral Norm used elsewhere in this article, differing mainly in how the norm is scaled by matrix dimensions rather than in the underlying geometry. It was first benchmarked at scale in the open-source Modded-NanoGPT speedrun repository. The optimizer has since been expanded through theoretical analyses, such as the Newton–Muon optimizer studied by Du and Su at UPenn, and scaled by Moonshot AI (the Kimi Team) to train the Moonlight 3B/16B Mixture-of-Experts (MoE) model across 5.7 trillion tokens.

    At its core, Muon operates through a multi-step update pipeline:

    • Momentum Collection: Gradient updates are first accumulated using standard Nesterov momentum.
    • Newton-Schulz Orthogonalization: Rather than computing an exact, computationally expensive Singular Value Decomposition (SVD) on GPUs, Muon applies an iterative Newton-Schulz Iteration (typically 5 steps). This quickly drives all singular values of the update matrix close to 1.0, generating an approximate Matrix Orthogonalization equivalent to the matrix-sign function and Polar Decomposition.
    • Aspect-Ratio Scaling: The orthogonalized update is rescaled according to matrix dimensions to preserve stable signal propagation across layers with varying shapes. Keller Jordan’s original implementation uses a ratio-based factor, sqrt(max(1, rows/cols)); the Moonlight paper (Kimi Team) proposes an explicit coefficient form instead, 0.2 * sqrt(max(A, B)), noting the two scalings are equivalent up to a global scale only when the matrices share the same second dimension, and can otherwise diverge.

    Muon is implemented as a hybrid optimizer. It is specialized strictly for 2D weight matrices in Hidden Layers and linear layers—such as Multi-Layer Perceptron (MLP) weights, attention projections in Transformer Architectures, and reshaped convolutional kernels. Everything else is optimized concurrently using a traditional optimizer like AdamW: genuinely 1D parameters such as Biases and Normalization Scales, plus the input embedding and output (unembedding) layers. Those embedding and output layers are typically 2D matrices themselves, but Muon’s authors exclude them for role-based reasons rather than dimensionality—embeddings on the theoretical grounds of the modular norm framework, and output layers on purely empirical grounds.

    This Matrix Orthogonalization scheme provides several key advantages over standard coordinate-wise optimizers:

    • Compute Efficiency and Convergence: Reported efficiency gains come from two distinct sources — Keller Jordan’s NanoGPT speedrun recorded roughly 1.35x faster wall-clock training time, while the Moonlight paper’s compute-optimal scaling law separately found up to 2x fewer training FLOPs versus AdamW to reach the same validation loss.
    • Lower Optimizer Memory Overhead: Muon tracks only a single first-moment momentum buffer per parameter and eliminates the second-moment buffer (v_t) required by AdamW, yielding notable VRAM savings on Hidden Layers.
    • Reducing Update Condition Number: By normalizing singular values, the optimizer prevents Weight Updates from collapsing into a small subset of high-variance singular directions—a near low-rank, high-condition-number pathology of raw gradients—better engaging the full representational capacity of weight matrices.
    • Empirical Performance: In addition to setting speedrun training records on single-node setups for CIFAR-10 and GPT-2 (small and medium), Muon has demonstrated stable scaling in multi-billion-parameter LLM pretraining. The Moonlight paper’s Distributed Muon implementation follows ZeRO-1 to partition optimizer states across the data-parallel dimension — with tensor parallelism usable only as a complementary, non-essential option — combined with adjusted weight decay.

    Sources: Muon: An optimizer for hidden layers in neural networks, GitHub – KellerJordan/Muon: Muon is an optimizer for hidden layers in neural networks · GitHub, Deriving Muon, medium.com/@md.abir1203/muon-optimizer-the-complete-guide-to-2-faster…, [2502.16982] Muon is Scalable for LLM Training

    Closing thoughts

    Taken together, these technical strengths illustrate how shifting from coordinate-wise updates to geometry-aware 2D matrix orthogonalization enables Muon to directly address the structural properties of hidden layers, yielding clear optimization advantages over standard first-order methods. Its hybrid architecture—applying efficient Newton-Schulz iterations to weight matrices while retaining AdamW for 1D parameters and embeddings—strikes a sensible balance between theoretical rigor and implementation practicality.

    Ultimately, with demonstrated FLOP reductions and lower memory overhead validated from single-node speedruns up to multi-trillion-token LLM pretraining, Muon establishes matrix-level spectral descent as a genuinely viable, high-efficiency paradigm for training modern neural networks.

    Frequently Asked Questions

    What is the Muon optimizer?

    Muon (MomentUm Orthogonalized by Newton-Schulz) is an optimization method designed for training hidden layers of neural networks by treating parameters as 2D matrices rather than flattened 1D vectors. It performs steepest descent with respect to the spectral norm to optimize the directional geometry and magnitude of weight updates.

    How does Muon achieve matrix orthogonalization efficiently?

    Instead of computing an expensive Singular Value Decomposition (SVD) on GPUs, Muon applies an iterative Newton-Schulz iteration, typically in 5 steps. This quickly drives update matrix singular values close to 1.0, providing an approximate orthogonalization equivalent to the matrix-sign function and polar decomposition.

    Why is Muon used as a hybrid optimizer alongside AdamW?

    Muon is specialized strictly for 2D weight matrices in hidden layers, such as MLP weights and attention projections. Everything else uses a traditional optimizer like AdamW instead: 1D parameters such as biases and normalization scales, plus the input and output embedding layers. Those embedding layers are typically 2D themselves, but are excluded from Muon for role-based reasons—not because they are non-2D.

    What are the primary benefits of Muon compared to AdamW?

    Muon reduces training compute versus AdamW — about 1.35x faster wall-clock time in Keller Jordan’s NanoGPT speedrun, and up to 2x fewer FLOPs under Moonlight’s compute-optimal scaling law — while also lowering VRAM usage by eliminating AdamW’s second-moment buffer. Additionally, it helps reduce the high condition number of raw weight updates by preventing them from collapsing into a small subset of high-variance singular directions.


    Keep reading: Training & Architectures

    All guides in this series: LLM Training & Architecture Guides: Fine-Tuning, RL & New Models