ICML 24 · 2024

MultiMax - Sparse and Multi-Modal Attention Learning

Y. Zhou, Mario Fritz, Margret Keuper

Why this publication matters

An attention mechanism needs to ignore irrelevant information without discarding several useful possibilities. MultiMax changes how attention weights are assigned to support both goals. Its evaluations show why this small mathematical component can matter across image and language tasks.

Abstract

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.

Abstract source ↗

Figures

Attention distributions illustrating sparsity and preservation of multiple modes. — **Figure 1.** We evaluate SoftMax, SparseMax, EntMax, EvSoftMax and MultiMax (using the parameters of a hidden layer MultiMax trained on ImageNet directly) functions on a series of example input points v ∈R3 and project the resulting distribution on a simplex ∆2. Informally, the interior of the simplex stands for trimodal distributions, the edges constitute the set of bimodal distributions, and the vertices are unimodal distributions. Notably, the above figures highlight the advantage of MultiMax’s multi-modality. EntMax, Sparsemax and SoftMax with small temperature (blue colored line) yield a (quasi) uni-modal distribution, which ignore the second largest entry. In contrary, SoftMax with higher temperatures (green and orange colored line) fails to ignore the negative entry. View in source ↗

**Figure 2.** Illustration of different reweighting functions in the twodimensional case. It can be seen clearly that MultiMax weigh the entries at small and large value ranges in a different manner, thus it does not suffer from the trade-off between sparse and multi-modal. View in source ↗

Cite this paper

Download .bib

@inproceedings{zhou2024multimaxsparseand49,
  title = {{MultiMax - Sparse and Multi-Modal Attention Learning}},
  author = {Y. Zhou and Mario Fritz and Margret Keuper},
  booktitle = {International Conference on Machine Learning},
  year = {2024},
  url = {https://openreview.net/forum?id=IC9UZ8lm25}
}

Figures and abstract are reproduced from the linked research sources. Credit remains with the authors and publishers.