# MultiMax: Sparse and Multi-Modal Attention Learning

arXiv: https://arxiv.org/abs/2406.01189 · HTML: https://mlsanity.com/paper/2406.01189/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 45 of its 63 references that are in the map.
- First posted on arXiv 3 Jun 2024.

## Builds on (most cited first, 30 of 45)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization](https://mlsanity.com/paper/1610.02391/) (2016)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [A Call for Clarity in Reporting BLEU Scores](https://mlsanity.com/paper/1804.08771/) (2018)
- [Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)](https://mlsanity.com/paper/1511.07289/) (2015)
- [Deep Networks with Stochastic Depth](https://mlsanity.com/paper/1603.09382/) (2016)
- [Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning](https://mlsanity.com/paper/1702.03118/) (2017)
- [Large Language Models Can Be Easily Distracted by Irrelevant Context](https://mlsanity.com/paper/2302.00093/) (2023)
- [Quantifying Attention Flow in Transformers](https://mlsanity.com/paper/2005.00928/) (2020)
- [Going deeper with Image Transformers](https://mlsanity.com/paper/2103.17239/) (2021)
- [DropEdge: Towards Deep Graph Convolutional Networks on Node Classification](https://mlsanity.com/paper/1907.10903/) (2019)
- [Adversarial Examples for Evaluating Reading Comprehension Systems](https://mlsanity.com/paper/1707.07328/) (2017)
- [Measuring and Relieving the Over-smoothing Problem for Graph Neural Networks from the Topological View](https://mlsanity.com/paper/1909.03211/) (2019)
- [Conditional Positional Encodings for Vision Transformers](https://mlsanity.com/paper/2102.10882/) (2021)
- [From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification](https://mlsanity.com/paper/1602.02068/) (2016)
- [Sparse Sequence-to-Sequence Models](https://mlsanity.com/paper/1905.05702/) (2019)
- [On the Properties of the Softmax Function with Application in Game Theory and Reinforcement Learning](https://mlsanity.com/paper/1704.00805/) (2017)
- [Anti-Oversmoothing in Deep Vision Transformers via the Fourier Domain Analysis: From Theory to Practice](https://mlsanity.com/paper/2203.05962/) (2022)
- [Structured Attention Networks](https://mlsanity.com/paper/1702.00887/) (2017)
- [System 2 Attention (is something you might need too)](https://mlsanity.com/paper/2311.11829/) (2023)
- [Bayesian Graph Neural Networks with Adaptive Connection Sampling](https://mlsanity.com/paper/2006.04064/) (2020)
- [Revisiting Over-smoothing in BERT from the Perspective of Graph](https://mlsanity.com/paper/2202.08625/) (2022)
- [KVT: k-NN Attention for Boosting Vision Transformers](https://mlsanity.com/paper/2106.00515/) (2021)
- [Understanding Anomaly Detection with Deep Invertible Networks through Hierarchies of Distributions and Features](https://mlsanity.com/paper/2006.10848/) (2020)

## Built on by (most cited first, 5 of 5)

- [Distributed LLMs and Multimodal Large Language Models: A Survey on Advances, Challenges, and Future Directions](https://mlsanity.com/paper/2503.16585/) (2025)
- [MOOSComp: Improving Lightweight Long-Context Compressor via Mitigating Over-Smoothing and Incorporating Outlier Scores](https://mlsanity.com/paper/2504.16786/) (2025)
- [SeedPolicy: Horizon Scaling via Self-Evolving Diffusion Policy for Robot Manipulation](https://mlsanity.com/paper/2603.05117/) (2026)
- [Accurate and Efficient Multi-Channel Time Series Forecasting via Sparse Attention Mechanism](https://mlsanity.com/paper/2603.18712/) (2026)
- [Sparse Attention for Dense Open-Vocabulary Prediction in CLIP](https://mlsanity.com/paper/2607.07135/) (2026)

## Abstract (arXiv)

SoftMax is a ubiquitous ingredient of modern machine learning algorithms. It maps an input vector onto a probability simplex and reweights the input by concentrating the probability mass at large entries. Yet, as a smooth approximation to the Argmax function, a significant amount of probability mass is distributed to other, residual entries, leading to poor interpretability and noise. Although sparsity can be achieved by a family of SoftMax variants, they often require an alternative loss function and do not preserve multi-modality. We show that this trade-off between multi-modality and sparsity limits the expressivity of SoftMax as well as its variants. We provide a solution to this tension between objectives by proposing a piece-wise differentiable function, termed MultiMax, which adaptively modulates the output distribution according to input entry range. Through comprehensive analysis and evaluation, we show that MultiMax successfully produces a distribution that supresses irrelevant entries while preserving multimodality, with benefits in image classification, language modeling and machine translation. The code is available at https://github.com/ZhouYuxuanYX/MultiMax.