Research map / Architectures and efficiency / Attention and transformers
Improving Transformers with Dynamically Composable Multi-Head Attention
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 DCMHA, which dynamically composes attention score and weight matrices across heads with input-dependent low-rank and gating maps, yielding DCFormer models that outperform Transformer baselines on Pile language modeling and downstream tasks.
Problem
Standard multi-head attention has independent heads, causing low-rank bottleneck of attention score matrices and head redundancy; prior head-collaboration methods are either static, surface-level, or computationally costly.
Approach
A Compose function that linearly mixes attention score (pre-softmax) and weight (post-softmax) matrices across heads, with input-dependent low-rank-plus-diagonal maps computed from queries and keys via a row-plus-column, low-rank-plus-diagonal tensor decomposition, keeping parameter and FLOP overhead small.
Key results
- DCPythia-6.9B achieves 5.95 Pile ppl vs 6.29 for Pythia-6.9B and 6.01 for Pythia-12B, and 56.7 vs 56.5 average zero-shot accuracy against Pythia-12B.
- DCFormer-834M matches Transformer loss at 1.87x compute and DCFormer++-834M matches Transformer++ at 1.67x compute.
- Training throughput of DCFormer++ is 74.5% to 89.2% of Transformer++ across 2.8B to 33B; inference is 87.9% to 95.1% of Transformer++ at the same window.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace MHA with DCMHA | Pythia-2.8B · Pile validation | perplexity: 6.63 -> 6.36 | Helped |
| replace MHA with DCMHA | Pythia-6.9B · Zero-shot downstream NLP tasks | average accuracy: 55.1 -> 56.7 | Helped |
| replace MHA with DCMHA | Pythia-6.9B · Synthetic dataset | accuracy: 31.9% -> 39.0% | Helped |
| replace MHA with DCMHA | ViT-S/16 · ImageNet-1k, Epoch 300 | Top-1 accuracy: 79.8 -> 80.1 | Helped |
| add dynamic projection and gating instead of static projection (ablation) | Transformer++ 405M · language modeling | perplexity: 11.17 -> 10.83 | Helped |
| add dynamic projection (ablation) | Transformer++ 405M · language modeling | perplexity: 11.68 -> 10.95 | Helped |
| add dynamic gating (ablation) | Transformer++ 405M · language modeling | perplexity: 11.68 -> 11.31 | Helped |
| add query-wise composition (ablation) | Transformer++ 405M · language modeling | perplexity: 11.68 -> 10.89 | Helped |
| add key-wise composition (ablation) | Transformer++ 405M · language modeling | perplexity: 11.68 -> 10.91 | Helped |
| add pre-composition on attention scores (ablation) | Transformer++ 405M · language modeling | perplexity: 11.68 -> 11.54 | Helped |
| add post-composition on attention weights (ablation) | Transformer++ 405M · language modeling | perplexity: 11.68 -> 11.05 | Helped |
| increase dynamic projection rank from 1 to 2 (ablation) | DCFormer++ 405M · language modeling | perplexity: 10.87 -> 10.83 | Helped |
Surprising findings
- Adapting pretrained Transformer checkpoints to DCFormer via finetuning gave no significant improvement, suggesting the architecture is not easily retrofitted.
- Post-compose on attention weights alone is much more effective than pre-compose on attention scores in ablations.
Limitations
- Training and inference overhead is larger than theoretical FLOP estimates, due to I/O-bound operations without custom kernels.
- Existing pretrained Transformers are difficult to adapt to DCFormer via continued pretraining.
Experimental setup
Datasets: Pile Flan Collection C4 ImageNet-1K synthetic entity-attribute reasoning dataset (888 examples, 74 tasks)
Benchmarks: Pile validation perplexity Flan perplexity LAMBADA PIQA WinoGrande ARC-Easy ARC-Challenge SciQ LogiQA BoolQ HellaSwag RACE
Metrics: perplexity accuracy top-1 accuracy top-5 accuracy training throughput inference speed
Baselines: Transformer Transformer++ Pythia-2.8B Pythia-6.9B Pythia-12B ViT-S/16 ViT-M/16 Talking-Heads Attention
Builds on: Talking-Heads Attention Transformer++ (RoPE + SwiGLU) Pythia QKNorm
Closest papers by meaning
- Memorization Capacity of Multi-Head Attention in Transformers (ICLR 2024 · similarity 0.87)
- MoH: Multi-Head Attention as Mixture-of-Head Attention (ICML 2025 · similarity 0.86)
- SAS: Simulated Attention Score (NeurIPS 2025 · similarity 0.86)
- Data-Adaptive Mahalanobis Metric Learning for Cross-Head Attention in Transformers (NeurIPS 2026 · similarity 0.86)
- SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention (NeurIPS 2024 · similarity 0.86)
- No Head Left Behind – Multi-Head Alignment Distillation for Transformers (AAAI 2024 · similarity 0.85)
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.