Research map / Learning theory and algorithms / Sample complexity
I/O Complexity of Attention, or How Optimal is FlashAttention?
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proves matching I/O complexity lower bounds for exact attention (including with fast matrix multiplication, via a new communication-complexity matrix compression argument), showing FlashAttention is optimal for cache size M ≥ d^2.
Problem
Whether FlashAttention's I/O complexity O(N^2 d^2 / M) is optimal for attention on a two-level memory hierarchy, for all cache sizes M, including when fast matrix multiplication is used.
Approach
Uses red-blue pebble game partition arguments for standard matrix multiplication, and for arbitrary (including fast matrix multiplication) algorithms reduces I/O to a one-way communication complexity matrix-entry compression problem, using Vandermonde matrices (large field) and binary BCH codes (binary inputs) to get linear-independence constraints; also gives an algorithm with O(N^2 d / sqrt(M))
Key results
- For standard matrix multiplication, attention I/O complexity is Theta(min(N^2 d/sqrt(M), N^2 d^2/M)), matching FlashAttention for M >= d^2.
- With fast matrix multiplication and field inputs (q > N), I/O complexity is Omega(min(N^2 d^2/M, N^2)), so FlashAttention is optimal for M >= d^2 even with FMM.
- For binary inputs the lower bound Omega(min(N^2 d^2/(M log^2 N), N^2)) is tight up to log^2 N factors.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Establish a matching lower bound for attention in the large-cache regime, including when fast matrix multiplication is allowed | FlashAttention · Attention, large cache (M ≥ d²), q > N | I/O complexity: \bigO{\frac{N^2 d^2}{M}}; \bigOmega{\min\left(\frac{N^2 d^2}{M}, N^2 \right)} | No clear effect |
| Use square tiling instead of FlashAttention | Attention, small cache (M < d²) · Attention, small cache | I/O complexity: \bigO{\frac{N^2 d}{\sqrt{M}}} vs. \bigO{\frac{N^2 d^2}{M}} | Helped |
| Restrict Q and K to binary matrices | Attention with any matrix multiplication algorithm · Attention, binary inputs | I/O complexity lower bound: \bigOmega{\min\left(\frac{N^2 d^2}{M \log^2 N}, N^2 \right)} | Mixed |
Surprising findings
- Attention for M < d^2 is equivalent in I/O cost to rectangular matrix multiplication, so the attention bound is no easier than matmul in that regime.
Limitations
- Tight bounds for small cache M < d^2 under fast matrix multiplication remain open.
- Field-based lower bounds assume entries in a finite field F_q with q > N; binary case is tight only up to log^2 N factors.
- Analysis is in an idealized two-level memory model without wall-clock measurements.
Experimental setup
Metrics: I/O complexity (asymptotic)
Baselines: FlashAttention
Builds on: FlashAttention red-blue pebble game Hong and Kung I/O model
Closest papers by meaning
- FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (ICLR 2024 · similarity 0.87)
- FlashMask: Efficient and Rich Mask Extension of FlashAttention (ICLR 2025 · similarity 0.86)
- MonarchAttention: Zero-Shot Conversion to Fast, Hardware-Aware Structured Attention (NeurIPS 2025 · similarity 0.85)
- FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (NeurIPS 2024 · similarity 0.85)
- Token Statistics Transformer: Linear-Time Attention via Variational Rate Reduction (ICLR 2025 · similarity 0.83)
- FlashBias: Fast Computation of Attention with Bias (NeurIPS 2025 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.