# Augmented KKT preconditioner for constrained neural training

- ID: 1299
- Canonical URL: https://synthcore.org/idea/1299/augmented-kkt-preconditioner-for-constrained-neural-training
- API JSON: https://synthcore.org/api/idea/1299.json
- API Markdown: https://synthcore.org/api/idea/1299.md
- Verification status: unverified
- Source: [arXiv:2607.22932](https://arxiv.org/abs/2607.22932)
- Category: optimization
- Solves: speedup, stability, scalability
- ML areas: optimizer, training, regularization
- Math tags: optimization, linear-algebra, numerical-analysis, convex-analysis
- Ratings: usefulness 6/10; difficulty 6/10; novelty 6/10

## Idea description

Replace a generic optimizer for equality-constrained neural-network updates with an augmented KKT solve that separately preconditions parameter and constraint variables. Add a positive semidefinite constraint penalty to the parameter block; this does not change an exactly feasible solution but makes the primal block better conditioned and the constraint Schur complement more robust.

## Mathematical statement

The paper augments the middle block of its saddle system as $-\eta\nu B_{1}^{T}M_{Q_{h}}^{-1}B_{1}$ while retaining the constraint equation $B_{1}(u_{h},\omega_{h})=0$, so the augmentation vanishes at a solution. Here $B_{1}$ is the discrete constraint operator, $M_{Q_{h}}$ is an SPD mass matrix defining the constraint-space metric, $\nu>0$ is a problem scale, and $\eta>0$ is the augmentation parameter. The resulting Schur block is $S_{U_{h},\eta}=B_{2}A^{-1}B_{2}^{T}+\eta\nu B_{1}^{T}M_{Q_{h}}^{-1}B_{1}=\nu(B_{2}M_{\Sigma_{h}^{+}}^{-1}B_{2}^{T}+\eta B_{1}^{T}M_{Q_{h}}^{-1}B_{1})$, where $A$ is the original primal SPD block, $B_{2}$ is the coupling operator, and $M_{\Sigma_{h}^{+}}$ is a mass-matrix approximation of $A/\nu$. For neural training, let $\theta\in\mathbb{R}^{n}$ be parameters, $c(\theta)\in\mathbb{R}^{m}$ equality constraints, $J=\partial c/\partial\theta$ their Jacobian, $H$ an SPD or damped Gauss--Newton approximation to the primal curvature, and $M_c$ an SPD constraint metric. At each linearized step, use the adapted KKT matrix $K_\eta=\begin{bmatrix}H+\eta J^{T}M_c^{-1}J & J^{T}\\J&0\end{bmatrix}$. The exact constrained Newton step is unchanged because the added term is zero whenever $J\Delta\theta=-c(\theta)$; the augmentation instead improves conditioning of the primal block. The paper's field-of-values requirement is $\inf_{\xi}\frac{(\xi,\mathcal{L}\bar{\mathcal{A}}_{\eta}\xi)_{\mathcal{H}}}{(\xi,\xi)_{\mathcal{H}}}\gtrsim1$ and $\sup_{\xi}\frac{\|\mathcal{L}\bar{\mathcal{A}}_{\eta}\xi\|_{\mathcal{H}}}{\|\xi\|_{\mathcal{H}}}\lesssim1$, where $\mathcal{L}$ is the block preconditioner and $\mathcal{H}$ is an SPD metric; in the adaptation these quantities diagnose stable preconditioned Krylov iterations.

## Key formulas

- $$\begin{bmatrix}A&B_{2}^{T}&0\\ B_{2}&-\eta\nu B_{1}^{T}M_{Q_{h}}^{-1}B_{1}&B_{1}^{T}\\ 0&B_{1}&0\end{bmatrix}\begin{bmatrix}\sigma_{h}\\ (u_{h},\omega_{h})\\ p_{h}\end{bmatrix}=\begin{bmatrix}0\\ g\\ 0\end{bmatrix},\qquad B_{1}(u_{h},\omega_{h})=0.$$
- $$S_{U_{h},\eta}=B_{2}A^{-1}B_{2}^{T}+\eta\nu B_{1}^{T}M_{Q_{h}}^{-1}B_{1}=\nu\left(B_{2}M_{\Sigma_{h}^{+}}^{-1}B_{2}^{T}+\eta B_{1}^{T}M_{Q_{h}}^{-1}B_{1}\right).$$
- $$K_{\eta}=\begin{bmatrix}H+\eta J^{T}M_{c}^{-1}J&J^{T}\\J&0\end{bmatrix},\qquad S_c=J(H+\eta J^{T}M_{c}^{-1}J)^{-1}J^{T}.$$
- $$\inf_{\xi}\frac{(\xi,\mathcal{L}\bar{\mathcal{A}}_{\eta}\xi)_{\mathcal{H}}}{(\xi,\xi)_{\mathcal{H}}}\gtrsim1,\qquad \sup_{\xi}\frac{\|\mathcal{L}\bar{\mathcal{A}}_{\eta}\xi\|_{\mathcal{H}}}{\|\xi\|_{\mathcal{H}}}\lesssim1.$$

## Implementation notes

Integrate this at the optimizer step for a model with explicit equality constraints, not inside the forward pass. Examples include enforcing conservation laws in a physics-informed network, preserving selected outputs during continual learning, or imposing adapter orthogonality. At iterate $\theta$, compute the constraint residual $r=c(\theta)$ and expose Jacobian-vector products for $J=\partial c/\partial\theta$. Form a damped curvature operator $H=G+\lambda I$, where $G$ is a Gauss--Newton, empirical-Fisher, or diagonal-curvature approximation and $\lambda>0$ is damping. Choose $M_c=\operatorname{diag}(\operatorname{mean}(J_i^2)+\epsilon)$ or a small constraint Gram matrix. Solve the adapted KKT system approximately with MINRES using block diagonal preconditioner $P=\operatorname{diag}(P_\theta,P_c)$, where $P_\theta\approx H+\eta J^TM_c^{-1}J$ and $P_c\approx JP_\theta^{-1}J^T$. Pseudocode: (1) evaluate $r$ and the loss gradient $g$; (2) define matrix-free products $v\mapsto Hv$, $v\mapsto Jv$, and $u\mapsto J^Tu$; (3) run MINRES on $K_\eta[\Delta\theta,\lambda]=[-g,-r]$; (4) update $\theta\leftarrow\theta+\alpha\Delta\theta$ and the multiplier estimate $\mu\leftarrow\mu+\alpha\lambda$; (5) refresh curvature and Jacobian statistics periodically. The paper's algebra supplies the augmentation and Schur structure; estimate $H$, $M_c$, and the Schur approximation empirically. Sweep $\eta\in\{0.1,1,10,100\}$ and monitor the constraint-conditioned spectrum or short Lanczos estimates. First test on a two-layer MLP solving constrained regression with 10--100 constraints, comparing Adam plus a penalty term and an unpreconditioned KKT solve at equal Jacobian-vector-product count. Success means fewer Krylov iterations, lower constraint violation at equal loss, and stable behavior as parameter count or constraint scaling increases.

## Disclaimer

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