Research map / Architectures and efficiency / Attention and transformers
Token Statistics Transformer: Linear-Time Attention via Variational Rate Reduction
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.
Proposes Token Statistics Self-Attention (TSSA), a linear-complexity attention operator derived from unrolling a variational form of the MCR² objective, and the Token Statistics Transformer (ToST) built by swapping it into XCiT, evaluated on ImageNet transfer, Long-Range Arena, and causal language modeling.
Problem
Standard self-attention has computation and memory that scale quadratically with the number of tokens, which limits scaling to long sequences and high-dimensional token counts.
Approach
Derives a variational upper bound on MCR² compression via concave spectral functions, then unrolls one gradient step to get an attention operator that computes token-level second-moment statistics of low-rank projections instead of pairwise token similarities, giving linear time and memory.
Key results
- On Long-Range Arena, ToST averages 59.90 versus 59.24 for the standard Transformer, with the best scores on Text, Retrieval, and Image.
- ToST-M reaches 80.3% ImageNet top-1 versus 80.5% for XCiT-S and 81.8% for ViT-B, with fewer parameters than XCiT-M and ViT-B.
- Causal ToST-Base achieves 0.60x time and 0.41x memory of GPT2-Base at 4k length, and 0.46x time and 0.24x memory at 8k length.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace self-attention with TSSA | ToST versus transformer baselines · LRA average | score: 59.24 -> 59.90 | Helped |
| Replace self-attention with TSSA | ToST versus transformer baselines · LRA Pathfinder | score: 77.05 -> 69.41 | Hurt |
| Replace self-attention with TSSA | ToST versus transformer baselines · LRA Image | score: 42.94 -> 46.62 | Helped |
| Replace self-attention with TSSA | ToST-M versus ViT-B · ImageNet-1k | top-1 accuracy: 81.8 -> 80.3 | Hurt |
| Replace self-attention with TSSA | ToST-M versus ViT-B · CIFAR100 | accuracy: 89.3 -> 84.5 | Hurt |
| Scale ToST from Base to Large | ToST causal language model · Causal language modeling average | cross entropy loss: 4.84 -> 4.02 | Helped |
| Use ToST-Base instead of GPT2-Base | Causal language model · Causal language modeling average | cross entropy loss: 4.26 -> 4.84 | Hurt |
| Replace MLP with ISTA (ablation) | ToST-S with L2 normalization · CIFAR10, 400 epochs | accuracy: 91.2 -> 87.4 | Hurt |
| Add L2 normalization when computing membership matrix (ablation) | ToST-S with MLP · CIFAR10, 400 epochs | accuracy: 85.2 -> 91.2 | Helped |
| Add L2 normalization when computing membership matrix (ablation) | ToST-S with ISTA · CIFAR10, 400 epochs | accuracy: 80.6 -> 87.4 | Helped |
| Replace MLP with ISTA (ablation) | ToST-S without normalization · CIFAR10, 400 epochs | accuracy: 85.2 -> 80.6 | Hurt |
Surprising findings
- Swapping self-attention for TSSA largely preserves language modeling performance, which the authors say calls into question whether pairwise similarity attention is critical to transformer success.
- ToST-Base has higher zero-shot loss than GPT2-Base on OpenWebText (3.20 vs 2.84), despite comparable parameter count.
Limitations
- Hyperparameters were not systematically tuned for causal language modeling, using GPT2-Base settings for all sizes.
- Does not outperform highly engineered architectures on ImageNet; gap to ViT-B and XCiT-M remains on several transfer tasks.
- Removes the LPI module and replaces CRATE's ISTA block with an MLP, departing from the pure white-box derivation.
Experimental setup
Datasets: ImageNet-1k CIFAR10 CIFAR100 Oxford Flowers-102 Oxford-IIIT-Pets Long-Range Arena OpenWebText Lambada Wikitext Penn Treebank
Benchmarks: ImageNet-1k ImageNet ReaL CIFAR10 CIFAR100 Oxford Flowers-102 Oxford-IIIT-Pets Long-Range Arena OpenWebText Lambada Wikitext PTB
Metrics: top-1 accuracy test accuracy zero-shot cross entropy loss time and memory relative usage
Baselines: ViT XCiT CRATE Reformer BigBird LinFormer Performer S4 GPT2-Base
Builds on: CRATE XCiT maximal coding rate reduction GPT-2
Closest papers by meaning
- Gated Linear Attention Transformers with Hardware-Efficient Training (ICML 2024 · similarity 0.88)
- Neural Attention Search (NeurIPS 2025 · similarity 0.88)
- Attention-Only Transformers via Unrolled Subspace Denoising (ICML 2025 · similarity 0.88)
- Frequency-Aware Token Reduction for Efficient Vision Transformer (NeurIPS 2025 · similarity 0.88)
- Theory, Analysis, and Best Practices for Sigmoid Self-Attention (ICLR 2025 · similarity 0.88)
- Delta Attention: Fast and Accurate Sparse Attention Inference by Delta Correction (NeurIPS 2025 · similarity 0.88)
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.