Hidden State — Module 11.1

How This Lesson Fits the Module & Volume

After embeddings and embedding space, tokens live as vectors. The Transformer stack (Vol. 10) transforms those vectors into hidden states—contextual representations that already mix information from the whole context window.

Hidden states are the bridge to logits, softmax, and every decoding strategy that follows. You cannot debug generation without knowing which vector is being projected into the vocabulary.

Learning Objectives

By the end of this lesson, students should be able to:

Define a hidden state as a contextual vector at a position and layer.
Distinguish token embeddings from post-attention hidden states.
Explain why the last-position hidden state drives next-token prediction.
Relate hidden states to residual blocks, LayerNorm, and MHSA from Vol. 10.
Read shapes such as (B, T, d_model) in a forward pass.
Connect hidden states to the LM head that produces logits.

Definition

A hidden state is the model’s internal vector representation of a token position after (or between) Transformer layers. For sequence length T and model width d_model, a layer outputs a tensor of shape (B, T, d_model). In causal LMs, the hidden state at the final context position is typically fed to the language-model head to score the next token.

From Embedding to Contextual Vector

Token IDs

Integer indices in the vocabulary.

Embeddings

Lookup + position; not yet contextual.

Layers

Attention + FFN mix context.

Hidden states

Ready for the LM head.

Object	Typical shape	What it encodes
Token embedding	(B, T, d_model)	Identity + position; little context yet
Layer-ℓ hidden state	(B, T, d_model)	Context mixed through ℓ blocks
Final hidden (last pos.)	(B, d_model)	Summary for next-token scoring
Logits (next lecture)	(B, V)	Unnormalized vocab scores

Which Position Matters for Generation?

Training (teacher forcing)

Every position predicts the next gold token.
Hidden states at all t are used.
Causal mask still blocks the future.

Autoregressive inference

Usually only the newest position is needed for the next step.
Past keys/values can be cached (see KV Cache).
Append the chosen token and repeat.

Bidirectional models

BERT-style models use all positions differently.
Covered in Module 11.2.
Not left-to-right next-token heads by default.

Code: Extract Final Hidden State

import torch
import torch.nn as nn

class TinyLM(nn.Module):
    def __init__(self, V=1000, d=64, n_layers=2):
        super().__init__()
        self.embed = nn.Embedding(V, d)
        self.blocks = nn.ModuleList([
            nn.TransformerEncoderLayer(d, nhead=4, batch_first=True)
            for _ in range(n_layers)
        ])
        self.lm_head = nn.Linear(d, V, bias=False)

    def forward(self, ids):
        # ids: (B, T)
        h = self.embed(ids)                 # (B, T, d)
        for block in self.blocks:
            h = block(h)                    # still (B, T, d)
        return h                            # hidden states

model = TinyLM()
ids = torch.randint(0, 1000, (2, 8))
h = model(ids)
last = h[:, -1, :]                          # (B, d) — drives next token
logits = model.lm_head(last)                # (B, V)
print(h.shape, last.shape, logits.shape)
# torch.Size([2, 8, 64]) torch.Size([2, 64]) torch.Size([2, 1000])

Strengths and Tradeoffs

Why hidden states matter

They are the reusable interface between backbone and task heads.
Probing / interpretability often inspects layer-wise hidden states.
Caching strategies store related K/V, not full recompute of past h.

Caveats

“Hidden state” is overloaded (RNN h_t vs Transformer h).
Which layer you read changes semantics a lot.
Last-token pooling is GPT-style; other tasks need other pools.

Common Misconception

“The embedding is the hidden state.” Embeddings are the input representation. After attention and FFNs, each position’s vector has absorbed context from other tokens. That contextual vector is the hidden state used for logits—not the raw embedding lookup.

Related module pages: Embedding, Embedding Space, Logits, Context Window, Next Token Prediction.

Knowledge Check

Short Answer: What typical shape do Transformer hidden states have? Answer: (B, T, d_model).
True/False: Token embeddings already include full bidirectional context. Answer: False—context is mixed in the layers.
Multiple Choice: For GPT-style next-token scoring we usually use: (a) the first position only, (b) the last context position, (c) the mean of random noise. Answer: (b).
Short Answer: What maps a final hidden state to vocabulary scores? Answer: The LM head (linear projection to V).
True/False: During teacher-forced training, every position’s hidden state can predict the next token. Answer: True.
Multiple Choice: Hidden states after attention are: (a) contextual, (b) always one-hot, (c) vocabulary IDs. Answer: (a).
Short Answer: Name one Vol. 10 mechanism that builds contextual hidden states. Answer: Self-attention / MHSA (or residual FFN blocks).
True/False: “Hidden state” always means an RNN cell state. Answer: False—Transformers use the term too.
Multiple Choice: KV cache stores primarily: (a) past keys and values, (b) raw images, (c) tokenizer JSON. Answer: (a).
Short Answer: What lecture comes next after hidden states in this module? Answer: Logits.

Key Takeaways

Hidden states are contextual vectors of shape (B, T, d_model).
They sit between embeddings and the LM head that produces logits.
Generation typically reads the last position’s hidden state each step.
Do not confuse raw embeddings with post-layer hidden states.
Next: Logits.

Trainer’s Guide

Hands-on idea: Print cosine similarity between position 0’s embedding and its layer-N hidden state on a short sentence; discuss how context changed the vector.

Discussion prompt: When would you pool all hidden states (mean/CLS) instead of using only the last token?

Recap: Hidden states are the Transformer’s contextual vectors; the last one feeds next-token scoring. Continue with Logits.