Research map / Learning theory and algorithms / Kernels and regression theory
Learning Linear Attention in Polynomial Time
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.
Shows that single-layer multi-head linear attention is strongly agnostically PAC-learnable in polynomial time by relaxing to a rank-H linear regression over cubic features, and gives a data-based certificate (min eigenvalue of second-moment matrix) for when all empirical risk minimizers compute the same function.
Problem
Whether single-layer Transformers with linear attention can be learned efficiently (polynomial time and samples) from data, and whether a learned model is uniquely determined by the data so it generalizes out of distribution.
Approach
Reformulate multi-head linear attention output as a linear function of a fixed cubic feature map over a d^2 x d^2 matrix W; solve least squares in polynomial time, then recover heads via SVD (at most d^2 heads). Identifiability is certified when the feature second-moment matrix is full rank.
Key results
- Algorithm returns MHLA with at most d^2 heads in time O(N d^4 n_max / epsilon) with sample complexity O(epsilon^-1 (d^4 + log(1/delta))) for strong agnostic PAC learning.
- For d=2 or 4, SGD-trained MHLA with d^2 heads performs best and is empirically comparable to the polynomial-time algorithm, approaching floating-point precision.
- On associative memory, distance in p-feature space to ground truth decreases as min eigenvalue of the data second-moment matrix increases; zero min eigenvalue yields spurious solutions.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase the number of heads, up to d^2 heads (ablation) | Multi-head linear attention trained with SGD · Synthetic random linear attention | performance | Helped |
| Use d^2 heads instead of fewer heads (ablation) | Multi-head linear attention trained with SGD · Synthetic random linear attention | performance relative to the polynomial-time algorithm | No clear effect |
| Increase the number of layers in single-head linear attention (ablation) | Multi-layer linear attention · Synthetic random linear attention | training loss convergence | Hurt |
| Add more layers (ablation) | Multi-layer linear attention · Synthetic random linear attention | performance | Hurt |
| Add more layers (ablation) | Multi-layer linear attention · Synthetic random linear attention | performance | Mixed |
| Replace single-layer linear attention with a full Transformer using softmax attention, MLP blocks, and layer normalization (ablation) | Single-layer linear attention · Synthetic random linear attention | performance | No clear effect |
| Train on identifiable data with minimum eigenvalue 0.06 rather than unidentifiable data with minimum eigenvalue 0 | MHLA trained with SGD · Associative lookup | distance to ground truth in p feature space: 0.06 vs. 0 | Helped |
| Increase the certificate value (minimum eigenvalue of the data covariance matrix) | Polynomial-time learning algorithm · Associative lookup | distance to ground truth in p feature space | Helped |
| Use unidentifiable data with minimum eigenvalue 0 instead of identifiable data | MHLA trained with SGD · Associative lookup | distance to ground truth in p feature space: 0.06 vs. 0 | Hurt |
Surprising findings
- Adding more layers can sometimes degrade performance of linear attention models, and the full softmax Transformer does not significantly outperform single-layer linear attention on the synthetic task.
Limitations
- Model class is single-layer linear attention, far simpler than modern LLMs and softmax Transformers.
- Learned model can have up to d^2 heads, and the guarantee relies on a relaxation to full-rank W.
- Identifiability certificate requires the second-moment matrix to be full rank, which fails for some natural data (e.g., unitary key-values).
Experimental setup
Datasets: synthetic linear attention network outputs associative memory synthetic data (Gaussian and random unitary key-values)
Benchmarks: synthetic random linear attention regression associative memory lookup learning universal Turing machine (theory only)
Metrics: training loss validation loss distance in p feature space min eigenvalue of second-moment matrix
Baselines: multi-layer linear attention softmax Transformer SGD-trained multi-head linear attention
Builds on: Transformer expressivity via Turing machine simulation PAC learning VC theory
Closest papers by meaning
- Transformers Learn Faster with Semantic Focus (NeurIPS 2025 · similarity 0.89)
- PolySketchFormer: Fast Transformers via Sketching Polynomial Kernels (ICML 2024 · similarity 0.88)
- Gated Linear Attention Transformers with Hardware-Efficient Training (ICML 2024 · similarity 0.88)
- Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical Functions (NeurIPS 2025 · similarity 0.88)
- kNN Attention Demystified: A Theoretical Exploration for Scalable Transformers (ICLR 2025 · similarity 0.87)
- Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators? (ICLR 2024 · similarity 0.87)
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.