Table of Contents
Fetching ...

Learning When Not to Attend Globally

Xuan Luo, Kailai Zhang, Xifeng Yan

TL;DR

This work tackles the quadratic cost of self-attention by introducing All-or-Here Attention (AHA), which attaches a binary router to each attention head to decide, per token, whether to use full attention or a local sliding-window span. The method yields a drop-in, hard-gated mechanism that, when trained with a joint loss and an $L1$ router-penalty, dramatically reduces full-attention usage while maintaining or even slightly exceeding Vanilla performance at moderate window sizes (e.g., $w=128$–$256$). Empirical results across multiple benchmarks reveal a long-tail distribution in context dependency: most tokens rely on local context, with only a sparse subset requiring global access; at $w=256$, over 90% of full-attention operations can be bypassed. This decoupling of local processing from global retrieval points to practical efficiency gains for inference in large language models, albeit with caveats about hardware speedups and the need for broader evaluations on additional models and attention variants.

Abstract

When reading books, humans focus primarily on the current page, flipping back to recap prior context only when necessary. Similarly, we demonstrate that Large Language Models (LLMs) can learn to dynamically determine when to attend to global context. We propose All-or-Here Attention (AHA), which utilizes a binary router per attention head to dynamically toggle between full attention and local sliding window attention for each token. Our results indicate that with a window size of 256 tokens, up to 93\% of the original full attention operations can be replaced by sliding window attention without performance loss. Furthermore, by evaluating AHA across various window sizes, we identify a long-tail distribution in context dependency, where the necessity for full attention decays rapidly as the local window expands. By decoupling local processing from global access, AHA reveals that full attention is largely redundant, and that efficient inference requires only on-demand access to the global context.

Learning When Not to Attend Globally

TL;DR

This work tackles the quadratic cost of self-attention by introducing All-or-Here Attention (AHA), which attaches a binary router to each attention head to decide, per token, whether to use full attention or a local sliding-window span. The method yields a drop-in, hard-gated mechanism that, when trained with a joint loss and an router-penalty, dramatically reduces full-attention usage while maintaining or even slightly exceeding Vanilla performance at moderate window sizes (e.g., ). Empirical results across multiple benchmarks reveal a long-tail distribution in context dependency: most tokens rely on local context, with only a sparse subset requiring global access; at , over 90% of full-attention operations can be bypassed. This decoupling of local processing from global retrieval points to practical efficiency gains for inference in large language models, albeit with caveats about hardware speedups and the need for broader evaluations on additional models and attention variants.

Abstract

When reading books, humans focus primarily on the current page, flipping back to recap prior context only when necessary. Similarly, we demonstrate that Large Language Models (LLMs) can learn to dynamically determine when to attend to global context. We propose All-or-Here Attention (AHA), which utilizes a binary router per attention head to dynamically toggle between full attention and local sliding window attention for each token. Our results indicate that with a window size of 256 tokens, up to 93\% of the original full attention operations can be replaced by sliding window attention without performance loss. Furthermore, by evaluating AHA across various window sizes, we identify a long-tail distribution in context dependency, where the necessity for full attention decays rapidly as the local window expands. By decoupling local processing from global access, AHA reveals that full attention is largely redundant, and that efficient inference requires only on-demand access to the global context.
Paper Structure (15 sections, 8 equations, 4 figures, 4 tables)

This paper contains 15 sections, 8 equations, 4 figures, 4 tables.

Figures (4)

  • Figure 1: Visualization of full average attention usage across tasks (local attention window size, $w=128$). The color gradient indicates the ratio of activating full attention, where dark blue represents 1.0 (Full Attention) and light blue represents 0.0 (Local Attention). This highlights that reliance on global context varies by task nature and token characteristics.
  • Figure 2: Overview of the All-or-Here Attention (AHA) architecture. A lightweight router computes importance scores for each head, generating binary gates that dynamically toggle between full causal attention and local sliding window attention.
  • Figure 3: Visualization of the average full attention usage $\mu_f$ across different layers and heads. Lighter colors indicate a higher frequency of triggering full attention.
  • Figure 4: Sorted average full attention usage for all attention heads. The data shows that a small minority of heads account for the majority of full attention usage.