Order the tail-end operations that turn the last residual stream into logits

**The correct order is**: final block output, final RMSNorm, unembedding projection, logits, softmax. Let's walk through each step. ### Step 1: residual stream output of the last block After the final transformer block runs, the residual stream has shape `(batch, seq_len, d_model)`. Each position carries a d_model-dimensional vector that encodes everything the model knows about that token in context. ### Step 2: final RMSNorm In a **pre-norm** architecture (every modern LLM), each transformer block applies its normalization *before* the attention and FFN sublayers. After the residual add at the end of each block, the residual stream is NOT normalized. After 80 blocks of accumulating un-normalized residual updates, the stream has grown in magnitude and needs a final normalization before any downstream operation can use it cleanly. The final RMSNorm (LayerNorm in older models) rescales each position's vector to unit RMS along d_model, then multiplies by a learned per-channel scale. The output is the same shape as the input but with magnitudes brought back into a healthy range. Without this step, the unembedding projection would see a residual stream with arbitrarily large magnitudes, and the resulting logits would be poorly calibrated. ### Step 3: unembedding projection The unembedding (also called the output head or LM head) is a single linear layer that maps each position's d_model vector to a vocab_size vector: ``` logits = x @ W_unembed.T ``` W_unembed has shape `(vocab_size, d_model)`. If the model uses **weight tying**, W_unembed is the same matrix as the input embedding table (with the transpose). Most modern LLMs (Llama, Mistral, Qwen, DeepSeek) do NOT tie weights at this scale; they use a separate unembedding matrix for capacity reasons. ### Step 4: logits The output of step 3 is the logits tensor, shape `(batch, seq_len, vocab_size)`. Each entry is an unnormalized score for one vocab token at one sequence position. ### Step 5: softmax Finally, softmax converts logits into a probability distribution over the vocabulary: ``` probs[i] = exp(logits[i]) / sum_j(exp(logits[j])) ``` During training, the softmax is fused into the cross-entropy loss for numerical stability. During inference, the softmax happens at sampling time (optionally with a temperature divisor before softmax: `logits / temperature`).

**The canonical tail-end order for a modern pre-norm LLM is**: final block output -> final RMSNorm -> unembedding projection -> logits -> softmax. ### Step 1: residual stream from the last block Shape `(batch, seq_len, d_model)`. After 32 to 96 transformer blocks of pre-norm processing, this is the model's accumulated belief about every token in context. In a pre-norm architecture, this stream has NOT been normalized since the last RMSNorm inside the last block (which was applied to the inputs of the last attention and FFN, not to the residual sum). ### Step 2: the final RMSNorm Why this step exists is the central insight of the question. Pre-norm models apply LayerNorm or RMSNorm *before* each sublayer, leaving the residual sum un-normalized: ``` x = x + Attn(RMSNorm(x)) x = x + FFN(RMSNorm(x)) ``` After the last block's last residual add, x has not been normalized. Its magnitude has grown through the stack as residual contributions accumulate. The final RMSNorm is the missing normalization that brings the stream back to unit RMS along d_model. In Llama, Mistral, Qwen, and DeepSeek configs you will see this as a single `model.norm` layer placed after the block stack and before the LM head. **Post-norm models** (Vaswani 2017 original, BERT, GPT-1) do not need this step because every block ends with a LayerNorm. The final residual stream is already normalized. ### Step 3: unembedding projection A linear layer mapping `d_model -> vocab_size`. Two implementation choices. - **Tied weights**: W_unembed.T is the same matrix as the input embedding. Saves vocab_size * d_model parameters. Used in GPT-2 and original-Llama. - **Untied weights**: W_unembed is a separate matrix. Costs vocab_size * d_model extra parameters (~400M for vocab=128k, d_model=8192) but typically improves quality. Most 2026 frontier LLMs untie at large scale. ### Step 4: logits Shape `(batch, seq_len, vocab_size)`.