Self-Attention Concept › The Role of Self-Attention in the Transformer Architecture

Self-attention does not operate in isolation. In a transformer, each self-attention layer is followed by a position-wise feed-forward network, layer normalization, and residual connections. Understanding how these components interact helps clarify what self-attention's specific contribution is.

The residual connection is particularly important for understanding self-attention's role. Rather than replacing the input representations entirely, self-attention adds to them. The output of a self-attention layer is added back to its input:

where $\text{SelfAttn}(\mathbf{x}_i, \mathbf{X})$ denotes the self-attention output for position $i$ given the full sequence $\mathbf{X}$. This means each token retains its original representation and supplements it with contextual information gathered from the sequence. The self-attention output is a correction or augmentation, not a replacement.

This residual formulation has important consequences for learning and for interpretation. It means the gradient always has a direct path from the output to any earlier representation, without passing through the attention operation. This prevents vanishing gradients and allows very deep networks to be trained. It also means that in the early layers of a transformer, the representations are still relatively close to the original embeddings, with self-attention providing small contextual corrections. In later layers, the residual additions accumulate, producing representations that are substantially richer than the original embeddings.

The feed-forward network that follows each self-attention layer plays a complementary role. While self-attention blends information across positions, the feed-forward network operates independently on each position, applying the same learned transformation to each token's representation. The feed-forward network introduces non-linearity that self-attention lacks, and it provides the capacity to perform complex transformations on individual token representations after they have been enriched with context.

Layer normalization, applied before or after each sub-layer (depending on the specific architecture variant), keeps the activations in a stable range throughout the network. This stabilizes training in deep transformers: without normalization, the accumulation of residual additions would cause the representations to grow in magnitude, eventually making the softmax in the attention mechanism saturate and gradients vanish.

Together, these components create a system where self-attention handles the cross-position information flow and the feed-forward network handles the per-position transformation. Self-attention is responsible for asking and answering the question "what information from the sequence is relevant to this position?" while the feed-forward network is responsible for asking "given this enriched representation, what feature should I extract or computation should I perform?"