Research map / Architectures and efficiency / Attention and transformers
Attention-Only Transformers via Unrolled Subspace Denoising
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 an attention-only transformer (AoT) obtained by unrolling multi-head subspace self-attention denoising under a mixture-of-low-rank-Gaussians token model, proves linear SNR improvement per layer, and shows near-baseline results on ImageNet and language/ICL tasks.
Problem
Standard transformers are empirically designed and hard to interpret, and some components (MLPs, value matrices, normalization) may be redundant; the paper asks whether a minimal, interpretable transformer with only attention and skip connections can match standard performance.
Approach
Model token representations as a mixture of noisy low-rank Gaussians, interpret multi-head subspace self-attention as a denoising step toward those subspaces, and unroll it into layers with only attention and skip connections (no MLP), with the projection matrices learned by backprop; each layer provably raises SNR at a linear rate.
Key results
- AoT-MSSA-V reaches 71.7% ImageNet top-1 with 22M parameters versus CRATE 79.5% at 39M; AoT-MHSA-V reaches 69.5% with 15M versus ViT 72.4% at 22M.
- AoT-MHSA-L Base (122M) gets LAMBADA zero-shot accuracy 0.38 and WikiText val loss 4.19, close to GPT-2 Base (0.40 and 4.13).
- Theorem shows each MSSA layer multiplies SNR by (1+eta*tau), i.e., linear SNR growth with depth under the noisy low-rank Gaussian model.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use attention-only AoT-MSSA-V instead of CRATE | ImageNet-1K classification · ImageNet-1K | Top-1 accuracy: 79.5\% -> 71.7\% | Hurt |
| Use attention-only AoT-MHSA-V instead of ViT | ImageNet-1K classification · ImageNet-1K | Top-1 accuracy: 72.4 \% -> 69.5\% | Hurt |
| Use AoT-MSSA-L Medium instead of AoT-MSSA-L Base | AoT-MSSA-L · LAMBADA | validation loss: 4.70 -> 4.47 | Helped |
| Use AoT-MHSA-L Base instead of AoT-MSSA-L Base | AoT language model · LAMBADA | validation loss: 4.70 -> 4.42 | Helped |
| Use AoT-MHSA-L Base instead of GPT-2 Base | Language modeling · PTB | validation loss: 5.75 -> 5.52 | Helped |
| Remove MLP layers from AoT (ablation) | Zero-shot language evaluation · LAMBADA | accuracy | No clear effect |
| Add MLP layers to the attention-only model (ablation) | AoT language model training · OpenWebText | training speed | Helped |
| Remove MLP layers (ablation) | AoT in-context learning model · Linear and sparse linear regression | normalized squared error | Mixed |
Surprising findings
- Adding MLP layers to AoT did not improve zero-shot performance, suggesting MLPs contribute only marginal gains.
- AoT-MSSA-L Base (102M) reports a LAMBADA validation loss of 4.70 and accuracy 0.25, underperforming MHSA variants.
Limitations
- Accuracy trails CRATE and ViT on ImageNet, and the paper states it does not aim for state of the art.
- Theory assumes orthogonal subspaces and a idealized Gaussian mixture; the analysis holds for the forward denoising step with known bases, not the learned backprop-trained model.
- Experiments are limited by compute, with ICL and language experiments on small models.
Experimental setup
Datasets: ImageNet-21K ImageNet-1K OpenWebText WikiText LAMBADA PTB CBT synthetic linear and sparse linear regression prompts
Benchmarks: ImageNet-1K WikiText LAMBADA PTB CBT
Metrics: top-1 accuracy validation loss zero-shot accuracy normalized squared error SNR
Baselines: CRATE ViT GPT-2 Base GPT-2 (ICL)
Builds on: CRATE sparse rate reduction GPT-2 nanoGPT
Closest papers by meaning
- Token Statistics Transformer: Linear-Time Attention via Variational Rate Reduction (ICLR 2025 · similarity 0.88)
- Xformer: Hybrid X-Shaped Transformer for Image Denoising (ICLR 2024 · similarity 0.87)
- Skip-Attention: Improving Vision Transformers by Paying Less Attention (ICLR 2024 · similarity 0.87)
- Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators? (ICLR 2024 · similarity 0.87)
- Transformer Meets Twicing: Harnessing Unattended Residual Information (ICLR 2025 · similarity 0.87)
- Rethinking Transformer-Based Blind-Spot Network for Self-Supervised Image Denoising (AAAI 2025 · 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.