Embedding and LM Head

Why are the input/output layers worth looking at separately?Copy Anchor Link

The heaviest computation in a Decoder-only Transformer happens in the blocks in the middle, but the input/output relationship of a language model is actually determined by the two ends:

$$\text{token ids} \xrightarrow{\text{embedding}} \text{hidden states} \xrightarrow{\text{LM Head}} \text{next-token logits}$$

If these two ends are not understood clearly, many later questions get tangled together: why the input is integers, why the output is not a token id, why the loss needs shifted targets, why the vocabulary size significantly affects the parameter count, and even why some models share the weights of the input embedding and the output LM Head.

The main thread of this chapter only covers the two ends, input and output: Embedding and LM Head.