Activation Functions › Activation Functions in Feed-Forward Networks vs Attention

Activation functions play a somewhat different role in different parts of a transformer compared to a feedforward network.

In a standard feedforward network, every hidden layer has an activation function, and the depth of the network means gradients must propagate through many activation functions. The choice of activation function is critical for gradient flow.

In a transformer, the self-attention computation itself is linear (the attention weights are applied to the values through a weighted sum), and the nonlinearity is concentrated in the feed-forward sub-layers that follow each attention block. These feed-forward layers are typically two-layer MLPs: a linear layer that expands the dimension, followed by an activation function, followed by a linear layer that contracts back to the model dimension. GELU is the standard choice for this nonlinearity in transformers.

The attention computation in transformers uses softmax, not a traditional activation function, to normalize the attention weights. Softmax is not a nonlinearity in the same sense as ReLU or GELU: it is applied across the sequence dimension to produce a probability distribution, not applied element-wise to introduce nonlinearity within a neuron. The feed-forward sub-layers are where the true element-wise nonlinearity lives in a transformer.

Some newer transformer variants (including LLaMA and Mistral) use Swish/SiLU with a gated structure in the feed-forward layers. This gated variant, called SwiGLU, computes two parallel linear projections of the input and multiplies them element-wise after applying SiLU to one of them. We will cover gated linear units and SwiGLU in detail when we reach the chapter on Transformer Blocks and Feed-Forward Networks.