Output and Gated Activations: Softmax, Sparsemax, GLU, and SIREN

7 minute read

Published:

TL;DR: The activation on your final layer is fixed by the loss you chose, not by taste. The commonest bug in the whole topic follows from that: 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.

TaskOutput activationWhat the loss needs
Binary classificationSigmoidOne probability in \((0, 1)\)
Multi-class classificationSoftmaxA vector on the probability simplex*
RegressionIdentityAn 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:

\[ \operatorname{softmax}(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}, \qquad \mathbf{z}\in\mathbb{R}^{K} \]

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]\).

StepComputationResult
Exponentiate\(e^{2.0},\ e^{1.0},\ e^{0.1}\)7.389, 2.718, 1.105
Sum7.389 + 2.718 + 1.10511.212
Normaliseeach divided by 11.2120.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:

T = 1.0
Softmax probabilities at the selected temperatureThree bars for classes C1, C2 and C3, from logits 2.0, 1.0 and 0.1. At T = 1 they read 66%, 24% and 10%. Lowering T concentrates probability on C1 until it approaches 100%; raising T flattens the bars toward 33% each. 100% 0% 66% 24% 10% C1 (z = 2.0) C2 (z = 1.0) C3 (z = 0.1)

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:

\[ \operatorname{sparsemax}(\mathbf{z}) = \arg\min_{\mathbf{p}\in\Delta^{K-1}} \lVert \mathbf{p}-\mathbf{z}\rVert_2^2 = \big[\,z_i - \tau(\mathbf{z})\,\big]_+ \]

The threshold \(\tau\) is set so the clipped values sum to one. (Entmax interpolates between the two.)

Softmax compared with sparsemax on identical logitsTwo bar charts built from logits 3, 2.4, 0, minus 1 and minus 2. Softmax gives 61.6, 33.8, 3.1, 1.1 and 0.4 percent, so every class keeps some mass. Sparsemax gives 80, 20, 0, 0 and 0 percent, so the last three classes are exactly zero. Softmax (dense) Sparsemax (sparse) 61.6%33.8% 3.1%1.1%0.4% 80%20% 000 C1C2C3 C4C5 C1C2C3 C4C5 logits z = [3, 2.4, 0, −1, −2]
Softmax leaves C3–C5 holding 3.1%, 1.1% and 0.4% — small, but never zero. Sparsemax zeroes them exactly, which is what sparse attention needs when some tokens must receive no weight. The trade-off: classes outside the support receive no gradient.

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:

\[ \operatorname{GLU}(\mathbf{x}) = \underbrace{(W_1\mathbf{x} + \mathbf{b}_1)}_{\text{content }\mathbf{a}} \;\odot\; \sigma\big(\underbrace{W_2\mathbf{x} + \mathbf{b}_2}_{\text{gate input }\mathbf{b}}\big) \]

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]\):

StepGLUPlain 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.

Key Insight — the parameter budget: A GLU block needs three matrices (\(W_1\), \(W_2\), and the down-projection) where a plain MLP needs two, so a like-for-like comparison has to shrink the hidden width to \(\tfrac{2}{3}\) of \(4d_{\text{model}}\). At \(d_{\text{model}} = 4096\) that is LLaMA's width of 11008, which brings SwiGLU back to ≈135M parameters against the GELU MLP's ≈134M — without the correction it would be ≈201M. The reported gains are real but modest, and Shazeer offers no theory for why Swish beats GELU as the gate.

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:

\[ \Phi_i(\mathbf{x}) = \sin\!\big(\omega_0 (W_i\mathbf{x} + \mathbf{b}_i)\big), \qquad \omega_0 = 30 \]

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

Three that show up constantly:
  1. Applying softmax before CrossEntropyLoss. That loss already applies log-softmax internally — hand it raw logits.
  2. Using sigmoid for mutually exclusive classes. Independent per-class probabilities will not sum to one; use softmax.
  3. 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 caseActivationWatch out for
Transformer feed-forwardSwiGLU or GeGLUShrink hidden width by 2/3 to keep parameters matched
Sparse attention weightsSparsemax or EntmaxNot in PyTorch core (pip install entmax); zeroed classes get no gradient
Implicit neural fieldsSIRENUseless without the \(\omega_0\) initialisation

Output activations answer to the loss. Everything inside answers to the task.

References

  1. Dauphin, Y. N., Fan, A., Auli, M., & Grangier, D. Language Modeling with Gated Convolutional Networks. ICML 2017.
  2. Shazeer, N. GLU Variants Improve Transformer. arXiv:2002.05202, 2020.
  3. Martins, A. & Astudillo, R. From Softmax to Sparsemax. ICML 2016.
  4. Peters, B., Niculae, V., & Martins, A. Sparse Sequence-to-Sequence Models. ACL 2019.
  5. Sitzmann, V., Martel, J., Bergman, A., Lindell, D., & Wetzstein, G. Implicit Neural Representations with Periodic Activation Functions. NeurIPS 2020.
  6. Guo, C., Pleiss, G., Sun, Y., & Weinberger, K. On Calibration of Modern Neural Networks. ICML 2017.