PaperlayerView exact revision on arXiv

Source-linked AI summary

Multi-Head Low-Rank Attention

Songtao Liu, Hongwu Peng, Zhiwei Zhang, Zhengyu Chen, Yue Guo

arXiv:2603.02188v1cs.LG

TL;DR

Long-context decoding is bottlenecked by repeated KV-cache movement, and MLA’s single latent head prevents efficient tensor-parallel sharding. MLRA makes latent states partitionable through independent low-rank branches, enabling 4-way TP decoding; it achieves strong language-modeling and reasoning results with about 2.8× speedup over MLA. The method’s scaling relies on an assumption that may not strictly hold during training.

  • Problem

    Long-context decoding repeatedly reloads the KV cache, while MLA’s single latent head cannot be sharded for tensor parallelism, limiting distributed-decoding efficiency.

  • Method

    MLRA decomposes MLA’s latent representation into independent low-rank branches whose outputs are summed, enabling native 4-way tensor parallelism and reducing per-device KV-cache loading.

  • Results

    MLRA-4 achieves state-of-the-art perplexity and zero-shot common-sense reasoning performance, while maintaining about 2.8× decoding speedup over MLA for long contexts.

  • Takeaways & Limitations

    MLRA provides a partitionable attention design for efficient 4-way tensor-parallel long-context decoding across sequences up to 2M tokens.

  • Takeaways & Limitations

    The scaling analysis assumes independently identically distributed zero-mean weights with common variance, but these conditions are not guaranteed during training.

Abstract

from arXiv · show

Long-context inference in large language models is bottlenecked by Key--Value (KV) cache loading during the decoding stage, where the sequential nature of generation requires repeatedly transferring the KV cache from off-chip High-Bandwidth Memory (HBM) to on-chip Static Random-Access Memory (SRAM) at each step. While Multi-Head Latent Attention (MLA) significantly reduces the total KV cache size, it suffers from a sharding bottleneck during distributed decoding via Tensor Parallelism (TP). Since its single latent head cannot be partitioned, each device is forced to redundantly load the complete KV cache for every token, consuming excessive memory traffic and diminishing TP benefits like weight sharding. In this work, we propose Multi-Head Low-Rank Attention (MLRA), which enables partitionable latent states for efficient 4-way TP decoding. Extensive experiments show that MLRA achieves state-of-the-art perplexity and downstream task performance, while also delivering a 2.8$\times$ decoding speedup over MLA. Code is available at https://github.com/SongtaoLiu0823/MLRA. Pretrained weights, along with the training and evaluation data, are available at https://huggingface.co/Soughing/MLRA.

1 INTRODUCTION

Long-context decoding is dominated by KV-cache data movement, while MLA reduces cache size but cannot shard its single latent head for tensor parallelism. MLRA addresses this limitation by decomposing the latent head into partitionable branches for 4-way TP decoding.

  • Motivation: KV-cache transfers, rather than computation, dominate long-context decoding latency because each sequential step reloads the cache from high-bandwidth memory.Long-context RAG and chain-of-thought increase the tokens processed at each decoding step.
  • Motivation: MLA compresses the KV cache into a 4.5d_h-per-token latent head and improves efficiency over MHA through query-side up-projection absorption.Its single latent head cannot be sharded across tensor-parallel devices.
  • Contribution: MLRA decomposes MLA’s latent head into four latent heads, independently up-projects them into NoPE keys and values, and sums the resulting attention outputs.The design enables native 4-way tensor parallelism and reduces per-device KV-cache loading.
  • Contribution: 13.672 average perplexity and 58.84% zero-shot common-sense accuracy make MLRA-4 the strongest reported model against MLA and GQA in the 2.9B experiments.The reported comparisons are 13.672 versus 13.727 and 14.139 for perplexity, and 58.84% versus 58.75% and 57.89% for accuracy.

2 BACKGROUND

MLA compresses query, key, and value representations into latent states, uses shared RoPE information, and absorbs up-projections to avoid materializing full KV tensors during decoding. Grouped Latent Attention extends this structure by splitting the latent representation into groups and summing block-wise projections.

  • Multi-Head Latent Attention: MLA derives latent query and KV states from hidden states, with KV and query latent dimensions much smaller than the full multi-head representation.The cited background states c ≪ h d_h for the latent dimensions.
  • Multi-Head Latent Attention: MLA caches C_KV and K_RoPE, then forms NoPE keys and values through up-projections before reshaping them into head-wise tensors.A shared RoPE key is combined with each head’s NoPE query and key to provide position-aware representations.
  • Efficient Decoding: During decoding, MLA absorbs up-projection weights into queries and uses shared latent KV states in an MQA-style attention computation.This avoids explicitly materializing the h heads of NoPE keys and values.
  • Efficient Decoding: MLA’s decoding procedure computes attention outputs from the shared latent states and then applies an output up-projection.The procedure is organized into query-side absorption, latent-KV decoding, and output up-projection.
  • Grouped Latent Attention: Grouped Latent Attention bisects MLA’s latent head into two groups, mapping each group to half of the attention heads.Its head-specific up-projections are decomposed into two row-wise blocks whose products are summed.
  • Grouped Latent Attention: GLA’s block decomposition expresses each head’s NoPE key and value as a sum of two block products formed from group-specific latent states and up-projection matrices.The same block construction is applied to both key and value projections.

3 MULTI-HEAD LOW-RANK ATTENTION

MLRA decomposes MLA’s latent head into independent low-rank branches whose attention outputs are summed, enabling 4-way tensor-parallel decoding. The design uses variance calibration and reduces KV-cache loading and decoding memory pressure relative to MLA and GLA-2.

  • MLRA shifts block summation from KV computation to attention outputs, treating each projection as an independent low-rank branch.
  • MLRA-2 uses two latent blocks, whereas MLRA-4 uses four blocks, with each MLRA-4 block serving all h heads.
  • 1.5dh per-device attention-logit space is achieved after absorption, compared with 4.5dh for MLA and 2.5dh for GLA-2.
  • Variance Calibration: Variance mismatch arises because RMSNorm is applied before up-projections, so scaling factors align latent-derived components with partial RoPE components.
  • Assumptions: The variance analysis assumes independently distributed, zero-mean weights and mutually uncorrelated attention outputs across latent blocks.The authors note that the weight assumptions may not hold strictly during training.
  • KV Cache: MLRA achieves 1.5dh per-device KV-cache loading with 4-way TP, while MLA and GLA-2 plateau at 4.5dh and 2.5dh, respectively.GQA and GTA require 8-way TP to reach 2dh and 1.5dh, respectively.
  • Arithmetic Intensity: MLRA-2 and MLRA-4 achieve arithmetic intensities of h and 2h, respectively, increasing the compute-to-memory ratio and shifting decoding toward a compute-limited regime.

4 EXPERIMENTS

Experiments evaluate MLRA against attention baselines across initialization, scaling, head-count, gating, quality, and decoding-efficiency settings. MLRA-4 achieves the strongest reported quality and decoding results while supporting efficient tensor-parallel inference.

  • Experimental Setup: All models are pretrained from scratch at 2.9B parameters on FineWeb-Edu-100B and evaluated on seven validation datasets plus seven common-sense reasoning benchmarks.The evaluation compares perplexity and zero-shot accuracy across broad language-modeling and downstream-task settings.
  • Initialization: Zero initialization outperforms N(0, σ = 0.02) for attention and feed-forward output projections, so subsequent experiments use zero initialization.The comparison is conducted across all models.
  • Scaling: Scaling improves convergence and lowers average perplexity for MLA, GLA-2, and MLRA-2, although MLRA-2 shows only a marginal gain.The largest improvements are reported for MLA and GLA-2.
  • Double Heads: Doubling attention heads increases loss and fails to reduce perplexity for GQA, MLA, and GLA-2, providing no measurable performance improvement.The models retain their default head counts in subsequent experiments.
  • Main Results: MLRA-4 achieves the best average perplexity at 13.672 versus MLA at 13.727 and the lowest FineWeb-Edu perplexity at 9.193.MLRA-4 also attains the highest average zero-shot accuracy across the common-sense reasoning tasks.
  • Gated Attention: Gating improves perplexity across all evaluated models; MLRA-4 remains best overall, while MLRA-2 is comparable to MLA at 13.651 versus 13.642.The gated comparison uses a constant parameter budget by reducing FFN intermediate size.
  • Decoding Efficiency: MLRA-4 consistently outperforms baselines from 128K to 2M tokens, achieving 1.05×–1.26× speedups over GQA and about 2.8× over MLA.The latency benchmark uses batch=1 on an NVIDIA H100 80GB GPU.
  • Decoding Efficiency: Decoding throughput is evaluated across sequence lengths from 1K to 16K tokens using eight H100 GPUs and tensor-parallel deployment.The supplied experiment description specifies the compared systems and fixed model dimensions.

5 CONCLUSION

MLRA is proposed as an attention mechanism with native 4-way tensor parallelism. At the 2.9B scale, it combines strong task performance with lower long-context decoding latency and higher throughput.

  • At the 2.9B scale, MLRA-4 achieves state-of-the-art perplexity and zero-shot common-sense reasoning performance.
  • MLRA achieves the lowest decoding latency for sequences up to 2M tokens.
  • MLRA-4 achieves the highest throughput across sequence lengths from 1K to 16K tokens with 4-way tensor parallelism.

B.2 ROTARY POSITION EMBEDDING

RoPE preserves translation equivariance for attention dot products because jointly shifting query and key positions leaves their relative displacement unchanged. However, arbitrary linear projections after RoPE generally break this property, while MQA and GQA reduce cached key-value heads to lower KV-cache costs.

  • Translation equivariance: RoPE preserves attention dot-product scores when both token positions are shifted by the same offset.The inner product remains unchanged under simultaneous translation of query and key positions.
  • Translation equivariance: The induced RoPE attention-score matrix is translation equivariant under simultaneous row and column shifts.This extends the dot-product property to the full attention-score matrix.
  • Effect of projections: Applying an arbitrary linear map after RoPE generally breaks translation equivariance.The property is preserved only under the specific constraint W_Q W_K^⊤ = I, which is difficult to enforce during training.
  • Attention heads and caching: Standard multi-head attention projects hidden states into queries, keys, and values, reshapes them by head, and caches keys and values during inference.Queries, keys, and values are separated into h heads of dimensionality d_h before caching the key and value tensors.
  • Attention heads and caching: MQA uses one shared key-value head, whereas GQA uses g shared key-value heads for h query heads and repeats them by r = h/g.Both methods reduce the number of key and value heads while retaining the full number of query heads.

C.3 MULTI-HEAD LATENT ATTENTION (MLA)

The section describes low-rank and grouped attention designs that reduce or factorize cached key-value states while analyzing translation equivariance under positional encoding. MLA uses shared compressed states, whereas MFA and grouped variants organize latent states differently across heads or groups.

  • MLA: MLA computes queries from hidden states, separates query heads, and uses a low-rank down-projection to obtain shared compressed key-value states.The compressed states and partial RoPE keys are cached, while full keys and values are reconstructed with learnable up-projections.
  • Translation equivariance: MLA’s attention score combines a position-independent term with a RoPE-dependent term that remains translation equivariant.The resulting property is termed semi-translation equivariance because MLA uses partial RoPE rather than full RoPE.
  • MFA: MFA uses h query heads but one shared key-value head, caching one key head and one value head before repeating them across query heads.Its translation-equivariance analysis is stated to be similar to MQA.
  • TPA: TPA represents each head’s keys and values as low-rank mixtures of component vectors and head-specific scalar coefficients.Inference caches component and coefficient tensors, then forms keys and values by linear combination; its RoPE scores preserve translation equivariance.
  • GLA: GLA divides the h attention heads into g groups, computes independent compressed states per group, and shares RoPE keys across groups.Each group has r = h/g heads, and the total KV cache size is dc + dR.

D LLAMA-3 ARCHITECTURE

The Llama-3 architecture applies attention and residual projection, followed by a gated MLP and another residual projection. The gated formulation uses elementwise nonlinearities and multiplication.

  • Attention block: The model computes attention outputs from hidden states, projects them back to the model dimension, and adds an attention residual.The attention output projection is applied before the residual update.
  • MLP block: A gated MLP block is applied after attention, using an elementwise nonlinearity and elementwise multiplication.The nonlinearity is described as a function such as SiLU in the ungated formulation and sigmoid in the gated formulation.
  • MLP block: The MLP output is projected back to the model dimension and added through a second residual connection.The output projection maps from the intermediate FFN dimension to the model dimension.

F.1 ARCHITECTURAL HYPERPARAMETERS FOR MAIN RESULTS

The main experiments use a 24-layer Llama-3-based architecture with 24 attention heads, hidden dimension 3072, head dimension 128, and FFN dimension 8192. Baseline configurations follow their original implementations, with method-specific settings reported in configuration tables.

  • Base architecture: The model uses 24 layers, 24 attention heads, hidden dimension d = 3072, head dimension dh = 128, and FFN dimension df = 8192.The architecture is largely based on Llama-3.2-3B but reduces the layer count from 28 to 24.
  • Attention configurations: MLA uses latent dimensions d′ and h = 0.5dh, while TPA uses ranks βq = 6 and βkv = 2.These are method-specific architectural settings for the main results.
  • Attention configurations: GQA and GTA use g = h/4 key-value heads in the main experiments.The detailed configurations for the evaluated baselines and proposed models are listed in Tables 7–17.

F.2 ARCHITECTURAL HYPERPARAMETERS FOR INITIALIZATION ABLATION STUDY

The initialization ablation isolates the effect of output-projection initialization by comparing zero initialization with Gaussian initialization while holding architecture and other hyperparameters fixed.

  • Initialization comparison: The ablation studies initialize the attention and FFN output projections either to zero or with N(0, σ = 0.02).The compared parameters are W_O,attn and W_O,mlp.
  • Experimental control: Architecture and all other hyperparameters remain identical to the main-result settings during the initialization comparison.This design is intended to isolate initialization strategy as the experimental variable.

F.3 ARCHITECTURAL HYPERPARAMETERS FOR SCALING ABLATION STUDY

The appendix evaluates scaling factors, attention-head counts, and gated attention through controlled architectural ablations with parameter parity and otherwise matched hyperparameters.

  • Scaling factors: Scaling ablations compare models with and without factors αq, αkv, and αattn applied to query latents, KV latents, and final attention outputs.The without condition fixes all three factors to 1 while keeping architecture and other hyperparameters unchanged.
  • Attention heads: Head-count ablations test whether increasing GQA, MLA, and GLA-2 to 48 attention heads improves performance while preserving KV-cache size.The FFN intermediate dimension is reduced to maintain parameter parity with the main results.
  • Gated attention: Gated-attention ablations add a gating mechanism to GQA, MLA, GLA-2, MLRA-2, and MLRA-4.The FFN intermediate dimension is proportionally decreased to offset the added gate and preserve parameter parity.

G ADDITIONAL EXPERIMENTAL RESULTS

Additional experiments examine initialization, scaling, and doubled-head settings using validation perplexity across seven datasets, alongside training-loss and architecture illustrations. The section also situates the work among KV-cache compression, low-rank approximation, attention systems, and linear-attention methods.

  • Validation perplexity: Validation perplexity is reported across Wikipedia, C4, Pile, RefinedWeb, Cosmopedia, FineWeb, and FineWeb-Edu for two output-projection initialization strategies.The compared strategies are zero initialization and Gaussian initialization N(0, σ = 0.02).
  • Validation perplexity: A separate validation-perplexity analysis compares models without and with scaling across the same seven datasets.The table caption identifies scaling as the specific comparison under study.
  • Validation perplexity: Another validation-perplexity analysis compares models with and without 2× attention heads across the same seven datasets.The comparison is explicitly framed around doubled attention heads.
  • Figures: The additional figures show training-loss curves for all models and illustrations of MLRA-2 and MLRA-4.Figures 7–9 respectively provide the training-loss overview and the two architecture illustrations.
  • Related work: Related work includes KV-cache compression, low-rank approximation, systems for attention, and linear attention.The cited approaches cover cache compression for pretrained models, low-dimensional representation compression, memory-efficient attention systems, and reformulated attention kernels.
Loading 2603.02188v1…