Research map / Architectures and efficiency / Attention and transformers
MultiMax: Sparse and Multi-Modal Attention Learning
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 MultiMax, a piecewise-differentiable, learnable replacement for SoftMax that modulates small and large input entries with separate temperatures to get both sparsity and multi-modality, improving ImageNet, WikiText-103 and IWSLT results.
Problem
SoftMax spreads probability mass onto irrelevant entries (poor sparsity), while existing sparse SoftMax variants need alternative losses, lose multi-modality, and have unstable training; a sparsity-multi-modality trade-off limits SoftMax expressivity.
Approach
Apply learnable piecewise (first- and second-order) modulation to the SoftMax input, using a low temperature for small entries (sparsity) and a high temperature for large entries (multi-modality), keeping full support and standard cross-entropy training.
Key results
- On ImageNet1K, DeiT-small with MultiMax in attention and output reaches 81.0% vs 80.4% SoftMax; DeiT-base reaches 82.6% vs 82.1%.
- On WikiText-103 with a 6-layer decoder, MultiMax in attention and output lowers test perplexity from 29.4 to 28.7.
- On IWSLT14 De-En, MultiMax scores 34.7 BLEU, on par with EvSoftMax (34.7) and above SoftMax (34.4); it adds ~40% training time in native PyTorch, ~10% with jit fusion.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace SoftMax with MultiMax in attention and output layers | DeiT-tiny · ImageNet-1k | accuracy (%): 72.8 -> 73.4 | Helped |
| Replace SoftMax with MultiMax in attention and output layers | DeiT-small · ImageNet-1k | accuracy (%): 80.4 -> 81.0 | Helped |
| Use Ev-SoftMax in attention instead of SoftMax | DeiT-small · ImageNet-1k | accuracy (%): 80.4 -> 80.0 | Hurt |
| Replace SoftMax with MultiMax in attention and output layers | DeiT-base · ImageNet-1k | accuracy (%): 82.1 -> 82.6 | Helped |
| Apply MultiMax in attention and output layers rather than Patch Diversification | DeiT-small · ImageNet-1k | accuracy (%): 81.2* vs 81.2 | No clear effect |
| Apply MultiMax in attention and output layers rather than AttnScale | DeiT-small · ImageNet-1k | accuracy (%): 80.9* -> 81.3 | Helped |
| Add MultiMax in attention only | 6-layer Transformer Decoder, 156M parameters · WikiText-103 test | perplexity: 29.4 -> 29.0 | Helped |
| Add MultiMax to the output layer in addition to attention | 6-layer Transformer Decoder, 156M parameters · WikiText-103 test | perplexity: 29.0 -> 28.7 | Helped |
| Replace SoftMax with MultiMax in attention | Transformer · IWSLT 2014 English-to-German | BLEU: 34.4 ± 0.07 -> 34.7 ± 0.07 | Helped |
| Add term (1), modulating smaller entries (ablation) | MultiMax on DeiT-small · ImageNet-1k | accuracy (%): 80.4 -> 80.6 | Helped |
| Add term (2), modulating larger entries, in addition to term (1) (ablation) | MultiMax on DeiT-small · ImageNet-1k | accuracy (%): 80.6 -> 80.7 | Helped |
| Add second-order terms to terms (1) and (2) (ablation) | MultiMax on DeiT-small · ImageNet-1k | accuracy (%): 80.7 -> 81.0 | Helped |
Surprising findings
- Ev-SoftMax training did not converge (accuracy below 10%) when used in the output layer on ImageNet.
- Learned MultiMax parameters show the first two attention layers are less sparse than SoftMax, contrary to deeper layers.
Limitations
- Native PyTorch implementation adds notable training overhead unless fused; no custom CUDA kernel.
- Several comparisons on ImageNet are single-run, and some baselines are not strictly comparable in training setup.
- Sparse SoftMax baselines were excluded from ImageNet classification due to training instability, limiting comparison.
Experimental setup
Datasets: ImageNet1K WikiText-103 IWSLT-2014 German-English
Benchmarks: ImageNet1K WikiText-103 IWSLT-2014 En-De
Metrics: top-1 accuracy perplexity BLEU patch similarity sparsity multi-modality
Baselines: SoftMax Top-k attention Ev-SoftMax SparseMax EntMax-1.5 Patch Diversification AttnScale
Builds on: SoftMax Sparsemax EntMax Ev-SoftMax DeiT
Closest papers by meaning
- Bridging the Divide: Reconsidering Softmax and Linear Attention (NeurIPS 2024 · similarity 0.88)
- Sparse Attention as Compact Kernel Regression (NeurIPS 2026 · similarity 0.87)
- Low-Rank Approximation for Sparse Attention in Multi-Modal LLMs (CVPR 2024 · similarity 0.87)
- Degrees of Freedom for Linear Attention: Distilling Softmax Attention with Optimal Feature Efficiency (NeurIPS 2025 · similarity 0.87)
- AdaSplash: Adaptive Sparse Flash Attention (ICML 2025 · similarity 0.87)
- Transformers Learn Faster with Semantic Focus (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.