Output and Gated Activations: Softmax, Sparsemax, GLU, and SIREN
Published:
CrossEntropyLoss already applies log-softmax, so calling softmax before it applies the normalisation twice and flattens your gradients. Inside the network the logic inverts — there the activation is a free design choice, and the modern answer is to make part of it learned.Part 3 of 3. Part 1 covers what activations do; Part 2 covers ReLU, GELU and SiLU. This page assumes you have met those.
The output layer is a contract with the loss
Hidden-layer activations shape what a network thinks. The output activation decides what it says, and must hand the loss exactly the object that loss was derived for.
| Task | Output activation | What the loss needs |
|---|---|---|
| Binary classification | Sigmoid | One probability in \((0, 1)\) |
| Multi-class classification | Softmax | A vector on the probability simplex* |
| Regression | Identity | An unconstrained real value |
*The probability simplex is just the set of valid probability vectors: all entries non-negative, summing to one. Softmax turns \(K\) logits into a point in it:
The subscript sits on the output: softmax maps a whole vector to a whole vector, so \(\operatorname{softmax}(z_i)\) is not a meaningful expression.
Worked example. A three-class classifier emits logits \(\mathbf{z} = [2.0,\ 1.0,\ 0.1]\).
| Step | Computation | Result |
|---|---|---|
| Exponentiate | \(e^{2.0},\ e^{1.0},\ e^{0.1}\) | 7.389, 2.718, 1.105 |
| Sum | 7.389 + 2.718 + 1.105 | 11.212 |
| Normalise | each divided by 11.212 | 0.659, 0.242, 0.099 |
These are normalised, not calibrated — trained networks are usually over-confident, which is why temperature scaling exists. Note 0.659 / 0.242 = 2.72: one unit of logit advantage multiplies the odds by \(e\).
Dividing logits by a temperature \(T\) before exponentiating rescales that sharpness. Drag the slider:
As T → 0 softmax approaches one-hot; as T → ∞ it approaches uniform. Distillation uses high T to expose the ratios between losing classes.
When you need exact zeros: sparsemax
Softmax never returns a zero — every class keeps some mass. Sparsemax instead takes the logit vector and finds the nearest valid probability vector to it in ordinary Euclidean distance. Because the nearest point often lies on an edge or corner of the simplex, entries genuinely hit zero:
The threshold \(\tau\) is set so the clipped values sum to one. (Entmax interpolates between the two.)
Sparsemax is often filed beside the elementwise shrinkage maps (SoftShrink, HardShrink), but those treat each coordinate independently with no normalisation, whereas sparsemax acts on the whole vector at once.
Gating: when part of the activation is learned
A ReLU asks one question per unit: should this value pass? A gated activation asks two projections to collaborate — one produces content, the other a per-channel gate that scales it:
Here \(\odot\) is the elementwise product. Swapping the gate nonlinearity gives the family: SwiGLU (SiLU), GeGLU (GELU), ReGLU (ReLU).
Worked example. With content \(\mathbf{a} = [1.2,\ -0.4,\ 0.8]\) and gate input \(\mathbf{b} = [2.1,\ -1.5,\ 0.3]\):
| Step | GLU | Plain linear |
|---|---|---|
| Gate | \(\sigma(\mathbf{b}) = [0.89,\ 0.18,\ 0.57]\) | — |
| Output | \(\mathbf{a} \odot \sigma(\mathbf{b}) = [1.07,\ -0.07,\ 0.46]\) | \([1.2,\ -0.4,\ 0.8]\) |
The second channel drops from −0.4 to −0.07 because its gate value is 0.18. Since \(\mathbf{b}\) is learned from the input, that suppression varies per example — a decision no plain linear layer can make.
Only the original sigmoid gate is bounded in \((0,1)\); SwiGLU, GeGLU and ReGLU gates are unbounded above, so they can amplify a channel, not just attenuate it. LLaMA, Mistral and Qwen use SwiGLU; Gemma uses GeGLU.
SIREN: when you need derivatives, not just values
Inside the network the activation is a free choice, and sometimes the task dictates it just as firmly as a loss does. To store an image or 3-D shape as a function from coordinates to values — an implicit neural representation — a ReLU network is a poor fit: it is piecewise linear, so its second derivative is zero almost everywhere. If you need the field’s curvature, that is fatal. SIREN uses a sine, whose derivatives are again sines:
The \(\omega_0\) factor is not cosmetic, and neither is its initialisation: hidden-layer weights are drawn from \(\mathcal{U}\!\left(-\sqrt{6/n}/\omega_0,\ +\sqrt{6/n}/\omega_0\right)\) for fan-in \(n\) (the first layer uses \(\mathcal{U}(-1/n,\ 1/n)\) instead). That keeps the pre-activation distribution stable with depth. Sine activations predate SIREN; this initialisation is what made deep ones trainable — implement \(\sin(\omega x)\) without it and the network will not converge.
Common mistakes
- Applying softmax before
CrossEntropyLoss. That loss already applies log-softmax internally — hand it raw logits. - Using sigmoid for mutually exclusive classes. Independent per-class probabilities will not sum to one; use softmax.
- Treating gate variants as drop-in swaps. Moving from SwiGLU to GeGLU changes optimisation behaviour, and changing the hidden width changes the parameter count.
Why the first one bites: the fused loss evaluates \(\log \sum_j e^{z_j} = m + \log \sum_j e^{z_j - m}\) with \(m = \max_j z_j\), which cannot overflow. Normalising yourself throws that away and normalises twice:
logits = model(x) # no softmax inside the model
loss = F.cross_entropy(logits, y) # applies log_softmax internally
probs = logits.softmax(-1) # only for reporting
Which one to reach for
The three output cases are in the table at the top. Beyond those:
| Use case | Activation | Watch out for |
|---|---|---|
| Transformer feed-forward | SwiGLU or GeGLU | Shrink hidden width by 2/3 to keep parameters matched |
| Sparse attention weights | Sparsemax or Entmax | Not in PyTorch core (pip install entmax); zeroed classes get no gradient |
| Implicit neural fields | SIREN | Useless without the \(\omega_0\) initialisation |
Output activations answer to the loss. Everything inside answers to the task.
References
- Dauphin, Y. N., Fan, A., Auli, M., & Grangier, D. Language Modeling with Gated Convolutional Networks. ICML 2017.
- Shazeer, N. GLU Variants Improve Transformer. arXiv:2002.05202, 2020.
- Martins, A. & Astudillo, R. From Softmax to Sparsemax. ICML 2016.
- Peters, B., Niculae, V., & Martins, A. Sparse Sequence-to-Sequence Models. ACL 2019.
- Sitzmann, V., Martel, J., Bergman, A., Lindell, D., & Wetzstein, G. Implicit Neural Representations with Periodic Activation Functions. NeurIPS 2020.
- Guo, C., Pleiss, G., Sun, Y., & Weinberger, K. On Calibration of Modern Neural Networks. ICML 2017.
