Normal-form block optimizer

(1) Integrate this into the optimizer rather than the forward network. Start with two or four coupled parameter blocks in a small Transformer, such as the attention output projection and MLP projection in one layer, or gain and bias parameters from two residual branches. Every $K$ steps, flatten each block gradient $v_i=\nabla_{p_i}L$, normalize it, and build the small Gram/coupling matrix $M^{\mathrm{ML}}_{ij}$ above. Use $D(p)=\mathrm{diag}(s_1,\ldots,s_m)$ with $s_i=\max(\|p_i\|,s_{\min})$; this is the ML analogue of the paper's scale factors $\mu^2,\mu,\lambda^2,\lambda$. (2) Pseudocode: compute gradients $v_i$; form $M_{ij}=\langle v_i,v_j\rangle/(\|v_i\|\|v_j\|+\varepsilon)+\rho\delta_{ij}$; compute the block forcing vector $B$ from the current gradients; solve $(M+\tau I)q=B$ by Cholesky; set $\Delta p=D^{-1}q$; apply the update and then decoupled weight decay. For a normal-form variant, maintain an exponential average $C_{ij}$ of off-diagonal block drift and replace $B_i$ by $B_i-\sum_{j\ne i}C_{ij}q_j$. (3) The paper supplies the rescaled-coordinate linear solve and the principle of removing first-order coupling; estimate $M$, $C$, and $B$ empirically from minibatch gradients. Clip eigenvalues of $M$ to $[10^{-3},10^2]$ and fall back to diagonal Adam when the condition number is too large. (4) First experiment: train a 6-layer small Transformer on TinyStories, comparing AdamW, block-diagonal Shampoo, and this optimizer at equal FLOPs. Measure loss oscillation, gradient cosine similarity between coupled blocks, condition number of $M$, update norms, and validation loss after a fixed number of updates. Success means lower update-norm variance and faster loss descent at equal compute, particularly when one branch is initialized at a substantially smaller scale.