Phase 1, Section 26: StoryLM-0 mean NLL

StoryLM-0 loss slice. Not a Phase 1 exit-criteria replacement. Not a new primitive. Uses the existing reduce_sum VJP plus a host scale.

1. What this section commits to

Section 15 returns per-token NLL [B, T] and seeds ones, so the cotangents are those of the sum. Common training reports the mean over the B · T positions.

This section locks mean NLL as existing arithmetic:

1. reduce_sum keepdims VJP — the restricted pass already implements it (softmax / broadcast backward graphs). This section writes the rule for a user-facing keepdims sum.
2. Mean as composition — keepdims-sum the [B, T] NLL over T then B, then multiply by host inv_n = 1 / (B · T).

There is no mean_nll keyword and no reduce_mean VJP. Tensor-literal payloads are discarded, so inv_n is a host tensor (same stance as Section 16 one).

The seed on the [1, 1] primal remains ones. That seed times inv_n is the mean scale.

2. Why not a primitive

reduce_sum, mul, and host scalars are already specified. A fused mean-NLL op would invent a reduction-axis ABI the numbered specs do not define. reduce_mean as a differentiated surface form is a different rule (broadcast(g) / n); it is out of scope.

3. reduce_sum VJP (keepdims)

Write y = reduce_keep(op=sum, axis=i) x with x, y the same rank and y.shape[i] = 1. Incoming cotangent g_y has y’s shape.

g_x = broadcast(shape(x), g_y)

Same-rank broadcast of a unit axis is already specified. keepdims=false (rank-dropping reduce) is not this slice: that VJP would emit a rank-changing broadcast.

Numeric isolated instance: x : [2, 3], axis 1, y : [2, 1]. Seed ones on y gives g_x = ones[2, 3].

4. Mean NLL composition

After the Section 15 per-token NLL nll : [B, T]:

over_t  = nll.reduce_keep(op=sum, axis=1)     // [B, 1]
over_bt = over_t.reduce_keep(op=sum, axis=0)  // [1, 1]
mean    = over_bt * inv_n                     // inv_n = 1/(B*T)

Numeric instance: B = 1, T = 4, inv_n = 0.25. reverseMode result (discrete ids omitted):

(mean, g_emb, g_mask, g_scale, g_wq, g_wk, g_wv, g_wo, g_lm, g_inv_n)

g_mask, g_scale, and g_inv_n are bound and unused. Each trainable cotangent equals the Section 15 sum cotangent divided by B · T.

The wrapper applies Section 11 SGD outside the differentiated function and returns (mean, emb', wq', wk', wv', wo', lm').

5. Surface instance

• ad_reduce_sum.kina — §3 wrapper fn; @main calls step.
• storylm0_mean_nll_step.kina — §4 decoder + mean + SGD.
• invalid_reduce_mean_ad.kina — reduce_keep(op=mean) in reverseMode is rejected.

No new tokens. Mean sits inside the differentiated function; SGD stays outside.

6. Tests

• ad_reduce_sum.kina — golden accepted by mlir-opt; step result is (y, g_x).
• storylm0_mean_nll_step.kina — golden accepted by mlir-opt; step result is the §4 tuple.
• invalid_reduce_mean_ad.kina — compile fails.
• scripts/validate_mean_nll.py compiles both positive examples to .so files and matches §3 / §4 (NumPy; the same algebra in PyTorch when present).

7. Out of scope

• A reduce_mean VJP, or keepdims=false reduce AD.
• Replacing the Section 24 / 25 train loops (they stay sum-NLL).
• Differentiating through the host scale as a learned parameter.
• Phase 1 exit criteria.

8. Open issues

1. reduce_mean VJP. broadcast(g) / n is not this slice.
2. Rank-dropping reduce AD. Unchanged.
3. Trained 32-token fixture. Specified in Section 27.
4. Prefill-T family / in-place write. Unchanged.