Research map / Architectures and efficiency / Attention and transformers
Customizing the Inductive Biases of Softmax Attention using Structured Matrices
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.
Replaces the low-rank attention scoring bilinear form with high-rank structured matrices (Bilinear BTT, Bilinear MLR) and uses contiguous MLR to impose distance-dependent compute, improving in-context regression, OpenWebText language modeling, and time-series forecasting.
Problem
Standard softmax attention's low-rank scoring function creates an information bottleneck for intrinsically high-dimensional inputs, and its single scoring function for all token pairs lacks a distance-dependent compute bias for local structure.
Approach
Replace the low-rank query-key bilinear form with structured matrices (BTT and contiguous Multi-Level Low Rank) that are high-rank yet FLOP-efficient, and use hierarchical MLR levels so that nearby token pairs get more compute per interaction than distant ones; also generalize both families into MLBTC.
Key results
- On in-context linear regression, Bilinear BTT with 8 heads learns the task at model width D=256, while standard 8-head attention fails even at D=512 and 1-head attention needs D=512.
- On OpenWebText with 6-layer transformers, 8-level MLR attention achieves lower validation loss than standard attention across widths D in {256, 384, 512, 768} when controlling for compute.
- On ETTh1 forecasting, MLR attention outperforms standard attention in MAE by around 1% as the time horizon grows.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace low-rank attention scoring with Bilinear MLR | Transformer · In-context linear regression | regression error | Helped |
| Replace low-rank attention scoring with Bilinear BTT | Transformer · In-context linear regression | regression error | Helped |
| Use 8-level MLR attention instead of standard attention | Transformer · OpenWebText | validation loss | Helped |
| Use MLR attention instead of sliding window attention | Transformer · OpenWebText | validation loss | Helped |
| Use MLR attention instead of Global + SWA | Transformer · OpenWebText | validation loss | Helped |
| Replace standard attention with MLR attention | Chronos · Time-series validation | validation loss | Helped |
| Use MLR attention with rank distribution 48|16 | Transformer · ETTh1 oil temperature forecasting | MAE | Helped |
| Use MLR attention with rank distribution 40|16|4|4 | Transformer · ETTh1 oil temperature forecasting | MAE | Helped |
| Increase the time horizon | MLR attention on ETTh1 · ETTh1 oil temperature forecasting | MAE | Helped |
Surprising findings
- Standard 8-head attention fails to learn the in-context regression task even at model width D=512, while full-rank 1-head attention succeeds.
Limitations
- Implementations are not optimized, making wall-clock time somewhat slower than standard attention.
- MLBTC full-rank scoring function is left to future work.
- Language modeling experiments use small 6-layer models with character-level tokenization and sequence length 1024.
Experimental setup
Datasets: OpenWebText ETT (ETTh1) Chronos data mixture (synthetic kernel data and real data) synthetic in-context linear regression data
Benchmarks: in-context linear regression (synthetic) OpenWebText validation loss ETTh1 oil temperature forecasting
Metrics: regression squared error validation loss mean absolute error (MAE) FLOPs
Baselines: standard multi-head attention sliding window attention Global + SWA Chronos with standard attention
Builds on: Monarch Butterfly matrices Multi-Level Low Rank matrices maximal update parameterization Chronos
Closest papers by meaning
- SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking (NeurIPS 2026 · similarity 0.87)
- Mind the Gap: a Spectral Analysis of Rank Collapse and Signal Propagation in Attention Layers (ICML 2025 · similarity 0.87)
- In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention (ICML 2025 · similarity 0.87)
- Positional LSH: Binary Block Matrix Approximation for Attention with Linear Biases (NeurIPS 2026 · similarity 0.86)
- MVA: Linear Attention with High-order Query-Keys Integration and Multi-level Vocabulary Decomposition (ICML 2025 · similarity 0.86)
- Efficient Low Rank Attention for Long-Context Inference in Large Language Models (NeurIPS 2025 · similarity 0.86)
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.