Train a tiny HSTU-inspired recommender: data, loss and top-k

Close the loop even for a tiny model

The notebook needs neither private logs nor a paid GPU. It generates controlled behavior patterns, trains a single-head model with U/Q/K/V, SiLU aggregation, gating and residuals, then produces next-item logits. Synthetic data and simplified temporal structure teach mechanics; they do not establish production impact.

IDs [B,T] become embeddings [B,T,D], the block preserves that shape, and an output head produces [B,T,V]. Every valid position can supervise next-item prediction. Padding target zero is excluded with ignore_index=0.

One backward pass is not learning evidence

Fix the seed and train deterministic sequences while plotting loss. Inspect gradient norms and compare top-k before and after training. Then test retained sequences with different starts and report actual metrics. Success on an easy synthetic rule only establishes learning within that controlled task.

# Uses model / inputs / targets defined in the complete notebook.
logits = model(inputs)
loss = torch.nn.functional.cross_entropy(
    logits.flatten(0,1), targets.flatten(), ignore_index=0)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()

This is a training-step fragment. The notebook supplies model definitions, data, optimizer, loop, plots and assertions.

Three essential correctness checks

Changing a future token must not alter earlier outputs: test causality. Replacing padding embeddings with nonzero random values must not affect legal outputs through masked padding. Exclude reserved ID zero before top-k evaluation.

Separate falling training loss from improving validation. A training-set-only result is an implementation sanity check, not generalization. Even the notebook's held-out synthetic patterns are not a real-world independent test set.

Check your understanding

Question: Loss is low but recommendations always return ID zero. What do you inspect first?

Answer

Check padding supervision, valid targets, candidate filtering and indexing of the final valid position. You may be reading a trailing padding position instead of the sequence's real endpoint.

Actual output from this run

Actual training curve for the 3,538-parameter teaching model. Synthetic convergence does not establish industrial recommendation quality.

Primary sources and further reading

• Official generative-recommenders implementation
• PyTorch cross-entropy

Sources checked on 2026-09-10. Teaching examples are not production benchmarks; confirm APIs and model support against the linked version.

Actual output appears after running.