Continuous Function

Attention & Transformers

Tiny LM Training Loop

Trace one tiny causal-LM training step as a state transition: which values are read, which are derived, and which must actually change if learning happened.

status: reviewimportance: criticaldifficulty 4/5math: undergraduateread: 14mlive demo
Back to Attention & TransformersNext: Training Loss Curves and Debugging (review)

Concept Structure

Tiny LM Training Loop

01Intuition

Start with the picture, metaphor, or geometric mechanism.

02Math

Make the objects explicit and connect them with notation.

03Code

Mirror the equations with runnable implementation details.

04Interactive Demo

Manipulate the mechanism and watch the idea respond.

5prerequisites
3next concepts
2related links

Learner Contract

What this page should let you do.

You are here becauseTrace one tiny causal-LM training step as a state transition: which values are read, which are derived, and which must actually change if learning happened.

This Attention & Transformers concept is the current idea: keep the same invariant visible across Intuition, Math, Code, Interactive Demo.

By the end4/4 sections ready | runnable code expected | live demo

Explain the mechanism, trace the main notation, and test one prediction in the live demo.

Do this firstIntuition

Read the intuition before the notation; the math should name a mechanism you already felt.

Test the linkManipulate one control and predict the visible change.Then continue to Training Loss Curves and Debugging (review)

Claim/source review status

Substantive review recorded

1/1 claims have bounded review metadata; still check caveats and source scope.Metadata-derived; review may be AI-assisted. Not a human certification.
Claims1/1 reviewed
Sources4 cited
Codeattached
Demolive
Reviewed2026-07-02
Updatedpage 2026-07-02

Learning item flow

4/4 sections readyAsk about thisResearch room
ConceptTiny LM Training LoopAttention & TransformersEquationTiny LM Training Loop equation 1Exact equationCodeTiny LM Training Loop runnable code 1Exact runnable codeDemoTiny LM Training Loop interactive demoInteractive demoClaimA tiny causal-LM training step reads token blocks and shifted labels,...Exact claim checkSourceStanford CS336 Assignment 1: BasicsExact source
ConceptTiny LM Training LoopAttention & Transformers
4 sources attachedLocal snapshot ready
concept:attention-transformers/tiny-lm-training-loop
CodenearbyPredictbefore revealRoomitem handoff
01IntuitionStart with the picture, metaphor, or geometric mechanism.02MathMake the objects explicit and connect them with notation.03CodeMirror the equations with runnable implementation details.04Interactive DemoManipulate the mechanism and watch the idea respond.
01

01

Intuition

Build the mental picture first so the rest of the page has something to attach to.

Section prompt

A tiny language-model training loop is not mainly a graph of boxes. It is a state transition.

Some objects are read: the token stream, the sampled input block, and the shifted labels. Some objects are derived for this step: logits and loss. Some objects must mutate if learning actually happened: gradient buffers, model parameters, and optimizer state. Validation should read the current model without adding training gradients. A checkpoint serializes enough state to make the run explainable later.

That separation matters because a loss number can fool you. You can compute a loss forever without updating the model. You can update parameters but accidentally validate in training mode. You can save weights but silently reset optimizer moments. The tiny loop teaches the contract before scale makes the bug expensive.

This page is deliberately modest. The runnable model is a tiny embedding-plus-linear LM, not a benchmark Transformer. You can replace the model with the decoder-only forward pass from the previous page; the loop contract stays the same.

02

02

Math

Translate the story into symbols, assumptions, and a derivation you can inspect.

Section prompt
Equation 1s0,s1,…,sN−1,si∈{0,…,V−1}.s_0, s_1, \ldots, s_{N-1}, \qquad s_i \in \{0,\ldots,V-1\}.Equation 2X[b,t]=sab+t,Y[b,t]=sab+t+1,X[b,t] = s_{a_b+t}, \qquad Y[b,t] = s_{a_b+t+1},

Let a tokenized corpus be a sequence

s0,s1,…,sN−1,si∈{0,…,V−1}.s_0, s_1, \ldots, s_{N-1}, \qquad s_i \in \{0,\ldots,V-1\}.

For batch size B and context length T, sample start indices a_b. The causal-LM batch is

X[b,t]=sab+t,Y[b,t]=sab+t+1,X[b,t] = s_{a_b+t}, \qquad Y[b,t] = s_{a_b+t+1},

so both X and Y have shape (B,T). The model produces logits

Z=fθ(X),Z∈RB×T×V.Z = f_\theta(X), \qquad Z \in \mathbb{R}^{B \times T \times V}.

Flattening the batch and time axes gives BT classification problems over the vocabulary. The training loss is

L(θ)=1BT∑b=1B∑t=1T−log⁡exp⁡Z[b,t,Y[b,t]]∑v=1Vexp⁡Z[b,t,v].L(\theta) = \frac{1}{BT} \sum_{b=1}^{B}\sum_{t=1}^{T} -\log \frac{\exp Z[b,t,Y[b,t]]}{\sum_{v=1}^{V}\exp Z[b,t,v]}.

The backward pass computes

gt=∇θL(θt).g_t = \nabla_\theta L(\theta_t).

An optimizer step is a state transition:

(θt+1,ut+1)=OptStep⁡(θt,ut,gt),(\theta_{t+1}, u_{t+1}) = \operatorname{OptStep}(\theta_t, u_t, g_t),

where u_t is optimizer state such as momentum or Adam moments. A correct training step mutates gradients, parameters, and optimizer state. It does not mutate the dataset or validation set.

Validation estimates a separate quantity such as

L^val(θt+1)=1∣V∣∑(X,Y)∈VCE⁡(fθt+1(X),Y),\hat{L}_{\mathrm{val}}(\theta_{t+1}) = \frac{1}{|\mathcal V|} \sum_{(X,Y)\in\mathcal V} \operatorname{CE}(f_{\theta_{t+1}}(X),Y),

under evaluation behavior and without gradient construction.

A resumable checkpoint is therefore not only theta. At minimum for this toy loop it records

Ct=(θt, ut, t, c, mt),C_t = (\theta_t,\ u_t,\ t,\ c,\ m_t),

where c is the run config and m_t is the metric record.

03

03

Code

Keep the implementation aligned with the notation so the algorithm is legible.

Section prompt
Runnable code 1import numpy as np rng = np.random.default_rng(0) ids = np.array([0, 1, 2, 3, 0, 4, 2, 3, 0,...python
import numpy as np

rng = np.random.default_rng(0)
ids = np.array([0, 1, 2, 3, 0, 4, 2, 3, 0, 5, 2, 3, 0, 6, 2, 3])
V, T, B, d, lr = 8, 4, 2, 16, 0.3
E, W = rng.normal(size=(V, d)) * 0.1, rng.normal(size=(d, V)) * 0.1
mE, mW = np.zeros_like(E), np.zeros_like(W)

def get_batch(split):
    data = ids[:10] if split == "train" else ids[10:]
    starts = rng.integers(0, len(data) - T, size=B)
    x = np.stack([data[i:i + T] for i in starts])
    y = np.stack([data[i + 1:i + T + 1] for i in starts])
    return x, y

def lm_step(split, train):
    global E, W, mE, mW
    x, y = get_batch(split)
    h = E[x]                         # (B, T, d)
    logits = h @ W                   # (B, T, V)
    z = logits.reshape(-1, V)
    targets = y.reshape(-1)
    p = np.exp(z - z.max(axis=1, keepdims=True))
    p /= p.sum(axis=1, keepdims=True)
    loss = -np.log(p[np.arange(len(targets)), targets]).mean()
    if train:
        dz = p
        dz[np.arange(len(targets)), targets] -= 1
        dz /= len(targets)
        dlogits = dz.reshape(B, T, V)
        dW = np.einsum("btd,btv->dv", h, dlogits)
        dE = np.zeros_like(E)
        np.add.at(dE, x, dlogits @ W.T)
        mE[:], mW[:] = 0.9 * mE + dE, 0.9 * mW + dW
        E[:], W[:] = E - lr * mE, W - lr * mW
    return loss

for step in range(30):
    train_loss = lm_step("train", train=True)
    if step % 10 == 0:
        print(step, round(train_loss, 3), round(lm_step("val", train=False), 3))

ckpt = {"step": step, "params": (E.copy(), W.copy()),
        "optimizer": (mE.copy(), mW.copy()), "config": {"B": B, "T": T, "V": V}}
04

04

Interactive Demo

Use direct manipulation to connect the explanation to a moving system.

Section prompt

The lab below treats one training step as a state ledger.

Before reveal, predict which state must mutate during one correct training step. After reveal, inspect what was only read, what was derived temporarily, what changed because learning happened, and what was serialized for resume.

The key invariant is: data is read, logits and loss are derived, gradients/parameters/optimizer state mutate, validation reads, and checkpoints serialize.

Live Concept Demo

Explore Tiny LM Training Loop

The stage is code-native and interactive. Use it to test the explanation against the mechanism.

difficulty 4/5undergraduatecode-aligned
Demo Prediction Checkpoint

Manipulate one control and predict the visible change.

Commit to what Tiny LM Training Loop should make visible before reading the result.

After The First Pass

Turn the concept into an inspected object.

Once the invariant is visible in the intuition, math, code, and demo, use these panels to inspect the mechanism visually, check source support, practice the idea, and attach a grounded research question.

Mechanism Storyboard

See the idea move before the page explains it

Trace one tiny causal-LM training step as a state transition: which values are read, which are derived, and which must actually change if learning happened.

Prediction open01 / Intuition
Prediction lens

Start with the picture, metaphor, or geometric mechanism.

Commit first

Before reading further, choose the kind of change Tiny LM Training Loop should make visible.

Visual Inquiry

Make the image answer a mathematical question

Trace one tiny causal-LM training step as a state transition: which values are read, which are derived, and which must actually change if learning happened.

4/4 stages readyLive demo connected
Prediction

Which visible object should carry the first intuition?

Commit first

Pick the cue that should make Tiny LM Training Loop easier to reason about before the page gives the answer.

Source Grounding

Canonical references for the mechanism on this page.

course-notes · 2025Stanford CS336 Assignment 1: BasicsStanford CS336 course staff

Source for language-model-from-scratch implementation contracts: sampled input/label blocks, Transformer LM logits, cross-entropy, gradient clipping, and minimal LM training ambition.

Open source
book · 2026Dive into Deep Learning: Language ModelsZhang, Lipton, Li, and Smola

Source for sampling input and shifted target sequences in minibatches, next-token prediction, cross-entropy averaged over sequence tokens, and perplexity as a language-model quality measure.

Open source
documentation · 2026PyTorch Tutorials: Optimizing Model ParametersPyTorch

Source for train and test loop structure, zero_grad, loss.backward, optimizer.step, model.train(), model.eval(), and torch.no_grad().

Open source
documentation · 2026PyTorch Tutorials: Saving and Loading ModelsPyTorch

Source for model and optimizer state_dicts and checkpoint dictionaries for resuming training.

Open source

Claim Review

Trace one tiny causal-LM training step as a state transition: which values are read, which are derived, and which must actually change if learning happened.

Status1 substantive review recorded

Claims without a substantive review badge still need exact source-support review.

Sources4 references

cs336-assignment1-basics, d2l-language-modeling, pytorch-optimization-loop, pytorch-saving-loading

Local checks4 local checks

Use equations, runnable code, and demos to check whether the source support is operational.

Substantively reviewedA tiny causal-LM training step reads token blocks and shifted labels, derives logits/loss, mutates gradients, parameters, and optimizer state, validates without gradients, and checkpoints resumable state.Claim metadata: source checked

CS336 and D2L support LM input/label block and logits/loss contracts; PyTorch supports zero/backward/step, eval/no_grad validation, and checkpoint state_dict mechanics.

Sources: Stanford CS336 Assignment 1: Basics, Dive into Deep Learning: Language Models, PyTorch Tutorials: Optimizing Model Parameters, PyTorch Tutorials: Saving and Loading ModelsThis is a local state-transition witness over a toy corpus. It does not claim benchmark quality, full Transformer training, tokenizer quality, distributed training, exact reproducibility across hardware, or production experiment tracking.A bounded review summary is present; still check caveats and exact reference scope.

Checked CS336 course/Assignment 1 adapters, D2L language-model sequence partitioning/perplexity material, PyTorch optimization-loop docs, and PyTorch checkpoint docs. The page remains review-status pending GPT Pro publication critique.

Reviewer: codex-local-source-audit; reviewed 2026-07-02
source-span-cs336-assignment1-basicssource-span-d2l-language-modelingsource-span-pytorch-optimization-loopsource-span-pytorch-saving-loadingEquation 1Equation 2Runnable code 1Interactive demo

Source support candidates

course-notes 2025Stanford CS336 Assignment 1: Basics

Source for language-model-from-scratch implementation contracts: sampled input/label blocks, Transformer LM logits, cross-entropy, gradient clipping, and minimal LM training ambition.

book 2026Dive into Deep Learning: Language Models

Source for sampling input and shifted target sequences in minibatches, next-token prediction, cross-entropy averaged over sequence tokens, and perplexity as a language-model quality measure.

documentation 2026PyTorch Tutorials: Optimizing Model Parameters

Source for train and test loop structure, zero_grad, loss.backward, optimizer.step, model.train(), model.eval(), and torch.no_grad().

documentation 2026PyTorch Tutorials: Saving and Loading Models

Source for model and optimizer state_dicts and checkpoint dictionaries for resuming training.

Mechanism checks

Equation 1
s0,s1,…,sN−1,si∈{0,…,V−1}.s_0, s_1, \ldots, s_{N-1}, \qquad s_i \in \{0,\ldots,V-1\}.
Equation 2
X[b,t]=sab+t,Y[b,t]=sab+t+1,X[b,t] = s_{a_b+t}, \qquad Y[b,t] = s_{a_b+t+1},
Runnable code 1import numpy as np rng = np.random.default_rng(0) ids = np.array([0, 1, 2, 3, 0, 4, 2, 3, 0,...Demo stateLive mechanism probe

Practice Loop

Try the idea before it explains itself

Trace one tiny causal-LM training step as a state transition: which values are read, which are derived, and which must actually change if learning happened.

Readiness0/3 checks ready
Predict

Before touching the demo, predict one visible change that should happen in Tiny LM Training Loop.

Hint 1

Reveal when your model needs a nudge.

Hint 2

Reveal when your model needs a nudge.

Hint 3

Reveal when your model needs a nudge.

Grounded research drawerClose
ConceptTiny LM Training LoopAttention & Transformers
Runnable code comparisonTiny LM Training Loop runnable code 1dW = np.einsum("btd,btv->dv", h, dlogits)Prediction before revealTiny LM Training Loop interactive demoManipulate one control and predict the visible change.
Grounded room questionWhat is the smallest example that makes Tiny LM Training Loop click without losing the math?Local snapshot ready

Research Room

Attach the question to a claim, equation, code, or demo

Pick the concept, equation, source, runnable code, claim, misconception, or demo state before asking for help. The handoff keeps that page item in context.
Next local actionNo local draft saved yet

Open the draft below to save one note and next action in this browser.

conceptAttention & Transformers

Tiny LM Training Loop

Attached question

What is the smallest example that makes Tiny LM Training Loop click without losing the math?

Local action draftNo local draft saved yetExpand only when ready to capture one local next action
Local action draft

This draft stays in this browser, attached to the selected learning item.

No local draft saved.
Evidence to inspect
  • References to inspect: attached references on this page.
  • Definition, prerequisite, and contrast concept links
  • The equation or runnable code that makes the concept operational
  • One demo state that shows the invariant instead of a slogan
What would resolve this
  • The learner can state the mechanism in their own words
  • The learner can name the prerequisite that would repair confusion
  • The learner can predict how the mechanism changes under one perturbation
Grounded AI handoff

I am working in Continuous Function's research reading room. Object: concept - Tiny LM Training Loop Selected item key: recorded for copy. Context: Attention & Transformers Page anchor: recorded for copy. Open question: What is the smallest example that makes Tiny LM Training Loop click without losing the math? Evidence to inspect: - References to inspect: attached references on this page. - Definition, prerequisite, and contrast concept links - The equation or runnable code that makes the concept operational - One demo state that shows the invariant instead of a slogan What would resolve this: - The learner can state the mechanism in their own words - The learner can name the prerequisite that would repair confusion - The learner can predict how the mechanism changes under one perturbation Answer as a careful research tutor: stay source-grounded, separate verified evidence from assumptions, name the relevant math objects, and end with one next action.

View it in context
concept/concept-notebook/attention-transformers/tiny-lm-training-loop concept:attention-transformers/tiny-lm-training-loop