Escaping the Variance Trap: Jacobian-Free Dynamics for Root-Finding Bilevel Optimization

Research Question and Approach

Many machine learning tasks—such as entropy tuning in reinforcement learning, stabilizing GANs, and hyperparameter optimization—are fundamentally root-finding problems where the goal is to satisfy an equilibrium condition ($h=0$). Standard approaches often force these into a minimization framework by minimizing the squared residual ($|h|^2$). The authors identify this as the "Variance Trap": because the gradient of the squared residual involves the implicit Jacobian, stochastic noise in the lower-level estimation is amplified, leading to instability and slow convergence.

To address this, the authors formalize these tasks as Root-Finding Bilevel Optimization (RF-BO). Instead of minimizing a loss, they propose using Two-Time-Scale Stochastic Approximation (TTSA) to update the upper-level parameters directly using the residual map. This Jacobian-free approach bypasses the need to compute or approximate complex implicit Jacobians or Hessian inverses, which are the primary sources of noise amplification.

The authors provide a theoretical diagnosis of the Variance Trap, proving that squared-residual minimization inherently amplifies noise proportional to the residual magnitude and the Jacobian condition number. In contrast, their proposed RF-TTSA algorithm maintains bounded update variance. They establish non-asymptotic convergence rates for this method under both strongly convex and Polyak-Łojasiewicz (PL) conditions, even in the presence of Markovian noise.

Empirical results demonstrate significant performance gains across diverse tasks:

• Generative Modeling: Achieved an 11.1% improvement in quality for WGAN-GP training.
• Representation Learning: Realized a 2.6% top-1 accuracy gain in SimCLR.
• Control: In non-linear ODE control tasks, the framework converged 17x faster than implicit-gradient baselines, which often failed to reach the target equilibrium due to vanishing gradients in saturation regions.