Watch a transformer learn 2-digit addition from scratch in your browser. Type two numbers and see it predict the sum digit by digit. Built with tensorgrad (autograd + WebGPU).
---
format: typebulb/v1
name: Transformer
---
**code.tsx**
```tsx
import presetWind3 from '@unocss/preset-wind3'
import { App, Component, div, h1, button, span, p, a, svg, g, rect, path, line, text, circle, animate, formField, inputRange, type VElement } from 'domeleon'
import { UnoThemeManager, type ThemeProxy } from 'domeleon/unocss'
import { inputNumber } from 'domeleon/maskito'
import {
Module, compile, checkWebGPU, lr, Linear, LayerNorm, crossEntropy, capture, singleFlight,
add, mul, sum, swapAxes,
relu, matmul, embedding, arange,
softmaxCausal, splitHeads, mergeHeads,
type Tensor,
type CompiledTraining, type CompiledForward,
} from 'tensorgrad'
// ========== MODEL / TRAINING ==========
// ---------- Constants ----------
// Vocab: digits 0..9 are tokens 0..9; '+' = 10, '=' = 11.
const VOCAB = 12
const TOK_PLUS = 10
const TOK_EQ = 11
const N_OP_DIGITS = 2
const N_RESULT_DIGITS = 3
// Sequence layout: [d d + d d = r r r], result digits reversed (LSB first) so
// the model can do left-to-right carry.
const SEQ_LEN = N_OP_DIGITS + 1 + N_OP_DIGITS + 1 + N_RESULT_DIGITS // 9
const RESULT_START = N_OP_DIGITS + 1 + N_OP_DIGITS + 1 // 6
const D_MODEL = 64
const N_LAYERS = 3
const N_HEADS = 4
const D_HEAD = D_MODEL / N_HEADS
const BATCH_SIZE = 128
// Training steps per UI yield. The loop yields to requestAnimationFrame once per
// BURST, not once per step — so the step rate isn't pinned to the ~60 Hz frame
// clock (which capped throughput at ~60·BATCH_SIZE ≈ 7.7k ex/s no matter how
// fast the GPU finished a step, leaving it idle most of each frame). Bursting
// lets training run at GPU speed between frames while the UI still refreshes
// every frame. Tune down if the page feels less responsive.
const STEPS_PER_YIELD = 16
const T_LEN = SEQ_LEN - 1
const LR_SCHEDULE = lr.linear({ peak: 0.005, final: 0.0005, steps: 1500 })
const range = (n: number) => Array.from({ length: n }, (_, i) => i)
const clamp = (n: number, min: number, max: number) => Math.min(Math.max(n, min), max)
// 6-token addition prompt: [a_tens, a_ones, +, b_tens, b_ones, =].
const promptTokens = (a: number, b: number): number[] => [
Math.floor(a / 10), a % 10, TOK_PLUS,
Math.floor(b / 10), b % 10, TOK_EQ,
]
// Deterministic train/test split by (a, b). ~20% held out.
const isTestPair = (a: number, b: number) => ((a * 100 + b) * 31 + 7) % 5 === 0
const N_TEST = 50
// ---------- Model hierarchy ----------
class Attention extends Module {
q = new Linear(D_MODEL, D_MODEL, { bias: false })
k = new Linear(D_MODEL, D_MODEL, { bias: false })
v = new Linear(D_MODEL, D_MODEL, { bias: false })
o = new Linear(D_MODEL, D_MODEL, { bias: false })
}
class MLP extends Module {
up = new Linear(D_MODEL, 4 * D_MODEL)
down = new Linear(4 * D_MODEL, D_MODEL)
}
class Block extends Module {
ln1 = new LayerNorm(D_MODEL)
attn = new Attention()
ln2 = new LayerNorm(D_MODEL)
mlp = new MLP()
}
class Transformer extends Module {
tok_emb: Tensor
pos_emb: Tensor
layers: Block[]
lnf = new LayerNorm(D_MODEL)
constructor() {
super()
this.tok_emb = this.param([VOCAB, D_MODEL])
this.pos_emb = this.param([SEQ_LEN, D_MODEL])
this.layers = range(N_LAYERS).map(() => new Block())
}
}
// ---------- Forward pass + loss ----------
// Capture sites (`capture(name, t)`) make intermediates readable post-run via
// `result.captures.get(name)`. Tensorgrad always reads captures back alongside
// the loss, so leaving viz captures in the training graph would mean ~7 MB
// GPU→main per step. `cap` gates them by trace mode: off during the training
// compile, on for the inference compile.
let captureMode: 'on' | 'off' = 'on'
const cap = (name: string, t: Tensor): Tensor => captureMode === 'on' ? capture(name, t) : t
function attentionFwd(p: Attention, x: Tensor, layerIdx: number): Tensor {
const q = cap(`q.${layerIdx}`, splitHeads(p.q.fwd(x), N_HEADS))
const k = cap(`k.${layerIdx}`, splitHeads(p.k.fwd(x), N_HEADS))
const v = cap(`v.${layerIdx}`, splitHeads(p.v.fwd(x), N_HEADS))
const scores = mul(matmul(q, swapAxes(k, -1, -2)), 1 / Math.sqrt(D_HEAD))
const attn = cap(`attn.${layerIdx}`, softmaxCausal(scores))
return p.o.fwd(mergeHeads(matmul(attn, v)))
}
// MLP is inlined so the post-ReLU hidden state can be captured.
function blockFwd(p: Block, x: Tensor, layerIdx: number): Tensor {
const x1 = add(x, attentionFwd(p.attn, p.ln1.fwd(x), layerIdx))
const xn2 = p.ln2.fwd(x1)
const h1 = cap(`mlp_hidden.${layerIdx}`, relu(p.mlp.up.fwd(xn2)))
return add(x1, p.mlp.down.fwd(h1))
}
function modelFwd(p: Transformer, tokens: Tensor): Tensor {
const tokE = embedding(p.tok_emb, tokens) // [B, T, D]
const posE = embedding(p.pos_emb, arange(T_LEN)) // [T, D] — broadcasts over batch
let x = add(tokE, posE)
for (let i = 0; i < p.layers.length; i++) {
cap(`residual.${i}`, x)
x = blockFwd(p.layers[i]!, x, i)
}
cap(`residual.${N_LAYERS}`, x)
const xnf = p.lnf.fwd(x)
// Tied embedding head: logits = xnf @ tok_emb^T.
return matmul(xnf, swapAxes(p.tok_emb, -1, -2))
}
// `mask` is float32 [T] with 1s on result-digit positions and 0s elsewhere —
// loss is only computed over the digits the model has to predict.
function lossFn(
p: Transformer,
{ tokens, targets, mask }: { tokens: Tensor; targets: Tensor; mask: Tensor },
): Tensor {
const ce = crossEntropy(modelFwd(p, tokens), targets, { reduction: 'none' }) // [B, T] of -log p(target)
return mul(sum(mul(ce, mask)), 1 / (BATCH_SIZE * N_RESULT_DIGITS))
}
// Inference forward extends modelFwd with logit-lens captures: at every
// residual depth, project the residual through the final LN + tied unembed.
// The final-layer logit-lens IS the model's prediction, so it's the output.
function predictFn(p: Transformer, { tokens }: { tokens: Tensor }): Tensor {
const tokE = embedding(p.tok_emb, tokens)
const posE = embedding(p.pos_emb, arange(T_LEN))
const tokEmbT = swapAxes(p.tok_emb, -1, -2)
const lensAt = (r: Tensor) => matmul(p.lnf.fwd(r), tokEmbT)
let x = add(tokE, posE)
for (let i = 0; i < p.layers.length; i++) {
capture(`residual.${i}`, x)
capture(`logit_lens.${i}`, lensAt(x))
x = blockFwd(p.layers[i]!, x, i)
}
capture(`residual.${N_LAYERS}`, x)
return capture(`logit_lens.${N_LAYERS}`, lensAt(x))
}
// Loss masks out non-result positions.
const RESULT_MASK = Float32Array.from({ length: T_LEN }, (_, t) => t >= RESULT_START - 1 ? 1 : 0)
// Excludes test pairs. Encodes with reversed result digits.
function makeBatch(): { tokens: Int32Array; targets: Int32Array } {
const tokens = new Int32Array(BATCH_SIZE * T_LEN)
const targets = new Int32Array(BATCH_SIZE * T_LEN)
for (let bi = 0; bi < BATCH_SIZE; bi++) {
let a, c
do { a = Math.floor(Math.random() * 100); c = Math.floor(Math.random() * 100) } while (isTestPair(a, c))
const sum = a + c
const seq = [
...promptTokens(a, c),
sum % 10, Math.floor(sum / 10) % 10, Math.floor(sum / 100),
]
for (let t = 0; t < T_LEN; t++) {
tokens[bi * T_LEN + t] = seq[t]!
targets[bi * T_LEN + t] = seq[t + 1]!
}
}
return { tokens, targets }
}
// ---------- Inference: forward-only graph at B=1 ----------
const HIDDEN = 4 * D_MODEL
// Stable softmax over a slice of a flat array.
function softmaxNumeric(arr: ArrayLike<number>, offset = 0, len = arr.length - offset): number[] {
let mx = -Infinity
for (let i = 0; i < len; i++) if (arr[offset + i]! > mx) mx = arr[offset + i]!
const probs = new Array<number>(len)
let sum = 0
for (let i = 0; i < len; i++) { probs[i] = Math.exp(arr[offset + i]! - mx); sum += probs[i]! }
for (let i = 0; i < len; i++) probs[i] = probs[i]! / sum
return probs
}
function argmaxNumeric(arr: number[]): number {
let best = 0
for (let i = 1; i < arr.length; i++) if (arr[i]! > arr[best]!) best = i
return best
}
type AttnMap = { data: number[]; T: number; q: number[]; k: number[]; v: number[] }
type LatticeCell = { topToken: number; topProb: number; probs: number[]; residual: number[] }
type InsideStep = { inputTokens: number[]; generated: number; lattice: LatticeCell[][]; mlpHiddens: number[][][] }
// Autoregressive generation. `infer` shares param buffers with training, so
// every step's update is visible on the next call.
async function predictAddition(
infer: CompiledForward,
a: number,
b: number,
opts: { collectViz?: boolean } = {},
): Promise<{
generated: number[]
generatedProbs: number[][]
attnMaps: AttnMap[][]
inside: InsideStep[]
}> {
// collectViz=false skips lattice/attention assembly for accuracy-only probes.
const collectViz = opts.collectViz !== false
const prefix = promptTokens(a, b)
const generated: number[] = []
const generatedProbs: number[][] = []
const attnMapsPerStep: AttnMap[][] = []
const insideSteps: InsideStep[] = []
const tokensBuf = new Int32Array(T_LEN)
for (let step = 0; step < N_RESULT_DIGITS; step++) {
const realLen = prefix.length + generated.length // 6, 7, or 8
tokensBuf.fill(0)
for (let i = 0; i < prefix.length; i++) tokensBuf[i] = prefix[i]!
for (let i = 0; i < generated.length; i++) tokensBuf[prefix.length + i] = generated[i]!
const inputTokens = prefix.concat(generated) // before this step's push
const r = await infer.run({ tokens: tokensBuf })
if (r.kind !== 'completed') return { generated, generatedProbs, attnMaps: attnMapsPerStep, inside: insideSteps }
const logitsAll = r.output
const captures = r.captures
// Restrict to digits 0..9 — '+' and '=' aren't valid at result positions.
const probs = softmaxNumeric(logitsAll, (realLen - 1) * VOCAB, 10)
const best = argmaxNumeric(probs)
generated.push(best)
generatedProbs.push(probs)
if (!collectViz) continue
const stepAttnMaps: AttnMap[] = []
for (let l = 0; l < N_LAYERS; l++) {
const attnH = captures.perHead(`attn.${l}`)
const qH = captures.perHead(`q.${l}`)
const kH = captures.perHead(`k.${l}`)
const vH = captures.perHead(`v.${l}`)
for (let h = 0; h < N_HEADS; h++) {
// attnH[h] is [T_LEN, T_LEN] flat; copy out the realLen×realLen block.
const attn: number[] = new Array(realLen * realLen)
for (let i = 0; i < realLen; i++) for (let j = 0; j < realLen; j++) {
attn[i * realLen + j] = attnH[h]![i * T_LEN + j]!
}
stepAttnMaps.push({
data: attn,
T: realLen,
q: Array.from(qH[h]!.subarray(0, realLen * D_HEAD)),
k: Array.from(kH[h]!.subarray(0, realLen * D_HEAD)),
v: Array.from(vH[h]!.subarray(0, realLen * D_HEAD)),
})
}
}
attnMapsPerStep.push(stepAttnMaps)
const stepLattice: LatticeCell[][] = []
for (let l = 0; l <= N_LAYERS; l++) {
const ll = captures.get(`logit_lens.${l}`) // [1, T, V] flat
const r = captures.get(`residual.${l}`) // [1, T, D] flat
const cells: LatticeCell[] = []
for (let pos = 0; pos < realLen; pos++) {
const ps = softmaxNumeric(ll, pos * VOCAB, VOCAB)
const topI = argmaxNumeric(ps)
cells.push({
topToken: topI,
topProb: ps[topI]!,
probs: ps,
residual: Array.from(r.subarray(pos * D_MODEL, (pos + 1) * D_MODEL)),
})
}
stepLattice.push(cells)
}
const stepMlp: number[][][] = []
for (let l = 0; l < N_LAYERS; l++) {
const m = captures.get(`mlp_hidden.${l}`) // [1, T, HIDDEN] flat
const perPos: number[][] = []
for (let pos = 0; pos < realLen; pos++) {
perPos.push(Array.from(m.subarray(pos * HIDDEN, (pos + 1) * HIDDEN)))
}
stepMlp.push(perPos)
}
insideSteps.push({ inputTokens, generated: best, lattice: stepLattice, mlpHiddens: stepMlp })
}
return { generated, generatedProbs, attnMaps: attnMapsPerStep, inside: insideSteps }
}
// Sibling-facing surface of Model — what UI panels read. Excludes training
// internals and Root-only stats (avgLoss, isRunning, paramCount).
interface IModel {
trainStep: number
predictedDigits: number[]
digitProbs: number[][]
attnMaps: AttnMap[][]
inside: InsideStep[]
recentFailures: { a: number; b: number; got: number; want: number }[]
embeddingHistory: { step: number; tokEmb: Float32Array }[]
getTokEmbSnapshot(): Float32Array | null
getMlpWeights(layerIdx: number): { w1: Float32Array; w2: Float32Array; b2: Float32Array } | null
getAttnWeights(layerIdx: number): { wQ: Float32Array; wK: Float32Array; wV: Float32Array; wO: Float32Array } | null
refreshPrediction(): Promise<void>
}
// ---------- Model ----------
class Model extends Component implements IModel {
status = 'Initializing...'
isReady = false
isRunning = false
trainStep = 0
examplesSeen = 0
avgLoss = 0
examplesPerSec = 0
testAcc = 0
predictedDigits: number[] = [] // LSB-first: [ones, tens, hundreds]
digitProbs: number[][] = []
attnMaps: AttnMap[][] = []
inside: InsideStep[] = []
lossHistory: { step: number; value: number }[] = []
accHistory: { step: number; value: number }[] = []
recentFailures: { a: number; b: number; got: number; want: number }[] = []
embeddingHistory: { step: number; tokEmb: Float32Array }[] = []
// CPU mirror of params, kept in sync from GPU for viz panels.
#params: Record<string, Float32Array> = {}
#train: CompiledTraining<Transformer> | null = null
// #infer is polymorphic over batch dim — first run() at B=1 (predictAddition)
// and at B=N_TEST (#writeDiagnostic) each trigger a sibling compile + cache.
#infer: CompiledForward | null = null
#compilePromise: Promise<void> | null = null
#trainingActive = false
#lossWindow: number[] = []
#stepTimes: number[] = []
#lastStepTs = 0
#diagnosticBusy = false
get root() { return this.ctx.root as any as IRoot }
constructor() {
super()
this.isReady = true
this.status = 'Ready (will compile WGSL on first run)'
}
// refreshPrediction needs this.root.ctx wired, which only happens at attach time.
// The WebGPU gate lives here, not the constructor: checkWebGPU is async
// because it probes for a real adapter, catching browsers that expose the
// API but hand out none (Chrome on Linux without its flags).
override async onAttached() {
const gpu = await checkWebGPU()
if (!gpu.ok) {
this.isReady = false
this.status = gpu.message
this.update()
return
}
this.refreshPrediction()
if (!this.isRunning) this.toggleRun()
}
get paramCount(): number {
return Object.values(this.#params).reduce((a, b) => a + b.length, 0)
}
getTokEmbSnapshot(): Float32Array | null {
const t = this.#params['tok_emb']
return t ? new Float32Array(t) : null
}
getMlpWeights(layerIdx: number): { w1: Float32Array; w2: Float32Array; b2: Float32Array } | null {
const lp = `layers.${layerIdx}.mlp`
const w1 = this.#params[`${lp}.up.W`]
const w2 = this.#params[`${lp}.down.W`]
const b2 = this.#params[`${lp}.down.b`]
return w1 && w2 && b2 ? { w1, w2, b2 } : null
}
getAttnWeights(layerIdx: number): { wQ: Float32Array; wK: Float32Array; wV: Float32Array; wO: Float32Array } | null {
const lp = `layers.${layerIdx}.attn`
const wQ = this.#params[`${lp}.q.W`]
const wK = this.#params[`${lp}.k.W`]
const wV = this.#params[`${lp}.v.W`]
const wO = this.#params[`${lp}.o.W`]
return wQ && wK && wV && wO ? { wQ, wK, wV, wO } : null
}
// Idempotent — concurrent callers await the same in-flight compile.
#ensureCompiled(): Promise<void> {
if (this.#train) return Promise.resolve()
if (this.#compilePromise) return this.#compilePromise
this.#compilePromise = this.#compile()
return this.#compilePromise
}
async #compile(): Promise<void> {
this.status = 'Compiling WGSL kernels...'
this.update()
const t0 = performance.now()
const model = new Transformer()
// Trace the training graph without captures — they'd be readback every step.
captureMode = 'off'
this.#train = await compile({
model,
loss: lossFn,
optimizer: {
kind: 'adamw',
lr: LR_SCHEDULE,
// Without weight decay, train loss bottoms out around 0.9 and held-out
// accuracy stays near 0% — the model memorizes.
weightDecay: 0.01,
clipGradNorm: 1.0,
},
inputs: {
tokens: { shape: [BATCH_SIZE, T_LEN], dtype: 'i32' },
targets: { shape: [BATCH_SIZE, T_LEN], dtype: 'i32' },
mask: [T_LEN],
},
})
captureMode = 'on'
this.#infer = await this.#train.attach({
forward: predictFn,
inputs: { tokens: { shape: [null, T_LEN], dtype: 'i32' } },
})
this.#params = await this.#train.downloadParams()
const compileMs = performance.now() - t0
this.status = `Ready (${this.#train.kernels.length} kernels, ${compileMs.toFixed(0)} ms)`
this.update()
}
async reset() {
// In-place: no recompile. Inference siblings see new params via shared buffers.
if (this.#train) {
await this.#train.reset()
this.#params = await this.#train.downloadParams()
} else {
this.#params = {}
}
this.isReady = true
this.status = this.#train ? `Ready` : 'Ready (will compile WGSL on first run)'
this.trainStep = 0
this.examplesSeen = 0
this.avgLoss = 0
this.testAcc = 0
this.#lossWindow = []
this.#stepTimes = []
this.#lastStepTs = 0
this.predictedDigits = []
this.digitProbs = []
this.attnMaps = []
this.inside = []
this.lossHistory = []
this.accHistory = []
this.recentFailures = []
this.embeddingHistory = []
this.update()
this.refreshPrediction()
}
toggleRun() {
this.isRunning = !this.isRunning
if (this.isRunning && !this.#trainingActive) this.runTrainingLoop()
this.update()
}
async trainOneStep() {
await this.#ensureCompiled()
if (!this.#train) return
const { tokens, targets } = makeBatch()
let lossVal: number
try {
const r = await this.#train.step({ tokens, targets, mask: RESULT_MASK })
if (r.kind !== 'completed') return
lossVal = r.loss
} catch (e: any) {
this.status = `Step error: ${e?.message ?? e}`
this.isRunning = false
this.update()
return
}
this.trainStep++
this.examplesSeen += BATCH_SIZE
if (!Number.isFinite(lossVal)) {
this.status = 'Loss diverged (NaN/Inf)'
this.isRunning = false
this.update()
return
}
this.#lossWindow.push(lossVal)
if (this.#lossWindow.length > 50) this.#lossWindow.shift()
this.avgLoss = this.#lossWindow.reduce((a, b) => a + b, 0) / this.#lossWindow.length
this.lossHistory.push({ step: this.trainStep, value: this.avgLoss })
if (this.lossHistory.length > 600) this.lossHistory.shift()
const now = performance.now()
if (this.#lastStepTs > 0) {
const dt = now - this.#lastStepTs
this.#stepTimes.push(dt)
if (this.#stepTimes.length > 30) this.#stepTimes.shift()
const avgDt = this.#stepTimes.reduce((a, b) => a + b, 0) / this.#stepTimes.length
this.examplesPerSec = avgDt > 0 ? BATCH_SIZE * 1000 / avgDt : 0
}
this.#lastStepTs = now
}
async runTrainingLoop() {
this.#trainingActive = true
let lastUiUpdate = 0
let lastPredictUpdate = 0
let lastDiagnostic = 0
while (this.isRunning) {
// Run a burst of steps at GPU speed, then yield ONCE for the UI. Yielding
// per-step (the old code) pinned the step rate to the ~60 Hz frame clock.
for (let i = 0; i < STEPS_PER_YIELD && this.isRunning; i++) {
await this.trainOneStep()
}
await new Promise<void>(r => requestAnimationFrame(() => r()))
const now = performance.now()
if (now - lastUiUpdate > 100) {
lastUiUpdate = now
this.update()
}
if (now - lastPredictUpdate > 500) {
lastPredictUpdate = now
this.refreshPrediction()
}
if (now - lastDiagnostic > 3000) {
lastDiagnostic = now
this.#writeDiagnostic()
}
}
// The last trainOneStep may have completed *after* isRunning went false (we
// were awaiting it when pause was clicked); flush its state changes.
this.update()
this.#trainingActive = false
}
async #syncParamsFromGpu() {
if (!this.#train) return
try {
this.#params = await this.#train.downloadParams()
} catch { /* next sync will retry */ }
}
// Latest-wins coalescing: rapid input changes drop stale predictions in flight.
// Sync CPU mirror for viz panels; inference itself reads from training's
// GPU buffers, no readback needed for the forward pass.
#refreshFlight = singleFlight(async (_: void) => {
if (!this.#infer) return null
await this.#syncParamsFromGpu()
return predictAddition(this.#infer, this.root.inputPanel.operandA, this.root.inputPanel.operandB)
})
async refreshPrediction() {
const r = await this.#refreshFlight()
if (r.kind === 'aborted' || !r.value) return
this.predictedDigits = r.value.generated
this.digitProbs = r.value.generatedProbs
this.attnMaps = r.value.attnMaps
this.inside = r.value.inside
this.update()
}
// Batched eval: N_TEST pairs decoded in parallel — 3 forward passes total
// (one per result digit) vs 150 for a per-pair loop.
async #writeDiagnostic() {
if (this.#diagnosticBusy) return
const evalCompiled = this.#infer
if (!evalCompiled) return
this.#diagnosticBusy = true
try {
await this.#syncParamsFromGpu()
const testPairs: { a: number; b: number }[] = []
let attempts = 0
while (testPairs.length < N_TEST && attempts < 5000) {
attempts++
const a = Math.floor(Math.random() * 100), b = Math.floor(Math.random() * 100)
if (isTestPair(a, b)) testPairs.push({ a, b })
}
// Defensive pad if isTestPair fails to yield N_TEST hits in 5000 tries.
while (testPairs.length < N_TEST) testPairs.push(testPairs[0]!)
const tokensBuf = new Int32Array(N_TEST * T_LEN)
for (let row = 0; row < N_TEST; row++) {
const { a, b } = testPairs[row]!
const off = row * T_LEN
const prompt = promptTokens(a, b)
for (let i = 0; i < prompt.length; i++) tokensBuf[off + i] = prompt[i]!
}
const generated: number[][] = Array.from({ length: N_TEST }, () => [])
for (let step = 0; step < N_RESULT_DIGITS; step++) {
const r = await evalCompiled.run({ tokens: tokensBuf }) // [N_TEST, T_LEN, VOCAB]
if (r.kind !== 'completed') return
const logits = r.output
const realLen = RESULT_START + step // 6, 7, 8
const readPos = realLen - 1
for (let row = 0; row < N_TEST; row++) {
const offset = (row * T_LEN + readPos) * VOCAB
let best = 0
for (let v = 1; v < 10; v++) {
if (logits[offset + v]! > logits[offset + best]!) best = v
}
generated[row]!.push(best)
if (step < N_RESULT_DIGITS - 1) tokensBuf[row * T_LEN + realLen] = best
}
}
let correct = 0
const failures: { a: number; b: number; got: number; want: number }[] = []
for (let row = 0; row < N_TEST; row++) {
const { a, b } = testPairs[row]!
const got = decodeReversedTokens(generated[row]!)
const want = a + b
if (got === want) correct++
else if (failures.length < 5) failures.push({ a, b, got, want })
}
const acc = correct / N_TEST
this.testAcc = acc
this.accHistory.push({ step: this.trainStep, value: acc })
if (this.accHistory.length > 200) this.accHistory.shift()
this.recentFailures = failures
this.embeddingHistory.push({ step: this.trainStep, tokEmb: new Float32Array(this.#params['tok_emb']) })
if (this.embeddingHistory.length > 120) this.embeddingHistory.shift()
this.update()
} catch (e: any) {
// reset() mid-dispatch surfaces as AbortError; next tick retries.
if (e?.name === 'AbortError') return
throw e
} finally {
this.#diagnosticBusy = false
}
}
}
// ========== UI ==========
// ---------- UI constants ----------
const STRIP_CELL_H = 20
const TOKEN_LABELS: string[] = ['0','1','2','3','4','5','6','7','8','9','+','=']
// LSB-first to match generation order.
const RESULT_DIGIT_NAMES = ['ones', 'tens', 'hundreds'] as const
const bareStageName = (s: string) => s.replace(/^After /, '')
// Sequence layout during prediction: hundreds digit is being generated at the
// final step, so it never appears as a captured position in attention.
const POSITION_LABELS: string[] = ['A tens', 'A ones', '+', 'B tens', 'B ones', '=', 'ones', 'tens']
// ---------- UI interfaces ----------
interface IRoot {
isNarrow: boolean
model: IModel
inputPanel: IInputPanel
}
interface IInputPanel {
operandA: number
operandB: number
}
// Layer index 0 = embedding, 1..N_LAYERS = layer 1..N. Default falls back to
// the bottom-right cell.
interface IInsidePanel {
effectiveSelection: { layer: number; pos: number }
}
// ---------- Theme ----------
const lightTheme = {
colors: {
primary: 'rgb(0, 128, 0)',
primaryStrong: 'rgb(0, 100, 0)',
accent: 'rgb(147, 51, 234)',
background: 'rgb(255, 255, 255)',
surface: 'rgb(245, 245, 245)',
text: 'rgb(95, 95, 95)',
textMuted: 'rgb(85, 85, 85)',
hoverBg: 'rgb(250, 250, 250)',
border: 'rgb(220, 220, 220)',
attn: 'rgb(70, 140, 230)',
mlp: 'rgb(180, 100, 0)',
error: 'rgb(220, 38, 38)'
}
}
const darkTheme = {
colors: {
primary: 'rgb(0, 200, 0)',
primaryStrong: 'rgb(0, 255, 0)',
accent: 'rgb(192, 132, 252)',
background: 'rgb(30, 30, 30)',
surface: 'rgb(40, 40, 40)',
text: 'rgb(245, 245, 245)',
textMuted: 'rgb(170, 170, 170)',
hoverBg: 'rgb(35, 35, 35)',
border: 'rgb(55, 55, 55)',
attn: 'rgb(70, 140, 230)',
mlp: 'rgb(220, 140, 0)',
error: 'rgb(239, 68, 68)'
}
}
const globalUnoCss = (theme: ThemeProxy<typeof lightTheme>) => ({
'body': `m-0 bg-${theme.colors.background} text-${theme.colors.text} font-sans antialiased overflow-x-hidden`,
'button': `transition-colors duration-200 ease-in-out`,
'a': `text-${theme.colors.primary} underline hover:text-${theme.colors.primaryStrong} transition-colors`,
'*': `box-border`
})
const getInitialTheme = () => {
const hostTheme = document.documentElement.getAttribute('data-theme')
if (hostTheme === 'light' || hostTheme === 'dark') return hostTheme
return window.matchMedia('(prefers-color-scheme: dark)').matches ? 'dark' : 'light'
}
const themeMgr = new UnoThemeManager({
id: 'app',
themes: { light: lightTheme, dark: darkTheme },
initialTheme: getInitialTheme(),
unoCssConfig: { presets: [presetWind3()] },
globalUnoCss
})
const styles = themeMgr.styles('xfm', (theme) => {
const { text, textMuted, primary, primaryStrong, surface, hoverBg, background, border, accent, error } = theme.colors
return {
layout: `flex flex-col min-h-screen w-full`,
header: `flex-none pt-3 bg-${background}`,
headerContainer: `max-w-[1300px] mx-auto px-4 md:px-6 flex flex-col items-center text-center md:items-start md:text-left`,
headerTitle: `text-${text} text-2xl font-bold tracking-tight my-2`,
content: `flex-1`,
controlStripBar: `w-full max-w-[1300px] mx-auto px-0 md:px-6`,
controlStripContent: `w-full bg-${surface} px-2 py-2 rounded-none md:rounded-lg`,
controlsRow: `w-full flex flex-wrap items-center justify-center md:justify-start gap-x-4 gap-y-2 mt-1`,
statsRow: `flex flex-wrap items-center justify-center gap-x-3 gap-y-1 sm:gap-x-8`,
controlGroup: `flex items-center gap-1 sm:gap-6`,
mainContainer: `max-w-[1300px] mx-auto px-0 md:px-6 py-4`,
labelText: `text-xs text-${textMuted} font-bold`,
valueText: `font-bold text-${text} whitespace-nowrap`,
iconBtn: `bg-transparent border-none p-0 cursor-pointer text-${primary} hover:text-${primaryStrong} transition-colors outline-none flex items-center justify-center`,
iconBtnPlay: `w-12 h-12`,
iconBtnSmall: `w-8 h-8`,
sectionTitle: `text-${text} text-lg font-semibold mb-3`,
body: `text-${textMuted} text-base`,
bodyText: `text-${text} text-base`,
bodyMono: `text-${text} text-base font-mono`,
panel: `bg-${surface} rounded-none md:rounded-lg p-4`,
btn: `px-3 py-1.5 rounded-md border border-solid border-${border} bg-${background} text-${text} text-xs cursor-pointer hover:bg-${hoverBg} transition-colors`,
waitContainer: `flex flex-col items-center justify-center min-h-screen`,
waitText: `text-${text} mb-4`,
waitSpinner: `mt-4 animate-bounce text-2xl`,
successHigh: `text-${primary}`,
successLow: `text-${error}`,
additionRow: `flex items-baseline gap-2 my-3`,
numberInput: `w-20 px-2 py-1 text-2xl text-center font-mono rounded-md border border-solid border-${border} bg-${background} text-${text} focus:outline-none focus:ring-2 focus:ring-${primary}`,
bigOp: `text-3xl font-mono text-${text}`,
bigResult: `text-3xl font-mono font-bold`,
presetRow: `flex flex-wrap gap-2 mt-3`,
digitProbs: `mt-3 flex flex-col gap-2`,
attnGrid: `grid gap-3`,
attnHeadCell: `flex flex-col items-center gap-1`,
attnHeadLabel: `text-${textMuted} text-[11px]`,
mainGrid: `grid gap-4`,
chartBox: `w-full h-auto block bg-${background} rounded-md`,
explainerContainer: `text-${textMuted}`,
detailPanel: `mt-5 p-3 rounded-md bg-${background} border border-${border} block max-w-full box-border`,
detailPanelTight: `mt-4 p-3 rounded-md bg-${background} border border-${border} block max-w-full box-border`,
monoLabelTiny: `text-[11px] font-mono text-${textMuted}`,
panelTitle: `text-[13px] text-${text} font-bold mb-2`,
tabBtn: `px-3 py-1.5 text-sm cursor-pointer rounded-md transition-colors`,
tabBtnActive: `bg-${primary} text-white border-none font-bold`,
tabBtnInactive: `bg-${background} text-${textMuted} border border-solid border-${border} font-normal hover:bg-${hoverBg}`,
subTabBtn: `px-2 py-1.5 text-sm cursor-pointer bg-transparent border-0 border-b-2 border-solid transition-colors rounded-none`,
subTabBtnActive: `text-${text} font-bold border-${primary}`,
subTabBtnInactive: `text-${textMuted} border-transparent hover:text-${text}`,
statCol: `flex flex-col items-center`,
statValueRow: `flex items-center gap-1.5`,
latticeCell: `flex items-center justify-center font-mono text-[13px] font-bold cursor-pointer box-border text-${text}`,
stripRow: `flex items-center gap-2 mb-0.5`,
stripRowLabel: `text-[11px] font-mono text-right pr-1 flex-shrink-0`,
rowLabel: `text-[11px] pr-2 text-right text-${textMuted}`,
pickMark: `text-[10px] h-3 text-${text}`,
vocabLabel: `text-[10px] mt-0.5 text-${textMuted}`,
softmaxPosLabel: `text-[9px] font-mono mt-0.5 text-${textMuted}`,
sparklineBox: `w-20 h-3.5 block flex-shrink-0 bg-${background} mt-0.5`,
freshSwatch: `w-3.5 h-3.5 bg-${accent}/40 border border-${accent}`,
attnDrillTitle: `text-sm font-mono text-${text}`,
mutedText: `text-${textMuted}`,
credit: `text-[0.85em] opacity-65 mt-1 mb-0 text-center`,
attnDrillHeatmapBox: `flex-shrink-0 max-w-[180px]`,
badge: `flex items-center justify-center rounded-full bg-${background}`,
}
})
const icon = (d: string) => svg({ viewBox: '0 0 24 24', fill: 'currentColor', class: 'w-full h-full' }, path({ d }))
const iconPlay = () => icon('M8 5v14l11-7z')
const iconPause = () => icon('M6 19h4V5H6v14zm8-14v14h4V5h-4z')
const iconReset = () => icon('M12 5V1L7 6l5 5V7c3.31 0 6 2.69 6 6s-2.69 6-6 6-6-2.69-6-6H4c0 4.42 3.58 8 8 8s8-3.58 8-8-3.58-8-8-8z')
// ---------- UI helpers ----------
function decodeReversedTokens(digits: number[]): number {
let n = 0
for (let i = digits.length - 1; i >= 0; i--) n = n * 10 + digits[i]
return n
}
// LSB-first: [ones, tens, hundreds]. Inverse of decodeReversedTokens.
function splitResultDigits(n: number): number[] {
return [n % 10, Math.floor(n / 10) % 10, Math.floor(n / 100)]
}
function maxAbs(values: ArrayLike<number>): number {
let m = 0
for (let i = 0; i < values.length; i++) {
const a = Math.abs(values[i])
if (a > m) m = a
}
return m
}
// Green for positive, blue for negative, intensity = |v| / maxVal.
function makeDivergingStripColor(maxVal: number): (v: number) => string {
const bg = themeMgr.theme.colors.surface.rawValue
const pos = themeMgr.theme.colors.primary.rawValue
const neg = themeMgr.theme.colors.attn.rawValue
if (maxVal === 0) return () => bg
return (v) => {
const t = Math.min(1, Math.abs(v) / maxVal)
return v >= 0 ? lerpRgb(bg, pos, t) : lerpRgb(bg, neg, t)
}
}
// 1D array → colored strip cells. cellW omitted = flex-grow.
function stripCells(
values: ArrayLike<number>,
color: (v: number) => string,
cellH: number,
tooltip: (i: number, v: number) => string,
cellW?: number
) {
return range(values.length).map(i => div({
class: cellW != null ? 'box-border' : 'flex-1 min-w-0 box-border',
style: cellW != null
? { width: cellW + 'px', height: cellH + 'px', backgroundColor: color(values[i]) }
: { height: cellH + 'px', backgroundColor: color(values[i]) },
title: tooltip(i, values[i])
}))
}
// Probability bar row with truth/pick highlighting.
// bar height/opacity = prob; ▼ marks pick;
// green bar = truth, red bar = picked-but-wrong, gray = neither.
// When confidently wrong, the truth bar can never reach green's max brightness
// (bounded by 1 - picked_prob) — presence of red disambiguates.
function probBarsView(opts: {
probs: number[]
picked: number
truth: number
labels: string[]
maxBarH: number
}) {
const { probs, picked, truth, labels, maxBarH } = opts
const primary = themeMgr.theme.colors.primary.css
const error = themeMgr.theme.colors.error.css
const textCol = themeMgr.theme.colors.text.css
// Frame height = maxBarH + 12 so the ▼ marker can ride on top of a max-height bar.
return div({ class: 'flex gap-1' },
probs.map((p, i) => {
const isPicked = i === picked
const isTruth = i === truth
const heightPx = Math.max(2, Math.round(p * maxBarH))
const bg = isTruth ? primary : (isPicked ? error : textCol)
return div({ class: 'flex flex-col items-center flex-1 min-w-0' },
div({ class: 'w-full flex flex-col items-center justify-end',
style: { height: (maxBarH + 12) + 'px' } },
div({ class: styles.pickMark }, isPicked ? '▼' : ''),
div({ class: 'w-full rounded-sm', style: {
height: heightPx + 'px',
backgroundColor: bg,
opacity: String(0.2 + 0.8 * p)
}})
),
div({ class: styles.vocabLabel }, labels[i])
)
})
)
}
function lerpRgb(from: string, to: string, t: number): string {
const parse = (s: string): [number, number, number] => {
const m = s.match(/\d+/g)
return m ? [Number(m[0]), Number(m[1]), Number(m[2])] : [128, 128, 128]
}
const [r1, g1, b1] = parse(from)
const [r2, g2, b2] = parse(to)
const k = clamp(t, 0, 1)
return `rgb(${Math.round(r1 + (r2 - r1) * k)},${Math.round(g1 + (g2 - g1) * k)},${Math.round(b1 + (b2 - b1) * k)})`
}
function pcaBasis(X: number[][]): { mean: number[]; v1: number[]; v2: number[] } {
const N = X.length, D = X[0].length
const mean = new Array(D).fill(0)
for (const row of X) for (let j = 0; j < D; j++) mean[j] += row[j]
for (let j = 0; j < D; j++) mean[j] /= N
const centered = X.map(row => row.map((v, j) => v - mean[j]))
const cov: number[][] = Array.from({ length: D }, () => new Array(D).fill(0))
for (let i = 0; i < D; i++) {
for (let j = i; j < D; j++) {
let s = 0
for (let k = 0; k < N; k++) s += centered[k][i] * centered[k][j]
cov[i][j] = cov[j][i] = s / N
}
}
const v1 = powerIter(cov)
const Mv1 = mvm(cov, v1)
const lambda1 = v1.reduce((s, x, i) => s + x * Mv1[i], 0)
for (let i = 0; i < D; i++) for (let j = 0; j < D; j++) cov[i][j] -= lambda1 * v1[i] * v1[j]
const v2 = powerIter(cov)
return { mean, v1, v2 }
}
function projectThroughBasis(X: number[][], basis: { mean: number[]; v1: number[]; v2: number[] }): [number, number][] {
return X.map(row => {
let p1 = 0, p2 = 0
for (let i = 0; i < row.length; i++) {
const c = row[i] - basis.mean[i]
p1 += c * basis.v1[i]
p2 += c * basis.v2[i]
}
return [p1, p2] as [number, number]
})
}
function powerIter(M: number[][], iters = 80): number[] {
const D = M.length
let v = new Array(D).fill(0).map(() => Math.random() - 0.5)
let n = Math.sqrt(v.reduce((s, x) => s + x * x, 0))
for (let i = 0; i < D; i++) v[i] /= n
for (let it = 0; it < iters; it++) {
const Mv = mvm(M, v)
n = Math.sqrt(Mv.reduce((s, x) => s + x * x, 0))
if (n < 1e-12) return v
for (let i = 0; i < D; i++) v[i] = Mv[i] / n
}
return v
}
function mvm(M: number[][], v: number[]): number[] {
return M.map(row => {
let s = 0
for (let i = 0; i < v.length; i++) s += row[i] * v[i]
return s
})
}
// Top-N entries by |w| as (i, j, w) triples sorted descending.
function pickTopWeights(w: Float32Array, rows: number, cols: number, topN: number): { i: number; j: number; w: number }[] {
const all: { i: number; j: number; w: number }[] = []
for (let i = 0; i < rows; i++) {
for (let j = 0; j < cols; j++) {
all.push({ i, j, w: w[i * cols + j] })
}
}
all.sort((a, b) => Math.abs(b.w) - Math.abs(a.w))
return all.slice(0, topN)
}
function countFiring(vals: ArrayLike<number>): number {
let n = 0
for (let i = 0; i < vals.length; i++) if (vals[i] > 0) n++
return n
}
// ---------- Attention panel ----------
type AttnDetailTab = 'project' | 'score' | 'output'
class AttentionPanel extends Component {
selectedAttnCell: { layer: number; head: number; qPos: number; kPos: number } | null = null
// Default to the last step; sticky across prediction refreshes.
activeAttnStep = N_RESULT_DIGITS - 1
// Sticky across cell selections.
attnDetailTab: AttnDetailTab = 'project'
get root() { return this.ctx.root as any as IRoot }
get currentStepAttnMaps(): AttnMap[] {
return this.root.model.attnMaps[this.activeAttnStep] ?? []
}
// -1 = "all positions fresh, no cache yet" (first step); otherwise the last
// position, which is the only K/V freshly computed this step.
get freshColIdx(): number {
const T_len = this.currentStepAttnMaps[0]?.T ?? 0
return this.activeAttnStep === 0 ? -1 : T_len - 1
}
view() {
if (this.selectedAttnCell !== null) {
return this.attnDrillDownView()
}
return div(
this.kvCacheStepNav(),
div({ class: [styles.body, 'mb-2'] }, 'Each row shows where that position is looking; brighter cells = more attention. Click any cell to drill in and see the math behind it (Q, K, V).'),
this.attnGridView(),
this.freshColumnHint()
)
}
// K/V are recomputed every step (vs cached in real LLM inference). Caption
// calls out which column would be the "fresh" one in a cached impl.
kvCacheStepNav() {
if (this.root.model.attnMaps.length === 0) return div()
const numSteps = this.root.model.attnMaps.length
const stepNames = RESULT_DIGIT_NAMES
const active = this.activeAttnStep
const T_len = this.currentStepAttnMaps[0]?.T ?? 0
const freshIdx = this.freshColIdx
const freshLabel = freshIdx >= 0 ? POSITION_LABELS[freshIdx] : null
const caption = active === 0
? `Generation Step 1/${numSteps}: the model sees the prompt (${T_len} tokens). It's about to generate the ${stepNames[active]} digit.`
: `Generation Step ${active + 1}/${numSteps}: the model has just generated the ${stepNames[active - 1]} digit (position ${freshIdx} = "${freshLabel}"). Only this new position needed fresh K and V values this step; positions 0–${freshIdx - 1} already had theirs from earlier steps (real LLMs save and reuse those past values — this is the K/V cache).`
return div({ class: 'mb-4' },
div({ class: 'flex items-center gap-3 flex-wrap mb-2' },
range(numSteps).map(s => {
const isActive = s === active
return button({
class: [styles.subTabBtn, isActive ? styles.subTabBtnActive : styles.subTabBtnInactive],
onClick: () => {
if (this.activeAttnStep === s) return
this.activeAttnStep = s
this.selectedAttnCell = null
this.update()
}
}, stepNames[s])
})
),
div({ class: styles.body }, caption)
)
}
freshColumnHint() {
if (this.root.model.attnMaps.length === 0) return div()
const freshIdx = this.freshColIdx
if (freshIdx < 0) return div()
const freshLabel = POSITION_LABELS[freshIdx]
return div({ class: [styles.labelText, 'mt-2 flex items-center justify-center gap-1.5'] },
div({ class: styles.freshSwatch }),
`= column "${freshLabel}" — K/V freshly computed at this step.`
)
}
attnDrillDownView() {
const sel = this.selectedAttnCell!
const idx = sel.layer * N_HEADS + sel.head
const am = this.currentStepAttnMaps[idx]
if (!am) {
this.selectedAttnCell = null
return div()
}
const qLabel = POSITION_LABELS[sel.qPos] ?? `pos ${sel.qPos}`
const kLabel = POSITION_LABELS[sel.kPos] ?? `pos ${sel.kPos}`
return div(
div({ class: 'flex items-center gap-3 mb-3 flex-wrap' },
button({
class: styles.btn,
onClick: () => { this.selectedAttnCell = null; this.update() }
}, '← Back to all heads'),
div({ class: styles.attnDrillTitle }, `Layer ${sel.layer + 1} · Head ${sel.head + 1}`)
),
div({ class: ['flex gap-4 items-center mb-1', this.root.isNarrow ? 'flex-col' : 'flex-row'] },
div({ class: styles.attnDrillHeatmapBox },
this.heatmapView(am, sel.layer, sel.head)
),
div({ class: 'flex-1 min-w-0' },
div({ class: styles.panelTitle },
`Attention weight: query at "${qLabel}" → key at "${kLabel}"`
),
div({ class: styles.body },
`Attention lets each position selectively read context from itself and earlier positions. Each position acts as three things: a query ("what am I looking for?"), a key ("what do I offer?"), and a value (the payload to share). Every query weights each available key, and those weights blend the values into the query's read.`
),
div({ class: [styles.body, 'mt-2'] },
`Every cell in this heatmap is one such weight, post-softmax. You clicked the cell where the query at "${qLabel}" attends to the key at "${kLabel}".`
)
)
),
this.attnDetailView()
)
}
attnDetailView() {
const sel = this.selectedAttnCell
if (!sel || this.currentStepAttnMaps.length === 0) return div()
const idx = sel.layer * N_HEADS + sel.head
const am = this.currentStepAttnMaps[idx]
if (!am) return div()
const T_len = am.T
if (sel.qPos >= T_len || sel.kPos >= T_len) return div()
const qVec = am.q.slice(sel.qPos * D_HEAD, (sel.qPos + 1) * D_HEAD)
const kVec = am.k.slice(sel.kPos * D_HEAD, (sel.kPos + 1) * D_HEAD)
const vVec = am.v.slice(sel.kPos * D_HEAD, (sel.kPos + 1) * D_HEAD)
const products = qVec.map((q, i) => q * kVec[i])
const dotProduct = products.reduce((a, b) => a + b, 0)
const score = dotProduct / Math.sqrt(D_HEAD)
const qLabel = POSITION_LABELS[sel.qPos] ?? `pos ${sel.qPos}`
const kLabel = POSITION_LABELS[sel.kPos] ?? `pos ${sel.kPos}`
// Causal mask: weights run 0..qPos.
const rowWeights: number[] = []
for (let j = 0; j <= sel.qPos; j++) rowWeights.push(am.data[sel.qPos * T_len + j])
const headOutput: number[] = new Array(D_HEAD).fill(0)
for (let j = 0; j <= sel.qPos; j++) {
const w = rowWeights[j]
for (let d = 0; d < D_HEAD; d++) {
headOutput[d] += w * am.v[j * D_HEAD + d]
}
}
const stripColor = makeDivergingStripColor(maxAbs([...products, ...headOutput]))
const cellH = this.root.isNarrow ? 18 : 24
const labelW = this.root.isNarrow ? 100 : 140
const renderStrip = (vec: number[], label?: string) => {
const cells = div({ class: 'flex gap-0 flex-1 min-w-0' },
stripCells(vec, stripColor, cellH, (i, v) => `dim ${i}: ${v.toFixed(3)}`)
)
return label
? div({ class: 'flex items-center gap-2 mb-1' },
div({ class: [styles.monoLabelTiny, styles.stripRowLabel], style: { width: labelW + 'px' } }, label),
cells)
: div({ class: 'flex mb-1' }, cells)
}
const inside = this.root.model.inside[this.activeAttnStep]
const attnW = this.root.model.getAttnWeights(sel.layer)
// Per-head slice of the layer's [D_MODEL, D_MODEL] projection matrix.
// Q/K/V take a column slice → [D_MODEL, D_HEAD]; W_o takes a row slice → [D_HEAD, D_MODEL].
const colSliceHead = (W: Float32Array): Float32Array => {
const out = new Float32Array(D_MODEL * D_HEAD)
for (let i = 0; i < D_MODEL; i++) {
for (let j = 0; j < D_HEAD; j++) {
out[i * D_HEAD + j] = W[i * D_MODEL + sel.head * D_HEAD + j]
}
}
return out
}
const rowSliceHead = (W: Float32Array): Float32Array => {
const out = new Float32Array(D_HEAD * D_MODEL)
for (let i = 0; i < D_HEAD; i++) {
for (let j = 0; j < D_MODEL; j++) {
out[i * D_MODEL + j] = W[(sel.head * D_HEAD + i) * D_MODEL + j]
}
}
return out
}
// Per-head contribution to residual_out at qPos = head_output @ W_o_slice.
// The full residual sums all heads' contributions; this drilldown shows only one.
let headContrib: Float32Array | null = null
let wOSlice: Float32Array | null = null
if (attnW) {
wOSlice = rowSliceHead(attnW.wO)
headContrib = new Float32Array(D_MODEL)
for (let i = 0; i < D_MODEL; i++) {
let s = 0
for (let d = 0; d < D_HEAD; d++) s += headOutput[d] * wOSlice[d * D_MODEL + i]
headContrib[i] = s
}
}
const projectSection = inside && attnW
? div({ class: 'mb-3' },
div({ class: [styles.body, 'mb-3 text-center'] },
`First, project each position's residual to produce its Q, K, V. For this pair: Q from the residual at "${qLabel}", K and V from the residual at "${kLabel}". Residual on top, projected vector below.`
),
this.renderProjectionWires(inside.lattice[sel.layer][sel.qPos].residual, qVec, colSliceHead(attnW.wQ), D_MODEL, D_HEAD, `Residual at "${qLabel}"`, `Query at "${qLabel}"`),
this.renderProjectionWires(inside.lattice[sel.layer][sel.kPos].residual, kVec, colSliceHead(attnW.wK), D_MODEL, D_HEAD, `Residual at "${kLabel}"`, `Key at "${kLabel}"`),
this.renderProjectionWires(inside.lattice[sel.layer][sel.kPos].residual, vVec, colSliceHead(attnW.wV), D_MODEL, D_HEAD, `Residual at "${kLabel}" (shared with K)`, `Value at "${kLabel}"`)
)
: div()
const scoreSection = div(
div({ class: [styles.body, 'mb-2 text-center'] },
`Next, measure how well every query matches every key (one matmul). Below, we zoom into your selected pair: Q at "${qLabel}" against K at "${kLabel}".`
),
div({ class: [styles.body, 'mb-1.5 text-center'] },
`The ${D_HEAD} cells below are the element-wise product of Q (at "${qLabel}") and K (at "${kLabel}") — one cell per dimension, each cell = Q[d] × K[d]. Bright = strong agreement (score up); pale cells contribute little either way.`
),
renderStrip(products),
div({ class: [styles.body, 'mt-2.5 text-center'] },
`Adding these ${D_HEAD} cells gives ${dotProduct.toFixed(2)}. Divided by √${D_HEAD}, that's this pair's raw score: ${score.toFixed(3)}.`
),
div({ class: [styles.body, 'mt-3.5 mb-2 text-center'] },
`Softmax across this row's scores (one per attended key — ${sel.qPos + 1} here, due to causal mask) → attention weights. The one for the selected key "${kLabel}" is `,
span({ style: { color: themeMgr.theme.colors.accent.css } }, '['),
span({ style: {
fontWeight: 'bold',
color: lerpRgb(themeMgr.theme.colors.surface.rawValue, themeMgr.theme.colors.primary.rawValue, rowWeights[sel.kPos])
} }, `${(rowWeights[sel.kPos] * 100).toFixed(1)}%`),
span({ style: { color: themeMgr.theme.colors.accent.css } }, ']'),
` of query "${qLabel}" 's attention.`
),
this.renderSoftmaxBars(rowWeights, sel.qPos, sel.kPos)
)
const outputSection = div(
div({ class: [styles.body, 'mb-2 text-center'] },
`Finally, compute the weighted sum of the V vectors using the row's attention weights (one matmul). The result is a ${D_HEAD}-dim vector — this head's output. All heads' outputs are concatenated and projected by W_o into a ${D_MODEL}-dim vector, then added to the residual stream. Below: this head's slice of that projection.`
),
headContrib && wOSlice
? this.renderProjectionWires(headOutput, headContrib, wOSlice, D_HEAD, D_MODEL, `V blend at "${qLabel}"`, `This head's output at "${qLabel}"`)
: renderStrip(headOutput, `This head's output at "${qLabel}"`)
)
const tabs: { id: AttnDetailTab; label: string }[] = [
{ id: 'project', label: 'Project →' },
{ id: 'score', label: 'Score →' },
{ id: 'output', label: 'Attend' },
]
const tabBar = div({ class: 'flex gap-3 mb-3' },
tabs.map(t => button({
class: [styles.subTabBtn, this.attnDetailTab === t.id ? styles.subTabBtnActive : styles.subTabBtnInactive],
onClick: () => { this.attnDetailTab = t.id; this.update() }
}, t.label))
)
return div({ class: styles.detailPanelTight },
tabBar,
this.attnDetailTab === 'project' ? projectSection
: this.attnDetailTab === 'score' ? scoreSection
: outputSection
)
}
// Bar chart of the row's softmaxed attention weights. The selected key's bar
// gets a purple outline matching the heatmap cell highlight.
renderSoftmaxBars(rowWeights: number[], qPos: number, kPos: number): VElement {
const textCol = themeMgr.theme.colors.text.rawValue
const primary = themeMgr.theme.colors.primary.rawValue
const accentCss = themeMgr.theme.colors.accent.css
return div({ class: 'flex items-end gap-0.5 justify-center', style: { minHeight: '54px' } },
range(qPos + 1).map(j => {
const w = rowWeights[j]
const isSelKey = j === kPos
const heightPx = Math.max(2, Math.round(w * 45))
return div({ class: 'flex flex-col items-center w-8' },
div({ class: 'box-content', style: {
width: '22px',
height: heightPx + 'px',
backgroundColor: isSelKey ? primary : textCol,
opacity: String(0.3 + 0.7 * w),
border: isSelKey ? `2px solid ${accentCss}` : 'none'
} }),
div({ class: styles.softmaxPosLabel }, POSITION_LABELS[j] ?? `${j}`)
)
})
)
}
// Source-strip → top-N weight fan → dest-strip diagram. Caller pre-slices the
// per-head weight matrix to [srcN, dstN]; this method is direction-agnostic.
renderProjectionWires(
srcVec: ArrayLike<number>,
dstVec: ArrayLike<number>,
headW: Float32Array,
srcN: number,
dstN: number,
srcLabel: string,
dstLabel: string,
): VElement {
const W = 700, H = 110
const PAD_X = 24
const STRIP_X0 = PAD_X
const STRIP_W = W - 2 * PAD_X
const ROW_H = STRIP_CELL_H
const Y_TOP = 24
const Y_BOT = 76
const xCenter = (j: number, n: number) => STRIP_X0 + STRIP_W * (j + 0.5) / n
const xEdge = (j: number, n: number) => STRIP_X0 + STRIP_W * j / n
const TOP_N = 80
const wTop = pickTopWeights(headW, srcN, dstN, TOP_N)
const wColor = makeDivergingStripColor(Math.abs(wTop[0]?.w ?? 1e-6))
const renderRow = (vals: ArrayLike<number>, n: number, y: number): VElement[] => {
const color = makeDivergingStripColor(maxAbs(vals) || 1)
const cellW = STRIP_W / n
const cells: VElement[] = []
for (let i = 0; i < n; i++) {
cells.push(rect({
// +0.5 closes hairline gaps from non-integer cellW.
x: xEdge(i, n), y, width: cellW + 0.5, height: ROW_H,
fill: color(vals[i])
}))
}
return cells
}
const fanLines = wTop.map(c => line({
x1: xCenter(c.i, srcN), y1: Y_TOP + ROW_H,
x2: xCenter(c.j, dstN), y2: Y_BOT,
stroke: wColor(c.w), strokeWidth: 0.6
}))
const topCells = renderRow(srcVec, srcN, Y_TOP)
const botCells = renderRow(dstVec, dstN, Y_BOT)
const labelColor = themeMgr.theme.colors.textMuted.css
const labels: VElement[] = [
text({ x: PAD_X, y: 16, fill: labelColor, fontSize: '11', textAnchor: 'start' }, srcLabel),
text({ x: PAD_X, y: Y_BOT + ROW_H + 14, fill: labelColor, fontSize: '11', textAnchor: 'start' }, dstLabel)
]
return svg({
viewBox: `0 0 ${W} ${H}`,
width: '100%',
class: 'block max-w-full mb-6',
preserveAspectRatio: 'xMidYMid meet'
}, ...fanLines, ...topCells, ...botCells, ...labels)
}
attnGridView() {
const stepMaps = this.currentStepAttnMaps
if (stepMaps.length === 0) {
return div({ class: [styles.body, 'py-6'] }, 'No attention to show yet.')
}
return div({ class: styles.attnGrid, style: { gridTemplateColumns: `repeat(${N_HEADS}, minmax(0, 1fr))` } },
range(N_LAYERS * N_HEADS).map(idx => {
const layer = Math.floor(idx / N_HEADS)
const head = idx % N_HEADS
const attn = stepMaps[idx]
const T = attn.T
let bestPos = 0
let bestWeight = -1
for (let j = 0; j <= T - 1; j++) {
const w = attn.data[(T - 1) * T + j]
if (w > bestWeight) { bestWeight = w; bestPos = j }
}
const label = POSITION_LABELS[bestPos] ?? `pos ${bestPos}`
return div({ class: styles.attnHeadCell },
div({ class: styles.attnHeadLabel }, `L${layer + 1} H${head + 1} → ${label} (${(bestWeight * 100).toFixed(0)}%)`),
this.heatmapView(attn, layer, head)
)
})
)
}
heatmapView(attn: AttnMap, layer: number, head: number) {
const T_len = attn.T
const labels: string[] = promptTokens(this.root.inputPanel.operandA, this.root.inputPanel.operandB)
.map(t => TOKEN_LABELS[t])
for (let i = 0; i < this.root.model.predictedDigits.length && labels.length < T_len; i++) {
labels.push(String(this.root.model.predictedDigits[i]))
}
while (labels.length < T_len) labels.push('?')
const cellSize = Math.max(18, Math.min(32, Math.floor(220 / T_len)))
const labelSize = 18
const totalSize = T_len * cellSize + labelSize
const cellBg = themeMgr.theme.colors.surface.rawValue
const cellHi = themeMgr.theme.colors.primary.rawValue
const cellColor = (a: number) => lerpRgb(cellBg, cellHi, a)
const freshCol = this.freshColIdx
const freshColor = themeMgr.theme.colors.accent.css // purple, matches selection ring
const cells: VElement[] = []
if (freshCol >= 0) {
cells.push(rect({
x: labelSize + freshCol * cellSize - 1, y: labelSize - 1,
width: cellSize + 1, height: T_len * cellSize + 1,
fill: freshColor, opacity: '0.18'
}))
}
for (let j = 0; j < T_len; j++) {
const isFresh = j === freshCol
cells.push(text({
x: labelSize + j * cellSize + cellSize / 2,
y: labelSize - 4,
textAnchor: 'middle',
fontSize: '10',
fontFamily: 'monospace',
fontWeight: isFresh ? 'bold' as any : 'normal' as any,
fill: isFresh ? freshColor : themeMgr.theme.colors.textMuted.css
}, labels[j]))
}
for (let i = 0; i < T_len; i++) {
cells.push(text({
x: labelSize - 4,
y: labelSize + i * cellSize + cellSize / 2 + 3,
textAnchor: 'end',
fontSize: '10',
fontFamily: 'monospace',
fill: themeMgr.theme.colors.textMuted.css
}, labels[i]))
}
const sel = this.selectedAttnCell
const selHere = sel && sel.layer === layer && sel.head === head
for (let i = 0; i < T_len; i++) {
for (let j = 0; j <= i; j++) {
const a = attn.data[i * T_len + j]
const isSel = selHere && sel!.qPos === i && sel!.kPos === j
cells.push(rect({
x: labelSize + j * cellSize, y: labelSize + i * cellSize,
width: cellSize - 1, height: cellSize - 1,
fill: cellColor(a),
stroke: isSel ? themeMgr.theme.colors.accent.rawValue : 'none',
strokeWidth: isSel ? '2' : '0',
class: 'cursor-pointer',
onClick: () => {
this.selectedAttnCell = { layer, head, qPos: i, kPos: j }
this.update()
}
}))
}
}
return svg({
viewBox: `0 0 ${totalSize} ${totalSize}`,
preserveAspectRatio: 'xMidYMid meet',
class: 'w-full h-auto block',
style: { maxWidth: totalSize + 'px' }
}, cells)
}
}
// ---------- Inside-the-model panel ----------
class InsidePanel extends Component implements IInsidePanel {
selectedCell: { layer: number; pos: number } | null = null
// Sticky across cell selections.
cellDetailTab: 'stages' | 'wires' | 'vocab' = 'stages'
wiresPanel = new WiresPanel()
get root() { return this.ctx.root as any as IRoot }
// Defaults to bottom-right cell (drives the next token).
get effectiveSelection(): { layer: number; pos: number } {
const inside = this.root.model.inside
const stagesLen = N_LAYERS + 1 // embedding + N_LAYERS
if (inside.length === 0) return { layer: stagesLen - 1, pos: 0 }
const step = inside[inside.length - 1]
const T = step.inputTokens.length
const sel = this.selectedCell ?? { layer: stagesLen - 1, pos: T - 1 }
return {
layer: clamp(sel.layer, 0, stagesLen - 1),
pos: clamp(sel.pos, 0, T - 1)
}
}
view() {
if (this.root.model.inside.length === 0) {
return div({ class: styles.body }, 'Waiting for prediction...')
}
// Last step's lattice subsumes all earlier ones — causal masking means
// residuals at positions 0..N depend only on inputs 0..N.
const step = this.root.model.inside[this.root.model.inside.length - 1]
return div(
div({ class: [styles.body, 'mb-2 text-center'] },
`One forward pass producing the final result. Each cell shows the model's most-likely next-token prediction at that (depth, position), with brightness = confidence. The three circled digits in the bottom row are the digits the model produced. Click any cell for its full distribution.`
),
this.latticeView(step)
)
}
latticeView(step: InsideStep) {
const T = step.inputTokens.length
const stageNames = this.root.isNarrow
? ['Emb', ...range(N_LAYERS).map(l => `L${l + 1}`)]
: ['After embedding', ...range(N_LAYERS).map(l => `After layer ${l + 1}`)]
const cellW = this.root.isNarrow ? 28 : 38
const cellH = this.root.isNarrow ? 24 : 28
const labelColW = this.root.isNarrow ? 44 : 110
const headerH = 28
const cellBg = themeMgr.theme.colors.surface.rawValue
const cellHi = themeMgr.theme.colors.primary.rawValue
const cellColor = (p: number) => lerpRgb(cellBg, cellHi, p)
const badgeSize = cellH - 6
const trueResultDigits = splitResultDigits(this.root.inputPanel.operandA + this.root.inputPanel.operandB)
const sel = this.selectedCell ?? { layer: stageNames.length - 1, pos: T - 1 }
const selLayer = clamp(sel.layer, 0, stageNames.length - 1)
const selPos = clamp(sel.pos, 0, T - 1)
return div({ class: 'overflow-auto' },
div({ class: 'mx-auto w-fit' },
div({ class: 'flex items-end mb-0.5 gap-1', style: { height: headerH + 'px' } },
div({ style: { width: labelColW + 'px' } }),
range(T).map(pos => {
// Positions > T - N_RESULT_DIGITS are the model's own previous outputs fed back.
const isPredInput = pos > T - N_RESULT_DIGITS
let colorClass = styles.mutedText
if (isPredInput) {
const idx = pos - (T - N_RESULT_DIGITS + 1)
colorClass = step.inputTokens[pos] === trueResultDigits[idx] ? styles.successHigh : styles.successLow
}
return div({
class: ['text-center text-[11px] font-mono', isPredInput ? 'font-bold' : '', colorClass],
style: { width: cellW + 'px' }
}, `${pos}\n` + TOKEN_LABELS[step.inputTokens[pos]])
})
),
stageNames.map((stageName, layerIdx) => {
const cells = step.lattice[layerIdx]
return div({ class: 'flex items-center mb-0.5 gap-1' },
div({ class: styles.rowLabel, style: { width: labelColW + 'px' } }, stageName),
range(T).map(pos => {
const cell = cells[pos]
const isFinalRow = layerIdx === stageNames.length - 1
// Final-layer cells at the last N_RESULT_DIGITS positions — each drove a generated digit.
const isGenerationCell = isFinalRow && pos >= T - N_RESULT_DIGITS
const isSelected = layerIdx === selLayer && pos === selPos
const accent = themeMgr.theme.colors.accent.css
const shadows: string[] = []
if (isSelected) shadows.push(`inset 0 0 0 3px ${accent}`)
if (shadows.length === 0) shadows.push(`inset 0 0 0 1px ${themeMgr.theme.colors.border.css}`)
let cellChild: any = TOKEN_LABELS[cell.topToken]
if (isGenerationCell) {
const idx = pos - (T - N_RESULT_DIGITS)
const isCorrect = cell.topToken === trueResultDigits[idx]
cellChild = span({
class: [styles.badge, isCorrect ? styles.successHigh : styles.successLow],
style: { width: badgeSize + 'px', height: badgeSize + 'px' }
}, TOKEN_LABELS[cell.topToken])
}
return div({
class: styles.latticeCell,
style: {
width: cellW + 'px',
height: cellH + 'px',
backgroundColor: cellColor(cell.topProb),
boxShadow: shadows.join(', ')
},
title: `${TOKEN_LABELS[cell.topToken]} ${(cell.topProb * 100).toFixed(0)}%`,
onClick: () => { this.selectedCell = { layer: layerIdx, pos }; this.update() }
}, cellChild)
})
)
})
),
this.cellDetailView(step, selLayer, selPos, stageNames)
)
}
cellDetailView(step: InsideStep, layerIdx: number, pos: number, stageNames: string[]) {
const T = step.inputTokens.length
// At result positions, inputTokens[pos+1] is the model's own prior prediction
// (possibly wrong) — use the ground-truth digit instead.
const correctToken = pos >= T - N_RESULT_DIGITS
? splitResultDigits(this.root.inputPanel.operandA + this.root.inputPanel.operandB)[pos - (T - N_RESULT_DIGITS)]
: step.inputTokens[pos + 1]
const subTabs: { id: 'stages' | 'wires' | 'vocab'; label: string }[] = [
{ id: 'stages', label: 'Residual' },
{ id: 'wires', label: 'MLP' },
{ id: 'vocab', label: 'Projection' }
]
return div({ class: styles.detailPanel },
div({ class: [styles.body, 'mb-2 text-center'] },
`Position ${pos} (token "${TOKEN_LABELS[step.inputTokens[pos]]}") · selected stage: after ${bareStageName(stageNames[layerIdx])} · correct next token at this position is "${TOKEN_LABELS[correctToken]}"`
),
div({ class: 'flex gap-x-3 gap-y-0 mb-3 flex-wrap justify-center' },
subTabs.map(t => button({
class: [styles.subTabBtn, this.cellDetailTab === t.id ? styles.subTabBtnActive : styles.subTabBtnInactive],
onClick: () => { this.cellDetailTab = t.id; this.update() }
}, t.label))
),
this.cellDetailTab === 'wires' ? this.wiresPanel.view() :
this.cellDetailTab === 'vocab' ? this.vocabView(step.lattice[layerIdx][pos], correctToken, stageNames[layerIdx]) :
this.stagesView(step, layerIdx, pos, stageNames)
)
}
// Globally normalized across stages so layer-to-layer magnitude differences are visible.
stagesView(step: InsideStep, layerIdx: number, pos: number, stageNames: string[]): VElement {
const allStageResiduals = step.lattice.map(stageCells => stageCells[pos].residual)
const stageMaxes = allStageResiduals.map(maxAbs)
const stripColor = makeDivergingStripColor(Math.max(...stageMaxes))
const stripLabelW = this.root.isNarrow ? 86 : 190
const primaryCss = themeMgr.theme.colors.primary.css
return div({ class: 'mb-4' },
div({ class: [styles.body, 'mb-1.5 text-center'] },
`Residual stream at this position, across all stages (block boundaries). After attention heads and MLPs have added into it.`
),
allStageResiduals.map((sRes, sIdx) => {
const isSel = sIdx === layerIdx
return div({ class: styles.stripRow },
div({
class: [styles.stripRowLabel, isSel ? styles.successHigh : styles.mutedText, isSel ? 'font-bold' : ''],
style: { width: stripLabelW + 'px' }
}, `${stageNames[sIdx]} max=${stageMaxes[sIdx].toFixed(2)}`),
div({
class: 'flex gap-0 flex-1 min-w-0 box-border',
style: { boxShadow: isSel ? `0 0 0 2px ${primaryCss}` : 'none' }
},
stripCells(sRes, stripColor, STRIP_CELL_H, (i, v) => `${stageNames[sIdx]} · channel ${i}: ${v.toFixed(3)}`)
)
)
})
)
}
vocabView(cell: LatticeCell, correctToken: number, stageName: string): VElement {
const topIdx = cell.probs.reduce((best, p, i) => p > cell.probs[best] ? i : best, 0)
return div(
div({ class: [styles.body, 'mb-1 text-center'] }, `Vocabulary projection after ${bareStageName(stageName)} (logit lens):`),
probBarsView({
probs: cell.probs,
picked: topIdx,
truth: correctToken,
labels: range(VOCAB).map(v => TOKEN_LABELS[v]),
maxBarH: 70
})
)
}
}
// ---------- Embeddings panel ----------
class TokenEmbeddingsPanel extends Component {
selectedEmbToken: number | null = null
embeddingScrubIdx: number | null = null // null = live latest; otherwise index into history
embeddingPlaying = false
#playTimer: number | null = null
get root() { return this.ctx.root as any as IRoot }
get embeddingHistory() { return this.root.model.embeddingHistory }
reset() {
this.embeddingScrubIdx = null
this.selectedEmbToken = null
this.embeddingPlaying = false
if (this.#playTimer != null) { clearInterval(this.#playTimer); this.#playTimer = null }
}
// null = live (latest frame).
get scrubValue(): number {
const len = this.embeddingHistory.length
if (len === 0) return 0
return this.embeddingScrubIdx == null ? len - 1 : this.embeddingScrubIdx
}
set scrubValue(v: number) {
const len = this.embeddingHistory.length
if (len === 0) return
this.#stopEmbeddingPlay()
this.embeddingScrubIdx = v >= len - 1 ? null : Math.max(0, v)
}
view() {
const flat = this.#currentEmbeddingFlat()
if (!flat || flat.length === 0) {
return div({ class: styles.body }, 'Waiting for first snapshot...')
}
const displayedRows = this.#rowsFromFlat(flat)
// Axes locked to LATEST snapshot so they don't wobble while scrubbing.
const latestFlat = this.embeddingHistory.length > 0
? this.embeddingHistory[this.embeddingHistory.length - 1].tokEmb
: flat
const latestRows = this.#rowsFromFlat(latestFlat)
const basis = pcaBasis(latestRows)
const latestPoints = projectThroughBasis(latestRows, basis)
const allXs = latestPoints.map(p => p[0])
const allYs = latestPoints.map(p => p[1])
const xRange = Math.max(...allXs) - Math.min(...allXs)
const yRange = Math.max(...allYs) - Math.min(...allYs)
// Pad 10% so early frames (clustered near center) have room.
const pad = 0.1
const minX = Math.min(...allXs) - xRange * pad
const maxX = Math.max(...allXs) + xRange * pad
const minY = Math.min(...allYs) - yRange * pad
const maxY = Math.max(...allYs) + yRange * pad
const points = projectThroughBasis(displayedRows, basis)
const W = 480, H = 360, plotPad = 30
const sx = (x: number) => plotPad + (x - minX) / Math.max(1e-9, maxX - minX) * (W - 2 * plotPad)
const sy = (y: number) => H - plotPad - (y - minY) / Math.max(1e-9, maxY - minY) * (H - 2 * plotPad)
const labelColor = (v: number) => v < 10 ? themeMgr.theme.colors.primary.rawValue : themeMgr.theme.colors.text.rawValue
const embStripColor = makeDivergingStripColor(
Math.max(...displayedRows.map(r => maxAbs(r)))
)
const sel = this.selectedEmbToken
const histLen = this.embeddingHistory.length
const isLive = this.embeddingScrubIdx == null
const displayStep = isLive ? this.root.model.trainStep : this.embeddingHistory[this.embeddingScrubIdx!].step
const sliderMax = histLen > 0 ? histLen - 1 : 0
return div(
div({ class: [styles.body, 'mb-2'] }, 'Each token\'s 64-dim embedding, shown in 2D — the plane that spreads the 12 tokens as far apart as possible (PCA). As training progresses, the digit tokens walk into a circle in numerical order — 0, 1, 2, …, 9, with 9 next to 0 — and the model has discovered the cyclic structure of mod-10 arithmetic. The "+" and "=" tokens sit near the center. The axes are locked to the latest snapshot, so when you scrub backward through training history the projection doesn\'t wobble — you watch the circle actually form.'),
div({ class: [styles.body, 'mb-2'] }, 'Reminiscent of the "grokking" paper (Power et al.): small transformers spontaneously discovering algebraic structure (here, the mod-10 cycle of digits) often after a long memorization plateau before a sudden breakthrough. Click any token to see its raw 64-dim embedding row.'),
div({ class: styles.chartBox },
svg({
viewBox: `0 0 ${W} ${H}`,
preserveAspectRatio: 'xMidYMid meet',
class: 'block mx-auto h-auto',
style: { maxWidth: W + 'px', width: '100%' }
},
points.map((p, i) => text({
x: sx(p[0]), y: sy(p[1]) + 5,
textAnchor: 'middle',
fontSize: i < 10 ? '20' : '18',
fontFamily: 'monospace',
fontWeight: 'bold',
fill: labelColor(i),
stroke: i === sel ? themeMgr.theme.colors.accent.css : 'none',
strokeWidth: i === sel ? '1' : '0',
class: 'cursor-pointer',
onClick: () => {
this.selectedEmbToken = this.selectedEmbToken === i ? null : i
this.update()
}
}, TOKEN_LABELS[i]))
)
),
sel != null && displayedRows[sel] ? div({ class: 'mt-3' },
div({ class: [styles.body, 'mb-1'] },
`Token "${TOKEN_LABELS[sel]}" — its ${D_MODEL}-dim row from tok_emb. These are the numbers PCA projects to the (x,y) point above. Click another token to switch, or click "${TOKEN_LABELS[sel]}" again to hide.`
),
div({ class: 'flex gap-0' },
stripCells(displayedRows[sel], embStripColor, STRIP_CELL_H, (i, v) => `dim ${i}: ${v.toFixed(3)}`)
)
) : div(),
div({ class: 'flex items-center gap-3 mt-2.5' },
button({
class: [styles.iconBtn, styles.iconBtnSmall],
onClick: () => this.#toggleEmbeddingPlay(),
title: this.embeddingPlaying ? 'Pause' : 'Play formation animation'
}, this.embeddingPlaying ? iconPause() : iconPlay()),
button({
class: styles.btn,
onClick: () => { this.embeddingScrubIdx = null; this.update() }
}, isLive ? 'Live' : 'Jump to live'),
div({ class: 'flex-1' },
formField({
target: this,
prop: () => this.scrubValue,
inputFn: inputRange,
label: '',
inputProps: {
attrs: { min: 0, max: sliderMax, step: 1, style: { width: '100%' } }
}
})
),
div({ class: [styles.monoLabelTiny, 'text-right'], style: { minWidth: '120px' } },
`step ${displayStep.toLocaleString()} ${isLive ? '(live)' : `· ${this.embeddingScrubIdx! + 1}/${histLen}`}`
)
)
)
}
#toggleEmbeddingPlay() {
if (this.embeddingPlaying) { this.#stopEmbeddingPlay(); return }
if (this.embeddingHistory.length < 2) return
this.embeddingPlaying = true
this.embeddingScrubIdx = 0
this.#playTimer = (setInterval(() => {
const len = this.embeddingHistory.length
if (len === 0) { this.#stopEmbeddingPlay(); return }
const cur = this.embeddingScrubIdx ?? len - 1
const next = cur + 1
if (next >= len) {
this.embeddingScrubIdx = null
this.#stopEmbeddingPlay()
return
}
this.embeddingScrubIdx = next
this.update()
}, 90) as unknown) as number
this.update()
}
#stopEmbeddingPlay() {
if (this.#playTimer != null) { clearInterval(this.#playTimer); this.#playTimer = null }
this.embeddingPlaying = false
this.update()
}
#currentEmbeddingFlat(): Float32Array | null {
if (this.embeddingScrubIdx != null && this.embeddingHistory[this.embeddingScrubIdx]) {
return this.embeddingHistory[this.embeddingScrubIdx].tokEmb
}
return this.root.model.getTokEmbSnapshot()
}
#rowsFromFlat(flat: Float32Array): number[][] {
const rows: number[][] = []
for (let i = 0; i < VOCAB; i++) {
const r: number[] = []
for (let j = 0; j < D_MODEL; j++) r.push(flat[i * D_MODEL + j])
rows.push(r)
}
return rows
}
}
// ---------- Input panel ----------
class InputPanel extends Component implements IInputPanel {
// Setters clamp to 0..99 and trigger a prediction refresh.
_operandA = 27
_operandB = 45
get operandA() { return this._operandA }
set operandA(v: number) { this._operandA = clamp(Math.floor(v), 0, 99); this.root.model.refreshPrediction() }
get operandB() { return this._operandB }
set operandB(v: number) { this._operandB = clamp(Math.floor(v), 0, 99); this.root.model.refreshPrediction() }
get root() { return this.ctx.root as any as IRoot }
get model() { return this.root.model }
setPreset(a: number, b: number) {
// Set storage directly to avoid two refreshes via the setters.
this._operandA = clamp(Math.floor(a), 0, 99)
this._operandB = clamp(Math.floor(b), 0, 99)
this.model.refreshPrediction()
this.update()
}
pickRandomPair(predicate: (a: number, b: number) => boolean, maxAttempts: number) {
let a = 0, b = 0
for (let attempt = 0; attempt < maxAttempts; attempt++) {
a = Math.floor(Math.random() * 100); b = Math.floor(Math.random() * 100)
if (predicate(a, b)) break
}
this.setPreset(a, b)
}
pickRandomTraining() { this.pickRandomPair((a, b) => !isTestPair(a, b), 200) }
pickRandomHeldOut() { this.pickRandomPair(isTestPair, 5000) }
pickRandomFailure() {
if (this.model.recentFailures.length === 0) return
const f = this.model.recentFailures[Math.floor(Math.random() * this.model.recentFailures.length)]
this.setPreset(f.a, f.b)
}
view() {
const trueDigits = splitResultDigits(this._operandA + this._operandB)
const predDigits = this.model.predictedDigits
const ready = predDigits.length === N_RESULT_DIGITS
// Hide the hundreds slot for sums < 100.
const showHundreds = ready && (predDigits[2] !== 0 || trueDigits[2] !== 0)
const places = showHundreds ? [2, 1, 0] : [1, 0]
const resultChildren = ready
? places.map(place => span({
class: predDigits[place] === trueDigits[place] ? styles.successHigh : styles.successLow
}, String(predDigits[place])))
: '?'
return div({ class: styles.panel },
div({ class: styles.sectionTitle }, 'Try an addition'),
div({ class: [styles.body, 'mb-2'] }, 'Type two numbers (0-99). Watch the model predict the sum, digit by digit.'),
div({ class: styles.additionRow },
formField({
target: this,
prop: () => this.operandA,
inputFn: inputNumber,
label: '',
inputProps: {
numberParams: { min: 0, max: 99, maximumFractionDigits: 0 },
attrs: { class: styles.numberInput }
}
}),
span({ class: styles.bigOp }, '+'),
formField({
target: this,
prop: () => this.operandB,
inputFn: inputNumber,
label: '',
inputProps: {
numberParams: { min: 0, max: 99, maximumFractionDigits: 0 },
attrs: { class: styles.numberInput }
}
}),
span({ class: styles.bigOp }, '='),
span({ class: styles.bigResult }, resultChildren)
),
div({ class: [styles.body, 'mb-2'] },
`True answer: ${this._operandA + this._operandB} ${isTestPair(this._operandA, this._operandB) ? '(held out from training)' : '(seen in training)'}`
),
div({ class: styles.presetRow },
button({ class: styles.btn, onClick: () => this.pickRandomTraining() }, 'Random in-training'),
button({ class: styles.btn, onClick: () => this.pickRandomHeldOut() }, 'Random held-out'),
button({ class: styles.btn, onClick: () => this.pickRandomFailure(), title: this.model.recentFailures.length === 0 ? 'No recent failures yet — train a bit first' : 'Pick a held-out pair the model is currently getting wrong' }, 'Random failure')
),
div({ class: [styles.sectionTitle, 'mt-6'] }, 'Per-digit prediction'),
this.digitProbsView()
)
}
digitProbsView() {
if (this.model.digitProbs.length === 0) return div({ class: styles.body }, 'No prediction yet.')
const trueDigits = splitResultDigits(this._operandA + this._operandB)
const labels = range(10).map(d => String(d))
return div({ class: styles.digitProbs },
this.model.digitProbs.map((probs, step) => {
const placeName = RESULT_DIGIT_NAMES[step] ?? `pos${step}`
const picked = this.model.predictedDigits[step]
return div(
probBarsView({
probs,
picked,
truth: trueDigits[step],
labels,
maxBarH: 50
}),
div({ class: [styles.bodyText, 'font-bold text-center mt-0.5'] }, placeName)
)
})
)
}
}
// ---------- Wires panel ----------
// MLP drilldown for InsidePanel's selected cell: residual-in → hidden ReLU →
// residual-out strips, connected by the strongest mlp.up.W / mlp.down.W weights.
class WiresPanel extends Component {
get root() { return this.ctx.root as any as IRoot }
get parent() { return this.ctx.parent as any as IInsidePanel }
view() {
const sel = this.parent.effectiveSelection
if (sel.layer === 0) {
return div({ class: [styles.body, 'text-center my-6'] },
'No MLP at the embedding stage — embeddings are just a lookup. Click any cell in the layer 1, 2, or 3 rows above to see that layer\'s MLP wires.'
)
}
const mlpLayerIdx = sel.layer - 1
const pos = sel.pos
const weights = this.root.model.getMlpWeights(mlpLayerIdx)
if (!weights) return div({ class: styles.body }, 'Initializing weights…')
const { w1, w2, b2 } = weights
const step = this.root.model.inside.at(-1) ?? null
const hidden = step?.mlpHiddens[mlpLayerIdx]?.[pos] ?? null
const afterResidual = step?.lattice[mlpLayerIdx + 1]?.[pos]?.residual ?? null
// residualOut = mlp.down.W · hidden + b2 (MLP contribution to the residual stream)
// residualIn = afterResidual − residualOut
let residualIn: Float32Array | null = null
let residualOut: Float32Array | null = null
if (hidden && afterResidual) {
residualOut = new Float32Array(D_MODEL)
for (let i = 0; i < D_MODEL; i++) {
let s = b2[i]
for (let j = 0; j < HIDDEN; j++) s += hidden[j] * w2[j * D_MODEL + i]
residualOut[i] = s
}
residualIn = new Float32Array(D_MODEL)
for (let i = 0; i < D_MODEL; i++) residualIn[i] = afterResidual[i] - residualOut[i]
}
const firingNote = hidden
? ` (${countFiring(hidden)}/${HIDDEN} firing — color = activation magnitude; never blue, ReLU forbids it)`
: ''
return div(
div({ class: [styles.body, 'mb-1 text-center'] },
`Inside layer ${mlpLayerIdx + 1}'s MLP at this position. ReLU on the hidden layer is where the transformer gets most of its non-linearity. The hidden strip is never blue — ReLU forbids negatives.`
),
this.renderSvg(w1, w2, residualIn, hidden, residualOut)
)
}
renderSvg(
w1: Float32Array,
w2: Float32Array,
residualIn: ArrayLike<number> | null,
hidden: ArrayLike<number> | null,
residualOut: ArrayLike<number> | null
) {
const W = 700, H = 222
const PAD_X = 24
const STRIP_X0 = PAD_X
const STRIP_W = W - 2 * PAD_X
const ROW_H = STRIP_CELL_H
const Y_TOP = 20
const STRIP_Y0 = 100
const STRIP_Y1 = STRIP_Y0 + ROW_H
const Y_BOT = 184
const xEdge = (j: number, n: number) => STRIP_X0 + STRIP_W * j / n
const xCenter = (j: number, n: number) => STRIP_X0 + STRIP_W * (j + 0.5) / n
const TOP_N = 300
const w1Top = pickTopWeights(w1, D_MODEL, HIDDEN, TOP_N)
const w2Top = pickTopWeights(w2, HIDDEN, D_MODEL, TOP_N)
// w*Top is sorted by descending |w|, so [0] is the matrix's global max.
const w1Color = makeDivergingStripColor(Math.abs(w1Top[0]?.w ?? 1e-6))
const w2Color = makeDivergingStripColor(Math.abs(w2Top[0]?.w ?? 1e-6))
const cellBg = themeMgr.theme.colors.surface.rawValue
const renderRow = (vals: ArrayLike<number> | null, n: number, y: number): VElement[] => {
const color = vals ? makeDivergingStripColor(maxAbs(vals) || 1) : null
const cellW = STRIP_W / n
const cells: VElement[] = []
for (let i = 0; i < n; i++) {
cells.push(rect({
// +0.5 prevents hairline gaps between cells at non-integer cellW.
x: xEdge(i, n), y, width: cellW + 0.5, height: ROW_H,
fill: color && vals ? color(vals[i]) : cellBg
}))
}
return cells
}
const drawFan = (
top: { i: number; j: number; w: number }[],
srcN: number, srcY: number, dstN: number, dstY: number,
color: (w: number) => string
): VElement[] => top.map(c => line({
x1: xCenter(c.i, srcN), y1: srcY,
x2: xCenter(c.j, dstN), y2: dstY,
stroke: color(c.w), strokeWidth: 0.6
}))
const lines = [
...drawFan(w1Top, D_MODEL, Y_TOP + ROW_H, HIDDEN, STRIP_Y0, w1Color),
...drawFan(w2Top, HIDDEN, STRIP_Y1, D_MODEL, Y_BOT, w2Color)
]
const topCells = renderRow(residualIn, D_MODEL, Y_TOP)
const midCells = renderRow(hidden, HIDDEN, STRIP_Y0)
const botCells = renderRow(residualOut, D_MODEL, Y_BOT)
const labelColor = themeMgr.theme.colors.textMuted.css
const textColor = themeMgr.theme.colors.text.css
const stripCenterY = STRIP_Y0 + ROW_H / 2 + 4 // +4 ≈ baseline offset for 11px text
const labels: VElement[] = [
text({ x: PAD_X, y: 14, fill: labelColor, fontSize: '11', textAnchor: 'start' }, `Residual in (${D_MODEL} channels) — flows into MLP`),
text({ x: W - PAD_X, y: 14, fill: labelColor, fontSize: '11', textAnchor: 'end' }, 'mlp.up.W wires below ↓'),
text({ x: W / 2, y: stripCenterY, fill: textColor, fontSize: '11', textAnchor: 'middle', fontWeight: 'bold' as any }, `Hidden ReLU (${HIDDEN} neurons)`),
text({ x: PAD_X, y: Y_BOT + ROW_H + 14, fill: labelColor, fontSize: '11', textAnchor: 'start' }, `MLP output (${D_MODEL} channels) — added into the residual stream`),
text({ x: W - PAD_X, y: Y_BOT + ROW_H + 14, fill: labelColor, fontSize: '11', textAnchor: 'end' }, '↑ mlp.down.W wires above')
]
return svg({
viewBox: `0 0 ${W} ${H}`,
width: '100%',
class: 'block max-w-full',
preserveAspectRatio: 'xMidYMid meet'
}, ...lines, ...topCells, ...midCells, ...botCells, ...labels)
}
}
// ---------- Explainer panel ----------
type ExplainerTopic = 'training' | 'architecture' | 'notes'
class ExplainerPanel extends Component {
get root() { return this.ctx.root as any as IRoot }
selectedTopic: ExplainerTopic = 'training'
view() {
const topics: { id: ExplainerTopic; label: string }[] = [
{ id: 'training', label: 'Training' },
{ id: 'architecture', label: 'Architecture' },
{ id: 'notes', label: 'Notes' }
]
return div({ class: styles.explainerContainer },
div({ class: 'flex gap-x-3 gap-y-0 mb-3 flex-wrap' },
topics.map(t => button({
class: [styles.subTabBtn, this.selectedTopic === t.id ? styles.subTabBtnActive : styles.subTabBtnInactive],
onClick: () => { this.selectedTopic = t.id; this.update() }
}, t.label))
),
this.topicView(),
this.liveTransformerDiagram()
)
}
topicView() {
switch (this.selectedTopic) {
case 'training': return this.trainingView()
case 'architecture': return this.architectureView()
case 'notes': return this.notesView()
}
}
trainingView() {
return div(
p('A small transformer is being trained from scratch in your browser to add 2-digit numbers. The vocabulary is just digits 0-9, "+", and "=". Each training example is one addition: "27+45=270". (The answer 72 is padded to 3 digits and reversed — the model generates the result one digit at a time, units first, so carries flow naturally left-to-right.)'),
p('At each training step, the model tries to predict the sum. How wrong those predictions are is the signal that drives every adjustment the model makes. If you\'re new to backprop and optimization, watch this ',
a({ href: 'https://typebulb.com/u/samples/xor-x-ray/full', target: '_blank' }, 'XOR X-ray'),
' demo, which has a simple neural network that you can see learn in real-time.'),
p('In our model, we use cross-entropy to measure wrongness, and AdamW for the optimizer. The loss is masked to the result-digit positions only — the model is graded on getting the answer right, not on predicting the operands. We hold out 20% of (a, b) pairs from training. Held-out accuracy measures real generalization: can the model add numbers it has never seen?')
)
}
architectureView() {
return div(
p(`What flows through every channel in the diagram below is a ${D_MODEL}-dim vector — a single point in ${D_MODEL}-dimensional space, where the model encodes information as *directions*. At the bottom, the residual starts where token and position vectors merge at ⊕ (so the model sees order); ${D_MODEL} dimensions has enough room to keep them distinguishable downstream. For the foundations — what a tensor is, what these additions are doing here — see this `, a({ href: 'https://typebulb.com/u/samples/tensors/full', target: '_blank' }, 'interactive Tensors tutorial'), '.'),
p(`The residual stream is the vertical channel: at every position, it runs upward through all ${N_LAYERS} layers. Every block in every layer reads from it and writes back into it at ⊕. In the diagram, it's the green vertical line at each position. At the top, the same token-embedding matrix transposed turns the final residual into a prediction.`),
p('The K/V stream is the horizontal channel: at each layer, K and V at every position are made available to all later positions in that same layer. Attention reads from it; each position\'s MLP reads only its own residual. In the diagram, the K/V bus is the horizontal line under each layer; each purple K/V circle writes to it, each attention block reads from it. Causal flow runs left-to-right, but the work is parallel — what RNNs do step by step, transformers do in one pass.')
)
}
notesView() {
return div(
p('Even the most sophisticated transformers follow this basic model. They do, however, have many refinements. For example, token positions are usually fed into the attention heads as rotations to pairs of K and Q projections (RoPE). In this example, we keep it simple, and positions are simply embedded with tokens at the start of each residual stream.'),
p('Attention is almost all linear algebra, aside from softmax. The bulk of the nonlinearity in the model lives in the MLPs.'),
p('See ', a({ href: 'https://tinyurl.com/44ayrzfp', target: '_blank' }, 'this transformer diagrammed by nn-dna'), ', which generates architecture diagrams from plain-English descriptions of neural networks.'),
p('Co-built with Claude Opus 4.7; inspired by @repligate / j⧉nus\'s "How Information Flows Through Transformers".')
)
}
// In-SVG so it travels with the diagram in screenshots; textLength+lengthAdjust
// forces consistent inter-item gaps despite per-glyph width variance.
diagramLegendSvg(W: number, H: number): VElement[] {
const muted = themeMgr.theme.colors.textMuted.css
const c = themeMgr.theme.colors
const y = H - 10
const swatchGap = 6, itemGap = 20, charW = 6, fontSize = '11'
type Shape = { width: number; render: (fill: string, x: number) => VElement }
const SHAPES: Record<string, Shape> = {
circle: { width: 10, render: (f, x) => circle({ cx: x + 5, cy: y - 4, r: '5', fill: f, stroke: 'none' }) },
square: { width: 10, render: (f, x) => rect({ x, y: y - 9, width: '10', height: '10', fill: f, stroke: 'none', rx: '1.5' }) },
line: { width: 20, render: (f, x) => rect({ x, y: y - 5, width: '20', height: '2', fill: f, stroke: 'none' }) }
}
const items: { kind: string; color: string; label: string }[] = [
{ kind: 'circle', color: c.accent.css, label: 'K/V Computation' },
{ kind: 'square', color: c.attn.css, label: 'Attention Computation' },
{ kind: 'square', color: c.mlp.css, label: 'MLP Computation' },
{ kind: 'line', color: c.primary.css, label: 'Residual Stream' },
{ kind: 'line', color: c.accent.css, label: 'K/V Stream' }
]
const labelW = (s: string) => s.length * charW
const itemW = (it: typeof items[0]) => SHAPES[it.kind].width + swatchGap + labelW(it.label)
const totalW = items.reduce((s, it) => s + itemW(it), 0) + itemGap * (items.length - 1)
const out: VElement[] = []
let x = (W - totalW) / 2
for (const it of items) {
const shape = SHAPES[it.kind]
out.push(shape.render(it.color, x))
out.push(text({
x: x + shape.width + swatchGap, y,
textAnchor: 'start', fontSize, fontFamily: 'sans-serif',
fill: muted,
textLength: String(labelW(it.label)),
lengthAdjust: 'spacingAndGlyphs' as any
}, it.label))
x += itemW(it) + itemGap
}
return out
}
// Wiring follows j⧉nus's "How Information Flows Through Transformers" — per-node
// K/V circle + attention block + MLP block, with a horizontal K/V bus per layer.
liveTransformerDiagram() {
const a = this.root.inputPanel.operandA
const b = this.root.inputPanel.operandB
const predicted = this.root.model.predictedDigits // [ones, tens, hundreds] when ready
const trueResultDigits = splitResultDigits(a + b)
// Position 6 = predicted ones fed back, position 7 = predicted tens fed back.
const prompt = promptTokens(a, b)
const inputTok = [...prompt, predicted[0] ?? -1, predicted[1] ?? -1]
const groundTruthNext = prompt.slice(1)
const tokLabel = (t: number) => (t < 0 || t >= TOKEN_LABELS.length) ? '?' : TOKEN_LABELS[t]
const N_POS = 8
const N_ROWS = N_LAYERS + 2 // head + N_LAYERS layers + embed
const cellW = 80
const nodeH = 92 // layer-node height; tuned so the K/V→attn stub reads as a distinct segment
const embedNodeH = 50 // embedding row: tok + pos lookups feed ⊕
const headNodeH = 50 // head row: lnf + tied unembed projection (residual → vocab)
const padL = 80, padR = 20, padT = 14, padB = 32 // stable frame; padB fits the in-SVG legend
// Content shift inside the stable frame. Reducing padL would shrink W and
// (with `w-full` sizing) make the whole SVG appear zoomed in; instead, we
// keep W/H constant and offset content (and the legend, via <g>) here.
const shiftX = -8
const shiftY = -10
const rowHeights = [headNodeH, nodeH, nodeH, nodeH, embedNodeH]
const W = padL + cellW * N_POS + padR
const H = padT + rowHeights.reduce((s, h) => s + h, 0) + padB
const xAt = (p: number) => padL + shiftX + cellW * (p + 0.5)
const yTop = (rowFromTop: number) =>
padT + shiftY + rowHeights.slice(0, rowFromTop).reduce((s, h) => s + h, 0)
// Head at top, embedding at bottom; layers in between.
const rowHead = 0
const rowEmbed = N_ROWS - 1
const layerRowsTopDown = [1, 2, 3]
const rowLabels = ['Unembed', 'Layer 3', 'Layer 2', 'Layer 1', 'Embed']
const residualColor = themeMgr.theme.colors.primary.css
const errorColor = themeMgr.theme.colors.error.css
const kvColor = themeMgr.theme.colors.accent.css
const attnColor = themeMgr.theme.colors.attn.css
const mlpColor = themeMgr.theme.colors.mlp.css
const mutedCol = themeMgr.theme.colors.textMuted.css
const textCol = themeMgr.theme.colors.text.css
const borderCol = themeMgr.theme.colors.border.css
const nodeFill = themeMgr.theme.colors.surface.css
const bgColor = themeMgr.theme.colors.background.css
const elements: VElement[] = []
// ===== Helpers =====
// All draw helpers push into `elements`.
// Marching dashes. Caller must orient the path so draw direction == flow
// direction (the residual is drawn bottom-to-top because flow is upward).
// 4 solid / 2 gap, 6 px/sec — short cycle so stubs (~6px) show a full dash.
const flowPath = (d: string, color: string, strokeWidth: number = 1.5) => {
elements.push(path({
d, stroke: color, strokeWidth: String(strokeWidth), fill: 'none',
strokeDashArray: '4 2'
}, animate({
attributeName: 'stroke-dashoffset',
from: '0', to: '-6',
dur: '1s',
repeatCount: 'indefinite'
})))
}
// Triangle arrowhead at (x, y) pointing in `dir`. Axial = wing offset along
// the arrow axis; perp = wing spread perpendicular to it.
const arrowhead = (
x: number, y: number, dir: 'up' | 'down' | 'left' | 'right',
color: string, axial = 3.5, perp = 2.5,
) => {
const d =
dir === 'right' ? `M ${x - axial} ${y - perp} L ${x} ${y} L ${x - axial} ${y + perp} Z` :
dir === 'left' ? `M ${x + axial} ${y - perp} L ${x} ${y} L ${x + axial} ${y + perp} Z` :
dir === 'up' ? `M ${x - perp} ${y + axial} L ${x} ${y} L ${x + perp} ${y + axial} Z` :
`M ${x - perp} ${y - axial} L ${x} ${y} L ${x + perp} ${y - axial} Z`
elements.push(path({ d, fill: color, stroke: 'none' }))
}
// Flow line ending in an arrowhead at (x2, y2). Horizontal or vertical only.
const drawArrow = (x1: number, y1: number, x2: number, y2: number, color: string) => {
flowPath(`M ${x1} ${y1} L ${x2} ${y2}`, color)
const dir = y1 === y2 ? (x2 > x1 ? 'right' : 'left') : (y2 > y1 ? 'down' : 'up')
arrowhead(x2, y2, dir, color)
}
// Filled with bg so the residual line is masked inside the circle — otherwise
// the vertical residual stroke overlaps the cross arm and ⊕ reads as ⊖.
const drawMerge = (mcx: number, mcy: number) => {
elements.push(circle({
cx: mcx, cy: mcy, r: '4',
fill: bgColor, stroke: residualColor, strokeWidth: '1.5'
}))
elements.push(path({
d: `M ${mcx - 2.5} ${mcy} L ${mcx + 2.5} ${mcy} M ${mcx} ${mcy - 2.5} L ${mcx} ${mcy + 2.5}`,
stroke: residualColor, strokeWidth: '1.2', fill: 'none'
}))
}
const drawBlock = (cx: number, nodeY: number, topY: number, color: string) => {
elements.push(rect({
x: blockX(cx), y: nodeY + topY, width: String(blockW), height: String(blockH),
fill: color, stroke: 'none', rx: '2'
}))
}
// ===== Drawing =====
// Residual: drawn bottom-to-top so dashes march upward (flow direction).
// Starts at the ⊕ inside the embed row (tok + pos merge); ends inside the
// head row entering the unembed box from below (which then displays the
// model's prediction at that position).
const embedMergeY = 7 // ⊕ y-offset within embed row
const headArrowY = 37 // arrowhead tip y-offset within head row (box bottom)
for (let p = 0; p < N_POS; p++) {
flowPath(
`M ${xAt(p)} ${yTop(rowEmbed) + embedMergeY} L ${xAt(p)} ${yTop(rowHead) + headArrowY + 6}`,
residualColor, 2
)
arrowhead(xAt(p), yTop(rowHead) + headArrowY, 'up', residualColor, 6, 4)
}
// Each layer node = attn block (bottom) + MLP block (top), each with a pre-LN
// side trip: residual splits, branch enters block, block output ⊕ skip = new
// residual. K/V circle sits in the attn side trip on the K/V bus.
const blockW = 18
const blockH = 12
const blockX = (cx: number) => cx + 20
const blockCx = (cx: number) => cx + 29 // K/V circle x
// Per-node vertical layout (y from nodeY = top of row). All other per-node
// y-coordinates derive from these.
const layerY = {
mlpTop: 14, // ⊕ MLP merge / MLP block top edge
mlpBottom: 26, // MLP block bottom = MLP-block split point
attnTop: 34, // ⊕ attn merge / attn block top edge
attnBottom: 46, // attn block bottom = top of K/V→attn stub
kvStubBottom: 58, // bottom of K/V→attn stub
kvCenter: 66, // K/V circle center / K/V bus y
kvBottom: 71, // K/V circle bottom (kvCenter + r=5)
attnSplit: 84, // residual → K/V branch starts below the bus
}
const busOffsetY = layerY.kvCenter
// Bus runs only between K/V_0 and K/V_{N-1} — pos 0 has no left bus neighbor,
// pos N-1 has no right one. Single mid-row arrowhead shows causal direction.
for (const rowL of layerRowsTopDown) {
const busY = yTop(rowL) + busOffsetY
const bcxFirst = blockCx(xAt(0))
const bcxLast = blockCx(xAt(N_POS - 1))
flowPath(
`M ${bcxFirst} ${busY} L ${bcxLast} ${busY}`,
kvColor, 2
)
// Arrow placed mid-bus (not at endpoint) so it doesn't suggest continuation.
const arrowTip = (blockCx(xAt(3)) + blockCx(xAt(4))) / 2
arrowhead(arrowTip, busY, 'right', kvColor, 10, 5)
}
// Per-node internals (layer rows only). Two non-obvious geometric facts:
// - Residual at cx visibly crosses the bus at kvCenter but doesn't
// interact — only the K/V circle (at bcx) is on the bus.
// - attnSplit is BELOW the bus so the Manhattan branch up to K/V doesn't
// overlap the prior K/V circle's bus segment.
for (const rowL of layerRowsTopDown) {
const nodeY = yTop(rowL)
for (let p = 0; p < N_POS; p++) {
const cx = xAt(p)
const bcx = blockCx(cx)
drawBlock(cx, nodeY, layerY.mlpTop, mlpColor)
// Block arrows inset 2px from the corners so they don't appear to come from the corner points.
drawArrow(cx, nodeY + layerY.mlpBottom - 2, blockX(cx), nodeY + layerY.mlpBottom - 2, residualColor)
drawArrow(blockX(cx), nodeY + layerY.mlpTop + 2, cx + 4, nodeY + layerY.mlpTop + 2, residualColor)
drawMerge(cx, nodeY + layerY.mlpTop + 2)
drawBlock(cx, nodeY, layerY.attnTop, attnColor)
// Attn-input Manhattan branch routes BELOW the bus to avoid overlapping
// the prior position's bus segment.
flowPath(
`M ${cx} ${nodeY + layerY.attnSplit} L ${bcx} ${nodeY + layerY.attnSplit} L ${bcx} ${nodeY + layerY.kvBottom}`,
residualColor
)
arrowhead(bcx, nodeY + layerY.kvBottom, 'up', residualColor)
drawArrow(bcx, nodeY + layerY.kvStubBottom, bcx, nodeY + layerY.attnBottom, kvColor)
drawArrow(blockX(cx), nodeY + layerY.attnTop + 2, cx + 4, nodeY + layerY.attnTop + 2, residualColor)
drawMerge(cx, nodeY + layerY.attnTop + 2)
// K/V circle drawn last so it sits on top of the arrow tip + bus line.
elements.push(circle({
cx: bcx, cy: nodeY + layerY.kvCenter, r: '5',
fill: kvColor, stroke: 'none'
}))
}
}
// Boxes for embed (tok + pos) and unembed: same size, hold actual values.
// Stroke is `border` (low-contrast palette tone) so the cell outlines don't
// outshine the digits inside — including green digits on correct predictions.
const lookupW = 24
const lookupH = 22
const lookupStroke = borderCol
const lookupYHead = 15
// Embed boxes shifted down so the arrow's vertical leg equals its
// horizontal leg (perfectly square turn into ⊕).
const lookupYEmbed = embedMergeY + 10 // boxCx (14) − ⊕ radius (4)
// Embed row: tok and pos lookups (off-axis) feed ⊕ on the residual line.
// The ⊕ is where positional info enters the model. Positions 6/7 hold the
// model's prior ones/tens prediction fed back (colored by correctness).
for (let p = 0; p < N_POS; p++) {
const nodeY = yTop(rowEmbed)
const cx = xAt(p)
const mergeY = nodeY + embedMergeY
const boxY = nodeY + lookupYEmbed
const tok = inputTok[p]
const isResultInput = p >= 6
const resultInputWrong = isResultInput && tok >= 0 && tok !== trueResultDigits[p - 6]
const tokColor =
!isResultInput ? textCol :
tok < 0 ? mutedCol :
resultInputWrong ? errorColor :
residualColor
// tok box (left of center)
elements.push(rect({
x: cx - 26, y: boxY, width: String(lookupW), height: String(lookupH),
fill: nodeFill, stroke: lookupStroke, strokeWidth: '1.5', rx: '2'
}))
elements.push(text({
x: cx - 14, y: boxY + 16,
textAnchor: 'middle', fontSize: '15', fontFamily: 'monospace', fontWeight: 'bold',
fill: tokColor
}, tokLabel(tok)))
// pos box (right of center)
elements.push(rect({
x: cx + 2, y: boxY, width: String(lookupW), height: String(lookupH),
fill: nodeFill, stroke: lookupStroke, strokeWidth: '1.5', rx: '2'
}))
elements.push(text({
x: cx + 14, y: boxY + 16,
textAnchor: 'middle', fontSize: '15', fontFamily: 'monospace',
fill: mutedCol
}, String(p)))
// Branches up to ⊕. Arrowheads end at the circle edge (r=4).
flowPath(`M ${cx - 14} ${boxY} L ${cx - 14} ${mergeY} L ${cx - 4} ${mergeY}`, residualColor)
arrowhead(cx - 4, mergeY, 'right', residualColor)
flowPath(`M ${cx + 14} ${boxY} L ${cx + 14} ${mergeY} L ${cx + 4} ${mergeY}`, residualColor)
arrowhead(cx + 4, mergeY, 'left', residualColor)
drawMerge(cx, mergeY)
}
// Head row: residual is consumed by the unembed projection (matmul with
// tok_emb^T, preceded by final LN). Box shows the model's prediction at
// each position — positions 5/6/7 are the ones/tens/hundreds result
// digits (green/red by correctness); 0–4 are untrained next-operand
// guesses, shown muted.
for (let p = 0; p < N_POS; p++) {
const nodeY = yTop(rowHead)
const cx = xAt(p)
const boxY = nodeY + lookupYHead
let label = '?'
let isPred = false
let isWrong = false
if (p < 5) {
label = tokLabel(groundTruthNext[p])
} else {
const idx = p - 5
const t = predicted[idx]
if (t !== undefined) {
label = tokLabel(t)
isPred = true
isWrong = t !== trueResultDigits[idx]
}
}
const fillColor = isPred ? (isWrong ? errorColor : residualColor) : mutedCol
elements.push(rect({
x: cx - 12, y: boxY, width: String(lookupW), height: String(lookupH),
fill: nodeFill, stroke: lookupStroke, strokeWidth: '1.5', rx: '2'
}))
elements.push(text({
x: cx, y: boxY + 16,
textAnchor: 'middle', fontSize: '15', fontFamily: 'monospace', fontWeight: 'bold',
fill: fillColor
}, label))
}
for (let i = 0; i < N_ROWS; i++) {
// Embed row's boxes are shifted down 2px (square-arrow alignment);
// keep the row label tracking with them.
const extra = i === rowEmbed ? (lookupYEmbed - lookupYHead) : 0
elements.push(text({
x: padL + shiftX - 4, y: yTop(i) + rowHeights[i] / 2 + 4 + extra,
textAnchor: 'end', fontSize: '11', fontFamily: 'monospace',
fill: mutedCol
}, rowLabels[i]))
}
elements.push(g({ transform: `translate(${shiftX} ${shiftY})` }, this.diagramLegendSvg(W, H)))
return svg({
viewBox: `0 0 ${W} ${H}`,
preserveAspectRatio: 'xMidYMid meet',
class: styles.chartBox
}, elements)
}
}
// ---------- Root ----------
class Root extends Component implements IRoot {
activeTab: 'explainer' | 'attention' | 'inside' | 'embeddings' = 'explainer'
// Child components must be public fields — domeleon enumerates them; # fields are invisible.
model = new Model()
inputPanel = new InputPanel()
attentionPanel = new AttentionPanel()
insidePanel = new InsidePanel()
tokenEmbeddingsPanel = new TokenEmbeddingsPanel()
explainerPanel = new ExplainerPanel()
#layout: 'wide' | 'narrow' = window.innerWidth >= 900 ? 'wide' : 'narrow'
get isNarrow() { return this.#layout === 'narrow' }
constructor() {
super()
new MutationObserver(() => {
const t = document.documentElement.getAttribute('data-theme')
if ((t === 'light' || t === 'dark') && t !== themeMgr.themeName) {
themeMgr.themeName = t
this.update()
}
}).observe(document.documentElement, { attributes: true, attributeFilter: ['data-theme'] })
window.addEventListener('resize', () => {
this.#layout = window.innerWidth >= 900 ? 'wide' : 'narrow'
this.update()
})
}
reset() {
this.tokenEmbeddingsPanel.reset()
this.model.reset()
}
view() {
if (!this.model.isReady) {
return div({ class: styles.waitContainer },
div({ class: styles.waitText }, this.model.status),
div({ class: styles.waitSpinner }, '⏳')
)
}
return div({ class: styles.layout },
div({ class: styles.header },
div({ class: styles.headerContainer },
h1({ class: styles.headerTitle }, 'Transformer learns addition'),
p({ class: [styles.body, 'mt-0 mb-5'] }, `A small transformer (${N_LAYERS} layers, ${N_HEADS} heads, ~${Math.round(this.model.paramCount / 1000)}K params) is being trained from scratch in your browser to add 2-digit numbers.`)
)
),
div({ class: styles.controlStripBar },
div({ class: styles.controlStripContent },
this.controlsContent()
)
),
div({ class: styles.content },
div({ class: styles.mainContainer },
this.mainContent()
)
)
)
}
controlsContent() {
const stat = (label: string, body: VElement) =>
div({ class: styles.statCol },
div({ class: [styles.labelText, 'whitespace-nowrap'] }, label),
body
)
const valueText = (s: string) => div({ class: styles.valueText }, s)
const withSpark = (body: VElement, spark: VElement) =>
div({ class: styles.statValueRow }, body, spark)
return div({ class: styles.controlsRow },
div({ class: styles.statsRow },
div({ class: styles.controlGroup },
button({
onClick: () => this.model.toggleRun(),
class: [styles.iconBtn, styles.iconBtnPlay]
}, this.model.isRunning ? iconPause() : iconPlay()),
button({
onClick: () => this.reset(),
class: [styles.iconBtn, styles.iconBtnSmall],
title: 'Reset model'
}, iconReset())
),
stat('Steps', valueText(this.model.trainStep.toLocaleString())),
stat('Examples', valueText(this.model.examplesSeen.toLocaleString())),
stat('Ex/s', valueText(this.model.examplesPerSec.toFixed(0))),
stat('Batch', valueText(String(BATCH_SIZE))),
div({ class: 'flex items-center gap-x-3 sm:gap-x-8' },
stat('Masked loss', withSpark(
valueText(this.model.avgLoss.toFixed(3)),
this.sparkline(this.model.lossHistory, themeMgr.theme.colors.text.rawValue, 'Loss', Math.log(VOCAB))
)),
stat('Held-out acc', withSpark(
div({ class: styles.valueText },
span({ class: this.model.testAcc > 0.9 ? styles.successHigh : styles.successLow }, (this.model.testAcc * 100).toFixed(0) + '%')
),
this.sparkline(this.model.accHistory, themeMgr.theme.colors.primary.rawValue, 'Held-out acc', 1)
))
)
)
)
}
mainContent() {
const gridTemplateColumns = this.isNarrow ? '1fr' : 'minmax(360px, 1fr) 2fr'
return div({ class: styles.mainGrid, style: { gridTemplateColumns } }, this.inputPanel.view(), this.tabsPanel())
}
tabsPanel() {
const tabs: { id: 'explainer' | 'attention' | 'inside' | 'embeddings'; label: string }[] = [
{ id: 'explainer', label: 'Explainer' },
{ id: 'attention', label: 'Attention' },
{ id: 'inside', label: 'Blocks' },
{ id: 'embeddings', label: 'Token embeddings' }
]
return div({ class: styles.panel },
div({ class: 'flex gap-1.5 mb-3 flex-wrap' },
tabs.map(t => button({
class: [styles.tabBtn, this.activeTab === t.id ? styles.tabBtnActive : styles.tabBtnInactive],
onClick: () => { this.activeTab = t.id; this.update() }
}, t.label))
),
this.activeTab === 'attention' ? this.attentionPanel.view() :
this.activeTab === 'inside' ? this.insidePanel.view() :
this.activeTab === 'embeddings' ? this.tokenEmbeddingsPanel.view() :
this.explainerPanel.view()
)
}
sparkline(data: { step: number; value: number }[], color: string, label: string, maxV: number) {
const W = 80, H = 14, pad = 1
if (data.length < 2) {
return svg({ viewBox: `0 0 ${W} ${H}`, preserveAspectRatio: 'none', class: styles.sparklineBox })
}
const minS = data[0].step
const maxS = data[data.length - 1].step
const sx = (s: number) => (s - minS) / Math.max(1, maxS - minS) * W
const sy = (v: number) => (H - pad) - Math.max(0, Math.min(1, v / maxV)) * (H - 2 * pad)
const d = data.map((p, i) => `${i === 0 ? 'M' : 'L'}${sx(p.step).toFixed(1)},${sy(p.value).toFixed(1)}`).join(' ')
return svg({
viewBox: `0 0 ${W} ${H}`,
preserveAspectRatio: 'none',
class: styles.sparklineBox,
title: `${label} · y: 0–${maxV.toFixed(2)} · steps ${minS.toLocaleString()}–${maxS.toLocaleString()}`
},
path({ d, stroke: color, strokeWidth: '1.5', fill: 'none' })
)
}
}
new App({
root: new Root(),
id: 'app',
cssAdapter: themeMgr.unoCssAdapter
})
```
**index.html**
```html
<div id="app"></div>
```
**config.json**
```json
{
"dependencies": {
"domeleon": "^0.6.3",
"@unocss/preset-wind3": "^66.5.3",
"tensorgrad": "^0.4.6"
},
"description": "Watch a transformer learn 2-digit addition from scratch in your browser. Type two numbers and see it predict the sum digit by digit. Built with tensorgrad (autograd + WebGPU)."
}
```