Linear Attention & Efficient Transformers
1. The Quadratic Bottleneck
Standard self-attention computes:
The attention matrix requires compute and memory.
1.1 Complexity Comparison
| Method | Time | Memory | Sequential |
|---|---|---|---|
| Standard Attention | |||
| Linear Attention | |||
| Sparse Attention |
2. Standard vs. Linear Attention Comparison
<svg viewBox="0 0 800 400" xmlns="http://www.w3.org/2000/svg">
<rect width="800" height="400" fill="#f8fafc"/>
{/* Title */}
<text x="400" y="30" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="16" font-weight="bold">Standard vs. Linear Attention</text>
{/* Standard Attention */}
<g transform="translate(50, 50)">
<text x="175" y="0" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="12" font-weight="bold">Standard Attention</text>
{/* Formula */}
<g transform="translate(0, 10)">
<rect width="350" height="40" rx="5" fill="#f8fafc" stroke="#ec4899"/>
<text x="175" y="25" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="10">Attn = softmax(QKᵀ/√d) · V</text>
</g>
{/* Matrix multiplication visualization */}
<g transform="translate(0, 60)">
{/* Q matrix */}
<rect x="0" y="0" width="60" height="120" rx="3" fill="#3b82f6"/>
<text x="30" y="65" text-anchor="middle" fill="white" font-family="monospace" font-size="10" font-weight="bold">Q</text>
<text x="30" y="135" text-anchor="middle" fill="#60a5fa" font-family="monospace" font-size="8">n × d</text>
{/* K matrix */}
<rect x="80" y="0" width="60" height="120" rx="3" fill="#a855f7"/>
<text x="110" y="65" text-anchor="middle" fill="white" font-family="monospace" font-size="10" font-weight="bold">Kᵀ</text>
<text x="110" y="135" text-anchor="middle" fill="#c084fc" font-family="monospace" font-size="8">d × n</text>
{/* Attention matrix */}
<rect x="160" y="20" width="80" height="80" rx="3" fill="#f59e0b"/>
<text x="200" y="65" text-anchor="middle" fill="white" font-family="monospace" font-size="10" font-weight="bold">A</text>
<text x="200" y="115" text-anchor="middle" fill="#fbbf24" font-family="monospace" font-size="8">n × n</text>
{/* V matrix */}
<rect x="260" y="0" width="60" height="120" rx="3" fill="#22c55e"/>
<text x="290" y="65" text-anchor="middle" fill="white" font-family="monospace" font-size="10" font-weight="bold">V</text>
<text x="290" y="135" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="8">n × d</text>
{/* Operation symbols */}
<text x="70" y="65" text-anchor="middle" fill="#64748b" font-family="monospace" font-size="14">×</text>
<text x="145" y="65" text-anchor="middle" fill="#64748b" font-family="monospace" font-size="14">→</text>
<text x="245" y="65" text-anchor="middle" fill="#64748b" font-family="monospace" font-size="14">×</text>
</g>
{/* Complexity */}
<text x="175" y="220" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="11" font-weight="bold">Complexity: O(n²d)</text>
<text x="175" y="240" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">Attention matrix is n×n</text>
<text x="175" y="255" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">Quadratic in sequence length</text>
</g>
{/* Linear Attention */}
<g transform="translate(450, 50)">
<text x="175" y="0" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="12" font-weight="bold">Linear Attention</text>
{/* Formula */}
<g transform="translate(0, 10)">
<rect width="350" height="40" rx="5" fill="#f8fafc" stroke="#22c55e"/>
<text x="175" y="25" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="10">Attn = φ(Q)(φ(K)ᵀV) / φ(Q)Σφ(K)ᵀ</text>
</g>
{/* Kernel trick visualization */}
<g transform="translate(0, 60)">
{/* φ(K) matrix */}
<rect x="0" y="0" width="60" height="80" rx="3" fill="#a855f7"/>
<text x="30" y="45" text-anchor="middle" fill="white" font-family="monospace" font-size="9" font-weight="bold">φ(K)ᵀ</text>
<text x="30" y="95" text-anchor="middle" fill="#c084fc" font-family="monospace" font-size="8">d × n</text>
{/* V matrix */}
<rect x="80" y="0" width="60" height="80" rx="3" fill="#22c55e"/>
<text x="110" y="45" text-anchor="middle" fill="white" font-family="monospace" font-size="9" font-weight="bold">V</text>
<text x="110" y="95" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="8">n × d</text>
{/* KV matrix */}
<rect x="160" y="10" width="60" height="60" rx="3" fill="#f59e0b"/>
<text x="190" y="45" text-anchor="middle" fill="white" font-family="monospace" font-size="9" font-weight="bold">KV</text>
<text x="190" y="85" text-anchor="middle" fill="#fbbf24" font-family="monospace" font-size="8">d × d</text>
{/* φ(Q) matrix */}
<rect x="240" y="0" width="60" height="80" rx="3" fill="#3b82f6"/>
<text x="270" y="45" text-anchor="middle" fill="white" font-family="monospace" font-size="9" font-weight="bold">φ(Q)</text>
<text x="270" y="95" text-anchor="middle" fill="#60a5fa" font-family="monospace" font-size="8">n × d</text>
{/* Operation symbols */}
<text x="70" y="45" text-anchor="middle" fill="#64748b" font-family="monospace" font-size="14">×</text>
<text x="145" y="45" text-anchor="middle" fill="#64748b" font-family="monospace" font-size="14">→</text>
<text x="225" y="45" text-anchor="middle" fill="#64748b" font-family="monospace" font-size="14">×</text>
</g>
{/* Complexity */}
<text x="175" y="220" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="11" font-weight="bold">Complexity: O(nd²)</text>
<text x="175" y="240" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">Kernel trick avoids n×n matrix</text>
<text x="175" y="255" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">Linear in sequence length (if d ≪ n)</text>
</g>
{/* Key insight */}
<g transform="translate(50, 310)">
<rect width="700" height="70" rx="8" fill="#f8fafc" stroke="#334155"/>
<text x="350" y="20" text-anchor="middle" fill="#f59e0b" font-family="monospace" font-size="11" font-weight="bold">Key Insight</text>
<text x="350" y="40" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="10">Standard: (QKᵀ)V — computes n×n attention first, then multiplies by V</text>
<text x="350" y="55" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="10">Linear: Q(KᵀV) — computes d×d matrix first, then multiplies by Q</text>
</g>
</svg>
3. Kernel Attention
3.1 The Kernel Trick
Replace softmax with a kernel feature map :
The attention becomes:
3.2 Associative Property
The key insight is the associative property of matrix multiplication:
Computing first gives instead of .
3.3 Common Kernel Functions
| Kernel | Properties | |
|---|---|---|
| ELU+1 | Simple, non-negative | |
| Random Fourier | Approximates RBF kernel | |
| Performer | Unbiased estimate | |
| ReLU | Simple, sparse |
4. Performer (Choromanski et al., 2021)
4.1 FAVOR+ (Fast Attention Via positive Orthogonal Random features)
The FAVOR+ algorithm provides an unbiased estimate of softmax attention:
where:
with .
4.2 Unbiased Estimate Property
The variance decreases as increases.
4.3 Positive Random Features
To ensure non-negativity (required for attention weights), use positive random features:
5. Performer Random Features
<svg viewBox="0 0 800 400" xmlns="http://www.w3.org/2000/svg">
<rect width="800" height="400" fill="#f8fafc"/>
{/* Title */}
<text x="400" y="30" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="16" font-weight="bold">Performer: FAVOR+ Random Features</text>
{/* Input space */}
<g transform="translate(50, 60)">
<text x="175" y="0" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="11">Input Space (d dimensions)</text>
{/* Q and K vectors */}
<g transform="translate(0, 15)">
<rect width="80" height="100" rx="5" fill="#f8fafc" stroke="#3b82f6" stroke-width="2"/>
<text x="40" y="25" text-anchor="middle" fill="#60a5fa" font-family="monospace" font-size="10" font-weight="bold">q</text>
<text x="40" y="50" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">∈ ℝᵈ</text>
<text x="40" y="70" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">Query</text>
</g>
<g transform="translate(100, 15)">
<rect width="80" height="100" rx="5" fill="#f8fafc" stroke="#a855f7" stroke-width="2"/>
<text x="40" y="25" text-anchor="middle" fill="#c084fc" font-family="monospace" font-size="10" font-weight="bold">k</text>
<text x="40" y="50" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">∈ ℝᵈ</text>
<text x="40" y="70" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">Key</text>
</g>
{/* Softmax kernel */}
<text x="85" y="140" text-anchor="middle" fill="#f59e0b" font-family="monospace" font-size="10">K(q,k) = exp(qᵀk)</text>
</g>
{/* Random features */}
<g transform="translate(300, 60)">
<text x="100" y="0" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="11">Random Features (m dimensions)</text>
{/* φ(q) */}
<g transform="translate(0, 15)">
<rect width="80" height="100" rx="5" fill="#f8fafc" stroke="#22c55e" stroke-width="2"/>
<text x="40" y="25" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="10" font-weight="bold">φ(q)</text>
<text x="40" y="50" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">∈ ℝᵐ</text>
<text x="40" y="70" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">φ(q)ᵀφ(k)</text>
<text x="40" y="85" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">≈ K(q,k)</text>
</g>
{/* φ(k) */}
<g transform="translate(100, 15)">
<rect width="80" height="100" rx="5" fill="#f8fafc" stroke="#f59e0b" stroke-width="2"/>
<text x="40" y="25" text-anchor="middle" fill="#fbbf24" font-family="monospace" font-size="10" font-weight="bold">φ(k)</text>
<text x="40" y="50" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">∈ ℝᵐ</text>
<text x="40" y="70" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">Random Fourier</text>
<text x="40" y="85" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">features</text>
</g>
</g>
{/* Feature map formula */}
<g transform="translate(50, 200)">
<rect width="700" height="100" rx="8" fill="#f8fafc" stroke="#334155"/>
<text x="350" y="20" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="12" font-weight="bold">FAVOR+ Feature Map</text>
<text x="350" y="45" text-anchor="middle" fill="#60a5fa" font-family="monospace" font-size="10">φ(x) = (e⁻‖x‖²/2 / √m) · [e^(ω₁ᵀx), e^(ω₂ᵀx), ..., e^(ωₘᵀx)]</text>
<text x="350" y="70" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="10">Property: E[φ(x)ᵀφ(y)] = softmax(xᵀy) (Unbiased estimate!)</text>
<text x="350" y="90" text-anchor="middle" fill="#fbbf24" font-family="monospace" font-size="9">ωᵢ ~ N(0, I_d) are random projection vectors</text>
</g>
{/* Complexity comparison */}
<g transform="translate(50, 320)">
<text x="350" y="0" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="11">Complexity Reduction</text>
<rect x="0" y="15" width="200" height="50" rx="5" fill="#ec4899" opacity="0.8"/>
<text x="100" y="35" text-anchor="middle" fill="white" font-family="monospace" font-size="10">Standard: O(n²d)</text>
<text x="100" y="52" text-anchor="middle" fill="white" font-family="monospace" font-size="8">n=4096, d=128: 2.1×10⁹ FLOPs</text>
<rect x="250" y="15" width="200" height="50" rx="5" fill="#22c55e" opacity="0.8"/>
<text x="350" y="35" text-anchor="middle" fill="white" font-family="monospace" font-size="10">Performer: O(nmd)</text>
<text x="350" y="67" text-anchor="middle" fill="white" font-family="monospace" font-size="8">n=4096, m=256, d=128: 1.3×10⁸ FLOPs</text>
<text x="550" y="45" fill="#22c55e" font-family="monospace" font-size="11" font-weight="bold">~16× faster!</text>
</g>
</svg>
6. Linear Attention Formulation
6.1 Core Linear Attention
6.2 Causal Linear Attention
For autoregressive models, maintain a running state:
This gives per-token complexity.
6.3 Linear Attention Recurrence
7. Variants of Linear Attention
7.1 Random Feature Attention (RFA)
Uses random Fourier features for the RBF kernel:
where and .
7.2 cosFormer (Qin et al., 2022)
Adds position-dependent re-weighting:
where is a cosine-based re-weighting function.
7.3 RWKV (Peng et al., 2023)
Combines linear attention with RNN-style recurrence:
8. Attention Complexity O(n)
8.1 When is Linear Attention Effective?
Linear attention is most effective when:
- (hidden dim much smaller than sequence length)
- Long-range dependencies are important
- Memory is the bottleneck
8.2 Trade-offs
| Aspect | Standard | Linear |
|---|---|---|
| Quality | Higher | Lower |
| Speed (short seq) | Faster | Slower |
| Speed (long seq) | Slower | Faster |
| Memory | ||
| Hardware efficiency | Better | Worse |
8.3 Hybrid Approaches
Combine linear and standard attention:
9. Training Linear Attention Models
9.1 Initialization
Linear attention models require careful initialization:
with chosen to match the kernel bandwidth.
9.2 Kernel Bandwidth
The RBF kernel bandwidth controls the trade-off:
- Small : Sharp attention, more local
- Large : Smooth attention, more global
9.3 Stable Training
Use numerically stable softmax approximation:
10. Practical Considerations
10.1 Hardware Efficiency
Linear attention is often slower in practice due to:
- Poor GPU utilization for matrix multiplication
- Higher memory bandwidth requirements
- Less optimized CUDA kernels
10.2 When to Use Linear Attention
| Scenario | Recommendation |
|---|---|
| Standard attention | |
| Consider linear attention | |
| Memory limited | Linear attention |
| Quality critical | Standard attention |
10.3 Future Directions
- Better kernel functions: Learnable feature maps
- Hybrid architectures: Combine linear and standard attention
- Hardware-aware design: Optimize for specific accelerators
- Training stability: Better initialization and normalization
Linear attention provides a promising path to scaling Transformers to very long sequences, though challenges remain in matching the quality of standard attention while maintaining computational efficiency.