Tag: Attention Mechanism

  • 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

  • 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