# Normal-form block optimizer

- ID: 2391
- Canonical URL: https://synthcore.org/idea/2391/normal-form-block-optimizer
- API JSON: https://synthcore.org/api/idea/2391.json
- API Markdown: https://synthcore.org/api/idea/2391.md
- Verification status: unverified
- Source: [arXiv:2608.16186](https://arxiv.org/abs/2608.16186)
- Category: dynamics
- Solves: stability, speedup, accuracy
- ML areas: optimizer, training-dynamics, fine-tuning
- Math tags: dynamical-systems, linear-algebra, optimization
- Ratings: usefulness 5/10; difficulty 5/10; novelty 4/10

## Idea description

Replace raw updates of strongly coupled parameter blocks by updates in rescaled, approximately normal-form coordinates. The optimizer estimates the local coupling matrix between block directions, solves a small modulation system for transformed velocities, and optionally subtracts predictable first-order cross-block drift.

## Mathematical statement

The paper derives a modulation system $Mq'=B$, where $q=(\mu^{2}\zeta',\,\mu\mu',\,\lambda^{2}\theta',\,\lambda\lambda')^{\mathsf T}$ contains rescaled parameter velocities, $M=(M_{ij})$ is obtained by differentiating orthogonality constraints, and $B$ contains residual forcing. A representative constraint is $\langle i\phi_{R\mu}e^{i\zeta}\Lambda W_{\mu},g\rangle=0$, where $g$ is the perturbation, $\phi_{R\mu}$ is a cutoff, $\zeta$ is a phase, $\mu$ is a scale, and $\Lambda W_{\mu}$ is a scaling mode. The modified system replaces the raw coordinates by $(\tilde\zeta,\tilde\mu,\theta,\tilde\lambda)$ and has $\tilde M\tilde q'=\tilde B$, with $\tilde q=(\mu^{2}\tilde\zeta',\,\mu\tilde\mu',\,\tilde\lambda^{2}\tilde\theta',\,\lambda\tilde\lambda')^{\mathsf T}$. The modified coordinates remove linear terms that otherwise make the modulation system poorly coercive. The paper shows that small forcing, such as $|\tilde B_1|+|\tilde B_2|\ll e^{-5|t|/4}$, $|\tilde B_3|\lesssim e^{-5|t|/2}$, and $|\tilde B_4|\lesssim e^{-2|t|}$, yields correspondingly small transformed velocities after inversion of $\tilde M$. The ML adaptation defines $v_i$ as gradients for coupled parameter blocks and uses their regularized normalized Gram matrix as the coupling matrix.

## Key formulas

- $$\begin{pmatrix}M_{11}&M_{12}&M_{13}&M_{14}\\M_{21}&M_{22}&M_{23}&M_{24}\\M_{31}&M_{32}&M_{33}&M_{34}\\M_{41}&M_{42}&M_{43}&M_{44}\end{pmatrix}\begin{pmatrix}\mu^{2}\zeta^{\prime}\\\mu\mu^{\prime}\\\lambda^{2}\theta^{\prime}\\\lambda\lambda^{\prime}\end{pmatrix}=\begin{pmatrix}B_{1}\\B_{2}\\B_{3}\\B_{4}\end{pmatrix}.$$
- $$\left\langle i\phi_{R\mu(t)}e^{i\zeta(t)}\Lambda W_{\mu(t)},g(t)\right\rangle=0.$$
- $$\begin{pmatrix}\tilde M_{11}&\tilde M_{12}&\tilde M_{13}&\tilde M_{14}\\\tilde M_{21}&\tilde M_{22}&\tilde M_{23}&\tilde M_{24}\\\tilde M_{31}&\tilde M_{32}&\tilde M_{33}&\tilde M_{34}\\\tilde M_{41}&\tilde M_{42}&\tilde M_{43}&\tilde M_{44}\end{pmatrix}\begin{pmatrix}\mu^{2}\tilde\zeta^{\prime}\\\mu\tilde\mu^{\prime}\\\tilde\lambda^{2}\tilde\theta^{\prime}\\\lambda\tilde\lambda^{\prime}\end{pmatrix}=\begin{pmatrix}\tilde B_{1}\\\tilde B_{2}\\\tilde B_{3}\\\tilde B_{4}\end{pmatrix}.$$
- $$M^{\mathrm{ML}}_{ij}=\frac{\langle v_i,v_j\rangle}{\|v_i\|\|v_j\|+\varepsilon}+\rho\,\delta_{ij},\qquad q=D(p)\,\Delta p,\qquad \Delta p=D(p)^{-1}(M^{\mathrm{ML}})^{-1}B.$$

## Implementation notes

(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.

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.
