Skip to content

Commit adbe67a

Browse files
OutisLiclaude
andcommitted
feat(dpa4): fold the focus competition norm into the fused value kernel
The fused SO(2) value-path training operator previously required an identity competition norm; a real norm (the focus entry of edge_norm) dropped the block to the unfused paths at roughly twice the training memory. The per-focus RMS norm now runs inside the operator: its learnable scales enter as the norm_scale input and follow the same input/gradient/second-gradient pattern as the projection weight. The forward folds the scales into the head projection and rescales the logit by the inverse RMS of the gate row in the same lane-strided pass; the backward pushes the logit gradient through the norm Jacobian in the per-edge competition kernel and contracts the scale gradient on the host; the second order derives the logit cotangent, the parameter curvature and the Jacobian's own gate dependence in closed form, with the head-Hessian term mapped onto the rotation operands by one rotation backward. The identity norm keeps the original code paths. Verified by the parity suite (19 cases, including the float64 structural bound on the norm shapes) and a Neo benchmark: the norm-enabled block now matches the film-only fused step time within run-to-run noise at 43% less peak memory than the unfused fallback. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent d92b1ac commit adbe67a

5 files changed

Lines changed: 479 additions & 145 deletions

File tree

0 commit comments

Comments
 (0)