---
title: Modified Softmax Functions
url: https://www.emergentmind.com/topics/modified-softmax-functions
type: topic
---

# Modified Softmax Functions

Modified softmax functions constitute a rapidly expanding class of alternatives to the canonical softmax activation. These modifications address diverse deficiencies of the standard function, including gradient vanishing, sparsity control, robustness to label noise, computational cost, and statistical learning rates. Contemporary research has produced a wide spectrum of variants, ranging from theoretically motivated geometric reinterpretations to engineering-driven numerical approximations and sparsification strategies. Below, the principal directions and technical features of modified softmax functions are surveyed.

## 1. Geometric and Clustering Perspectives

A rigorous formal equivalence exists between softmax-based classifiers and $k$-means clustering in the transformed feature space. Specifically, given any feedforward neural network with a final linear layer $W\in\mathbb{R}^{d\times c}$ and penultimate mapping $f_p\colon\mathbb{R}^n\to\mathbb{R}^d$, the softmax classifier partitions the $f_p$-space into Voronoi cells with equidistant centroids $Z_{\cdot k}=W_{\cdot k}+v$, and performs nearest-centroid decoding:
\[
\arg\max_{k}\;f_p(x)^\top W_{\cdot k}
\;=\;
\arg\min_{k}\;\|f_p(x)-Z_{\cdot k}\|^2
\]
Under this interpretation, the standard softmax loss encourages “cone-based” decisions, disregarding within-cone distances. The “Gauss” (or centroid-based) alternative replaces the softmax with a distribution explicitly decaying with squared Euclidean distance to a learned centroid $\mu_k$ per class:
\[
P(y=k\mid x) =
\frac{\exp(-\|f_p(x)-\mu_k\|^2)}{\sum_j \exp(-\|f_p(x)-\mu_j\|^2)}
\]
This replacement enforces tighter mapping of samples to their respective centroids, raising the robustness threshold for adversarial perturbations via direct control over the Lipschitz continuity of the penultimate mapping and increasing inter-centroid distance. Empirically, the “Gauss” activation delivers 6–7× improved resistance to one-pixel attacks over vanilla softmax, with comparable clean accuracy and better confidence calibration [2001.01987].

## 2. Gradient-Stabilized Softmax Modifications

Standard softmax exhibits severe gradient vanishing when input elements attain extreme values, as $\frac{\partial \alpha_j}{\partial x_j} = \alpha_j(1-\alpha_j)\to0$ for $\alpha_j\approx 0$ or $1$. Self-Adjust Softmax (SA-Softmax) mitigates this by introducing a multiplicative self-modulating term,
\[
\text{SA-Softmax:}\quad \beta_j = x_j \cdot \mathrm{softmax}(x)_j
\]
and a normalized variant,
\[
\tilde\beta_j = \frac{x_j - \min(\min_k x_k, 0)}{\max(0, \max_k x_k) - \min(\min_k x_k, 0)} \cdot \mathrm{softmax}(x)_j
\]
The Jacobian of SA-Softmax contains an additive identity $\mathrm{diag}(\alpha)$, boosting the smallest singular value and preventing gradient collapse even in saturated regimes. Integrated into transformer attention mechanisms, SA-Softmax yields consistent perplexity reductions (1–3%) and nontrivial accuracy and BLEU improvements across classification and translation tasks (up to 6.3 pp/1.0 BLEU, respectively). The normalized variant tends to stabilize training further [2502.18277].

## 3. Structural and Computational Modifications for Efficiency

Efficient realization of softmax in privacy-preserving computation, such as homomorphic encryption (HE), motivates approximations that avoid max and division operations. MGF-softmax replaces the partition function with a moment-generating-function- (MGF-) based exponential shift:
\[
\text{softmax}_{\mathrm{MGF}^{(m)}}(x)_i = \exp \left( x_i 
-\sum_{j=1}^m \frac{\hat\kappa_j}{j!} 
- \ln n \right),\quad \hat\kappa_j: \text{sample cumulants}
\]
This yields multiplicative depth reductions $O(k+6)$ (vs $O(8k+9)$ for baseline), eliminating bootstrapping for typical input sizes, and matches the accuracy of plaintext softmax within 1%, as validated on LLaMA-3.2B (NLU) and ViT/DeiT (vision) [2602.01621]. Other variants in the MPC/HE literature, such as ReLU-based output normalization, trade off accuracy and only offer modest speed-ups except in shal-low nets, and are not recommended for multi-layer models [2011.11202].

The online normalizer approach fuses the maximum-finding and sum-accumulation passes in the standard “safe” softmax via running recurrences, reducing memory transfers per element by $\sim$33% and achieving up to $1.3\times$ speed-ups for large-batch computations on GPUs. If paired with TopK selection, memory access can be reduced up to $5\times$ [1805.02867].

## 4. Sparsity-Inducing and Selective Variants

Classical softmax assigns strictly positive probabilities to all classes, which can hinder convergence speed and model discrimination for high-dimensional outputs. Multiple modified functions have been proposed to achieve sparsity and dynamic support selection:

- **r-softmax:** Enforces a fixed fraction $r$ of exactly zero outputs by thresholding logits at the empirical quantile, yielding exactly $rn$ zeros. The mechanism is controlled via a quantile-based shift before softmax, offering differentiability nearly everywhere and adaptability by coupling $r$ to auxiliary networks. r-softmax outperforms other sparse mappings and produces higher F1 in multi-label settings, and yields improvements when used for attention in transformers [2304.05243].
- **Sparse-softmax:** Masks all but the $k$ largest logits, exponentiates and normalizes over that support. The choice of $k$ controls both margin and learning focus: for $d\gg1$, sparse-softmax reduces the required margin from $\log(d-1)$ to $\log(k-1)$ in cross-entropy. It yields faster convergence and improved macro/micro F1 in high-dimensional text and sequence models for moderate $k$ (e.g., $k=20$ with $d\sim 100$) [2112.12433].
- **Input/Feature Transformed Gating:** In mixture-of-experts models with softmax gating, transforming the input $X\mapsto M(X)$ before gating removes pathological expert–gate PDE couplings that otherwise slow parameter estimation rates below any polynomial. Choices such as $M(X) = \tanh(X), \cos(X), X^m (m\geq3)$ restore parametric estimation and improve convergence in multinomial logistic MoEs [2310.14188].

## 5. Alternative Basis Functions and Spherical/Taylor Families

The exponential in softmax can be replaced by other normalizers for various optimization or architectural benefits.

- **Spherical-family losses:** Use quadratic-based normalizers, e.g., log-spherical softmax $\propto o_k^2+\epsilon$ and log-Taylor softmax is based on the second-order Taylor expansion of the exponential, $1+o_k+\frac12 o_k^2$. These normalizers permit $O(d^2)$ weight updates (vs $O(dD)$ for softmax) and excel for low $D$, but are outperformed by softmax for very large output dimensions [1511.05042].
- **Taylor and Soft-Margin Taylor Softmax:** Explicitly replace $\exp(z)$ with finite/infinite Taylor expansions in both the normalizer and the gradient. The SM-Taylor variant adds a margin penalty $m$ to promote class separability. On benchmarks, SM-Taylor with small $n$ (e.g., $n$=2 for MNIST/CIFAR-10) always meets or exceeds softmax accuracy [2011.11538].
- **Periodic Softmax Alternatives:** In attention mechanisms where dot-product scores are not normally distributed and softmax gradients vanish, replacing $\exp(x)$ by periodic functions such as $\exp(\sin(x))$ or phase-shifted $\sin^2(x)$ yields oscillatory but bounded gradients, preventing vanishing and yielding up to $4$pp accuracy gains in deep self-attention models. Pre-normalization and careful phase design are required for stability [2108.07153].

## 6. Noise-Robust and Distributional Modulation Variants

Noisy labels degrade standard loss functions unless outputs are close to one-hot or symmetric. $\epsilon$-Softmax augments the largest softmax output by a constant $m$ and divides by $m+1$, guaranteeing proximity to the one-hot set and providing explicit excess-risk bounds under asymmetric noise:
\[
f_\epsilon(h(x))_k = \frac{p_k + m\, \delta_{k=t}}{1+m}
\]
where $t = \arg\max p$. The excess risk decays as $O(1/\sqrt{m})$, and empirical results on CIFAR-10/100 and large-scale datasets show state-of-the-art robustness across all noise regimes, with negligible cost [2508.02387].

Invertible modifications, as in the “softmax$_{++}$” mapping, build reparameterizable distributions on the simplex (for variational inference), allowing the use of Gaussian sources with explicit Jacobians and closed-form KL divergences, outperforming Gumbel-Softmax in likelihood and gradient variance [1912.09588].

## 7. Sampled and Reinforcement Learning Operators

For large $C$, sampled softmax with kernel-based importance sampling replaces the exact partition function by a mini-batch estimate; adaptive kernels tracking $e^{o_i}$ reduce the bias inherent to uniform sampling (by orders of magnitude), achieving close-to-exact performance with tens or hundreds of samples [1712.00527].

In sequential decision-making, “mellowmax” is a quasi-arithmetic mean operator
\[
\mathrm{mellowmax}_\omega(X) = \frac{1}{\omega} \log \frac{1}{n} \sum_{i=1}^n e^{\omega x_i}
\]
which, unlike the Boltzmann operator, is a non-expansion under $\|\cdot\|_\infty$ and guarantees unique fixed-point convergence in planning and SARSA. Policies derived via state-dependent temperature maximize expected value plus entropy with moment-matching constraints, yielding stable and performant RL updates [1612.05628].

---

In summary, modified softmax functions span a continuum from geometric reinterpretations and gradient stabilization, through sparsification and noise-robustification, to architecture-specific and efficiency-driven transformations. These advances have been directly integrated into modern deep learning pipelines for classification, language modeling, vision, transformers, multi-label learning, and variational inference, producing gains in empirical accuracy, scalability, robustness, and theoretical convergence.

Source: https://www.emergentmind.com/topics/modified-softmax-functions