13 · Map

End-to-end map and catalogue

The stacked forward equation, the train path, the infer path, Jacobians, and what each piece is for.

Stacked forward map

x₁:Ttokens

H⁽⁰⁾embed + RoPE

H⁽ⁿ⁾n blocks

Zlogits

p₁:Tsoftmax

Forward map. Training attaches CE → Jacobians → AdamW after p. Inference samples from p.

x_{1 : T} H^{(ℓ - 1)} H^{(n)} W_{E}, pos/RoPE H^{(0)} Norm, Attn, + Z^{(ℓ)} Norm, MLP, + H^{(ℓ)} Norm, W_{U} Z softmax p_{1 : T}

(13.1)

Train path

p_{1 : T}, x_{2 : T + 1} CE L reverse-mode AD / Jacobians \nabla_{θ} L AdamW θ_{new}

(13.2)

Infer path

z_{t} τ, top- k / p p ~_{t} sample x_{t + 1}

(13.3)

Jacobian catalogue

Reverse-mode AD never materializes most of these as matrices. It computes the vector-Jacobian products. Written out, they are the whole backward pass.

Map	Jacobian / VJP
softmax	diag(a) − aaᵀ
CE + softmax	p − y
Y = XW	∂L/∂X = G Wᵀ, ∂L/∂W = Xᵀ G
RMSNorm	r⁻¹ (I − uuᵀ / (d r²))
LayerNorm	same, after P = I − 11ᵀ/d
residual H ↦ H + F(H)	I + J_Fᵀ
RoPE	Rᵀ = R⁻¹
SiLU	σ(z) + z σ(z)(1 − σ(z))
embedding lookup	scatter into row x_t

What each piece is for

The whole machine is one differentiable map from token IDs to next-token distributions. Training differentiates that map and walks θ downhill on cross-entropy. Use-time throws the derivative away and only evaluates the map.

Piece	Job
Tokenizer / BPE	Discrete interface to text. Frozen. Defines the units of loss.
Embedding + position	Put tokens into a space where order exists
Residual stream	Common currency every layer reads and writes
Norm	Control scale so depth does not explode
Causal attention	Move information from past tokens to the current one
√dₕ scale	Keep scores O(1) so softmax does not saturate
GQA / MQA	Share KV heads so the cache, not the matmul, stays small
MLP / SwiGLU	Per-token computation and stored features
Unembed + softmax	Turn vectors into a distribution over V
Cross-entropy = MLE	Scalar “how wrong was the next-token guess”
Entropy + KL	Floor versus avoidable error
Softmax–CE gradient p − y	Error signal into logits
Jacobians / backprop	Route that scalar error to every weight
Skip-connection Jacobian I	Make deep nets trainable
AdamW + schedule + clip	Stable updates at billion-parameter scale
C ≈ 6ND / Chinchilla E	Compute unit and the fitted floor
Validation loss	Measure generalization; zero is not the target
Sampling + KV cache	Use the trained pθ without changing it