Skip to content

Log-Linear Attention

ArXiv: 2506.04761

🎯 Pitch

Log-Linear Attention introduces a novel attention mechanism that bridges the gap between the efficiency of linear RNN/state-space models and the expressive power of softmax attention. By replacing a single fixed-size hidden state with a logarithmically growing, hierarchy-organized set of states, it enables sub-quadratic (O(T log T)) training and logarithmic-time/space decoding, all while retaining parallel-friendly matrix multiplications. This innovation makes it possible to train and deploy models on much longer contexts with improved recall and retrieval abilities—crucial for advancing language models in tasks like long-form understanding, retrieval, and summarization—without sacrificing computational efficiency.

1. Executive Summary (2–3 sentences)

This paper introduces log-linear attention, a middle ground between efficient linear-time RNN-style models and expressive softmax attention. It replaces a single fixed-size state with a logarithmically growing set of states organized by a Fenwick-tree hierarchy, yielding O(T log T) training time, O(T) training memory, and O(log T) decoding time and memory (Sections 3–3.2, Fig. 1–3, Eq. 3–4). The framework is general and is instantiated for Mamba-2 and Gated DeltaNet, improving long-context recall and retrieval on several benchmarks while retaining matmul-friendly parallelism (Sec. 3.3, Fig. 4, Tables 2–4, 6–7).

2. Context and Motivation

  • Problem the paper addresses
  • Standard softmax attention is accurate and hardware-friendly to train via matrix multiplications, but its compute scales quadratically with sequence length and memory scales linearly, limiting long-context use (Sec. 1–2).

  • Linear attention and state-space models (SSMs) achieve linear-time and constant-memory decoding by using a fixed-size hidden state, but this fixed capacity fundamentally limits context recall and associative retrieval (Sec. 1; Table 1; limitations summarized with references [2, 61, 34] in Sec. 1–2).

  • Why it matters

  • Long-context capabilities are increasingly needed (retrieval, summarization, code, scientific text). A model that preserves the training efficiency of matmul-heavy implementations yet expands effective memory could unlock longer, more reliable context use (Sec. 1–2).

  • Prior approaches and where they fall short

  • Softmax attention: accurate but O(T^2) compute and O(T) memory (Table 1).
  • Linear attention and modern linear RNNs/SSMs (e.g., RetNet, Mamba/Mamba-2, DeltaNet/GLA): linear-time decoding and chunk-parallelizable training, but rely on a fixed-size state (Sec. 2; Table 1). This limits associative recall and degrades with longer contexts (Sec. 1–2; refs [2, 61, 34]).

  • Long convolution models and other structured patterns (e.g., Toeplitz, Hyena): sometimes O(T log T) time but often O(T) memory at inference, or require distillation back to RNNs (Sec. 2).

  • Positioning of this work

  • The paper proposes log-linear attention, which maintains O(T log T) training and O(log T) decoding memory/time by expanding the number of states logarithmically with context length instead of keeping a single fixed state (Sec. 3–3.2). It keeps training matmul-friendly and can be “dropped in” on top of linear-attention variants (e.g., Mamba-2, Gated DeltaNet) by only modifying the masking structure M in a unifying formulation O = (A ⊙ M) V (Eq. 1; Sec. 2–3.3, Table 1).

3. Technical Approach

At a high level, the paper reframes many efficient sequence layers through a common equation (Eq. 1): - P = A ⊙ M, O = P V - A captures the query–key interactions (e.g., QK^T or more structured forms). - M is a lower-triangular, causal mask that encodes how much past positions influence the current token. Its structure determines compute/memory efficiency (Sec. 2).

The core idea is to replace the simple causal mask used in linear attention with a hierarchical, data-dependent mask so each time step only consults O(log T) “memories” at different temporal scales.

Step-by-step mechanism

1) From linear attention to a hierarchical view - Baseline linear attention (without feature maps/normalization) in parallel form is O = (QK^T ⊙ M) V with M_ij = 1{i ≤ j} (Sec. 2). - Linear attention’s recurrent form keeps a single matrix-valued state S_t with updates S_t = S_{t-1} + v_t k_t^T, output o_t = S_t q_t (Sec. 2). Gated versions add decay α_t to forget older content, still using a single state (Eq. 2; Sec. 2).

2) Fenwick-tree memory partitioning (how O(log T) states arise) - A Fenwick tree (also known as a binary indexed tree) supports fast prefix sums by partitioning [0, t) into O(log t) power-of-two segments. The paper uses this to split the prefix context for each time t into disjoint “buckets” at multiple temporal scales (Sec. 3.1; Fig. 2). - Implementation detail: - lssb(t) is the index of the least significant set bit of t; it guides how to greedily subtract largest powers of two to tile the prefix (Sec. 3.1). - Buckets B_t^{(ℓ)} have size 2^{ℓ-1} for ℓ ≥ 1 and a sentinel bucket B_t^{(0)} of size 1 (Sec. 3.1). - The model maintains one hidden state per level S_t^{(ℓ)} = ∑_{s ∈ B_t^{(ℓ)}} v_s k_s^T and mixes them with learned weights λ_t^{(ℓ)} ≥ 0: - Output: o_t = ∑_{ℓ=0}^{L-1} λ_t^{(ℓ)} q_t^T S_t^{(ℓ)} (Eq. 3). These λ_t^{(ℓ)} are predicted from the input, letting the model emphasize recent or coarse-scale context adaptively (Sec. 3.1).

3) O(log T)-space/time decoding (how to update the states efficiently) - Online (token-by-token) updates use a Fenwick-tree style recurrence (Sec. 3.1): - Insert v_t k_t^T into the finest level ℓ=0. - Zero all levels 0 < ℓ ≤ lssb(t). - Merge levels 0..lssb(t) into the next coarser level ℓ = lssb(t)+1. - Keep coarser levels above this unchanged. - This ensures only O(log T) states are stored and updated per time step, giving O(log T) time and memory per-step decoding (Sec. 3.1; recurrence under “Memory-efficient decoding”).

4) Parallel, matmul-friendly training (how to keep GPUs happy) - The same computation can be written in parallel form: - O = (QK^T ⊙ M^H) V where M^H_{t,s} = λ_t^{ℓ(t,s)} if s ≤ t, else 0; ℓ(t,s) is the level that token s belongs to for query time t (Eq. 4). - The mask M^H is a special lower-triangular hierarchical matrix. The paper connects it to the H-matrix family—specifically a quasi-HODLR (“hierarchically off-diagonal low-rank”) structure—which enables efficient algorithms (Sec. 3.1 and Appendix B; Fig. 3 left). - Chunkwise parallel training algorithm (Sec. 3.2; Fig. 3 right; Algorithm 1): - Split the sequence into chunks of length C (as in prior linear attention training). - Decompose M^H into a block-diagonal part D (handles intra-chunk interactions) plus a sum of inter-chunk low-rank pieces M^{(ℓ)} for ℓ ≥ 1 (Sec. 3.2). - Intra-chunk: multiply dense lower-triangular blocks (cost O(T·C)). - Inter-chunk: for each level, use existing linear-RNN/SSM “state passing” primitives to scan across chunks in O(T) per level; with O(log(T/C)) levels, total O(T log T) (Sec. 3.2; line 5–10 of Algorithm 1). - This is a “chunkwise parallel scan” (a hierarchical generalization of parallel prefix-sum) that remains matmul-rich and hardware-friendly (Sec. 3.2).

Why this approach over alternatives - A single fixed state (linear attention/SSMs) cannot, in principle, store high-fidelity multi-item associations over long contexts (Sec. 1 with [2]); growing the number of states logarithmically is a principled middle ground between O(1) and O(T). - Fenwick partitions naturally bias toward high resolution for recent tokens, which often matter most for language modeling and retrieval (Sec. 3.1 and remark). - The hierarchical mask preserves training parallelism and leads to O(T log T) complexity with good GPU utilization (Sec. 3.2–3.4, Fig. 4).

Applications to existing models (how to “plug in”) - Under the unified O = (A ⊙ M) V view, keep each model’s A and compose its original mask (e.g., gated/semi-separable M^S) with the hierarchical mask M^H: - Log-Linear Mamba-2: O = (QK^T ⊙ M^S ⊙ M^H) V (Sec. 3.3). - Log-Linear Gated DeltaNet: O = ([QK^T ⊙ L]·[I + KK^T ⊙ (L − I)]^{-1} ⊙ M^S ⊙ M^H) V where the [·]^{-1} term arises from the delta-rule/Householder formulation (Sec. 2 and 3.3; see “DeltaNet parallel form” and Gated DeltaNet recurrence). - Generalization to more expressive linear RNNs with matrix-valued transitions and weights is provided via “SSS tensors” and “H tensors” (Appendix A), showing the framework is not restricted to scalar gates.

4. Key Insights and Innovations

  • Hierarchical memory with O(log T) states per position
  • Novelty: Replaces a single fixed state with a logarithmically growing multi-scale set of states using a Fenwick tree partition (Sec. 3.1; Fig. 1–2; Eq. 3). This explicitly encodes multiple temporal resolutions.

  • Significance: Expands representational capacity for recall while keeping per-step decoding cost O(log T), much smaller than full attention (Sec. 3.1).

  • A quasi-H-matrix masking perspective that preserves parallel training

  • Novelty: The hierarchical mask M^H is a “quasi-H” matrix connecting HODLR/HSS ideas to attention masks (Appendix B; remark in Sec. 3.1; Fig. 3 left). It is crafted to admit O(log T) inference and O(T log T) training with matmul-rich computation.

  • Significance: Provides a principled structured-matrix view that explains both efficiency and how to compose with existing linear-time models (Sec. 3.2–3.3).

  • Chunkwise parallel scan for hierarchical state passing

  • Novelty: Extends chunked parallelization used by linear attention to a hierarchical setting, decomposing M^H into block-diagonal and low-rank inter-chunk pieces and invoking scans O(log T) times (Sec. 3.2; Algorithm 1; Fig. 3 right).

  • Significance: Keeps the hardware advantages of matmul-heavy training while achieving log-linear runtime.

  • Drop-in generalization for linear-time architectures (Mamba-2, Gated DeltaNet)

  • Novelty: Compose existing models’ masks with M^H without changing their A terms (Sec. 3.3).

  • Significance: Demonstrates the framework’s versatility and practical impact, showing improved recall/long-context use in multiple tasks (Tables 2–4, 6–7; Fig. 6–7).

  • Practical kernel design

  • Novelty: A Triton implementation that fuses work across levels and analytically unifies gradients for K/V across levels, improving training speed over a naïve multi-level approach (Sec. 3.4 and Appendix C).
  • Significance: Evidence that theory translates to actual speed benefits on GPUs (Fig. 4).

5. Experimental Analysis

Evaluation setup - Implementations and efficiency - Custom Triton kernels; level fusion; unified backward computation (Sec. 3.4; Appendix C). - Throughput and runtime measured across sequence lengths up to 131k on H100 GPUs, comparing Log-Linear Mamba-2 to FlashAttention-2 and Mamba-2 (Fig. 4).

  • Language modeling pretraining
  • 50B tokens on Long-Data-Collections, 16k context; 21-layer models around 700–800M parameters with a 24-layer Transformer (778M) included for parameter-matched comparison (Sec. 4.2).

  • Metrics: WikiText perplexity, LAMBADA perplexity/accuracy, and zero-shot commonsense tasks (Table 2).

  • Recall/long-context tasks

  • Synthetic associative recall (MQAR) with 256-token sequences, varying key-value pair counts (Sec. 4.1). Results summarized in Table 6.
  • RULER Needle-In-A-Haystack (NIAH): single-needle and multi-needle variants at 4k–16k (Fig. 7; Table 7).
  • In-context retrieval: SWDE, SQuAD, FDA, TriviaQA, DROP, NQ at lengths 512–16k (Table 3).
  • LongBench (14 long-context tasks): QA, summarization, few-shot, and code (Table 4).
  • Per-position loss on Book3 to visualize long-context utilization (Fig. 6).

Main quantitative results

  • Efficiency (training kernels)
  • Fig. 4 shows:
    • “Log-Linear Mamba-2 (naive)” is slower due to repeated primitive calls; “Log-Linear Mamba-2” with fused kernels narrows the gap and “outperforms FlashAttention-2 (forward+backward) at sequence lengths beyond 8K.” The caption also notes a throughput dip at 131k due to gradient checkpointing.

  • Conclusion: With engineering care, the O(T log T) algorithm can be competitive in wall-clock time for long contexts, not just asymptotics (Sec. 3.4; Fig. 4).

  • Synthetic associative recall (MQAR, Table 6)

  • Log-linear variants improve or match baselines at most sizes:
    • Mamba-2: accuracy at dim=64 improves from 89.6% to 92.9%.
    • Gated DeltaNet: matches ≥99% at dim=64; at dim=32 improves from 79.0% to 84.4%.
  • Quote: > “Training was early stopped when accuracy exceeded 99%.” (Table 6 note)

  • Interpretation: Expanding memory capacity helps even strong linear RNNs on associative recall.

  • Language modeling (Table 2)

  • Log-Linear Mamba-2:
    • LAMBADA perplexity improves from 24.14 to 21.86; LAMBADA accuracy from 36.2% to 37.0%.
  • Log-Linear Gated DeltaNet:
    • WikiText perplexity 21.73 → 21.44; LAMBADA perplexity 19.71 → 18.08; “Avg.” commonsense accuracy 45.0 → 45.6.

  • Quote: > “Notably, [Log-Linear Gated DeltaNet] also outperforms a layer-matched Transformer across all metrics and a parameter-matched Transformer on half of them.” (Table 2 paragraph)

  • Per-position loss (Fig. 6)

  • Both log-linear variants lower (smoothed) loss compared to their linear counterparts across positions up to 16k, indicating better long-range context use. Log-Linear Gated DeltaNet tracks the 24-layer Transformer more closely than its linear variant.

  • Needle-In-A-Haystack (Fig. 7; Table 7)

  • Strong gains for log-linear Mamba-2 at long contexts:
    • Single-needle pass-key retrieval at 16k: 21.6% → 72.4%.
    • Multi-key line retrieval at 8k: 18.6% → 39.8%.
  • Gated DeltaNet already strong in some single-needle tasks; log-linear improves several multi-needle metrics:
    • MK-NIAH-1 at 4k: 23.0% → 49.4%; multi-query at 16k: 7.2% → 9.8%.

  • Overall pattern matches the motivation: added hierarchical memory especially helps multi-item recall over long contexts.

  • In-context retrieval (Table 3)

  • Log-Linear Gated DeltaNet shows consistent improvements or parity over its linear version across tasks/lengths (e.g., FDA at 2048: 33.2% → 39.1%; SWDE at 2048: 27.2% → 35.3%).

  • Log-Linear Mamba-2 is mixed: improves on SQuAD and NQ but degrades on FDA at longer lengths.

  • LongBench (Table 4)

  • Log-Linear Mamba-2 and Log-Linear Gated DeltaNet each outperform their baselines on 8/14 tasks.
  • Results vary by task family; some summarization/code tasks remain challenging.

Do the experiments support the claims? - Yes, for the two central claims: - Efficiency: practical kernel shows competitive throughput at long lengths (Fig. 4). - Better long-context recall: improved MQAR, many NIAH metrics, lower per-position loss, and several retrieval/LongBench gains (Tables 6–7, 3–4, Fig. 6–7). - Nuance: - Improvements are stronger and more consistent for Gated DeltaNet than for Mamba-2 on some real-world retrieval tasks (Table 3). - A performance gap to Transformers persists on several benchmarks (Sec. 5 Discussion).

Ablations and robustness - The paper explores a “naive” vs. fused kernel (Fig. 4) and discusses a stronger “admissibility” H-matrix variant that marginally improves accuracy but is much slower (Appendix B.4). - Limited hyperparameter exploration for λ_t^{(ℓ)}; the authors note more tuning could help (Sec. 5 Discussion and Limitations).

6. Limitations and Trade-offs

  • Efficiency trade-off vs. linear RNNs/SSMs
  • Training time increases from O(T) (linear attention with chunking) to O(T log T). Decoding time/memory increase from O(1) to O(log T) (Table 1; Sec. 3.2).

  • While kernels are competitive at long lengths (Fig. 4), constant factors and engineering complexity are higher than linear attention (Sec. 5 Discussion).

  • Inductive bias from Fenwick-tree partitioning

  • Favoring fine granularity for recent tokens and coarse compression for distant ones may not fit all domains (Sec. 5 Discussion).

  • The strongly admissible H-matrix alternative that would utilize all hierarchical levels incurred up to 4× slowdown with marginal accuracy gains (Appendix B.4).

  • Mixed empirical gains

  • Some retrieval tasks show limited or negative changes for Log-Linear Mamba-2 (e.g., FDA at 16k in Table 3). There remains a gap to Transformers on several benchmarks (Sec. 5 Discussion; Table 2–4).

  • Complexity of implementation

  • Intra-chunk computation requires bespoke kernels; backward pass needs extra care for λ gradients; multiple levels add orchestration overhead (Sec. 5 Discussion; Appendix C).

  • Parameterization of λ_t^{(ℓ)}

  • The paper primarily uses scalar, per-level weights per head/time that are learned from inputs; richer parameterizations (e.g., matrix-valued Λ_t^{(ℓ)}) are discussed theoretically (Appendix A) but not explored empirically. The authors suggest tuning here could improve results (Sec. 5 Discussion).

7. Implications and Future Directions

  • How this changes the landscape

  • Demonstrates a principled way to move beyond fixed-size recurrent states without abandoning GPU-friendly, matmul-centric training. The hierarchical mask perspective (quasi-H matrices) could guide new efficient attention designs that sit between O(1) and O(T) state sizes (Sec. 3–3.2; Appendix B).

  • Practical applications

  • Long-context language modeling, retrieval-augmented systems, code understanding, and tasks needing multi-item recall over thousands of tokens. The O(log T) decoding memory/time makes server-side generation more scalable than full attention while improving recall over linear RNNs (Table 1; Fig. 6–7; Tables 3–4, 7).

  • Follow-up research directions

  • Learn the hierarchy: replace fixed Fenwick partitioning with learned multi-resolution segmentations while keeping parallelism (Sec. 5 Discussion).
  • Stronger hierarchical structures: revisit strongly admissible H-matrices with better kernels or approximations to reduce their current 4× overhead (Appendix B.4).
  • Richer λ (or Λ) parameterizations: explore matrix-valued, head-coupled, or content-aware level mixing (Appendix A). Combine with more expressive transition matrices (e.g., recent dynamic-state RNNs).
  • Hybridization with softmax attention: use log-linear layers for most tokens and sparse/full attention for rare, critical long-range interactions.
  • Broader domains: time series, audio, and scientific modeling where multi-scale interactions are natural; evaluate whether the recency bias helps or hurts.
  • Tooling: consolidate the Triton kernels into public libraries (the paper provides code at the linked repository; Sec. 1 footnote).

Overall, the paper offers a clear conceptual bridge between linear RNN efficiency and attention-like expressiveness via hierarchical states. Its empirical results, while mixed in places, show consistent gains on long-context recall tasks and competitive training efficiency, making it a promising direction for scalable long-sequence modeling.