PaperlayerView exact revision on arXiv

Source-linked AI summary

LoGo: Token-Level Dynamic Local-Global Attention

Yuqi Pan, Zheng Li, Bohao Tang, Zhen Qin, Guoqi Li

arXiv:2608.29539v1cs.CLcs.LG

TL;DR

Long-context attention is increasingly costly because standard Transformers assign uniform budgets despite varying token needs, and existing local-global hybrids usually allocate spans statically. LoGo learns token-level span allocation through coupled local/global branches, gated routing, and controlled global budgets. Experiments show preserved Transformer scaling, improved matched-budget performance including long-range retrieval, and interpretable allocation patterns.

  • Problem

    As context lengths scale, attention becomes a dominant computational cost, while standard Transformers allocate the same budget to tokens with different contextual demands and existing hybrids allocate spans statically.

  • Method

    LoGo gives every token local attention and uses a learned gate, adaptive threshold, progressive masking, and query-sparse kernels to selectively provide global attention under a target budget.

  • Results

    LoGo preserves Transformer scaling across 200M–3.3B parameters and improves full-attention and matched-budget static-hybrid baselines, with clear gains on long-range retrieval.

  • Takeaways & Limitations

    Learned token-level span allocation offers an effective and scalable approach to improving the long-context performance-compute trade-off.

Abstract

from arXiv · show

As context lengths scale, attention increasingly becomes a primary computational bottleneck in large language models. Standard Transformers remain powerful but computationally inefficient, as they allocate the same attention budget to every token regardless of its contextual demand. Existing local-global hybrids provide a more efficient alternative by mixing restricted- and full-context attention, but they typically allocate span statically across layers or heads. To address these limitations, we propose LoGo, a token-level dynamic local-global attention mechanism that uses attention span as a direct proxy for attention budget allocation. Each LoGo layer contains coupled local and global branches: all tokens receive efficient local attention over a restricted context window, while a learned gate activates global attention with full-context access only for tokens requiring long-range information. A threshold-based budget controller maintains a target global ratio without auxiliary losses, and a progressive masking schedule stabilizes training before sparse routing takes effect. We further implement query-sparse Triton kernels that convert reduced global-attention computation into practical speedups. Extensive experiments validate LoGo's effectiveness, showing that it preserves the scaling behavior of full-attention Transformers across model sizes. In controlled comparisons, LoGo improves over the full-attention Transformer and matched-budget static local-global hybrids, with clear gains on long-range retrieval. Analysis further shows that LoGo learns interpretable span allocation patterns. These results suggest that learned token-level span allocation is an effective and scalable way to improve the long-context performance-compute trade-off.

1 Introduction

LoGo addresses the rising cost of long-context attention by learning token-level local/global span allocation instead of assigning every token the same attention budget. It combines selective global routing with budget control and reports improved efficiency and long-range retrieval while preserving Transformer scaling.

  • As context lengths grow, attention becomes a dominant cost, while standard Transformers allocate computation uniformly despite differing token-level contextual demands.
  • Attention span serves as a proxy for budget because it controls both accessible context and computation cost, while long-range retrieval tokens may require global access.
  • Existing local-global hybrids allocate spans statically across layers or heads, forcing all tokens through the same pathway and potentially wasting global computation.
  • The research questions test whether learned token-level allocation can target long-range needs, preserve backbone scaling, and improve the performance-compute trade-off over static hybrids.
  • LoGo couples local and global branches, using a learned token gate and adaptive threshold to activate full-context attention only for selected queries under a target ratio.
  • LoGo preserves or improves standard-Transformer scaling from 200M to 3.3B parameters and improves matched-budget comparisons, especially on long-range retrieval.

2 Related Work

Prior work reduces long-context cost through conditional computation and local-global attention, but existing attention hybrids generally use static or handcrafted routing. LoGo is positioned as a learned token-level span-allocation alternative.

  • Conditional computation allocates token-adaptive budgets through experts, depth, adaptive head selection, and token-wise sparse attention.
  • Local-global hybrids combine restricted-span and full-context attention, typically mixing them statically across layers or heads.
  • Sequence-adaptive hybrids often use fixed patterns, while TransMamba and Oryx rely on handcrafted per-token budgets rather than learned routing.

3 Methods

LoGo allocates local and global attention dynamically at the token level, combining learned routing with budget control and query-sparse execution. Its design couples span-specific branches while progressively introducing sparse global routing during training.

  • LoGo Attention: LoGo uses global attention only for selected token queries, while all tokens receive local attention over a restricted causal window.The main experiments use a local window of w = 128; selected queries retain full-context access.
  • LoGo Attention: LoGo shares the main attention projections across branches and adds lightweight branch-specific transformations so local and global representations can specialize.The transformations are initialized as identities, and the output projection remains shared.
  • LoGo Attention: A learned scalar gate predicts each token’s preference for global span and provides the routing decision and fusion weight.Unselected tokens reduce to the local output, whereas selected tokens combine local and global outputs.
  • Budget Control and Progressive Masking: A threshold controller maintains the realized global activation ratio near a target without auxiliary balancing losses or modifying the language-modeling objective.When the realized ratio exceeds the target, the threshold increases; for example, ρ^(ℓ) = 0.5 activates global attention for half of query tokens per layer.
  • Budget Control and Progressive Masking: Progressive masking starts with dense global supervision and gradually introduces hard routing as the gate and both branches learn.The threshold is initialized so all tokens initially activate global attention, then moves toward the target operating point during the masking schedule.
  • Query-Sparse Implementation: Query-sparse kernels pack selected global queries and skip pruned query rows, preserving dense-attention semantics for selected rows while reducing computation.During decoding, unselected tokens skip the global call, although batched decoding may launch a dense call when any selected row exists.

4 Experiments

Experiments evaluate LoGo’s scaling, matched-budget performance, efficiency, ablations, robustness, hybrid compatibility, and learned span allocation. Across these tests, LoGo preserves scaling, improves the performance-compute trade-off, accelerates sparse attention, and develops interpretable routing patterns.

  • Scaling Property: LoGo consistently matches or improves the full-attention Transformer across model sizes from 200M to 3.3B parameters.Both models use the same MHA-SwiGLU architecture and training hyperparameters.
  • Comparison with Static Hybrid Paradigms: Under matched parameters and attention FLOPs, LoGo achieves the best overall language-modeling results against full-attention and static local-global baselines.The comparison uses four 1.5B attention paradigms with a matched global ratio of 0.5.
  • Recall: LoGo remains competitive on short-context recall and provides clear gains on long-range retrieval, with larger gains at longer sequence lengths.On RULER-32k and RULER-64k, LoGo obtains the best score on five of eight subtasks.
  • Efficiency: At a 0.5 attention budget, the query-sparse kernel reaches 1.99× speedup at 64k relative to the dense Triton baseline.The kernel’s speedups scale with sparsity, while comparison with dense CUDA shows an advantage only at high sparsity.
  • Ablation Study: Ablations identify lightweight branch decoupling, ContextNorm, weighted fusion, threshold-based budget control, and P-mask as important design choices.Fully sharing branches hurts performance, while fully decoupling them adds substantial parameters without improving recall.
  • Analysis: LoGo is robust across global ratios, can complement static hybrids, and learns domain- and token-dependent allocation patterns.The average global ratio stays close to the target 0.5, while long-range-demand tokens receive more global span and locally predictable tokens receive less.

5 Conclusion

LoGo dynamically allocates local and global attention at the token level, using attention span as a controlled proxy for attention budget. The method is designed to complement static hybrids and improves long-context performance-compute trade-offs while preserving Transformer scaling behavior.

  • 5 Conclusion: LoGo routes each token between local-only and local-plus-global computation using attention span as a direct proxy for attention budget.Every token receives local attention, while a learned gate selectively activates global attention under a controlled budget.
  • 5 Conclusion: LoGo preserves Transformer scaling behavior and improves over full-attention and static hybrid baselines under matched parameters and computation.The reported gains include clear improvements on long-range retrieval.
  • 5 Conclusion: LoGo differs from prior conditional-computation methods by learning span allocation end-to-end and evaluating it at pretrained scale under matched parameters and FLOPs.The design emphasizes long-context retrieval rather than only short-context quality or inference speed.
  • 5 Conclusion: Static inter-layer and intra-layer hybrids assign the same span pathway across tokens, whereas LoGo adapts global computation to token-level contextual demand.Inter-layer designs vary attention across layers, while intra-layer designs vary attention across heads.
  • 5 Conclusion: LoGo is complementary to static local-global hybrids because it keeps the attention form fixed while learning token-level span allocation under a controlled global-attention ratio.This enables direct comparisons with inter-layer and intra-layer hybrids.

B Efficient Implementation Details

LoGo implements dynamic global attention through selected-query sparsity while retaining local attention for all tokens. Its attention cost matches the dominant order of static hybrids with the same global fraction, but its dense key-value cache limits memory savings.

  • B Efficient Implementation Details: During prefill and training, selected queries are compacted while keys and values remain dense, preserving causal visibility through original sequence positions.The backward pass preserves the same row-sparse structure for query gradients and dense scheduling for key-value gradients.
  • B Efficient Implementation Details: During decoding, locally routed tokens can skip the global attention call, while globally routed tokens attend over the full global cache and fuse both outputs.The local branch attends to at most the most recent w cached tokens.
  • B Efficient Implementation Details: The conservative batched decoding policy skips global computation only when no query activates it; otherwise, it launches the dense kernel and masks inactive rows.This avoids per-step query-compaction overhead while capturing all-local decoding steps.
  • B Efficient Implementation Details: LoGo evaluates the global branch only for selected query rows, yielding total attention cost O(TwHdh + ρ(ℓ)T 2Hdh).The local branch processes all tokens, while the global branch is restricted by target global ratio ρ(ℓ).
  • B Efficient Implementation Details: LoGo does not by itself provide proportional memory savings because its key-value cache remains dense and linear in context length.It can be combined with memory-saving local-global designs by replacing their global-attention components with query-sparse LoGo layers.
  • B Efficient Implementation Details: Row sparsity gives all heads the same token-level routing decision, avoiding head-wise load imbalance while preserving compatibility with tensor and sequence parallelism.The dense key-value layout also preserves sequence order and causal structure.

C.1 Training Setup

The scaling and controlled-comparison experiments use matched Transformer backbones and standardized optimization settings. LoGo is trained with a fixed target global ratio, sliding-window size, and progressive masking initialization.

  • C.1 Training Setup: LoGo scaling runs use target global ratio ρ(ℓ) = 0.5, sliding-window size w = 128, and threshold update rate γ = 0.0005.The threshold is initialized according to a 20% P-mask duration.
  • C.1 Training Setup: Scaling experiments fix the tokens-per-parameter ratio at 20 and use a maximum sequence length of 8192.Batch size is 0.5M tokens for the first four scales and 1M tokens for the 3.3B scale.
  • C.1 Training Setup: The controlled 1.5B comparison trains all models for 100B tokens with maximum sequence length 8k and global batch size 4M tokens.Models share the MHA-SwiGLU backbone and differ only in their attention span pathway.
  • C.1 Training Setup: For a P-mask duration of S steps, the initial threshold is set to b(ℓ)_th,0 = 0.5 − γS; with 5000 steps, it equals -2.0.The 5000-step duration corresponds to 20% of the 100B-token pretraining schedule.

C.2 Evaluation Setup

The evaluation covers language modeling, zero-shot commonsense reasoning, short-context recall, and long-range needle-style retrieval. These tasks span both standard quality measures and context-dependent retrieval behavior.

  • C.2 Evaluation Setup: Language modeling is evaluated with training loss and perplexity on WikiText and Lambada.The benchmarks are reported as Wiki. and LMB.
  • C.2 Evaluation Setup: Zero-shot commonsense reasoning reports average accuracy across Lambada, PIQA, HellaSwag, WinoGrande, and ARC-Easy or ARC-Challenge.The evaluation aggregates performance over multiple commonsense reasoning tasks.
  • C.2 Evaluation Setup: Short-context recall includes FDA, SWDE, SQuAD, TriviaQA, Natural Questions, and DROP.The evaluation follows the setting of Arora et al.
  • C.2 Evaluation Setup: Long-range recall uses all needle-style RULER tasks, including three single-needle and five multi-needle variants.The multi-needle variants include MK-NIAH, MQ-NIAH, and MV-NIAH.

C.3 Additional Experimental Results

The additional results document LoGo’s evaluation after extending context to 32k and 128k tokens, covering language modeling, commonsense reasoning, and per-task long-context retrieval.

  • Evaluation scope: Tables 7–9 report LoGo results after context extension, including language modeling, commonsense reasoning, and RULER evaluations.Table 7 covers 32k context; Tables 8 and 9 provide per-task RULER results for 32k and 128k checkpoints.
  • Evaluation scope: Table 7 evaluates language modeling and commonsense reasoning after 32k context extension using global-attention ratio as the attention-budget measure.Loss and perplexity are lower-is-better metrics, while accuracy is higher-is-better.
  • Long-context retrieval: Tables 8 and 9 report needle-style RULER scores across progressively longer input lengths for 32k and 128k checkpoints.The 32k checkpoints are evaluated at 8k, 16k, and 32k inputs, while 128k checkpoints are evaluated at 32k and 64k inputs.

C.4.1 Additional Ablations

The additional ablations examine branch design, normalization, fusion, budget control, local/global ratios, and compatibility with memory-saving hybrids.

  • Architecture ablations: Fully sharing branch parameters hurts modeling and recall, while fully decoupling them increases parameters by about 20% without improving recall.The ablation supports lightweight branch decoupling rather than either complete sharing or complete separation.
  • Architecture ablations: Removing ContextNorm prevents proper convergence, showing that explicit output normalization is important when local and global branches have different variances.The result holds even when scale-and-offset transformation reduces branch discrepancy.
  • Architecture ablations: Weighted fusion outperforms additive fusion, supporting explicit competition between local and global spans.The ablation identifies fusion choice as a contributor to LoGo’s architecture behavior.
  • Architecture ablations: Threshold-based budget control outperforms auxiliary-loss control for modeling performance and training loss.Replacing threshold-based control degrades modeling performance and yields a higher training loss.
  • Compatibility with other hybrids: LoGo combined with a 1:1 inter-layer hybrid raises average recall to 65.6/66.9 versus 65.1 while further reducing compute.The 0.5-threshold variant also reaches 65.6 versus 64.7 for an equal-FLOPs 1:3 inter-layer hybrid.
  • Local/global ratios: The ratio-sweep tables compare LoGo with inter-layer hybrids under matched global-attention budgets for short-context recall and RULER performance.Table 10 additionally reports average recall excluding FDA because FDA varies noticeably with the global ratio.

C.4.2 Additional Analysis

Additional analyses show that LoGo maintains its target global ratio while allocating global attention non-uniformly across layers and tokens, with patterns linked to contextual structure.

  • Span allocation statistics: The average global ratio stays close to the target ratio of 0.5 across four domains, with mild domain variation.This confirms that the budget controller maintains its intended operating point.
  • Span allocation statistics: Global-span allocation varies by domain and layer: Wiki and Quest concentrate it near the final layer, GSM8K peaks in the first, and code uses the least.Figure 4 reports realized layer-wise ratios and token-wise distributions over 0–24 activating layers.
  • Word-level visualization: Figure 5 visualizes word-level span allocation, with warmer colors indicating more layers activating global attention.The visualization records global-span activation at the word level.
  • Word-level visualization: Repeated content triggers more global activation when the matching material lies outside the local window, while nearby repetition uses less.The analysis hypothesizes that anchor-token matches can lead subsequent tokens to rely on global retrieval.
  • Word-level visualization: After a context switch, global-span usage drops quickly and rises again when the passage resumes long-range matching.This pattern connects routing changes to the location of long-range contextual matches.
Loading 2608.29539v1…