Rediscover the Structure of an LLM

For long, I depicted the structure of a typical transformer-based LLM (a decoder-only model, for simplicity) like this:

sequence -> embedding -> {MHA -> residual+norm -> MHA -> residual+norm -> FFN -> linear -> residual -> softmax} -> ... -> {MHA -> residual+norm -> FFN -> linear + residual -> softmax} -> outputs

Where each transformer block appears to consist of MHA -> residual+norm -> MHA -> FFN -> linear -> softmax.

If you thought about it this way too, you were just as wrong as I was. There are barely any explanations that explicitly correct this misconception. Because the same figure from Attention Is All You Need appears everywhere, newcomers like me can easily misread the architecture and think in the wrong way for a long time.

Think Generally

Let’s think about the structure more generally. A typical neural network can be treated as a function that maps an input to an output, so its structure can be shown as:

input -> network -> output

The network first processes or projects the input, passes it through its main functional component (the backbone), and finally processes the result into the output required by the task. In short:

pre-processor -> backbone -> task head

These parts are built from basic arithmetic operations, are usually grouped into blocks for re-using. Greater depth enables a network to capture more complex relationships (Deep Residual Learning for Image Recognition). Stacking blocks makes the network deeper, forming its backbone:

backbone block -> backbone block -> ... -> backbone block

Overall, the structure looks like:

pre-processor -> backbone (block -> ... -> block) -> task head

For example, an object detector can be shown in the same way:

image -> input preprocessing -> backbone -> detection head -> predictions

LLM Structure

Likewise, the structure of a typical decoder-only LLM should be depicted as follows:

sequence -> embedding -> transformer -> ... -> transformer -> norm -> linear -> softmax -> logits

The embedding stage is the pre-processor. A decoder-only transformer block has one MHA -> residual+norm, since they don’t need the second one for cross attention with an encoder, hence MHA -> residual+norm -> FFN -> linear -> residual -> softmax. Multiple transformer blocks make up the backbone. And the final linear -> softmax stage is the LM head.