Research map / Architectures and efficiency / Attention and transformers
Wavy Transformer
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 attention-layer dynamics equal graph neural diffusion on a complete graph and proposes Wavy Transformer, a second-order wave-equation attention residual (plus velocity-aware LN/FFN) that mitigates over-smoothing across NLP, CV, and sparse-graph tasks.
Problem
Deep transformers suffer from over-smoothing (token representations converge across layers); the paper asks whether this can be understood via dissipative diffusion dynamics and mitigated with an intrinsic wave-based mechanism.
Approach
Interpret stacked attention residual updates as graph diffusion on a complete graph, then replace the first-order diffusive residual with a second-order wave (momentum) update, adding velocity-specific LayerNorm and FFN so the state-velocity relation is preserved under the chain rule.
Key results
- On GLUE, Mix (+Light Wave) reaches average 66.12 vs 64.13 for diffusion baseline (+1.99), with STS-B gain of +12.65.
- On ImageNet, DeiT-Ti with Diffusion + Light Wave reaches 73.09% top-1 vs 72.17% baseline (+0.92); CaiT-XXS-24 improves from 77.6 to 78.6 with Full Wave.
- On OGBN-Arxiv with 7 layers, DIFFormer diffusion collapses to 24.44% while + Light Wave reaches 66.73%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Mix diffusion with Full Wave residual | BERT-like pretraining model · BERT validation | PPL: 31.76 -> 29.00 | Helped |
| Use Full Wave residual instead of diffusion | BERT-like pretraining model · BERT validation | PPL: 31.76 -> 31.99 | Hurt |
| Mix diffusion with Light Wave residual | BERT-like model · GLUE average | score: 64.13 -> 66.12 (+1.99) | Helped |
| Mix diffusion with Light Wave residual | BERT-like model · GLUE STS-B | score: 52.11 -> 64.76 | Helped |
| Add Light Wave residual | DeiT-Ti · ImageNet-1K val | Top-1 accuracy: +0.92 | Helped |
| Add Full Wave residual | CaiT-XXS-24 · ImageNet-1K | Top-1 accuracy: +1.0 | Helped |
| Add Light Wave residual | DIFFormer · OGBN-Arxiv, 7 layers | accuracy: 24.44 ± 4.51 -> 66.73 ± 0.33 (+42.29) | Helped |
| Add Light Wave residual | DIFFormer · OGBN-Proteins, 5 layers | ROC-AUC: 69.42 ± 2.31 -> 80.14 ± 0.67 (+10.72) | Helped |
| Add Light Wave residual | DIFFormer · Cora, 20 layers, τ=0.2 | accuracy: 39.92 ± 5.14 -> 85.18 ± 0.48 (+45.26) | Helped |
| Remove velocity FFN (ablation) | Full Wave DeiT-Tiny · ImageNet-1K val, 150 epochs | Top-1 accuracy: 68.92 -> 68.39 (--0.53) | Hurt |
| Replace Diffusion with Light Wave residual | DIFFormer · OGBN-Arxiv, 3 layers | accuracy: 68.77 ± 0.28 -> 66.68 ± 0.66 (-2.09) | Hurt |
| Add Light Wave residual | DIFFormer · OGBN-Arxiv, 5 layers | accuracy: 61.72 ± 1.11 -> 67.06 ± 0.03 (+5.34) | Helped |
Surprising findings
- Wave residual alone sometimes underperforms the diffusion baseline on GLUE (e.g., Full Wave STS-B drops to 32.91 vs 52.11), and Light Wave regresses at depth 3 on Arxiv (-2.09).
- Full Wave roughly halves throughput (about 2600 to 1300 img/s on DeiT-Ti) and doubles training memory.
Limitations
- Full Wave variant incurs substantial inference slowdown and memory overhead.
- Causal-attention (GPT-2 fine-tuning) gains are negligible (about 0.06%) and not conclusively assessed.
- Some results show regressions and high variance (e.g., Citeseer depth 8, Cora at tau=0.5 depth 12).
Experimental setup
Datasets: Wikipedia BookCorpus GLUE SQuAD ImageNet-1K OGBN-Arxiv OGBN-Proteins Cora Citeseer PubMed OpenWebText
Benchmarks: GLUE ImageNet-1K OGBN-Arxiv OGBN-Proteins Cora Citeseer PubMed SQuAD
Metrics: perplexity MLM accuracy GLUE score top-1 accuracy ROC-AUC cosine similarity spectral gap throughput
Baselines: Diffusion (standard residual) DeiT-Ti DeiT-Ti + FeatScale CaiT-XXS-24 DIFFormer BERT
Builds on: graph neural diffusion DIFFormer DeiT BERT FeatScale
Closest papers by meaning
- Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape (ICML 2024 · similarity 0.88)
- Rare Text Semantics Were Always There in Your Diffusion Transformer (NeurIPS 2025 · similarity 0.87)
- Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis (NeurIPS 2024 · similarity 0.87)
- Graph Convolutions Enrich the Self-Attention in Transformers! (NeurIPS 2024 · similarity 0.87)
- Viewing Transformers Through the Lens of Long Convolutions Layers (ICML 2024 · similarity 0.86)
- Transformers Get Stable: An End-to-End Signal Propagation Theory for Language Models (ICML 2024 · 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.