Phase 1, Section 15: NLL ∘ decoder and multi-weight SGD

StoryLM-0 composition slice. Not a Phase 1 exit-criteria replacement. No new operator. Uses Sections 9–14.

1. What this section commits to

Section 14 differentiates the decoder graph and applies Section 11 SGD to the embedding table, with seed ones on the logits. Section 10 already differentiates

- logits.log_softmax(axis).take_along_axis(targets, axis)

This section composes those rules in one program:

1. The differentiated function is the decoder let-spine plus the Section 9 NLL tail. The primal is per-token NLL Tensor[F32, [B, T], @host].
2. Host-provided I64 token_ids and targets are discrete parameters (Section 14 §3). They do not appear in the reverseMode result tuple.
3. After handle … with reverseMode, Section 13 names every floating cotangent. Section 11 SGD is applied independently to each trainable weight.

There is no new primitive, IR constructor, or FFI.

2. What this section does not commit to

• A reduce_mean / reduce_sum of the [B, T] NLL. The seed remains ones of the primal (Section 10 / 12), so the cotangents are those of sum of per-token NLL, not mean. Mean NLL is Section 26.
• Next-token shift inside the graph. targets is a host tensor of the same shape as the sequence, not token_ids[t+1].
• SGD on causal_mask or attention_scale. Those are host constants. Their cotangents are still produced and must be bound; they are discarded.
• AdamW, momentum, weight decay, argmax, seeded sampling, KV-cache, or Phase 2 handler AD.
• Differentiating through LetTuple (Section 13).
• Rank-changing broadcast VJP (Section 14 open issue 7). Unit batch dims still use volume-preserving reshape.
• Device take_along_axis / ScatterAdd.

3. Differentiated function

The decoder body is the storylm0_ad_step.kina let-spine (every operator in its own let; 2D weights add a unit batch dim with reshape). After logits:

let log_probs = logits.log_softmax(axis=2) in
let picked = log_probs.take_along_axis(targets, axis=2) in
let nll = -picked in
nll

The body ends with a Var. Method chains are split the same way as Section 14.

Parameters, in order:

reverseMode result type (Section 12 flatten, Section 14 discrete filter):

(nll, g_emb, g_mask, g_scale, g_wq, g_wk, g_wv, g_wo, g_lm)

Seed ones on nll : [1, 4] means

g_logits[b, t, :] = softmax(logits[b, t, :]) - one_hot(targets[b, t])

which is the standard sum-NLL cotangent. No new VJP.

4. Multi-weight SGD

Section 11 already says several parameters update independently with the same rule and that this is written as one composition per tensor. The trainable set is

{token_embedding, wq, wk, wv, wo, lm_head}

Each update is

w' = w - lr.broadcast(shape(w)) * g_w

lr : Tensor[F32, [1, 1], @host] is broadcastable to every trainable shape under the existing Phase 0 broadcast rule. g_mask and g_scale are bound and unused.

Projection and the updates sit outside the differentiated function.

5. Train-step result

The wrapper fn returns a constructed flat tensor tuple (surface (e₀, e₁, …); Section 12 flatten; MLIR multi-result). It is not a new ABI. The numeric instance is

(nll, emb', wq', wk', wv', wo', lm')

All seven components are rank-2 F32 host tensors. The LLVM C wrapper is the existing ≥2-result packed-struct rule (Section 12 §6). @main of the closed program has the same flattened result list.

6. Surface instance

let storylm0_nll = fn(
  token_ids: tensor<I64, [1, 4], @host>,
  targets: tensor<I64, [1, 4], @host>,
  token_embedding: tensor<F32, [8, 4], @host>,
  causal_mask: tensor<F32, [1, 1, 4, 4], @host>,
  attention_scale: tensor<F32, [1, 1, 1, 1], @host>,
  wq: tensor<F32, [4, 4], @host>,
  wk: tensor<F32, [4, 4], @host>,
  wv: tensor<F32, [4, 4], @host>,
  wo: tensor<F32, [4, 4], @host>,
  lm_head: tensor<F32, [4, 8], @host>
) -> tensor<F32, [1, 4], @host> { … decoder …; nll } in
let step = fn(
  token_ids: tensor<I64, [1, 4], @host>,
  targets: tensor<I64, [1, 4], @host>,
  token_embedding: tensor<F32, [8, 4], @host>,
  causal_mask: tensor<F32, [1, 1, 4, 4], @host>,
  attention_scale: tensor<F32, [1, 1, 1, 1], @host>,
  wq: tensor<F32, [4, 4], @host>,
  wk: tensor<F32, [4, 4], @host>,
  wv: tensor<F32, [4, 4], @host>,
  wo: tensor<F32, [4, 4], @host>,
  lm_head: tensor<F32, [4, 8], @host>,
  lr: tensor<F32, [1, 1], @host>
) -> (
  tensor<F32, [1, 4], @host>,
  tensor<F32, [8, 4], @host>,
  tensor<F32, [4, 4], @host>,
  tensor<F32, [4, 4], @host>,
  tensor<F32, [4, 4], @host>,
  tensor<F32, [4, 4], @host>,
  tensor<F32, [4, 8], @host>
) {
  let (nll, g_emb, g_mask, g_scale, g_wq, g_wk, g_wv, g_wo, g_lm) =
    handle storylm0_nll(
      token_ids, targets, token_embedding, causal_mask, attention_scale,
      wq, wk, wv, wo, lm_head
    ) with reverseMode in
  let emb_next = token_embedding - lr.broadcast([8, 4]) * g_emb in
  let wq_next = wq - lr.broadcast([4, 4]) * g_wq in
  let wk_next = wk - lr.broadcast([4, 4]) * g_wk in
  let wv_next = wv - lr.broadcast([4, 4]) * g_wv in
  let wo_next = wo - lr.broadcast([4, 4]) * g_wo in
  let lm_next = lm_head - lr.broadcast([4, 8]) * g_lm in
  (nll, emb_next, wq_next, wk_next, wv_next, wo_next, lm_next)
}

storylm0_nll stays a pure tensor let-chain. The constructed tuple is only in step.

7. Tests

• storylm0_nll_step.kina — §6 program; golden accepted by mlir-opt; the step fn result is the seven-tensor tuple.
• After transform_ad, the loss function contains both ScatterAlongAxis (NLL) and ScatterAdd (embedding gather).
• scripts/validate_nll_decoder_train.py: injects patterned tokens, targets, weights, and lr; checks nll and every updated weight against NumPy (PyTorch when present).

8. Open issues

1. Mean NLL. Specified in Section 26.
2. In-graph next-token shift. Would need a specified slice or roll VJP.
3. AdamW / decode. AdamW is Section 16. Decoder-wide AdamW is Section 25. Deterministic argmax is Section 17. Seeded sampling is Section 18. KV-cache append is Section 19. The host generation loop is Section 20.
4. Device NLL / gather AD. Unchanged from Sections 9 and 14.