













Stanford CS336 Assignment 1 is titled Building a Transformer LM. It covers the main ideas behind a decoder-only language model: BPE tokenization, Transformer architecture, cross-entropy, SGD, AdamW, learning-rate schedules, gradient clipping, and decoding with temperature and top-p sampling.
This post is not a line-by-line translation of the handout, nor is it mainly about software engineering. Instead, it organizes the assignment as a review of the large-language-model concepts behind training and inference.
The central pipeline is:
text→tokens→Transformer→logits→loss→gradient update\text{text} \rightarrow \text{tokens} \rightarrow \text{Transformer} \rightarrow \text{logits} \rightarrow \text{loss} \rightarrow \text{gradient update}
During training, the model learns to predict the next token from the preceding tokens. During inference, it feeds its own predictions back into the context and generates a sequence one token at a time.
Suppose a text sequence has been tokenized as:
x=(x1,x2,…,xT)x=(x_1,x_2,\ldots,x_T)
An autoregressive language model decomposes the probability of the whole sequence using the chain rule:
p(x1,x2,…,xT)=∏t=1Tp(xt∣x1,…,xt−1)p(x_1,x_2,\ldots,x_T) =\prod_{t=1}^{T}p(x_t\mid x_1,\ldots,x_{t-1})
At every position, the model performs the same task: given a prefix, predict the next token.
Training maximizes the log-likelihood of the data:
maxθ∑t=1Tlogpθ(xt∣x<t)\max_\theta \sum_{t=1}^{T}\log p_\theta(x_t\mid x_{<t})
Equivalently, we minimize the negative log-likelihood, which becomes the next-token cross-entropy loss:
L(θ)=−∑t=1Tlogpθ(xt∣x<t)\mathcal{L}(\theta) =-\sum_{t=1}^{T}\log p_\theta(x_t\mid x_{<t})
The important point is that training does not require running the model separately for every position. One forward pass can produce predictions for all positions in parallel. The causal mask ensures that position tt cannot access future tokens.
During training, position tt receives the ground-truth token xtx_t and is trained to predict xt+1x_{t+1}:
input:x1x2x3x4\text{input}:\quad x_1\quad x_2\quad x_3\quad x_4
target:x2x3x4x5\text{target}:\quad x_2\quad x_3\quad x_4\quad x_5
This is called teacher forcing. The model always receives the correct history during training, rather than its own previous prediction.
Teacher forcing makes training highly parallelizable, but it creates exposure bias: during inference, the model must condition on tokens that it generated itself, including possible earlier mistakes.
Unicode assigns each character an abstract code point. For example:
1 | ord("s") |
Using Unicode code points directly as tokens would create a large and highly imbalanced vocabulary. CS336 instead encodes text as UTF-8 bytes before learning subword tokens.
A byte has a value between 00 and 255255, so the initial vocabulary contains only 256 possible byte tokens. UTF-8 can represent arbitrary Unicode text and remains compatible with ASCII.
The tradeoff is sequence length: one Unicode character may occupy several UTF-8 bytes. Byte-level tokenization almost eliminates unknown tokens, but it can produce longer sequences and therefore increase Transformer computation.
BPE, or Byte-Pair Encoding, is a compromise between byte-level and word-level tokenization.
Suppose a neighboring pair (A,B)(A,B) occurs frequently. BPE merges it into a new token ABAB:
[A,B,C,A,B]→[AB,C,AB][A,B,C,A,B]\rightarrow[AB,C,AB]
The algorithm starts with 256 byte tokens and repeatedly merges the most frequent adjacent pair. After MM merge operations, the vocabulary size is approximately:
∣V∣=256+M+∣Vspecial∣|V|=256+M+|V_{\text{special}}|
BPE is not learning word semantics directly. It is learning a useful compression scheme for frequent byte sequences. Common words may become a single token, while rare words can fall back to shorter subword or byte sequences.
Before counting byte pairs, the corpus is split into pre-tokens. This has two purposes:
The assignment uses a GPT-2-style regular expression and preserves leading spaces. As a result, " text" and "text" can be represented differently.
Special tokens create hard boundaries. For example, <|endoftext|> should remain one indivisible token and should not allow ordinary BPE merges to cross from one document into the next.
After BPE training, the tokenizer stores a vocabulary and an ordered list of merges. Encoding proceeds as follows:
Decoding performs the reverse operation: look up the byte string for every token ID, concatenate the bytes, and decode them as UTF-8. An arbitrary sequence of token IDs may not form valid UTF-8, so malformed bytes are usually replaced with the Unicode replacement character U+FFFD.
One distinction is worth remembering: BPE training learns merge rules from a corpus, while encoding applies an already learned set of rules to new text.
Given token IDs,
X∈NB×TX\in\mathbb{N}^{B\times T}
the token embedding maps them to dense vectors:
H∈RB×T×dmodelH\in\mathbb{R}^{B\times T\times d_{model}}
The hidden states pass through several Transformer blocks and are finally projected to vocabulary logits:
Z∈RB×T×∣V∣Z\in\mathbb{R}^{B\times T\times |V|}
The full model can be summarized as:
Token IDs→Embedding→Transformer Blocks→Final Norm→LM Head→Logits\text{Token IDs} \rightarrow \text{Embedding} \rightarrow \text{Transformer Blocks} \rightarrow \text{Final Norm} \rightarrow \text{LM Head} \rightarrow \text{Logits}
The model normally outputs logits rather than probabilities. During training, cross-entropy can be computed directly from logits. During inference, logits are converted into a probability distribution for decoding.
For hidden states XX, attention applies three learned linear transformations:
Q=XWQ,K=XWK,V=XWVQ=XW_Q,\qquad K=XW_K,\qquad V=XW_V
An intuitive interpretation is:
The dot product between a query and a key measures how relevant one position is to another. The values are then averaged according to those relevance scores.
The basic attention operation is:
Attention(Q,K,V)=softmax(QKTdhead+M)V\operatorname{Attention}(Q,K,V) =\operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_{head}}}+M\right)V
Why divide by dhead\sqrt{d_{head}}? If the components of QQ and KK have variance close to 1, the variance of their dot product grows with dheadd_{head}. Without scaling, large head dimensions can produce very large scores, causing softmax saturation and small gradients.
The scaling factor keeps the score distribution in a more manageable range.
An autoregressive language model cannot look at future tokens. The causal mask is:
Mij={0,j≤i−∞,j>iM_{ij}=\begin{cases} 0,&j\leq i\\ -\infty,&j>i \end{cases}
After adding this mask, future positions receive zero probability after softmax. Position ii can attend only to x1,…,xix_1,\ldots,x_i.
This is one of the most important constraints in language-model training. If the mask direction is reversed, the model can see the answer during training. The training loss may look excellent, while generation completely fails.
Multi-head attention splits the model dimension into hh smaller heads:
dhead=dmodelhd_{head}=\frac{d_{model}}{h}
Each head computes attention independently. The results are concatenated and passed through an output projection. Different heads can learn different types of dependencies, such as local syntax, long-range references, or positional relationships.
Without positional information, self-attention is permutation-equivariant: it has no way to tell that the order of the tokens has changed.
CS336 uses Rotary Position Embedding. RoPE treats every pair of dimensions as a two-dimensional plane and rotates query and key vectors by an angle determined by position:
R(θ)=[cosθ−sinθsinθcosθ]R(\theta)= \begin{bmatrix} \cos\theta&-\sin\theta\\ \sin\theta&\cos\theta \end{bmatrix}
For positions mm and nn:
qm′=Rmqm,kn′=Rnknq_m'=R_mq_m,\qquad k_n'=R_nk_n
The rotated inner product is:
(Rmq)T(Rnk)=qTRmTRnk(R_mq)^T(R_nk)=q^TR_m^TR_nk
Since rotation matrices satisfy RmTRn=Rn−mR_m^TR_n=R_{n-m}, the attention score depends on the relative distance n−mn-m.
This is the central intuition behind RoPE: by rotating queries and keys, their dot products naturally encode relative position.
RoPE is applied to QQ and KK, but not to VV. It changes how positions are matched; it does not alter the content carried by the values.
LayerNorm subtracts the mean and divides by the standard deviation. RMSNorm only normalizes the root mean square.
For a hidden vector a∈Rda\in\mathbb{R}^d:
RMS(a)=1d∑i=1dai2+ϵ\operatorname{RMS}(a)=\sqrt{\frac{1}{d}\sum_{i=1}^{d}a_i^2+\epsilon}
RMSNorm(a)=aRMS(a)⊙g\operatorname{RMSNorm}(a)=\frac{a}{\operatorname{RMS}(a)}\odot g
Here gg is a learned gain parameter. RMSNorm does not force the mean to zero; it mainly controls the scale of the hidden vector.
Why is normalization useful? As depth increases, the scale of activations and gradients can change substantially. Normalization gives each sub-layer a more stable input distribution, which makes optimization easier.
The pre-norm Transformer block used in the assignment is:
z=x+Attention(RMSNorm(x))z=x+\operatorname{Attention}(\operatorname{RMSNorm}(x))
y=z+FFN(RMSNorm(z))y=z+\operatorname{FFN}(\operatorname{RMSNorm}(z))
Normalization happens before each sub-layer, while the residual addition remains outside the sub-layer. The original Transformer is closer to the post-norm form:
z=Norm(x+Attention(x))z=\operatorname{Norm}(x+\operatorname{Attention}(x))
The key intuition behind pre-norm is that the residual stream provides a relatively direct path for both information and gradients. Each sub-layer learns an incremental correction to the stream instead of having to reconstruct a completely new representation.
This is why the assignment compares pre-norm and post-norm as an ablation. The location of normalization is not merely a formatting choice; it changes training stability and optimization behavior.
Attention mixes information across sequence positions. The feed-forward network applies the same nonlinear transformation independently to each position:
FFN(x)=W2 σ(W1x)\operatorname{FFN}(x)=W_2\,\sigma(W_1x)
Modern language models often use the gated SwiGLU variant:
SwiGLU(x)=W2(SiLU(W1x)⊙W3x)\operatorname{SwiGLU}(x) =W_2\left(\operatorname{SiLU}(W_1x)\odot W_3x\right)
where:
SiLU(x)=x⋅σ(x)\operatorname{SiLU}(x)=x\cdot\sigma(x)
W3xW_3x is the gate branch. It is multiplied element-wise with SiLU(W1x)\operatorname{SiLU}(W_1x), allowing the network to dynamically amplify or suppress features depending on the input.
SwiGLU uses three matrices, while a standard FFN uses two. To keep their parameter counts roughly comparable, the inner dimension of SwiGLU is often chosen near:
dff≈83dmodeld_{ff}\approx\frac{8}{3}d_{model}
The TinyStories configuration in the assignment uses dmodel=512d_{model}=512 and dff=1344d_{ff}=1344. The latter is close to 8/3×5128/3\times512 and is also divisible by 64, which is convenient for GPU hardware.
At one sequence position, the model produces vocabulary logits:
z=(z1,z2,…,z∣V∣)z=(z_1,z_2,\ldots,z_{|V|})
Softmax converts them into probabilities:
pi=ezi∑jezjp_i=\frac{e^{z_i}}{\sum_j e^{z_j}}
If the correct token has index yy, the cross-entropy loss is:
L=−logpy\mathcal{L}=-\log p_y
Substituting the softmax gives:
L=−zy+log∑jezj\mathcal{L} =-z_y+\log\sum_j e^{z_j}
The loss therefore has two effects: it rewards increasing the correct logit zyz_y, while the log-sum-exp term accounts for competition with every vocabulary item.
Let pp be the softmax probability vector and let eye_y be the one-hot vector for the correct class. Then:
∂L∂z=p−ey\frac{\partial\mathcal{L}}{\partial z}=p-e_y
This is the most important result for softmax cross-entropy.
If the model already assigns almost all probability to the correct class, p≈eyp\approx e_y and the gradient is close to zero. If it is highly confident but wrong, the gradient is large.
Directly computing ezie^{z_i} can overflow. Let m=maxizim=\max_i z_i. Then:
log∑iezi=m+log∑iezi−m\log\sum_i e^{z_i} =m+\log\sum_i e^{z_i-m}
Because zi−m≤0z_i-m\leq0, the exponentials are much safer to evaluate.
In practice, a cross-entropy implementation usually combines log_softmax and negative log-likelihood rather than materializing the full probability matrix first.
If the average per-token loss is L\mathcal{L}, perplexity is defined as:
PPL=eL\operatorname{PPL}=e^{\mathcal{L}}
It can be interpreted loosely as the effective number of equally likely choices the model faces at each position. Lower perplexity is generally better, but perplexities from different tokenizers or datasets are not directly comparable.
The basic gradient-descent update is:
θt+1=θt−αgt\theta_{t+1}=\theta_t-\alpha g_t
where α\alpha is the learning rate and gt=∇θL(θt)g_t=\nabla_\theta\mathcal{L}(\theta_t).
The learning rate controls the step size. If it is too small, training is slow. If it is too large, the loss may oscillate or diverge.
AdamW maintains an exponential moving average of the gradient and of the squared gradient:
mt=β1mt−1+(1−β1)gtm_t=\beta_1m_{t-1}+(1-\beta_1)g_t
vt=β2vt−1+(1−β2)gt2v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2
mtm_t tracks the direction of the gradient, while vtv_t tracks its scale. Since both states start at zero, the early estimates are biased toward zero. Adam corrects this bias using:
m^t=mt1−β1t,v^t=vt1−β2t\hat m_t=\frac{m_t}{1-\beta_1^t},\qquad \hat v_t=\frac{v_t}{1-\beta_2^t}
The adaptive update is:
θt←θt−αm^tv^t+ϵ\theta_t\leftarrow \theta_t-\alpha\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}
AdamW adds decoupled weight decay:
θt←θt−αλθt\theta_t\leftarrow \theta_t-\alpha\lambda\theta_t
Combining the two terms gives:
θt←θt−αm^tv^t+ϵ−αλθt\theta_t\leftarrow \theta_t -\alpha\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon} -\alpha\lambda\theta_t
Weight decay gently pulls parameters toward zero. The important distinction is that it is decoupled from the adaptive gradient update; it should not simply be added to the gradient and passed through Adam.
Typical LLM settings include β1=0.9\beta_1=0.9, β2=0.95\beta_2=0.95, and ϵ=10−8\epsilon=10^{-8}, although the best values depend on the model and the data.
Transformer training commonly uses a learning-rate schedule with a warmup phase followed by cosine decay.
During warmup:
αt=tTwαmax,t<Tw\alpha_t=\frac{t}{T_w}\alpha_{max},\qquad t<T_w
During cosine decay:
αt=αmin+12(1+cos(t−TwTc−Twπ))(αmax−αmin)\alpha_t=\alpha_{min} +\frac{1}{2}\left(1+\cos\left(\frac{t-T_w}{T_c-T_w}\pi\right) \right)(\alpha_{max}-\alpha_{min})
Here TwT_w is the end of warmup and TcT_c is the end of cosine decay. After TcT_c, the learning rate remains at αmin\alpha_{min}.
The intuition behind warmup is that the model has not yet reached a stable activation and gradient regime at the beginning of training. Starting immediately with a large learning rate may destabilize the optimization.
If a batch produces an unusually large gradient, we can scale the entire gradient vector so its L2 norm does not exceed MM:
g←g⋅min(1,M∥g∥2+ϵ)g\leftarrow g\cdot\min\left(1,\frac{M}{\lVert g\rVert_2+\epsilon}\right)
When ∥g∥2≤M\lVert g\rVert_2\leq M, the gradient is unchanged. Otherwise, every parameter gradient is multiplied by the same factor.
Gradient clipping does not make learning intrinsically faster. Its purpose is to prevent a small number of pathological batches from causing catastrophic parameter updates.
Ignoring biases, the main parameters in one Transformer block come from attention and the feed-forward network.
The QQ, KK, VV, and output projections together contribute approximately:
4dmodel24d_{model}^2
The three SwiGLU matrices contribute approximately:
3dmodeldff3d_{model}d_{ff}
For a model with LL layers, a rough parameter-count formula is:
P≈∣V∣dmodel+L(4dmodel2+3dmodeldff+2dmodel)+∣V∣dmodel+dmodelP\approx |V|d_{model} +L\left(4d_{model}^2+3d_{model}d_{ff}+2d_{model}\right) +|V|d_{model}+d_{model}
This includes the token embedding, each Transformer block, the final LM head, and the final RMSNorm. It assumes that the input embedding and output projection do not share weights. With weight tying, the two vocabulary-sized terms can be combined.
Self-attention computes QKTQK^T, whose sequence dimension has quadratic complexity in context length. The main compute of a Transformer block can be summarized as:
O(BTdmodel2+BT2dmodel)O\left(BT d_{model}^2+BT^2d_{model}\right)
The first term comes mainly from linear projections and the FFN. The second term comes from attention scores and the weighted sum over values.
For short sequences, the model dimension and FFN often dominate. For long contexts, the T2T^2 attention term becomes increasingly important.
If parameters, gradients, first moments, and second moments are all stored in float32, each parameter requires roughly:
4+4+4+4=16 bytes4+4+4+4=16\text{ bytes}
This excludes activation memory. Training therefore has several distinct memory costs:
Inference usually requires less memory because it does not need backward activations or optimizer states. Increasing batch size increases activation memory, while increasing context length increases both token computation and the quadratic attention cost.
Given a prefix x1:tx_{1:t}, the Transformer produces logits at every position. To generate the next token, we use only the final position:
v=TransformerLM(x1:t)tv=\operatorname{TransformerLM}(x_{1:t})_t
The logits are converted into a distribution:
p(xt+1=i∣x1:t)=evi∑jevjp(x_{t+1}=i\mid x_{1:t}) =\frac{e^{v_i}}{\sum_j e^{v_j}}
We sample xt+1x_{t+1}, append it to the prefix, and repeat until <|endoftext|> is generated or a maximum length is reached.
Training can compute all positions in a sequence in parallel. Naive generation is sequential because each new token depends on the previously generated result. This is the fundamental difference between training and autoregressive inference.
Greedy decoding selects the most likely token at every step:
xt+1=argmaxipix_{t+1}=\arg\max_i p_i
It is stable but can produce repetitive and overly conservative text.
Sampling draws from the probability distribution itself. It preserves multiple plausible continuations, but it can also select low-quality tokens, so the distribution is often modified before sampling.
Temperature τ\tau rescales the logits before softmax:
pi=softmax(v/τ)ip_i=\operatorname{softmax}(v/\tau)_i
Temperature does not change the model parameters. It changes only the sampling distribution used during inference.
Suppose the probabilities are sorted as q1≥q2≥⋯q_1\geq q_2\geq\cdots. Top-p sampling chooses the smallest candidate set V(p)V(p) whose cumulative probability reaches pp:
∑i∈V(p)qi≥p\sum_{i\in V(p)}q_i\geq p
It then renormalizes and samples only from this set:
P(i)={qi∑j∈V(p)qj,i∈V(p)0,i∉V(p)P(i)= \begin{cases} \dfrac{q_i}{\sum_{j\in V(p)}q_j},&i\in V(p)\\ 0,&i\notin V(p) \end{cases}
Top-p is adaptive. When the model is confident, only a few tokens are retained. When the model is uncertain, the candidate set becomes larger.
The assignment starts with TinyStories and later moves to OpenWebText.
TinyStories has a relatively simple distribution, so a small model can learn stable grammar and story patterns quickly. It is therefore useful for studying architecture and hyperparameters. OpenWebText is more varied and noisy, so the same model and compute budget usually produce higher loss and worse generations.
Loss values from different datasets should not be compared without context. Loss depends on data difficulty, tokenizer compression, vocabulary size, and sequence distribution.
Learning-rate experiments typically reveal three regimes:
The “edge of stability” intuition is that the best learning rate is often close to the largest rate that remains stable. A very conservative rate wastes the compute budget, while an aggressive rate may diverge.
Increasing batch size can improve hardware utilization and reduce the noise of the gradient estimate, but larger is not always better. When comparing batch sizes, we must specify whether we are holding the number of steps, the number of processed tokens, or wall-clock time fixed.
The assignment asks us to compare several architectural variants:
An ablation should not be judged only by its final loss. We should also inspect:
If the model architecture changes together with the parameter count, training token budget, or learning rate, the final difference cannot be attributed cleanly to the component being studied.
The main lesson of CS336 Assignment 1 is not a collection of APIs. It is an understanding of why a language model can be trained and how it generates text after training.
BPE maps open-ended Unicode text into a finite token vocabulary. The Transformer uses attention to aggregate context. RoPE provides positional information. The causal mask enforces autoregressive factorization. Cross-entropy turns next-token prediction into an optimization objective. AdamW, learning-rate schedules, and gradient clipping control how the parameters learn. Temperature and top-p determine how the model chooses among possible continuations during inference.
Together, these ideas explain the full process behind:
p(x)=∏tp(xt∣x<t)p(x)=\prod_t p(x_t\mid x_{<t})
A Transformer does not directly “understand” an entire document in one indivisible operation. At every position, it estimates a conditional distribution for the next token. With enough data, model capacity, and effective optimization, these local predictions can give rise to language understanding, knowledge recall, and increasingly complex reasoning behavior.
此内容由惯性聚合(RSS阅读器)自动聚合整理,仅供阅读参考。 原文来自 — 版权归原作者所有。