# Unpacking Softmax: How Temperature Drives Representation Collapse, Compression, and Generalization

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 44 references that are in the map.
- First posted on arXiv 2 Jun 2025.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Vision Transformers Need Registers](https://mlsanity.com/paper/2309.16588/) (2023)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [Rethinking the Value of Network Pruning](https://mlsanity.com/paper/1810.05270/) (2018)
- [Understanding Dimensional Collapse in Contrastive Self-supervised Learning](https://mlsanity.com/paper/2110.09348/) (2021)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [An Unconstrained Layer-Peeled Perspective on Neural Collapse](https://mlsanity.com/paper/2110.02796/) (2021)
- [Head2Toe: Utilizing Intermediate Representations for Better Transfer Learning](https://mlsanity.com/paper/2201.03529/) (2022)
- [Theory, Analysis, and Best Practices for Sigmoid Self-Attention](https://mlsanity.com/paper/2409.04431/) (2024)
- [A Study on ReLU and Softmax in Transformer](https://mlsanity.com/paper/2302.06461/) (2023)
- [Deep Neural Collapse Is Provably Optimal for the Deep Unconstrained Features Model](https://mlsanity.com/paper/2305.13165/) (2023)
- [Neural Collapse for Unconstrained Feature Model under Cross-entropy Loss with Imbalanced Data](https://mlsanity.com/paper/2309.09725/) (2023)
- [Linguistic Collapse: Neural Collapse in (Large) Language Models](https://mlsanity.com/paper/2405.17767/) (2024)
- [The Tunnel Effect: Building Data Representations in Deep Neural Networks](https://mlsanity.com/paper/2305.19753/) (2023)
- [Weight decay induces low-rank attention layers](https://mlsanity.com/paper/2410.23819/) (2024)
- [Linking Neural Collapse and L2 Normalization with Improved Out-of-Distribution Detection in Deep Neural Networks](https://mlsanity.com/paper/2209.08378/) (2022)
- [Stabilizing Contrastive RL: Techniques for Robotic Goal Reaching from Offline Data](https://mlsanity.com/paper/2306.03346/) (2023)
- [Neural Collapse in the Intermediate Hidden Layers of Classification Neural Networks](https://mlsanity.com/paper/2308.02760/) (2023)
- [Neural Collapse versus Low-rank Bias: Is Deep Neural Collapse Really Optimal?](https://mlsanity.com/paper/2405.14468/) (2024)
- [Eureka-Moments in Transformers: Multi-Step Tasks Reveal Softmax Induced Optimization Problems](https://mlsanity.com/paper/2310.12956/) (2023)
- [What Variables Affect Out-of-Distribution Generalization in Pretrained Models?](https://mlsanity.com/paper/2405.15018/) (2024)
- [The Impact of Geometric Complexity on Neural Collapse in Transfer Learning](https://mlsanity.com/paper/2405.15706/) (2024)

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

- [LUCID: Attention with Preconditioned Representations](https://mlsanity.com/paper/2602.10410/) (2026)
- [Over-Alignment vs Over-Fitting: The Role of Feature Learning Strength in Generalization](https://mlsanity.com/paper/2602.00827/) (2026)
- [How Much Rank Does LoRA Need? Rank-Error Bounds for Transformer Attention](https://mlsanity.com/paper/2608.26052/) (2026)

## Abstract (arXiv)

The softmax function is a fundamental building block of deep neural networks, commonly used to define output distributions in classification tasks or attention weights in transformer architectures. Despite its widespread use and proven effectiveness, its influence on learning dynamics and learned representations remains poorly understood, limiting our ability to optimize model behavior. In this paper, we study the pivotal role of the softmax function in shaping the model's representation. We introduce the concept of rank deficit bias - a phenomenon in which softmax-based deep networks find solutions of rank much lower than the number of classes. This bias depends on the softmax function's logits norm, which is implicitly influenced by hyperparameters or directly modified by softmax temperature. Furthermore, we demonstrate how to exploit the softmax dynamics to learn compressed representations or to enhance their performance on out-of-distribution data. We validate our findings across diverse architectures and real-world datasets, highlighting the broad applicability of temperature tuning in improving model performance. Our work provides new insights into the mechanisms of softmax, enabling better control over representation learning in deep neural networks.