Phase 1, Section 16: Functional AdamW

StoryLM-0 update slice. Not a Phase 1 exit-criteria replacement. Not a new primitive. Uses Phase 0 elementwise ops, pow, sqrt, and broadcast.

1. What this section commits to

Section 11 locks plain SGD. A training loop that matches common LLM practice still needs AdamW: exponential moving averages of the gradient and its square, bias correction, and decoupled weight decay.

This section locks AdamW as existing arithmetic. There is no adamw keyword, IR constructor, or FFI.

The formula is the PyTorch AdamW algebra (decoupled decay, bias-corrected first and second moments, eps added after the corrected second-moment square root). It is not AMSGrad, NAdam, or fused CUDA AdamW.

2. Why not a primitive

mul, add, sub, div, pow, sqrt, and broadcast are already specified. A fused adamw op would invent rounding, in-place updates, and a step clock that Phase 0 does not define. Functional AdamW is the composition below. Differentiating through the update is out of scope (same stance as Section 11).

3. State and scalars

w, g, m, and v have the same dtype, shape σ, and placement. All hyperparameters are host-provided ranked tensors. Bootstrap tensor-literal payloads are discarded, so this slice does not write numeric constants into .kina source.

[1, 1] is the same ranked scalar stand-in as Section 11. Rank-0 Tensor[T, []] is allowed by existing broadcast; it is not the numeric instance.

t is carried by the host. This function does not increment it. The numeric instance takes host-computed β1 ** t and β2 ** t because host pow currently emits linalg.powf, which this LLVM/MLIR pipeline does not lower to a .so. The algebra is still 1 - β ** t. The host passes powers for the current t and recomputes them for t+1 on the next call.

T is a floating dtype from the implemented frem set. Every scalar is broadcastable to σ. Placements match.

4. Operational meaning

Let ⊙ be pointwise multiply and / pointwise divide. After broadcasting every [1, 1] scalar to σ:

m'     = β1 ⊙ m + (1 - β1) ⊙ g
v'     = β2 ⊙ v + (1 - β2) ⊙ (g ⊙ g)
bc1    = 1 - β1 ** t
bc2    = 1 - β2 ** t
denom  = sqrt(v') / sqrt(bc2) + ε
decay  = 1 - lr ⊙ λ
w'     = (w ⊙ decay) - (lr / bc1) ⊙ m' / denom

** is the existing pointwise pow. sqrt is the existing unary. λ is wd.

No AMSGrad max, no gradient clipping, no in-place store, no schedule. Each call produces new tensors w', m', v'.

Several parameters update independently with the same rule. This slice does not add a tuple-map combinator; write one composition per tensor (Section 11).

5. Typing (derived)

Each line is existing elementwise / unary / broadcast typing. pow and sqrt require floating operands. After broadcast, mul / add / sub / div require equal shapes.

The wrapper result is a constructed flat tuple (Section 12 flatten):

(w', m', v') : (Tensor[T, σ], Tensor[T, σ], Tensor[T, σ])

6. Surface instance

let adamw_step = fn(
  w: tensor<F32, [4, 8], @host>,
  g: tensor<F32, [4, 8], @host>,
  m: tensor<F32, [4, 8], @host>,
  v: tensor<F32, [4, 8], @host>,
  lr: tensor<F32, [1, 1], @host>,
  beta1: tensor<F32, [1, 1], @host>,
  beta2: tensor<F32, [1, 1], @host>,
  eps: tensor<F32, [1, 1], @host>,
  wd: tensor<F32, [1, 1], @host>,
  beta1_pow: tensor<F32, [1, 1], @host>,
  beta2_pow: tensor<F32, [1, 1], @host>,
  one: tensor<F32, [1, 1], @host>
) -> (
  tensor<F32, [4, 8], @host>,
  tensor<F32, [4, 8], @host>,
  tensor<F32, [4, 8], @host>
) { … §4 lets …; (w_next, m_next, v_next) }

No new tokens. The update sits outside any reverseMode function. Naming g from a handle is Section 13; this section’s example takes g as an argument, like Section 11.

7. Tests

• adamw_step.kina — §6 program; golden accepted by mlir-opt; the step fn result is (tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>).
• scripts/validate_adamw.py compiles the example to a .so, injects patterned w / g / m / v and the §3 scalars (t = 3 as host β**t), and matches §4 (NumPy; the same algebra in PyTorch when present).

8. Out of scope

• Differentiating through the update.
• AMSGrad, Nesterov, Lion, Sophia, or fused kernels.
• A language-level step clock or schedule.
• argmax, seeded sampling, KV-cache. Decoder-wide composition is Section 25.

9. Open issues

1. Decoder AdamW. Specified in Section 25.
2. Host pow .so lowering. Replace linalg.powf with a math.powf generic (the float8 path already does this) so β ** t can live in-graph.
3. argmax / sampling / KV-cache / generate. argmax is Section 17. Seeded sampling is Section 18. KV-cache append is Section 19. The host generation loop is Section 20.
4. Mean NLL. Unchanged from Section 15.