Train your own embedding model
The training loop·30 min

Train to a target loss

0 / 1
# grade this stage (from your course directory)
npx theinference test --stage train-to-target --watch

Put it all together. For K steps: draw the next fixed-order batch, run forward → InfoNCE → backward → Adam, and watch the loss fall. Then print the final loss as JSON. This is the loop that actually trains your embedding model.

the training loop
for step in 0..K:
    batch  = nextBatch(seededOrder)
    Q, P   = encode(batch)            # forward
    S      = Q · Pᵀ
    loss   = infonce(S, τ)
    dS     = (softmax − onehot)/τ     # backward
    grads  = contrastiveBackward(...)
    adamStep(params, grads)           # update
print({ finalLoss })

What to build

Your task

Implement the training loop

Wire train(corpus, cfg) to run the loop from the frozen init over the seeded batch order, and make your_program.sh train print {"finalLoss": <number>}. The frozen config (seeds, steps, batch size, lr) is fixed so your final loss matches the reference exactly.

your_program.sh train
$ ./your_program.sh train
{"finalLoss":0.378758}
How it's graded

Process mode: stdout must be JSON matching fixtures/expected.json within atol 1e-4, rtol 1e-3. Because init, batch order, and the data are all seeded and the run is single-thread, your final loss equals the reference's frozen value — the small tolerance just absorbs benign float drift.

Sanity: the loss curve

The loss starts near ln(B) (uniform over B in-batch classes) and should fall well below it within a few hundred steps. If it sits flat at ln(B), you forgot the temperature scale (chapter 3) or flipped a gradient sign (chapter 4). A rising loss is almost always a sign error.