Power-based Partial Attention: Bridging Linear-Complexity and Full Attention

TL;DR

The paper addresses the quadratic bottleneck of standard self-attention by introducing power-based partial attention (PPA), a parameterized mechanism with complexity $O(L^{1+p})$ that smoothly interpolates between linear ($p=0$) and full ($p=1$) attention. It develops a causal masking scheme that unions incremental-stride attention with a sliding window and systematically studies performance as a function of the exponent $p$. Empirically, the authors observe an S-shaped transition in accuracy, with a rapid gain region around $p$ in the range 0.75–0.875, showing that sub-quadratic attention can approach full-attention performance after task-specific adaptation. The work highlights practical implications for efficiency-aware transformer design and motivates further hardware-optimized implementations and broader validation across tasks and model sizes.

Abstract

It is widely accepted from transformer research that "attention is all we need", but the amount of attention required has never been systematically quantified. Is quadratic $O(L^2)$ attention necessary, or is there a sub-quadratic attention mechanism that can achieve comparable performance? To answer this question, we introduce power-based partial attention (PPA), an attention mechanism of order $O(L^{1+p})$, where $0 \leq p \leq 1$, such that $p=0$ corresponds to sliding window attention with linear complexity, and $p=1$ corresponds to full attention. With this attention construction, we can explore how transformer architecture performance varies as a function of the attention scaling behavior controlled by $p$. The overall trend from our experiments shows an S-curve-like behavior where the performance transitions from sliding-window (linear-complexity) attention to full attention over a narrow window of $p$ values, and plateaus as $p$ approaches $1$. In our experiments, we show that there exists $0<p<1$ such that $O(L^{1+p})$ attention is sufficient to achieve similar results as $O(L^2)$ full attention.

Figures (3)

• Figure 1: Illustrations of attention masks for (a) fixed stride attention, (b) dynamic stride attention, (c) incremental stride attention, and (d) incremental stride combined with sliding window attention. Each row represents a query token position, and each column represents a key token position. In fixed stride attention, each token attends to previous tokens at fixed intervals (strides). In dynamic stride attention, each token attends to previous tokens at intervals that scale with the sequence length. In incremental stride attention (red cells), the stride increases incrementally for each token, allowing for more comprehensive coverage of previous tokens while maintaining sub-quadratic complexity. In (d), blue cells show the sliding window attention overlay that ensures local context visibility, which is combined with the incremental stride pattern (red cells where not overlapped by the sliding window) to form the complete power-based partial attention mechanism used in our experiments.
• Figure 2: Evaluation loss curves during fine-tuning for different PPA power values. The curves show a clear ordering by $p$, with higher $p$ consistently achieving lower evaluation loss. The most pronounced separation occurs between $p=0.5$ and $p=0.875$, foreshadowing the performance hierarchy observed in the final benchmark results.
• Figure 3: Performance of power-based partial attention (PPA) on MATH500 (left) and GSM8k (right) benchmarks as a function of the power parameter $p$ controlling attention complexity. Both benchmarks exhibit S-shaped curves, with rapid performance gains occurring in the transition region around $p \approx 0.75$--$0.875$. The horizontal reference lines indicate the performance of sliding window attention ($p=0$) and full attention ($p=1$), demonstrating that sub-quadratic attention mechanisms with intermediate $p$ values can achieve near-full-attention performance.