MultiMax

This is the official implementation of our ICML 2024 paper "MultiMax: Sparse and Multi-Modal Attention Learning".

Updates

[2026-07] MultiMax is now integrated into PaddleFleet, the large-scale distributed training framework by PaddlePaddle. Key highlights of the integration:

• Renamed modulation function to SegLU: To avoid naming collision with the existing SELU (Scaled Exponential Linear Unit) in PyTorch and deep learning literature, the MultiMax modulation function is renamed SegLU (Segmented Linear Unit) in PaddleFleet.
• Triton-fused SegLU kernel: SegLU is fused directly into the chunked cross-entropy kernel (via LigerKernel), eliminating the need to materialize the full [B, S, V] logits tensor and reducing peak activation memory.
• Fused LM head path: The LM head emits a 5-tuple (hidden_states, weight, bias, multimax_ranges, multimax_ts) so SegLU is applied inside the chunked CE kernel without a separate logits pass.

Best practices for integrating MultiMax into LLM pretraining

Optimizer: Use Adam (or AdamW) for MultiMax scalar parameters. Muon is designed for matrix-shaped weights (it applies Newton-Schulz orthogonalization) and is not applicable to scalar or 1-D parameters — the convention in Muon-based training is to route all non-matrix parameters to Adam.

Weight decay: Exclude MultiMax scalar parameters (multimax_ts, multimax_ranges) from the weight decay list. This follows the standard convention of not applying weight decay to bias and 1-D parameters (e.g., HuggingFace Transformers does this by default). If your training framework does not handle this automatically, add these parameters explicitly to the no-decay group in your optimizer parameter groups.

Compatibility with global gradient norm clipping: MultiMax scalar parameters are shared across the full vocabulary, so their gradients naturally accumulate over all vocabulary entries and their gradient magnitude scales as ~sqrt(V). Under global gradient norm clipping (ClipGradByGlobalNorm / clip_grad_norm_), this can cause frequent clipping that slows down backbone learning. Two solutions: (1) Normalize the MultiMax gradient norm by 1/sqrt(vocab_size) in the clip path — this brings it to a per-element scale comparable to the backbone, and since Adam is invariant to constant gradient scaling, the effective parameter updates are unchanged. (2) Place MultiMax parameters in a separate clip group with its own norm threshold, fully decoupled from the backbone gradient norm.

Improved multi-modality and sparsity

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 ∈ R^3$ 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.

Figure 2: Grad-CAM of Deit-small using SoftMax (the two images on the left) and MultiMax (the two images on the right). The MultiMax attention maps are better localized on the objects and are close to zero in most background regions, indicating sparsity at the attention level.

MultiMax is as efficient as SoftMax

As shown in Equation 4, the total amount of additional parameters for a 12 layer Transformer with 2nd-order MultiMax is just 8x12=96, because each order only contains 4 parameters, including t_b, t_d, b and d. Moreover, the modulation function σ(x) merely consists of cheap element-wise operations, i.e., multiplication with t_b and t_d, subtraction with b and d, two Max operations, addition of the two terms at each order as well as a residual addition. Thus a second-order MultiMax requires 7x2+1=15 extra Floating Point Operations (FLOPs) for a univariant input. For Deit-small model with input length of 256, hidden dimension of 384 and 12 layers, replacing MultiMax with SoftMax in all attention layers leads to 0.0168G extra FLOPs. It is only 0.37% of the original model’s 4.6G FLOPs.

In practice, customized layers often run much slower than the highly optimized built-in Pytorch layers. The performance gap between theory and practice is mainly because PyTorch framework is eagerly evaluated and thus brings additional memory access time and kernel launch time, please refer to https://residentmario.github.io/pytorch-training-performance-guide/jit.html for more details. As discussed in Appendix D, a native Pytorch implementation of MultiMax increases the training time of Deit-small on ImageNet by about 40% (0.1887 s/iteration vs 0.2641 s/iteration), while the increase in inference time is negligible (less than 2%). However, we are able to achieve a reduction from 40% (native Pytorch) to only 12.3% increase of training time (0.2120 s/iteration) by simply applying torch.jit.script decorator to fuse the pointwise operations of our MultiMax, following https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html. Notably, a fully optimized implementation of MultiMax in C++ or CUDA as done with Pytorch built-in layers might further reduce such a gap.

Model	Dataset	Config	Parameters	FLOPs	Training Speed
Deit-small (12 layers)	ImageNet-1K	SoftMax	22M	4.6 G	0.1887 second/iteration
		MultiMax	+0.09K(+0.0004%)	+168M (+0.37%)	0.2120 second/iteration (+12.3%)
Transformer-XL-Base (18 layers)	One Billion Words	SoftMax	460M	-	1.949 second/iteration
		MultiMax	+0.14K (+0.00003%)	-	2.035 second/iteration (+4.4%)

Implementation

• Implementation of MultiMax:

  • We implement the Max operator with 0 in Equation 6 as Pytorch built-in ReLU function
  • We apply torch.jit.script decorator to fuse the remaining elementwise operations in Equation 6, following the official documentation of TorchScript
  • We term the implementation of our modulation function as Segmented Linear Unit (SeLU)

• Main changes in vision_transformer.py:

  • The modulator function in Equation 6 of our paper is implemented in line 101.
  • The attention layer with MultiMax is implemented at line 133 by modulating the input to SoftMax via SeLU.
  • The output layer with MultiMax is implemented at line 324 in the same way.
  • We adopt Global Average Pooling (GAP) instead of Classification Token to aggregate the spatial information for our baseline model.

• Main changes in multihead_attention.py:

  • Include the multi_head_attention_forward function from the source code of torch.nn.functional.
  • The modulator function in Equation 6 of our paper is implemented in line 281.
  • The attention layer with MultiMax is implemented at line 666 by modulating the input to SoftMax via SeLU.

4. Main changes in transformer_decoder.py:
   • The modulator function in Equation 6 of our paper is implemented in line 38.
   • The output layer with MultiMax is implemented at line 438 by modulating the input to SoftMax via SeLU.

Training

Our experiment results are fully reproducible:

• Train a Vision Transformer with MultiMax

  1. Replace timm/models/vision_transformer.py with our provided vision_transformer.py
  2. Follow the training receipe of Deit to reproduce our experiment results

• Train a Language Transformer with MultiMax

  1. Replace fairseq/modules/multihead_attention.py with our provided multihead_attention.py to apply MultiMax in the attention layers.
  2. Replace fairseq/models/transformer/transformer_decoder.py with our provided transformer_decoder.py to apply MultiMax in the output layer.
  3. Follow the training receipe of fairseq/examples/language_model to reproduce our experiment results

Acknowledgements

This repo is based on Deit, timm, fairseq and pytorch.

Thanks to the original authors for their great work!

References

@InProceedings{pmlr-v235-zhou24g,
  title = 	 {{M}ulti{M}ax: Sparse and Multi-Modal Attention Learning},
  author =       {Zhou, Yuxuan and Fritz, Mario and Keuper, Margret},
  booktitle = 	 {Proceedings of the 41st International Conference on Machine Learning},
  pages = 	 {61897--61912},
  year = 	 {2024},
  editor = 	 {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix},
  volume = 	 {235},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {21--27 Jul},
  publisher =    {PMLR},
  url = 	 {https://proceedings.mlr.press/v235/zhou24g.html},
}