Category: AI/ML Concepts

Explainers of AI/ML research techniques — optimizers, architectures and post-training methods — grounded in the original papers and official technical sources.

  • 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

  • Multi-Head Latent Attention: Scaling LLM Serving Throughput via Low-Rank KV Compression

    Multi-Head Latent Attention: Scaling LLM Serving Throughput via Low-Rank KV Compression

    TL;DR

    • Multi-Head Latent Attention (MLA) overcomes serving memory bottlenecks by compressing Key-Value states into compact latent vectors via low-rank projections rather than relying on coarse head-sharing.

    • During inference, MLA absorbs key and value up-projection matrices into query and output projections; it pairs this with Decoupled RoPE so that position embeddings do not interfere with matrix absorption.

    • In DeepSeek-V2, MLA enables a 93.3% reduction in KV cache footprint compared to DeepSeek 67B. Together with sparse MoE computation (activating 21B out of 236B parameters), this contributes to a 5.76-fold boost in maximum generation throughput, while separate architectural ablations show MLA outperforms standard Multi-Head Attention.

    The Memory Wall: Why the KV Cache Constrains LLM Serving

    During Autoregressive Generation in Large Language Models (LLMs), Inference Throughput is fundamentally bottlenecked by the Key-Value (KV) Cache. In standard Multi-Head Attention (MHA), the model must retain the Key and Value tensors of all previously generated tokens across every layer. These cached states reside in High-Bandwidth Memory (HBM) and must be repeatedly transferred into on-chip memory at every decoding step to compute attention for each new token. This persistent, high-volume memory traffic creates severe HBM bandwidth contention, capping decoding throughput and strictly limiting both the serving batch size and the practical context window length.

    Prior attempts to alleviate this memory wall rely primarily on Multi-Query Attention (MQA) and Grouped-Query Attention (GQA). Both techniques reduce memory requirements by forcing multiple query heads to share a reduced number of KV heads. While this head-sharing structure lowers the absolute volume of data stored in HBM and transferred to on-chip memory, it introduces an architectural trade-off: sharing KV projections across disparate query heads compromises per-head representational diversity and impairs retrieval expressiveness.

    Consequently, LLM serving faces a structural tension between memory bandwidth efficiency and model capacity. Standard MHA preserves distinct per-head projections at the cost of high HBM bandwidth consumption, whereas MQA and GQA reduce cache overhead at the expense of representational richness. Overcoming this bottleneck requires moving beyond coarse head-sharing, addressing the memory wall via Low-Rank Joint Compression of the stored representations into a Compressed KV Latent Vector rather than sacrificing individual per-head projection capacity.

    Sources: [2405.04434] DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model, arxiv.org/pdf/2405.04434, vertexaisearch.cloud.google.com/grounding-api-redirect/AUZIYQF73W_INb…, Not Acceptable!, Paper page – TransMLA: Multi-head Latent Attention Is All You Need

    Empirical Gains: KV Footprint Reduction and Serving Speedups in Practice

    The practical efficacy of this low-rank compression strategy is clearly demonstrated in modern architectures. Compared to the dense DeepSeek 67B baseline, DeepSeek-V2 reduces the Key-Value (KV) Cache footprint by 93.3% and boosts maximum generation throughput up to 5.76 times. This throughput gain is a system-wide result: the KV cache reduction stems directly from MLA’s low-rank compression, while overall serving efficiency is further multiplied by the DeepSeekMoE architecture, which activates only 21B parameters out of 236B total parameters per token.

    Standard Multi-Head Attention (MHA) caches all uncompressed keys and values across every attention head, requiring 2 * n_h * d_h * l elements per token, where n_h represents the head count, d_h is the per-head dimension, and l denotes the layer count. Alternative mechanisms like Grouped-Query Attention (GQA) and Multi-Query Attention (MQA) scale memory down to 2 * n_g * d_h * l elements (for n_g groups) and 2 * d_h * l elements, respectively, but sacrifice modeling capability.

    In contrast, MLA projects keys and values into a single low-rank latent vector c_t^KV of dimension d_c. However, standard Rotary Position Embedding (RoPE) is position-sensitive and couples directly with keys and queries. If RoPE were applied directly to the content keys, the key up-projection matrix W^UK could not be absorbed into the query projection during inference because a token-dependent rotation matrix would sit between them, destroying associativity and forcing the system to recompute prefix keys. MLA resolves this by introducing Decoupled RoPE: the content keys k_t^C remain unrotated (enabling matrix absorption), while a separate, low-dimensional shared key k_t^R of dimension d_h^R carries the positional embeddings.

    During inference, MLA caches only the compressed latent vector c_t^KV alongside the decoupled key k_t^R, requiring (d_c + d_h^R) * l elements per token. In the DeepSeek-V2 configuration (n_h = 128, d_h = 128, d_c = 512 = 4 * d_h, and d_h^R = 64 = d_h / 2), the cache requirement is (d_c + d_h^R) * l = 4.5 * d_h * l elements. In terms of per-token cache storage footprint, this is equivalent to GQA configured with only 2.25 groups. In architectural evaluations, MLA nevertheless achieves stronger performance than full MHA.

    Furthermore, during inference, the key up-projection matrix W^UK is absorbed directly into the query projection W^Q, and the value up-projection matrix W^UV is absorbed into the output projection W^O. As a result, the full per-head keys and values never need to be materialized in memory for attention computation, greatly reducing HBM bandwidth traffic.

    Sources: DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model

    Closing thoughts

    These architectural efficiencies highlight a practical approach to addressing memory bottlenecks in LLM serving. Rather than relying entirely on coarse head-sharing mechanisms like GQA or MQA that compromise representational diversity, Multi-Head Latent Attention demonstrates that memory bandwidth contention can be significantly mitigated through low-rank joint compression.

    A key design element is decoupling positional embeddings from content projections. Because standard RoPE operations are position-dependent and prevent key up-projection matrices from being absorbed into queries, Decoupled RoPE isolates the positional stream. This allows the content stream to remain purely linear, enabling inference engines to absorb W^UK into W^Q and W^UV into W^O without materializing uncompressed per-head Key-Value states in memory.

    MLA does not eliminate all architectural trade-offs: practitioners must still select appropriate compression dimensions d_c, allocate cache for decoupled positional keys, and navigate specialized CUDA kernel implementations. Nevertheless, its empirical impact in DeepSeek-V2—where MLA’s 93.3% KV cache reduction relative to DeepSeek 67B combines with sparse MoE execution to deliver up to a 5.76-fold throughput improvement, alongside ablation results outperforming standard MHA—makes latent compression a compelling pattern for high-throughput serving systems.

    Frequently Asked Questions

    Why does the standard Key-Value (KV) cache create a bottleneck in LLM serving?

    In standard Multi-Head Attention, retaining uncompressed Key and Value tensors across all heads and layers requires repeatedly transferring massive volumes of data from High-Bandwidth Memory to on-chip memory at each decoding step. This creates severe memory bandwidth contention that restricts decoding throughput, serving batch sizes, and context window lengths.

    How does Multi-Head Latent Attention (MLA) differ from MQA and GQA?

    MQA and GQA reduce memory by forcing multiple query heads to share a reduced count of KV heads, which compromises representational diversity and retrieval expressiveness. In contrast, MLA compresses stored representations into a compact latent vector via low-rank joint compression without relying on coarse head-sharing.

    What does MLA store in the KV cache per token?

    MLA caches only the Compressed KV Latent Vector alongside a decoupled key carrying Decoupled Rotary Position Embedding (Decoupled RoPE). This requires (d_c + d_h^R) * l elements per token rather than materializing full uncompressed keys and values across every head.

    How does MLA optimize computation during inference?

    During inference, the key up-projection matrix is absorbed directly into the query projection, and the value up-projection matrix is absorbed into the output projection. Because Decoupled RoPE keeps the content stream position-free, this linear absorption avoids having to compute or materialize full per-head Key and Value tensors in memory during attention.

    What empirical throughput and memory gains were achieved by DeepSeek-V2 using MLA?

    Compared to DeepSeek 67B, DeepSeek-V2 reduced the KV cache footprint by 93.3% and increased maximum generation throughput up to 5.76 times. The KV cache reduction is driven by MLA, while the throughput gains reflect both the compressed cache and the sparse DeepSeekMoE architecture activating 21B out of 236B parameters. In terms of cache memory footprint per token, MLA is equivalent to GQA with just 2.25 groups, while separate paper ablations show it delivers stronger modeling performance than full Multi-Head Attention.


    Keep reading: LLM Inference & Serving

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

  • Differential Transformer: Canceling Context Noise with Differential Attention

    Differential Transformer: Canceling Context Noise with Differential Attention

    TL;DR

    • The Differential Transformer cancels attention noise and focuses attention by calculating scores as the difference between two separate softmax attention maps.
    • It achieves standard Transformer performance with approximately 35% fewer parameters or training tokens while reducing contextual hallucinations and enhancing long-context retrieval.
    • The DEX adaptation method enables existing pretrained language models to effectively incorporate key advantages of differential attention using minimal adaptation data (< 0.01%).

    Introduction to the Differential Transformer

    Developed by researchers from Microsoft Research and Tsinghua University—including Tianzhu Ye, Li Dong, Yuqing Xia, Yutao Sun, Yi Zhu, Gao Huang, and Furu Wei—the Differential Transformer (Diff Transformer) introduces a foundational architectural modification to standard Transformer models. Accepted as an Oral Presentation at ICLR 2025, the architecture addresses an inherent limitation in standard softmax attention: attention heads tend to overallocate probability weight to irrelevant context, which is often viewed as a consequence of softmax normalizing scores into strictly positive values summing to 1. In the authors’ motivation, this attention noise may relate to observed challenges such as contextual hallucinations, prompt-order sensitivity, and the “lost-in-the-middle” phenomenon.

    Drawing inspiration from noise-canceling headphones and differential amplifiers in electrical engineering, the Differential Transformer incorporates a Differential Attention Mechanism designed to cancel common-mode background noise and extraneous context. The architecture demonstrates several key advantages across language modeling tasks:

    • Improved Scaling Efficiency: The Diff Transformer matches the performance of standard Transformers while requiring approximately 35% fewer parameters or training tokens.
    • Long-Context Robustness: The model demonstrates stronger retrieval performance on needle-in-a-haystack tasks and long-context evaluations up to 64k tokens, showing less degradation across needle depths and context lengths.
    • Reduced Contextual Hallucinations: By being less distracted by irrelevant context, it decreases hallucination rates on text summarization and question-answering tasks.
    • In-Context Learning Stability: The architecture mitigates sensitivity to the permutation order of few-shot examples, offering higher stability during prompt reordering.
    • Quantization Readiness: Alongside noise cancellation, the architecture suppresses extreme activation outliers, facilitating cleaner post-training low-bit quantization as a distinct representational benefit.

    Sources: [2410.05258] Differential Transformer, Differential Transformer, microsoft.com/en-us/research/publication/differential-transformer, Microsoft’s Differential Transformer cancels attention noise in LLMs, OpenReview (paper discussion page)

    Cancelling Context Noise Through Attention Map Subtraction

    At the heart of these performance gains is the specific way Differential Attention operates. The Differential Attention Mechanism cancels attention noise by calculating attention scores as the difference between two separate softmax attention maps:

    DiffAttn(X) = (softmax(Q1 * K1^T / sqrt(d)) - lambda * softmax(Q2 * K2^T / sqrt(d))) * V
    

    where query and key projections are split into two distinct paths (Q1, K1 and Q2, K2) and λ is a learnable scalar parameter. In practice, λ is reparameterized via learnable vectors λ_q1, λ_k1, λ_q2, and λ_k2 (specifically computed as λ = exp(λ_q1 · λ_k1) – exp(λ_q2 · λ_k2) + λ_init, where λ_init is set according to the layer index). This subtraction directly cancels common-mode noise and counteracts the tendency of conventional architectures to overallocate attention to irrelevant context, amplifying focus on relevant tokens and promoting the emergence of sparse attention patterns.

    By canceling attention noise through this dual-map subtraction, the Diff Transformer achieves consistent empirical advantages across both scaling and downstream evaluations, outperforming standard Transformer baselines across various model sizes and training token volumes.

    Sources: Differential Transformer

    Follow-Up Analysis and Adapting Pretrained Models

    While the original Differential Transformer paper framed its mechanism primarily around noise cancellation, an independent follow-up study titled “Understanding Differential Transformer Unchains Pretrained Self-Attentions” (arXiv:2505.16333) by Chaerin Kong, Jiho Jang, and Nojun Kwak investigated the underlying factors behind its empirical success. The authors uncovered three key factors:

    • Enhanced expressivity via negative attention
    • Reduced redundancy among attention heads
    • Improved learning dynamics

    Because the standard Differential Transformer architecture requires training from scratch, which limits its direct use with existing pretrained models, the study proposed DEX, an efficient adaptation method designed to integrate the advantages of differential attention into pretrained models.

    Rather than computing two separate softmax maps, DEX reuses the original softmax attention scores and adds a lightweight differential operation on the output value matrix. This formulation allows pretrained language models to effectively incorporate key advantages of differential attention while remaining lightweight in training and inference, achieving significant performance gains across diverse benchmarks with minimal adaptation data (< 0.01%).

    Sources: Understanding Differential Transformer Unchains Pretrained Self-Attentions

    Closing thoughts

    In conclusion, by subtracting attention noise through two softmax maps, the Differential Transformer addresses a key limitation of standard softmax attention. Its empirical gains across scaling efficiency, long-context retrieval, and contextual hallucination reduction, alongside activation outlier suppression, highlight the advantages of differential attention representations. Furthermore, independent analysis reveals that benefits such as negative attention expressivity, reduced head redundancy, and improved learning dynamics also underpin its performance, inspiring practical adaptation methods like DEX for pretrained models. Differential attention thus stands as a grounded architectural refinement that delivers tangible, parameter- and data-efficient performance gains.

    Frequently Asked Questions

    What is the primary problem the Differential Transformer solves?

    The Differential Transformer addresses attention noise in standard softmax attention, where attention heads tend to overallocate probability weight to irrelevant context, often seen as a consequence of softmax normalizing scores into strictly positive values summing to 1.

    How does the Differential Attention Mechanism cancel context noise?

    It computes query and key projections across two separate paths, takes the difference between two distinct softmax attention maps weighted by a learnable parameter λ, and subtracts common-mode background noise to focus on salient tokens.

    What performance advantages does the Differential Transformer provide?

    It matches standard Transformer performance using approximately 35% fewer parameters or training tokens, improves long-context retrieval across various depths and lengths up to 64k tokens, lowers contextual hallucination rates in QA and summarization, stabilizes in-context learning against prompt reordering, and suppresses activation outliers for cleaner low-bit quantization.

    What is the DEX method and why is it needed?

    Because the original Differential Transformer architecture requires training from scratch, the DEX method was proposed in an independent follow-up study (arXiv:2505.16333 by Chaerin Kong, Jiho Jang, and Nojun Kwak) to adapt pretrained language models. Instead of computing two softmax maps, DEX reuses existing softmax attention scores and applies a lightweight differential operation on the output value matrix using minimal adaptation data (< 0.01%).


    Keep reading: Training & Architectures

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

  • 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

  • Byte Latent Transformer: Replacing Tokenizers with Dynamic Entropy-Based Patching

    Byte Latent Transformer: Replacing Tokenizers with Dynamic Entropy-Based Patching

    TL;DR

    • The Byte Latent Transformer introduces a tokenizer-free architecture that overcomes the noise brittleness and vocabulary biases of traditional subword tokenization.
    • It processes raw byte streams efficiently by dynamically segmenting bytes into patches based on next-byte entropy.
    • BLT matches or exceeds the performance of traditional tokenizer-based models like LLaMA at scales up to 8 billion parameters.

    Introduction: The Shift from Fixed Tokenizers to Byte-Level Modeling

    Traditional large language models depend heavily on heuristic Subword Tokenization algorithms, such as Byte-Pair Encoding (BPE), to segment text into static subword vocabularies. While widely adopted, fixed Subword Tokenization introduces fundamental structural vulnerabilities:

    • Brittleness to Noise and Manipulation: Models exhibit high sensitivity to typos, spelling errors, and character-level perturbations, while struggling on character-level manipulation and arithmetic tasks.
    • Vocabulary and Representation Issues: Fixed vocabularies lead to out-of-vocabulary fragmentation and create structural biases favoring English and high-resource scripts over low-resource languages.

    Operating directly on raw byte streams eliminates the need for tokenizers, but naive Byte-Level Language Model design introduces severe computational bottlenecks. Expanding text directly into raw byte sequences dramatically inflates sequence lengths, imposing extreme quadratic computational overhead on standard Transformer self-attention mechanisms.

    To overcome both the rigidity of subword tokenizers and the computational expense of naive byte sequences, researchers from Meta FAIR, the University of Washington, and the University of Chicago (Artidoro Pagnoni, Ram Pasunuru, Pedro Rodriguez, John Nguyen, Benjamin Muller, Margaret Li, Chunting Zhou, Lili Yu, Jason Weston, Luke Zettlemoyer, Gargi Ghosh, Mike Lewis, Ari Holtzman, and Srinivasan Iyer) introduced the Byte Latent Transformer (BLT). BLT represents the first compute- and FLOP-controlled scaling study demonstrating that a tokenizer-free Byte-Level Language Model can match or exceed the performance of standard tokenizer-based models, such as LLaMA baselines, across scales up to 8 billion parameters and trillions of training bytes.

    Sources: youtube.com/watch, Byte Latent Transformer: Patches Scale Better Than Tokens, [2412.09871] Byte Latent Transformer: Patches Scale Better Than Tokens, OpenReview (paper discussion page), Byte Latent Transformer: Patches Scale Better Than Tokens – ACL Anthology

    Mechanism: Dynamically Segmenting Patches via Next-Byte Entropy

    At the core of this breakthrough is the Byte Latent Transformer (BLT) architecture’s approach to sequence structure: raw bytes are segmented into dynamically sized patches through Dynamic Entropy-Based Patching, serving as the primary units of computation and bypassing the need for Subword Tokenization. Patch segmentation is driven by Next-Byte Entropy predicted by an Entropy Model, which enables the model to dynamically allocate more compute and model capacity where increased data complexity demands it.

    When the underlying data is predictable, Dynamic Entropy-Based Patching constructs longer patches. This entropy-based segmentation enhances both training and inference efficiency, allowing BLT to simultaneously scale model size and patch length under fixed inference costs.

    Sources: Byte Latent Transformer: Patches Scale Better Than Tokens

    Closing thoughts

    Ultimately, by tying sequence segmentation directly to informational entropy rather than static vocabulary heuristics, the Byte Latent Transformer successfully resolves the long-standing computational bottleneck of naive byte-level processing. In my view, the fact that BLT matches or exceeds standard baselines like LLaMA at scales up to 8 billion parameters demonstrates that the structural brittleness and vocabulary biases of BPE are no longer necessary compromises for scalable architectures. Dynamically scaling patch lengths according to data predictability proves to be a well-reasoned solution, allowing byte-level modeling to achieve character-level robustness while strictly maintaining compute and inference efficiency.

    Frequently Asked Questions

    What are the main drawbacks of traditional subword tokenization methods like BPE?

    Fixed subword tokenization introduces brittleness to noise, typos, and character-level perturbations, while causing out-of-vocabulary fragmentation. It also struggles with character manipulation tasks and creates structural biases favoring English and high-resource scripts over low-resource languages.

    Why is naive byte-level modeling computationally inefficient?

    Directly expanding text into raw byte sequences dramatically inflates sequence lengths. This imposes extreme quadratic computational overhead on standard Transformer self-attention mechanisms.

    How does the Byte Latent Transformer (BLT) process byte sequences efficiently?

    BLT uses Dynamic Entropy-Based Patching driven by next-byte entropy predicted by an Entropy Model. It constructs longer patches when data is predictable, dynamically allocating more compute and capacity where data complexity demands it.

    How does the performance of BLT compare to standard tokenizer-based models?

    BLT matches or exceeds the performance of standard tokenizer-based models, such as LLaMA baselines, across scales up to 8 billion parameters and trillions of training bytes.


    Keep reading: Training & Architectures

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

  • Titans: Learning to Memorize at Test Time with Neural Long-Term Memory

    Titans: Learning to Memorize at Test Time with Neural Long-Term Memory

    TL;DR

    • Titans introduces a dual-memory architecture combining attention for short-term dependencies with a Neural Long-Term Memory module for persistent historical context.
    • The framework resolves trade-offs between quadratic attention and lossy recurrent models through three variants: Memory as Context, Memory as Gate, and Memory as Layer.
    • Titans scales to context windows exceeding 2 million tokens and achieves superior performance across language, genomics, time series, and Needle-in-a-Haystack benchmarks.

    Introduction to Titans and Test-Time Neural Memory

    Titans introduces a dual-memory architectural paradigm that combines attention with a dedicated Neural Long-Term Memory module. In this framework, attention acts as short-term memory, leveraging its accurate token dependency modeling over a limited context window, while the Neural Long-Term Memory module acts as persistent memory by learning to memorize historical context. This Neural Long-Term Memory allows the system to retain and utilize long-past information to assist attention on the current context, while maintaining the advantages of fast parallelizable training and fast inference.

    Built upon the integration of these two complementary modules, the Titans family features three architectural variants designed to incorporate memory effectively: Memory as Context (MAC), Memory as Gate (MAG), and Memory as Layer (MAL). Across evaluations spanning language modeling, common-sense reasoning, genomics, and time series benchmarks, Titans demonstrates greater effectiveness than Transformers and modern linear recurrent models. Furthermore, the architecture scales to context window sizes exceeding 2M tokens while achieving higher accuracy on Needle-in-a-Haystack tasks relative to baseline models.

    Sources: Titans: Learning to Memorize at Test Time

    Architectural Paradigms: Comparing Attention and Recurrent Memory Trade-Offs

    To understand the motivation behind this architecture, it is helpful to examine modern sequence modeling, which has long been characterized by a fundamental tension between attention-based architectures and recurrent models. On one hand, Transformers excel at capturing precise, fine-grained, and direct dependencies between tokens across a context window. However, self-attention incurs a quadratic computational complexity of O(N^2) with respect to sequence length, resulting in extreme memory footprints and computational bottlenecks that severely restrict the size of usable context windows.

    In contrast, Recurrent Neural Networks (RNNs) and linear recurrent models offer linear computational complexity of O(N) and fast inference by compressing historical context into a fixed-size hidden state. However, this fixed-capacity compression introduces significant information loss and degrades associative recall over extensive contexts.

    Titans reconciles these conflicting paradigms by structuring memory hierarchically. Rather than relying entirely on quadratic self-attention or static hidden-vector compression, the architecture assigns short-term, high-resolution local dependencies to an attention mechanism while offloading past historical context to a dynamic Neural Long-Term Memory module. This design retains the fast inference and linear compute scaling of recurrent models without sacrificing the associative fidelity required for long-context sequence modeling.

    Sources: [2501.00663] Titans: Learning to Memorize at Test Time, youtube.com/watch, medium.com/@avd.sjsu/titans-learning-to-memorize-at-test-time-271e944…, OpenReview (paper discussion page), Paper page – Titans: Learning to Memorize at Test Time

    Closing thoughts

    Ultimately, by reframing sequence modeling as a hierarchical dual-memory system rather than a single compromised mechanism, Titans convincingly resolves the structural tension between quadratic attention bottlenecks and lossy recurrent compression. In our view, assigning persistent historical context to an active neural long-term memory while reserving standard attention for high-resolution local dependencies is the key architectural insight that allows the model to scale past two million tokens without degrading retrieval fidelity.

    Furthermore, the demonstrated gains across domains as varied as genomics, time series, and language benchmarks suggest that modular memory integration—whether via context, gating, or dedicated layers—is fundamentally more effective than forcing a uniform mechanism to handle both immediate and distant dependencies. In doing so, Titans proves that sequence models can achieve fast inference and linear scaling without sacrificing the associative accuracy essential for vast context windows.

    Frequently Asked Questions

    What is the dual-memory architecture used in Titans?

    Titans combines an attention mechanism that serves as short-term memory with a dedicated Neural Long-Term Memory module for persistent memory. This design allows the model to retain past information to support attention on the current context while maintaining fast parallelizable training and fast inference.

    What are the three architectural variants of Titans?

    The Titans family includes three architectural variants: Memory as Context (MAC), Memory as Gate (MAG), and Memory as Layer (MAL).

    How does Titans resolve the trade-offs between Transformers and recurrent models?

    Titans organizes memory hierarchically by assigning short-term, high-resolution local dependencies to attention while offloading past historical context to a Neural Long-Term Memory module. This enables linear compute scaling and fast inference without the information loss typical of static recurrent compression.

    What context window size can Titans achieve?

    Titans scales to context window sizes exceeding 2 million tokens while delivering higher accuracy on Needle-in-a-Haystack tasks compared to baseline models.


    Keep reading: Training & Architectures

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

  • Group Relative Policy Optimization for Efficient Reinforcement Learning in Language Models

    Group Relative Policy Optimization for Efficient Reinforcement Learning in Language Models

    TL;DR

    • Group Relative Policy Optimization (GRPO) eliminates the Critic Model from standard PPO architectures, cutting GPU VRAM usage nearly in half during LLM reinforcement learning.
    • Standard GRPO relies on static sampling and fixed rollouts, which can waste computational resources on easy prompts while under-training difficult reasoning tasks.
    • Combining GRPO with dynamic techniques like Prompt-GDRO and Rollout-GDRO directs compute toward hard tasks, significantly improving reasoning accuracy.

    Understanding Group Relative Policy Optimization (GRPO) in LLM Training

    Reinforcement learning (RL) training methods enhance the alignment and reasoning performance of Large Language Models (LLMs), specifically by improving their capacity to understand human intents, follow user instructions, and strengthen inferential processing.

    Across the broader LLM lifecycle, RL strategies are integrated across several phases: pre-training, alignment fine-tuning, and reinforced reasoning. In particular, RL approaches deployed during the reinforced reasoning phase act as a primary driver for advancing model reasoning limits, with significant focus placed on Reinforcement Learning with Verifiable Rewards (RLVR). Fine-tuning and evaluation within these training frameworks draw upon varied data sources and benchmarks, including human-annotated datasets, AI-assisted preference data, and program-verification-style corpora.

    Sources: Reinforcement Learning Meets Large Language Models: A Survey of Advancements and Applications Across the LLM Lifecycle

    Architectural Shift: Eliminating the Critic Model for Memory Efficiency

    To operationalize these reinforced reasoning phases effectively, attention has increasingly turned to optimizing trainer architectures. In standard Proximal Policy Optimization (PPO), reinforcement learning relies on an Actor-Critic architecture. Fine-tuning a Large Language Model (LLM) Actor Model under PPO requires training a separate Critic Model of a similar size. Running this dedicated value network nearly doubles GPU VRAM requirements and compute overhead during post-training.

    Group Relative Policy Optimization (GRPO) executes an architectural shift by completely eliminating the Critic Model. Rather than using a value network to predict output values, GRPO generates a group of outputs {o1, o2, …, oG} for each input prompt q under the current old policy pi_theta_old. Scores for each output are computed using a reward function—such as rule-based checks for answer correctness—and then standardized across the group to determine relative performance. Removing the Critic Model cuts VRAM usage nearly in half during reinforcement learning training. This direct reduction in memory overhead significantly decreases hardware costs and frees up GPU capacity, enabling models to be trained with higher batch sizes.

    Sources: [2402.03300] DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models, GRPO – Group Relative Policy Optimization: How DeepSeek Trains Reasoning Models, DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning, Implementing GRPO in TRL · Hugging Face, TRL – Transformers Reinforcement Learning · Hugging Face

    Impact on Model Performance, Throughput, and Optimization Limitations

    Despite these architectural memory savings, standard Group Relative Policy Optimization (GRPO) suffers from a structural optimization limitation rooted in static uniformity: it relies on uniform prompt sampling and a fixed number of rollouts per prompt. For heterogeneous, heavy-tailed reasoning datasets, this static approach creates inefficiencies by wasting computational resources on already-solved patterns while under-training the long tail of difficult problems.

    To overcome these compute and optimization bottlenecks, Multi-Adversary Group Distributionally Robust Optimization adapts the training distribution dynamically using an Online Difficulty Classifier that partitions prompts into pass@k difficulty groups: Prompt-GDRO employs an EMA-debiased multiplicative-weights bandit sampler to target the intensive difficulty margin and upweight persistently hard prompt groups without introducing frequency bias, while Rollout-GDRO uses a shadow-price controller guided by a variance-proxy analysis to reallocate rollouts across difficulty groups—aiming for a square-root optimal rollout allocation—to maximize gradient variance reduction on hard tasks under a fixed mean budget.

    Because Rollout-GDRO operates under a fixed mean budget, it dynamically reallocates compute resources in a compute-neutral manner rather than increasing total throughput demands. Qualitative evaluations show that this creates an emergent curriculum, shifting optimization resources toward the evolving reasoning frontier as the model learns. When validated on the DAPO 14.1k dataset using Qwen3-Base models across 1.7B, 4B, and 8B parameter scales, these optimization adjustments yield measurable performance improvements over the standard GRPO baseline, with Prompt-GDRO achieving an average relative gain of +10.6% in pass@8 accuracy and Rollout-GDRO achieving an average relative gain of +10.1% in pass@8 accuracy.

    Sources: Group Distributionally Robust Optimization-Driven Reinforcement Learning for LLM Reasoning

    Closing thoughts

    In summary, by eliminating the memory-heavy Critic Model, Group Relative Policy Optimization offers an invaluable structural solution to the severe VRAM bottlenecks inherent in standard PPO frameworks. However, raw memory efficiency alone falls short when static sampling wastes compute on already-solved problems while under-training the long tail of complex tasks.

    In my assessment, the true breakthrough lies in pairing this critic-free architecture with dynamic techniques like Prompt-GDRO and Rollout-GDRO, which establish an emergent curriculum by targeting compute specifically at the model’s evolving reasoning frontier. Ultimately, these findings demonstrate that maximizing reinforcement learning efficacy requires both stripping away architectural redundancy and intelligently shifting optimization resources toward hard, unresolved prompts.

    Frequently Asked Questions

    What is Group Relative Policy Optimization (GRPO)?

    Group Relative Policy Optimization (GRPO) is a reinforcement learning training method for language models that eliminates the separate Critic Model used in standard PPO architectures. Instead of relying on a value network, GRPO generates a group of outputs per prompt and standardizes their reward scores to measure relative performance.

    How does GRPO reduce memory overhead during training?

    Standard PPO requires training a separate Critic Model of similar size to the Actor Model, which nearly doubles GPU VRAM requirements. By completely removing the Critic Model, GRPO cuts VRAM usage nearly in half and allows models to train with higher batch sizes.

    What limitation exists in standard GRPO?

    Standard GRPO relies on uniform prompt sampling and a fixed number of rollouts per prompt. This static approach can waste computational resources on already-solved patterns while under-training difficult, long-tail reasoning problems.

    How do Prompt-GDRO and Rollout-GDRO address the limitations of GRPO?

    Prompt-GDRO upweights persistently hard prompt groups without frequency bias, while Rollout-GDRO dynamically reallocates rollouts across difficulty groups under a fixed mean budget. Together, they shift optimization resources toward the evolving reasoning frontier, improving accuracy over standard GRPO baselines.


    Keep reading: Training & Architectures

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

  • SimPO: Reference-Free Preference Optimization for Large Language Models

    SimPO: Reference-Free Preference Optimization for Large Language Models

    TL;DR

    • Developed by researchers at Princeton and UVA, SimPO is a reference-free preference optimization method that eliminates the need for a frozen reference model in LLM alignment.
    • By utilizing sequence-average log probabilities and a target reward margin, SimPO aligns training objectives with inference while reducing GPU memory usage and training overhead.
    • SimPO consistently outperforms DPO across standard benchmarks like AlpacaEval 2 and Arena-Hard while mitigating token verbosity and length bias.

    Introduction to SimPO: A Simple Reference-Free Approach

    SimPO (Simple Preference Optimization) is a reference-free preference optimization method designed to align large language models (LLMs) directly with human preferences. Introduced by researchers Yu Meng, Mengzhou Xia, and Danqi Chen from Princeton University and the University of Virginia, the approach was published in the paper SimPO: Simple Preference Optimization with a Reference-Free Reward (arXiv:2405.14734) and accepted as a NeurIPS 2024 paper. Official code implementations and open-source model checkpoints are hosted under the Princeton NLP organization.

    SimPO was motivated by the operational and theoretical limitations of Direct Preference Optimization (DPO). Although DPO simplified Reinforcement Learning from Human Feedback (RLHF) by training policy models directly on human preference pairs without a separate reward model, DPO still relies on a Reference Model—a frozen copy of the base supervised fine-tuned (SFT) model—to constrain policy drift using Kullback-Leibler (KL) divergence. This design introduces two key challenges:

    • Training-Inference Discrepancy: DPO calculates Implicit Reward based on a log-likelihood ratio relative to the frozen Reference Model. Because LLMs generate text during inference based directly on sequence log-probabilities rather than relative likelihood ratios, this ratio-based Implicit Reward creates a misalignment between the training objective and actual generation behavior.
    • Compute and Memory Inefficiency: Keeping a second frozen Reference Model in GPU memory throughout training requires extra forward passes, increasing VRAM usage and compute overhead.

    SimPO resolves these issues by eliminating the Reference Model entirely from the alignment objective. Removing the Reference Model reduces memory consumption and speeds up training, enabling researchers and practitioners to fine-tune larger models on smaller GPU clusters. Furthermore, SimPO aligns the training objective with inference generation while avoiding the token verbosity and length bias frequently observed in traditional RLHF and DPO implementations. Across standard alignment evaluations—such as AlpacaEval 2, Arena-Hard, and MT-Bench—SimPO consistently outperforms DPO and its variants across multiple model families, including Llama 3, Gemma 2, and Mistral.

    Sources: [2405.14734] SimPO: Simple Preference Optimization with a Reference-Free Reward, [논문리뷰] SimPO: Simple Preference Optimization with a Reference-Free Reward – 전생했더니 인공지능이었던 건에 대하여, OpenReview (paper discussion page), SimPO: Simple Preference Optimization with a Reference-Free Reward, Paper page – SimPO: Simple Preference Optimization with a Reference-Free Reward

    Algorithmic Design: Sequence-Average Probabilities and Target Margins

    To address these limitations, the core algorithmic design of SimPO relies on using the Average Log Probability of a sequence as its Implicit Reward formulation. By utilizing sequence Average Log Probability, the Implicit Reward directly aligns with the language model’s generation process. Crucially, this reward formulation eliminates the requirement for a Reference Model during training, making the optimization process both compute and memory efficient.

    Additionally, SimPO modifies the traditional Bradley-Terry objective by incorporating a target reward margin. This target margin explicitly enforces a larger margin between the Implicit Rewards of the winning and dispreferred losing responses, driving a distinct separation between candidate outputs to further enhance the algorithm’s performance.

    Sources: SimPO: Simple Preference Optimization with a Reference-Free Reward

    Closing thoughts

    Building on these algorithmic refinements, SimPO demonstrates that the architectural reliance on a frozen reference model was an unnecessary hurdle in preference optimization. By aligning the implicit reward directly with sequence-average log probabilities and enforcing a target margin, the algorithm logically resolves the mismatch between how models are trained and how they actually generate text. In my view, its primary strength lies in achieving this conceptual elegance alongside immediate practical gains—slashing GPU memory requirements while simultaneously curbing the length bias that plagues traditional RLHF. As evidenced by its superior performance across benchmarks like AlpacaEval 2 and Arena-Hard, SimPO proves that streamlining alignment mechanics yields a significantly more resource-efficient workflow without compromising on output quality.

    Frequently Asked Questions

    What is SimPO and who developed it?

    SimPO (Simple Preference Optimization) is a reference-free preference optimization method designed to align large language models directly with human preferences. It was introduced by researchers Yu Meng, Mengzhou Xia, and Danqi Chen from Princeton University and the University of Virginia.

    How does SimPO improve upon Direct Preference Optimization (DPO)?

    Unlike DPO, SimPO completely eliminates the need for a frozen reference model, which slashes GPU memory usage and speeds up training. It also resolves the discrepancy between training objectives and inference generation while curbing length bias.

    What are the core algorithmic components of SimPO?

    SimPO uses the sequence-average log probability of an output as its implicit reward formulation. It also modifies the Bradley-Terry objective by adding a target reward margin to enforce a clear separation between preferred and dispreferred responses.

    How does SimPO perform on standard language model benchmarks?

    SimPO consistently outperforms DPO and its variants across standard alignment evaluations, including AlpacaEval 2, Arena-Hard, and MT-Bench. These gains are demonstrated across multiple model families, such as Llama 3, Gemma 2, and Mistral.


    Keep reading: Training & Architectures

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

  • Test-Time Compute Scaling with Process Reward Models: Mechanisms, Strategies, and Trade-offs

    Test-Time Compute Scaling with Process Reward Models: Mechanisms, Strategies, and Trade-offs

    TL;DR

    • Process Reward Models evaluate each intermediate step in a Chain-of-Thought trajectory, catching reasoning errors at their point of origin to improve accuracy.
    • Step-level generative verification leverages explicit reasoning and code verification, allowing compact models like GenPRM to outperform much larger baselines.
    • Offline compute strategies like Sleep-time Compute pre-calculate context before queries arrive, significantly reducing test-time latency and operational costs.

    Understanding Process Reward Models and Step-Level Feedback

    Test-Time Compute Scaling (also known as Inference-Time Scaling) allocates additional computational resources, measured in FLOPs, during the generation process rather than relying solely on pre-training scale to improve response accuracy. Evaluating reasoning trajectories during inference depends heavily on the feedback granularity of the verifier model.

    Outcome Reward Models (ORMs) evaluate only the final output of a completed response. Because they score solely the terminal state, ORMs suffer from credit assignment problems—frequently rewarding a generation trajectory that happened to arrive at the correct final answer despite relying on flawed intermediate logic.

    Process Reward Models (PRMs) address this credit assignment challenge by introducing step-by-step process supervision. Instead of assessing only the end answer, PRMs evaluate each individual intermediate step in a Chain-of-Thought (CoT) trajectory. This step-level feedback provides dense, fine-grained scoring across the response, catching reasoning errors at their point of origin and preventing early mistakes from compounding as the trajectory progresses.

    Landmark research by Lightman et al. (2023) established that training reward models with step-by-step process supervision significantly outperforms outcome-based supervision. This concept underpins modern reasoning paradigms more broadly: models such as OpenAI’s o1 and o3 are widely understood to achieve their gains through reinforcement-learned internal reasoning rather than an external verifier, an approach closely related to the Implicit Test-Time Compute Scaling paradigm described below.

    Recent developments highlight two distinct architectural approaches to leveraging step-level feedback and compute during inference:

    • External Search with PRMs: Using separate, dedicated PRM verifier models to score intermediate steps and steer explicit search algorithms across trajectories.
    • Implicit Test-Time Compute Scaling via Reinforcement Learning (RL): Training policy models via RL to produce longer, self-correcting internal Chain-of-Thought (CoT) tokens directly inside a single inference pass without relying on external verifier loops.

    Sources: [2506.00027] From Mathematical Reasoning to Code: Generalization of Process Reward Models in Test-Time Scaling, self-supervised.cs.jhu.edu/fa2024/files/presentations/9-26-Reasoning-…, youtube.com/watch, [논문리뷰] LLM 1B 모델로 405B 모델 성능을? Test-Time Scaling!, medium.com/wix-engineering/explaining-openais-o1-breakthrough-the-rev…

    Why Step-Level Generative Verification Improves Reasoning

    Building on the External Search paradigm introduced above, step-level generative verification provides a critical advancement over standard PRM verifier designs. Traditional Process Reward Models (PRMs) face structural limitations due to their reliance on scalar value predictions without leveraging the generative capabilities of language models, which restricts their process supervision, generalization capabilities, and ability to leverage Test-Time Compute Scaling. Step-level generative verification overcomes these constraints by requiring the model to perform explicit Chain-of-Thought (CoT) reasoning alongside code verification before rendering a judgment for each individual reasoning step.

    In frameworks such as GenPRM, step-level generative verification utilizes Relative Progress Estimation (RPE) and a rationale synthesis framework that integrates code verification to produce high-quality process supervision labels and explicit rationale data. This structured step-level rationale generation allows process verification to unlock Test-Time Compute Scaling. Consequently, a 1.5B GenPRM model leverages Test-Time Compute Scaling to outperform GPT-4o, while a 7B GenPRM model surpasses Qwen2.5-Math-PRM-72B on ProcessBench using only 23K training instances from the MATH dataset. Beyond basic step evaluation, generating explicit reasoning and code verification for each step enables the verifier to serve as a critic model for policy model refinement, bridging the operational gap between PRMs and critic models.

    Sources: GenPRM: Scaling Test-Time Compute of Process Reward Models via Generative Reasoning

    Execution-Based Reward Models for Code Generation Search

    While generative verification strengthens step-level scoring, code generation presents a related but distinct challenge for search-based inference. Rather than scoring intermediate reasoning steps directly like a classic PRM, search and verification mechanisms in this domain frequently generate multiple complete candidate solutions and validate each one using execution results from model-generated unit tests as an outcome-based reward signal. However, because models can produce flawed tests with high confidence, these test-based reward signals are often unreliable, degrading the overall quality of candidate verification.

    Scaling the number of unit tests provides a positive correlation with reward signal quality, yielding higher performance benefits on more challenging problems. To balance reward precision with execution efficiency during search, a dynamic scaling mechanism can adaptively adjust the quantity of generated unit tests according to problem difficulty. Utilizing CodeRM-8B—a lightweight unit test generator designed for high-quality scaling—alongside dynamic unit test scaling delivers marked accuracy improvements across multiple benchmarks, including performance gains of 18.43% for Llama3-8B and 3.42% for GPT-4o-mini on HumanEval Plus.

    Sources: Dynamic Scaling of Unit Tests for Code Reward Modeling

    Latency Overhead and Operational Cost Consequences

    Although integrating search and verification mechanisms significantly boosts accuracy, scaling compute during inference imposes severe latency overhead and high operational inference costs during active user interactions. To mitigate these execution-time constraints, offline pre-computation approaches such as Sleep-time Compute can reallocate compute demands by allowing models to “think” offline about context prior to query presentation. By anticipating user queries and pre-computing useful quantities ahead of time, this approach significantly decreases the compute required during the actual test-time phase.

    Empirical evaluations on Stateful GSM-Symbolic and Stateful AIME show that Sleep-time Compute reduces the amount of Test-Time Compute Scaling needed to achieve equivalent accuracy by approximately 5x. Furthermore, scaling the amount of Sleep-time Compute allocated offline yields direct accuracy improvements, driving gains of up to 13% on Stateful GSM-Symbolic and 18% on Stateful AIME.

    To further lower per-query operational costs, Sleep-time Compute can be amortized across multiple requests sharing the same context. In Multi-Query GSM-Symbolic—an extension incorporating multiple related queries per context—amortizing pre-computed quantities across these queries reduces the average compute cost per query by 2.5x. Analysis indicates that the predictability of incoming user queries correlates directly with the effectiveness of Sleep-time Compute in reducing test-time latency and operational overhead. Case study evidence confirms the applicability of this strategy to complex, realistic workloads, including agentic software engineering (SWE) tasks.

    Sources: Sleep-time Compute: Beyond Inference Scaling at Test-time

    Closing thoughts

    Ultimately, the evidence demonstrates that shifting from coarse outcome evaluations to step-level generative verification allows compact models to achieve superior reasoning by systematically catching intermediate errors at their point of origin. However, because scaling search and verification during live inference introduces severe latency and operational cost bottlenecks, relying solely on real-time verifications remains practically constrained. By pairing dynamic step-level feedback with offline strategies like Sleep-time Compute, systems can pre-calculate context and amortize overhead across predictable queries to drastically reduce test-time costs. In my view, the true promise of test-time scaling lies in this precise synergy—balancing fine-grained process supervision with intelligent offline compute allocation to deliver high-accuracy reasoning efficiently.

    Frequently Asked Questions

    What is the main difference between Outcome Reward Models (ORMs) and Process Reward Models (PRMs)?

    Outcome Reward Models evaluate only the final output of a completed response, which can result in rewarding flawed intermediate logic. In contrast, Process Reward Models evaluate each individual intermediate step in a Chain-of-Thought trajectory to catch reasoning errors at their point of origin.

    How does step-level generative verification improve reasoning over traditional PRMs?

    Traditional PRMs rely on scalar value predictions without using the generative capabilities of language models. Step-level generative verification overcomes this limitation by requiring explicit Chain-of-Thought reasoning alongside code verification before judging each step.

    What are two main architectural approaches to using step-level feedback during inference?

    The two main approaches are External Search with PRMs, which uses dedicated verifiers to steer explicit search algorithms, and Implicit Test-Time Compute Scaling via RL, which trains policy models to generate self-correcting internal reasoning tokens in a single pass.

    How does Sleep-time Compute help mitigate test-time latency overhead?

    Sleep-time Compute reallocates compute demands by allowing models to pre-compute context offline before a user presents a query. This reduces the test-time compute required to achieve equivalent accuracy by approximately 5x and lowers per-query costs when amortized across multiple requests.


    Keep reading: Training & Architectures

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

  • Kolmogorov-Arnold Networks: Rethinking Neural Architecture via Learnable Spline Activations

    Kolmogorov-Arnold Networks: Rethinking Neural Architecture via Learnable Spline Activations

    TL;DR

    • Kolmogorov-Arnold Networks (KANs) serve as an alternative to Multi-Layer Perceptrons by placing learnable B-spline activation functions directly on network edges to model complex, non-monotonic dependencies.
    • KANs significantly enhance model interpretability in applications such as cognitive diagnosis while achieving competitive training efficiency through optimized implementations.
    • Architectural extensions like the Kurkova-Kolmogorov-Arnold Network outperform traditional MLPs and original KANs in function approximation and operator learning tasks.

    Introduction to Kolmogorov-Arnold Networks (KANs)

    Kolmogorov-Arnold Networks (KANs) were introduced in April 2024 by a team of researchers from MIT, Caltech, and Northeastern University (several affiliated with the NSF Institute for Artificial Intelligence and Fundamental Interactions): Ziming Liu, Yixuan Wang, Sachin Vaidya, Fabian Ruehle, James Halverson, Marin Soljačić, Thomas Y. Hou, and Max Tegmark. Accepted as an oral presentation at ICLR 2025, KANs offer a fundamentally different neural network architecture intended to serve as an alternative to standard Multi-Layer Perceptrons (MLPs).

    Sources: [2404.19756] KAN: Kolmogorov-Arnold Networks, ICLR Oral KAN: Kolmogorov–Arnold Networks, KAN: Kolmogorov–Arnold Networks, [2408.10205] KAN 2.0: Kolmogorov-Arnold Networks Meet Science, GitHub – KindXiaoming/pykan: Kolmogorov Arnold Networks · GitHub

    Mathematical Foundation and Structural Differences from MLPs

    At the core of this architectural alternative is the mathematical foundation provided by the Kolmogorov-Arnold Representation Theorem. One notable application of this foundation outside standard neural networks is in Fuzzy Cognitive Maps (FCMs), a graph-based reasoning framework where nodes represent factors and edges represent causal influence between them. Drawing upon this theorem, Kolmogorov-Arnold Fuzzy Cognitive Maps (KA-FCMs) redefine the causal transmission mechanism to overcome the limitations of standard FCMs. The standard FCM formulation relies on static scalar synaptic weights and monotonic activation functions, placing non-linearity at the nodes’ aggregation phase, which fundamentally constrains its ability to model non-monotonic causal dependencies. In contrast, the KA-FCM architecture replaces static scalar weights with learnable univariate B-spline activation functions located directly on the model edges. This shift moves the non-linearity from the node aggregation phase directly to the causal influence phase along the edges, enabling the modeling of arbitrary, non-monotonic causal relationships without increasing graph density or introducing hidden layers.

    Sources: Non-monotonic causal discovery with Kolmogorov-Arnold Fuzzy Cognitive Maps

    Model Interpretability and Efficiency Benefits

    Beyond these core mathematical properties, this shift in architecture provides significant advantages in model interpretability and efficiency. In domain-specific applications such as intelligent education and cognitive diagnosis, Kolmogorov-Arnold Networks (KANs) address the long-standing interpretability limitations of Multi-Layer Perceptrons (MLPs). Neural Cognitive Diagnosis Models rely on embeddings for students, exercises, and knowledge concepts to reveal proficiency for downstream recommendation tasks. While neural Cognitive Diagnosis Models typically outperform traditional models, their reliance on MLPs yields poor interpretability even when constrained by monotonicity assumptions. Replacing MLPs with KANs in cognitive diagnosis architectures—such as in the KAN2CD framework—enhances model interpretability through two structural designs: direct replacement of standard MLPs in existing neural Cognitive Diagnosis Models with KANs, and hierarchical combination of student, exercise, and concept embeddings through several distinct KANs before combining and learning their outputs within a unified KAN to generate final predictions. The learned structures of KANs allow these enhanced neural Cognitive Diagnosis Models to maintain the clear interpretability of traditional Cognitive Diagnosis Models while surpassing existing neural Cognitive Diagnosis Models in interpretability.

    Regarding computational efficiency, original KAN architectures present a challenge due to slow training speeds. However, modifying the implementation of original KANs accelerates training, yielding computational costs that are competitive with existing models. On four real-world datasets, efficient KAN-based models outperform traditional Cognitive Diagnosis Models and maintain a performance lead over existing neural Cognitive Diagnosis Models while retaining high interpretability and competitive training costs.

    Sources: Endowing Interpretability for Neural Cognitive Diagnosis by Efficient Kolmogorov-Arnold Networks

    Architectural Extensions and Learning Dynamics

    Building upon these interpretability and efficiency benefits, researchers have also introduced broader architectural extensions and examined their learning dynamics. Inspired by the Kolmogorov-Arnold Representation Theorem and Kurkova’s principle of using approximate representations, the Kurkova-Kolmogorov-Arnold Network (KKAN) is a two-block architecture that combines robust Multi-Layer Perceptron (MLP) based inner functions with flexible linear combinations of basis functions as outer functions. Proven to be a universal approximator, KKAN demonstrates versatility across scientific machine-learning applications, including function regression, Physics-Informed Machine Learning (PIML), and operator-learning frameworks. In benchmark results, KKANs outperform MLPs and original Kolmogorov-Arnold Networks (KANs) in function approximation and operator learning tasks, while achieving performance comparable to fully optimized MLPs for PIML.

    An analysis using information bottleneck theory provides insight into the geometric complexity and learning dynamics of these models, identifying three universal learning stages across all types of architectures: fitting, transition, and diffusion. Optimal generalization is achieved during the diffusion stage. Additionally, a strong correlation exists between geometric complexity and signal-to-noise ratio (SNR). To dynamically maintain a high SNR, self-scaled residual-based attention weights can be used, ensuring uniform convergence and prolonged learning.

    Sources: KKANs: Kurkova-Kolmogorov-Arnold Networks and Their Learning Dynamics

    Closing thoughts

    Taking these empirical and structural insights together, Kolmogorov-Arnold Networks demonstrate how relocating non-linear activation functions directly onto network edges via learnable B-splines provides a compelling alternative to traditional MLPs that fundamentally enhances how complex, non-monotonic dependencies are modeled. In my view, this paradigm shift successfully resolves the historical trade-off between interpretability and predictive power, as evidenced by clear explainability gains in cognitive diagnosis tasks and superior function approximation in hybrid extensions like KKANs. As optimized implementations continue to resolve early training inefficiencies and leverage dynamic controls like residual-based attention, KAN-based architectures demonstrate that rethinking foundational neural building blocks can yield significant gains in transparency without sacrificing performance or computational viability.

    Frequently Asked Questions

    What are Kolmogorov-Arnold Networks (KANs) and who introduced them?

    Kolmogorov-Arnold Networks (KANs) are an alternative neural network architecture to standard Multi-Layer Perceptrons (MLPs). They were introduced in April 2024 by a team of researchers from MIT, Caltech, and Northeastern University.

    How do KANs differ structurally from standard Multi-Layer Perceptrons?

    KANs replace static scalar weights with learnable univariate B-spline activation functions located directly on the model edges. This moves non-linearity from the node aggregation phase to the edges, enabling the model to represent arbitrary non-monotonic causal relationships.

    How do KANs address interpretability and efficiency challenges?

    KANs improve interpretability in domain-specific applications like cognitive diagnosis, matching the explainability of traditional models while maintaining neural performance. Additionally, modified KAN implementations accelerate training speeds to make computational costs competitive with existing models.

    What is the Kurkova-Kolmogorov-Arnold Network (KKAN)?

    The KKAN is a two-block architectural extension that combines MLP-based inner functions with flexible linear combinations of basis functions as outer functions. Proven to be a universal approximator, it outperforms standard MLPs and original KANs in function approximation and operator learning tasks.


    Keep reading: Training & Architectures

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