CS336 Assignment 1: Large Language Model Training and Inference

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{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.

1. What Does a Language Model Learn?

Autoregressive language modeling

Suppose a text sequence has been tokenized as:

$$x=(x_1,x_2,\ldots,x_T)$$

An autoregressive language model decomposes the probability of the whole sequence using the chain rule:

$$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_\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:

$$\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 $t$ cannot access future tokens.

Teacher forcing

During training, position $t$ receives the ground-truth token $x_t$ and is trained to predict $x_{t+1}$:

$$\text{input}:\quad x_1\quad x_2\quad x_3\quad x_4$$

$$\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.

2. Tokenization: Why Start with Byte-Level BPE?

Unicode, UTF-8, and byte vocabulary

Unicode assigns each character an abstract code point. For example:

ord("s")  # 115
ord("é")  # 233

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 $0$ and $255$, 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.

The compression idea behind BPE

BPE, or Byte-Pair Encoding, is a compromise between byte-level and word-level tokenization.

Suppose a neighboring pair $(A,B)$ occurs frequently. BPE merges it into a new token $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 $M$ merge operations, the vocabulary size is approximately:

$$|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.

Why pre-tokenization matters

Before counting byte pairs, the corpus is split into pre-tokens. This has two purposes:

1. It reduces the amount of pair counting that needs to be performed.
2. It prevents arbitrary merges across word-like or document boundaries.

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.

Encoding and decoding

After BPE training, the tokenizer stores a vocabulary and an ordered list of merges. Encoding proceeds as follows:

1. Pre-tokenize the input text.
2. Convert each pre-token into UTF-8 bytes.
3. Apply the learned merges in their training order.
4. Map the resulting byte strings to integer token IDs.

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.

3. The Overall Transformer Language Model

Given token IDs,

$$X\in\mathbb{N}^{B\times T}$$

the token embedding maps them to dense vectors:

$$H\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\in\mathbb{R}^{B\times T\times |V|}$$

The full model can be summarized as:

$$\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.

4. Self-Attention: How Does the Model Aggregate Context?

Query, Key, and Value

For hidden states $X$, attention applies three learned linear transformations:

$$Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V$$

An intuitive interpretation is:

• Query: what information is the current position looking for?
• Key: what information does each position offer?
• Value: what content is actually aggregated?

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.

Scaled dot-product attention

The basic attention operation is:

$$\operatorname{Attention}(Q,K,V) =\operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_{head}}}+M\right)V$$

Why divide by $\sqrt{d_{head}}$? If the components of $Q$ and $K$ have variance close to 1, the variance of their dot product grows with $d_{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.

Causal masking

An autoregressive language model cannot look at future tokens. The causal mask is:

$$M_{ij}=\begin{cases} 0,&j\leq i\\ -\infty,&j>i \end{cases}$$

After adding this mask, future positions receive zero probability after softmax. Position $i$ can attend only to $x_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

Multi-head attention splits the model dimension into $h$ smaller heads:

$$d_{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.

5. RoPE: Where Does Positional Information Come From?

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(\theta)= \begin{bmatrix} \cos\theta&-\sin\theta\\ \sin\theta&\cos\theta \end{bmatrix}$$

For positions $m$ and $n$:

$$q_m'=R_mq_m,\qquad k_n'=R_nk_n$$

The rotated inner product is:

$$(R_mq)^T(R_nk)=q^TR_m^TR_nk$$

Since rotation matrices satisfy $R_m^TR_n=R_{n-m}$, the attention score depends on the relative distance $n-m$.

This is the central intuition behind RoPE: by rotating queries and keys, their dot products naturally encode relative position.

RoPE is applied to $Q$ and $K$, but not to $V$. It changes how positions are matched; it does not alter the content carried by the values.

6. RMSNorm and Pre-Norm Transformers

RMSNorm

LayerNorm subtracts the mean and divides by the standard deviation. RMSNorm only normalizes the root mean square.

For a hidden vector $a\in\mathbb{R}^d$:

$$\operatorname{RMS}(a)=\sqrt{\frac{1}{d}\sum_{i=1}^{d}a_i^2+\epsilon}$$

$$\operatorname{RMSNorm}(a)=\frac{a}{\operatorname{RMS}(a)}\odot g$$

Here $g$ 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.

Pre-norm block

The pre-norm Transformer block used in the assignment is:

$$z=x+\operatorname{Attention}(\operatorname{RMSNorm}(x))$$

$$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=\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.

7. Feed-Forward Networks and SwiGLU

Attention mixes information across sequence positions. The feed-forward network applies the same nonlinear transformation independently to each position:

$$\operatorname{FFN}(x)=W_2\,\sigma(W_1x)$$

Modern language models often use the gated SwiGLU variant:

$$\operatorname{SwiGLU}(x) =W_2\left(\operatorname{SiLU}(W_1x)\odot W_3x\right)$$

where:

$$\operatorname{SiLU}(x)=x\cdot\sigma(x)$$

$W_3x$ is the gate branch. It is multiplied element-wise with $\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:

$$d_{ff}\approx\frac{8}{3}d_{model}$$

The TinyStories configuration in the assignment uses $d_{model}=512$ and $d_{ff}=1344$. The latter is close to $8/3\times512$ and is also divisible by 64, which is convenient for GPU hardware.

8. Cross-Entropy: Why Does the Loss Have This Form?

From logits to probabilities

At one sequence position, the model produces vocabulary logits:

$$z=(z_1,z_2,\ldots,z_{|V|})$$

Softmax converts them into probabilities:

$$p_i=\frac{e^{z_i}}{\sum_j e^{z_j}}$$

If the correct token has index $y$, the cross-entropy loss is:

$$\mathcal{L}=-\log p_y$$

Substituting the softmax gives:

$$\mathcal{L} =-z_y+\log\sum_j e^{z_j}$$

The loss therefore has two effects: it rewards increasing the correct logit $z_y$, while the log-sum-exp term accounts for competition with every vocabulary item.

The gradient is prediction minus target

Let $p$ be the softmax probability vector and let $e_y$ be the one-hot vector for the correct class. Then:

$$\frac{\partial\mathcal{L}}{\partial z}=p-e_y$$

This is the most important result for softmax cross-entropy.

• For the correct class, the gradient is $p_y-1$, usually negative, so gradient descent increases its logit.
• For an incorrect class, the gradient is $p_i$, so gradient descent decreases its logit.

If the model already assigns almost all probability to the correct class, $p\approx e_y$ and the gradient is close to zero. If it is highly confident but wrong, the gradient is large.

Numerical stability

Directly computing $e^{z_i}$ can overflow. Let $m=\max_i z_i$. Then:

$$\log\sum_i e^{z_i} =m+\log\sum_i e^{z_i-m}$$

Because $z_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.

Perplexity

If the average per-token loss is $\mathcal{L}$, perplexity is defined as:

$$\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.

9. Optimizers: How Do Parameters Learn?

SGD

The basic gradient-descent update is:

$$\theta_{t+1}=\theta_t-\alpha g_t$$

where $\alpha$ is the learning rate and $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 and its two moments

AdamW maintains an exponential moving average of the gradient and of the squared gradient:

$$m_t=\beta_1m_{t-1}+(1-\beta_1)g_t$$

$$v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2$$

$m_t$ tracks the direction of the gradient, while $v_t$ tracks its scale. Since both states start at zero, the early estimates are biased toward zero. Adam corrects this bias using:

$$\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:

$$\theta_t\leftarrow \theta_t-\alpha\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}$$

AdamW adds decoupled weight decay:

$$\theta_t\leftarrow \theta_t-\alpha\lambda\theta_t$$

Combining the two terms gives:

$$\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 $\beta_1=0.9$, $\beta_2=0.95$, and $\epsilon=10^{-8}$, although the best values depend on the model and the data.

Cosine learning-rate schedule

Transformer training commonly uses a learning-rate schedule with a warmup phase followed by cosine decay.

During warmup:

$$\alpha_t=\frac{t}{T_w}\alpha_{max},\qquad t<T_w$$

During cosine decay:

$$\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 $T_w$ is the end of warmup and $T_c$ is the end of cosine decay. After $T_c$, the learning rate remains at $\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.

Gradient clipping

If a batch produces an unusually large gradient, we can scale the entire gradient vector so its L2 norm does not exceed $M$:

$$g\leftarrow g\cdot\min\left(1,\frac{M}{\lVert g\rVert_2+\epsilon}\right)$$

When $\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.

10. Training Compute and Resource Intuition

Parameter count

Ignoring biases, the main parameters in one Transformer block come from attention and the feed-forward network.

The $Q$, $K$, $V$, and output projections together contribute approximately:

$$4d_{model}^2$$

The three SwiGLU matrices contribute approximately:

$$3d_{model}d_{ff}$$

For a model with $L$ layers, a rough parameter-count formula is:

$$P\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.

Attention as a scaling bottleneck

Self-attention computes $QK^T$, whose sequence dimension has quadratic complexity in context length. The main compute of a Transformer block can be summarized as:

$$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 $T^2$ attention term becomes increasingly important.

Why training uses more memory than inference

If parameters, gradients, first moments, and second moments are all stored in float32, each parameter requires roughly:

$$4+4+4+4=16\text{ bytes}$$

This excludes activation memory. Training therefore has several distinct memory costs:

• model parameters;
• gradients;
• optimizer states;
• intermediate activations needed for backpropagation.

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.

11. Inference: How Does the Model Generate Text?

Autoregressive decoding

Given a prefix $x_{1:t}$, the Transformer produces logits at every position. To generate the next token, we use only the final position:

$$v=\operatorname{TransformerLM}(x_{1:t})_t$$

The logits are converted into a distribution:

$$p(x_{t+1}=i\mid x_{1:t}) =\frac{e^{v_i}}{\sum_j e^{v_j}}$$

We sample $x_{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 and sampling

Greedy decoding selects the most likely token at every step:

$$x_{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

Temperature $\tau$ rescales the logits before softmax:

$$p_i=\operatorname{softmax}(v/\tau)_i$$

• $\tau<1$: a sharper distribution, closer to greedy decoding;
• $\tau=1$: the original model distribution;
• $\tau>1$: a flatter distribution with more randomness.

Temperature does not change the model parameters. It changes only the sampling distribution used during inference.

Top-p, or nucleus sampling

Suppose the probabilities are sorted as $q_1\geq q_2\geq\cdots$. Top-p sampling chooses the smallest candidate set $V(p)$ whose cumulative probability reaches $p$:

$$\sum_{i\in V(p)}q_i\geq p$$

It then renormalizes and samples only from this set:

$$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.

12. How Should We Read the Assignment Experiments?

TinyStories versus OpenWebText

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 and batch size

Learning-rate experiments typically reveal three regimes:

1. Too small: stable but slow improvement.
2. Appropriate: fast and stable loss reduction.
3. Too large: oscillation or divergence.

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.

What do the ablations tell us?

The assignment asks us to compare several architectural variants:

1. Remove RMSNorm and observe whether training becomes unstable.
2. Replace pre-norm with post-norm.
3. Remove RoPE and test whether the model can infer position from causal attention alone.
4. Replace SwiGLU with an ungated SiLU FFN while approximately matching parameter count.

An ablation should not be judged only by its final loss. We should also inspect:

• the speed of learning-curve descent;
• whether the run diverges;
• activation, gradient, and parameter norms;
• the coherence and repetition of generated text.

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.

13. Common Confusions

Logits, probabilities, and loss

• Logits are arbitrary real-valued outputs of the final linear layer.
• Softmax converts logits into a probability distribution over the vocabulary.
• Cross-entropy uses the probability assigned to the correct token.
• In practice, training usually computes cross-entropy directly from logits without explicitly materializing probabilities.

Training and inference

• Training uses teacher forcing and can process all positions in parallel.
• Inference is autoregressive: each generated token depends on the previous generated context.
• Training needs gradients, saved activations, and optimizer states.
• Inference needs only forward computation and a decoding strategy.

RoPE and causal masking

• RoPE provides positional information.
• The causal mask restricts which positions are visible.
• RoPE answers “where are these positions?”
• The causal mask answers “is the future allowed to be seen?”

RMSNorm and residual connections

• RMSNorm controls the scale of activations.
• Residual connections provide a direct information and gradient path.
• Pre-norm places normalization before attention or the FFN.
• Post-norm places normalization after residual addition.

Conclusion

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)=\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.

References

1. Stanford CS336, Assignment 1: Building a Transformer LM.
2. Vaswani et al., Attention Is All You Need, 2017.
3. Sennrich, Haddow, and Birch, Neural Machine Translation of Rare Words with Subword Units, 2016.
4. Zhang and Sennrich, Root Mean Square Layer Normalization, 2019.
5. Loshchilov and Hutter, Decoupled Weight Decay Regularization, 2019.
6. CS336 Assignment 1 repository