Zero-Temperature Softmax

• Zero-temperature softmax is a limiting case of the softmax function where T → 0, resulting in a hard, deterministic one-hot selection.
• It underpins hard attention, discrete sampling, and differentiable categorical relaxations, such as the Gumbel-Softmax, in neural models.
• Practical challenges include vanishing gradients, rank collapse, and dispersion in long sequences, necessitating careful temperature tuning.

Zero-temperature softmax refers to the limiting behavior of the temperature-controlled softmax function as the temperature parameter $T \to 0^+$. In this regime, the softmax outputs approach a hard one-hot assignment, selecting the argmax of the input logits deterministically. This concept is central to discrete decision-making, hard attention, categorical relaxations (such as the Gumbel-Softmax), and the analysis of neural network learning dynamics. The zero-temperature limit exposes both mathematical structure and practical issues in differentiation, optimization, and generalization.

1. Mathematical Formalism of Temperature-Scaled Softmax

The temperature-scaled softmax for input logits $z = (z_1, \ldots, z_c) \in \mathbb{R}^c$ and $T > 0$ is defined as

$$\operatorname{softmax}_T(z)_i = \frac{\exp(z_i/T)}{\sum_{j=1}^c \exp(z_j/T)}$$

for $i=1,\ldots, c$. This parameterization allows explicit control over the "sharpness" of the output distribution: $T \gg 1$ yields near-uniform assignments, while $T \ll 1$ produces sharper distributions. Equivalently,

$$\operatorname{softmax}_T(z) = \operatorname{softmax}_1(z/T),$$

so lowering $T$ amplifies the impact of logit differences (Masarczyk et al., 2 Jun 2025, Veličković et al., 2024, Lee-Jenkins, 28 Aug 2025).

At the variational level, softmax arises as the unique maximizer of the "free energy"

$\mathcal{F}(p) = \langle p, \phi \rangle + T H(p)$

over the probability simplex $z = (z_1, \ldots, z_c) \in \mathbb{R}^c$0, with $z = (z_1, \ldots, z_c) \in \mathbb{R}^c$1 the Shannon entropy. The Karush-Kuhn-Tucker (KKT) conditions yield the same temperature-scaled softmax expression for the optimum $z = (z_1, \ldots, z_c) \in \mathbb{R}^c$2 (Lee-Jenkins, 28 Aug 2025).

2. Zero-Temperature Limit: Hard Argmax and Discrete Distributions

As $z = (z_1, \ldots, z_c) \in \mathbb{R}^c$3, the softmax concentrates its entire probability mass on the maximal element of $z = (z_1, \ldots, z_c) \in \mathbb{R}^c$4. If $z = (z_1, \ldots, z_c) \in \mathbb{R}^c$5 is unique, then

$z = (z_1, \ldots, z_c) \in \mathbb{R}^c$6

This is a deterministic one-hot ("hard") decision selecting the argmax; non-maximal entries vanish exponentially fast in $z = (z_1, \ldots, z_c) \in \mathbb{R}^c$7: $z = (z_1, \ldots, z_c) \in \mathbb{R}^c$8 recovering the classical argmax rule. This limiting behavior underpins hard attention, discrete sampling, and deterministic routing in neural models (Masarczyk et al., 2 Jun 2025, Veličković et al., 2024, Lee-Jenkins, 28 Aug 2025, Jang et al., 2016).

3. Categorical and Gumbel-Softmax Relaxations

The Gumbel-Softmax ("Concrete") distribution facilitates differentiable relaxation of categorical sampling. For categorical probabilities $z = (z_1, \ldots, z_c) \in \mathbb{R}^c$9 and independent Gumbel noise $T > 0$0, softmax-based sampling is defined as

$T > 0$1

with temperature $T > 0$2. As $T > 0$3, $T > 0$4 converges to a one-hot vector at index $T > 0$5, recovering an exact categorical sample via the Gumbel-Max trick (Jang et al., 2016). For intermediate $T > 0$6, $T > 0$7 interpolates smoothly between uniform and categorical extremes.

The straight-through Gumbel-Softmax (ST-GS) estimator enforces hard argmax discretization in the forward pass ($T > 0$8), while using a continuous softmax with moderate temperature in the backward pass ($T > 0$9) for gradient computation. This decoupling improves both exactness and gradient fidelity in discrete optimization tasks, outperforming single-temperature estimators in autoencoding and generative settings (Shah et al., 2024).

4. Optimization, Gradient Behavior, and Training Dynamics at Zero Temperature

In softmax-based models trained via the cross-entropy loss, introducing inverse temperature $$\operatorname{softmax}_T(z)_i = \frac{\exp(z_i/T)}{\sum_{j=1}^c \exp(z_j/T)}$$0 yields

$$\operatorname{softmax}_T(z)_i = \frac{\exp(z_i/T)}{\sum_{j=1}^c \exp(z_j/T)}$$1

As $$\operatorname{softmax}_T(z)_i = \frac{\exp(z_i/T)}{\sum_{j=1}^c \exp(z_j/T)}$$2 ($$\operatorname{softmax}_T(z)_i = \frac{\exp(z_i/T)}{\sum_{j=1}^c \exp(z_j/T)}$$3), the model outputs approach a one-hot vector aligned with the maximal logit, but crucially:

• The gradient $$\operatorname{softmax}_T(z)_i = \frac{\exp(z_i/T)}{\sum_{j=1}^c \exp(z_j/T)}$$4 becomes exponentially small as $$\operatorname{softmax}_T(z)_i = \frac{\exp(z_i/T)}{\sum_{j=1}^c \exp(z_j/T)}$$5 (where $$\operatorname{softmax}_T(z)_i = \frac{\exp(z_i/T)}{\sum_{j=1}^c \exp(z_j/T)}$$6 is the logit gap), leading to vanishing updates.
• The network remains in the linear (NTK) regime for extended periods ($$\operatorname{softmax}_T(z)_i = \frac{\exp(z_i/T)}{\sum_{j=1}^c \exp(z_j/T)}$$7), impeding feature learning.
• Convergence becomes excessively slow, with model weights barely moving from initialization, and generalization degrades (Agarwala et al., 2020).

Careful annealing or tuning of $$\operatorname{softmax}_T(z)_i = \frac{\exp(z_i/T)}{\sum_{j=1}^c \exp(z_j/T)}$$8 is recommended, with optimal values architecture-dependent and typically in the range $$\operatorname{softmax}_T(z)_i = \frac{\exp(z_i/T)}{\sum_{j=1}^c \exp(z_j/T)}$$9.

5. Representation Theory, Rank Collapse, and Generalization Effects

At very low temperature, the softmax sharpens, and the network's learned representations (rows of the logit matrix across input samples) exhibit rank-deficit bias: numerical ranks much lower than the class count. Empirical findings indicate:

• At low $i=1,\ldots, c$0, hidden-layer matrix rank drops sharply, fewer layers contribute meaningfully to linear separability, and gradients also collapse in rank.
• Increasing logit norm (decreasing $i=1,\ldots, c$1) raises post-softmax matrix rank but, during training, high $i=1,\ldots, c$2 solutions often exhibit persistently low pre-softmax rank.
• Rank collapse impairs out-of-distribution generalization while sometimes improving compressed model performance (Masarczyk et al., 2 Jun 2025).

Explicit temperature tuning, architectural parameters (e.g., width, normalization), and annealing schedules can balance sharpness, representation diversity, and numerical stability. Low $i=1,\ldots, c$3 can require smaller learning rates and careful regularization.

6. Zero-Temperature Softmax in Attention and Sequence Models: Dispersion and Limitations

In attention mechanisms (e.g., Transformers), zero-temperature softmax yields exact selection of the token with maximal similarity (hard attention). For any fixed $i=1,\ldots, c$4, as sequence length $i=1,\ldots, c$5 grows, the maximal attention weight disperses: $i=1,\ldots, c$6 where $i=1,\ldots, c$7 bounds all logits. As $i=1,\ldots, c$8, $i=1,\ldots, c$9 for fixed $T \gg 1$0—sharpness is lost on long sequences. Adaptive temperatures based on entropy can partially restore sharpness at inference but cannot overcome the limitation for arbitrarily large $T \gg 1$1 (Veličković et al., 2024). Non-softmax attention mechanisms or explicit discrete routing are necessary for indefinitely sharp, scale-robust reasoning.

7. Dynamical Perspective: Replicator Dynamics and the Time-Rescaling Role of Temperature

The evolution of output probabilities under entropic mirror ascent with temperature $T \gg 1$2 yields, in continuous time, the replicator equation: $T \gg 1$3 where $T \gg 1$4 is the score vector. Temperature appears as a pure time rescaling parameter: $T \gg 1$5. As $T \gg 1$6, the trajectory in probability simplex space collapses exponentially fast onto the vertex corresponding to the argmax of $T \gg 1$7, with convergence rates determined by logit gaps. This interpretation provides a principled theoretical link between zero-temperature softmax and argmax selection as a dynamical system (Lee-Jenkins, 28 Aug 2025).

---

Key References:

• Gumbel-Softmax and categorical reparameterization: (Jang et al., 2016, Shah et al., 2024)
• Softmax temperature, rank collapse, generalization: (Masarczyk et al., 2 Jun 2025, Agarwala et al., 2020)
• Dispersion and adaptive temperature in attention: (Veličković et al., 2024)
• Variational, mirror descent, and replicator dynamics perspectives: (Lee-Jenkins, 28 Aug 2025)