Selective Adjoint Response Renormalization

Integrate this into the backward pass of a DEQ or another implicit layer, immediately after the forward equilibrium state $z^\star$ has been found and before parameter gradients are computed. The tensor to modify is the adjoint vector passed into vector-Jacobian products with $f_\theta(z^\star,x)$; do not alter the forward fixed-point solver in the first experiment. Use a matrix-free randomized SVD or Lanczos routine because explicitly materializing $K$ is infeasible for large hidden states. Pseudocode: (1) compute $z^\star$ with the existing fixed-point solver and compute $g=\nabla_z\mathcal L$; (2) define JVP/VJP routines for $Kq=q-J_fq$ and $K^\top q=q-J_f^\top q$; (3) estimate the smallest $r$ singular triplets $(\sigma_i,u_i,v_i)$ using Lanczos on $K^\top K$ or a randomized eigensolver; (4) calculate $a_i=v_i^\top g$; (5) mark mode $i$ critical when $\sigma_i<\tau_\sigma$ and $|a_i|/(\|g\|_2+\varepsilon)>\tau_g$; (6) form the low-rank corrected adjoint $\widetilde v=v_{\mathrm{base}}+\sum_{i\in C}a_i(1/\widetilde\sigma_i-1/\sigma_i)u_i$, where $v_{\mathrm{base}}$ is the ordinary iterative solve or the truncated SVD response; (7) feed $\widetilde v$ into the standard DEQ parameter-gradient VJP. Estimate singular values and vectors empirically; the paper supplies the response decomposition, while $\tau_\sigma$, $\tau_g$, rank $r$, and $\lambda_0$ are hyperparameters. Start with $\tau_\sigma$ equal to the 1st-5th percentile of estimated singular values and sweep $\tau_g\in\{0.01,0.05,0.1\}$. The first cheap experiment should use a 2-4 layer MLP DEQ on MNIST or CIFAR-10, comparing exact/standard implicit differentiation, global gradient clipping, Tikhonov damping, and selective renormalization at equal training steps. Measure loss variance, gradient norm spikes, failed linear solves, final accuracy, and wall-clock overhead. Success is fewer adjoint-solve failures and lower gradient variance with no more than 1-2% accuracy loss; on deliberately near-critical DEQs, target faster loss descent and substantially fewer exploding updates at equal FLOPs.