---
title: All-or-Here Attention (AHA) in Transformers
url: https://www.emergentmind.com/topics/all-or-here-attention-aha
type: topic
---

# All-or-Here Attention (AHA) in Transformers

All-or-Here Attention (AHA) is a dynamic architectural mechanism designed to drastically reduce the frequency of full global attention computations in Transformer-based large language models (LLMs) by enabling per-head, per-token routing between global (full) and local (sliding-window) attention. With a lightweight router and supervised fine-tuning, AHA achieves extensive sparsity in attention computations, replacing over 90% of global attention with local alternatives at no cost to downstream accuracy for practical window sizes. Empirical analysis reveals that the requirement for global attention exhibits a highly skewed distribution, with only a minority of tokens and heads necessitating access to the entire context, while the majority operate effectively within local context windows [2512.22562].

## 1. Formal Specification and Mechanisms

AHA extends the standard decoder-only Transformer architecture by incorporating a binary routing function, per layer and per head, that determines whether each token–head pair requires global or local context. Let $X\in\mathbb{R}^{n\times d}$ denote hidden states, $n$ sequence length, $d$ model dimensionality. Each attention head $h$ defines projections $W_Q, W_K, W_V \in \mathbb{R}^{d\times(d/m)}$:

- **Full Self-Attention (Global):**  
  $$
  \text{Attention}_\text{all}(q_t,\ K,\ V) = \text{Softmax}\left(\frac{q_t K^\top}{\sqrt{d_k}}\right)V,\quad q_t=X_t W_Q,\ d_k=\frac{d}{m}
  $$
- **Sliding-Window Attention (Local, window $w$):**  
  $$
  \text{Attention}_\text{here}(q_t, K_{t-w+1:t}, V_{t-w+1:t}) = \text{Softmax}\left(\frac{q_t K_{t-w+1:t}^\top}{\sqrt{d_k}}\right)V_{t-w+1:t}
  $$

The routing function is implemented as $S = \sigma(X W_\text{router}) \in (0,1)^{n\times m}$, where $W_\text{router}\in\mathbb{R}^{d\times m}$. Binarization follows, with $g_{t,h} = \mathbb{1}(s_{t,h}>\tau)$ for threshold $\tau=0.5$. Each head's output is
$$
a_{t,h} = g_{t,h}\ \text{Attention}_\text{all}(q_{t,h},K_{1:t,h},V_{1:t,h}) + (1-g_{t,h})\ \text{Attention}_\text{here}(q_{t,h},K_{t-w+1:t,h},V_{t-w+1:t,h})
$$

Gradients during backpropagation are passed via a Straight-Through Estimator (STE): $\partial\mathcal{L}/\partial s_{t,h} \approx \partial\mathcal{L}/\partial g_{t,h}$.

## 2. Architectural Integration and Modifications

AHA substitutes the standard full self-attention module in each Transformer (decoder-only) block with an AHA block. The router is applied as a linear projection immediately before standard Q/K/V computations. For each token and head, binary gating determines whether to perform a full global attention or a local sliding-window attention read. All other architectural components, including normalization, feed-forward modules, and residual pathways, remain unchanged. The system is designed for compatibility with established models (e.g., LLaMA, OLMo-2); implementation proceeds by initializing from a pre-trained checkpoint, inserting $W_\text{router}$, and continuing with supervised fine-tuning [2512.22562].

## 3. Training Objective and Regularization

The optimization objective is a sum of standard autoregressive language modeling loss and an $L_1$-based regularization on the router's output:
$$
\mathcal{L} = \mathcal{L}_\text{LM} + \lambda\,\mathcal{L}_\text{reg}
$$
where
- $\mathcal{L}_\text{LM}$: cross-entropy next-token prediction,
- $\mathcal{L}_\text{reg} = \frac{1}{L n m}\sum_{k=1}^L\sum_{t=1}^n\sum_{h=1}^m s_{t,h}^{(k)}$ penalizes global attention usage,
- $\lambda$ ($3\times10^{-4}$ by default) controls sparsity-performance trade-off,
- threshold $\tau = 0.5$.  
STE enables gradients to flow through discrete gate selection.

## 4. Empirical Evaluation and Performance

AHA was evaluated by fine-tuning OLMo-2-0425-1B-SFT for one epoch on TULU-v3 using AdamW optimization (learning rate $3\times 10^{-5}$, $\beta_1 = 0.9,\ \beta_2 = 0.95,\ $warmup=3\%, batch=128). Window sizes $w$ explored: 16, 32, 64, 128, 256.

Downstream evaluation included:
- Single-token: MMLU (Acc), HellaSwag (Acc_norm), CSQA (Acc)
- Multi-token: GSM8K (EM, 5-shot), MBPP (Pass@1, zero-shot), MultiNews (ROUGE, zero-shot)

Key quantitative results for $w=256$:
- Relative retained accuracy: 102.5% of baseline
- Average global attention usage $\mu_f \approx 6.7\%$—i.e., 93.3% of full attention operations supplanted by local attention
- At $w=128$, $\mu_f\approx11.6\%$; at $w=16$, $\mu_f\approx52.7\%$
- Even with extreme locality ($w=16$), model retains $\approx92.7\%$ baseline performance

## 5. Distributional Analysis of Context Dependency

Analysis reveals a pronounced “long-tail” in the necessity of full attention as a function of window size:
- As $w$ increases: $\mu_f$ drops from 52.7% (16) → 41.4% (32) → 28.1% (64) → 11.6% (128) → 6.7% (256)
- Most heads and tokens rarely invoke global attention, confirmed by heatmaps across layers/heads where only a minority display high global attention rates (“heavy-hitters”)

Representative per-head “full-attention gap” (number of tokens between global-attention activations):
- Layer 5, Head 2: gap $\approx$ 1 (always-on)
- Layer 12, Head 8: gap in tens–hundreds
- Layer 3, Head 6: gap in thousands

This suggests most model computations can be executed efficiently with only local context, with sporadic global reads required.

## 6. Inference Algorithm and Implementation

Efficient inference with AHA is enabled by jointly computing importance scores and per-head binary gating, followed by selective routing to either global or local attention:

```python
for each layer k in 1..L:
  X ← hidden_states
  S = sigmoid( X @ W_router^(k) )          # shape n×m
  G = (S > τ).float()                      # binary gates n×m
  for each head h in 1..m:
    Q_h = X @ W_Q^(k,h)
    K_h = X @ W_K^(k,h)
    V_h = X @ W_V^(k,h)
    for each token t in 1..n:
      if G[t,h]==1:
        # full causal attention
        A[t,h] = Attention( Q_h[t], K_h[1:t], V_h[1:t] )
      else:
        # sliding window attention
        start = max(1, t-w+1)
        A[t,h] = Attention( Q_h[t], K_h[start:t], V_h[start:t] )
  # Output aggregation
  Concat all A[:, h] across heads, project by W_O, add resid & FFN
```

During backpropagation, gradients propagate through the STE to train $W_\text{router}$ in an end-to-end manner.

## 7. Limitations and Potential Extensions

AHA exhibits several constraints and extension opportunities:
- **Sparsity–Performance Trade-off:** Extreme regularization (e.g., $\lambda = 1\times10^{-3}$, $\mu_f \approx 4.5\%$) can hurt accuracy (e.g., GSM8K EM drop from 0.4291 to 0.3730), while insufficient regularization yields excess global attention ($\lambda = 1\times10^{-4}$, $\mu_f \approx 53.5\%$) without accuracy benefit. $\lambda = 3\times10^{-4}$ achieves effective sparsity–retention balance.
- **System-Level Speedup:** While AHA achieves algorithmic sparsity, current hardware and attention kernels (e.g., FlashAttention) are optimized for static dense computation. Realizing wall-clock speedups is contingent on hardware-adaptive, dynamic computation support.
- **Scalability and Generalization:** No from-scratch pre-training is required for AHA deployment. Integration and extension possibilities include multi-choice routing (heads select among multiple window sizes or global), mixture-of-experts/spans routing, and layering AHA atop other efficient attention mechanisms (sparse, low-rank, index-based) [2512.22562].

AHA establishes a minimal, data-dependent conditional computation design for dynamic sparsity in large-scale Transformers, conclusively showing that full global attention is usually redundant except for a narrow subset of model computations.

Source: https://www.emergentmind.com/topics/all-or-here-attention-aha