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

2026-06-21

AI Summary

This work addresses the limitations of conventional bilevel optimization, which forcibly recasts stochastic root-finding problems as loss minimization, often leading to variance amplification and unstable convergence. The paper introduces the first formalization of bilevel optimization as a root-finding problem (RF-BO) and proposes a Jacobian-free algorithm that bypasses implicit Jacobian computations. Instead, it directly updates along the root error using two-timescale stochastic approximation (TTSA). Under Markovian noise, the method enjoys non-asymptotic convergence guarantees. Empirical results demonstrate its effectiveness: it improves top-1 accuracy by 2.6% in SimCLR, accelerates convergence in nonlinear ODE control by 17×, significantly enhances entropy stability in reinforcement learning, and boosts generative model quality by 11.1%.

Abstract

Many central machine learning tasks, from entropy tuning in reinforcement learning to equilibrating generative adversarial networks, are fundamentally stochastic root-finding problems rather than loss minimization. Yet, they are frequently forced into a minimization framework via squared residuals, introducing a critical flaw we identify as the Variance Trap. Standard bilevel minimization algorithms require estimating hypergradients involving implicit Jacobians; in stochastic settings, these terms act as noise amplifiers, destabilizing convergence. We formalize Root-Finding Bilevel Optimization (RF-BO) as a distinct problem class that bypasses this pathology. We propose a Jacobian-free solution using Two-Time-Scale Stochastic Approximation (TTSA) that updates directly along the root error, structurally avoiding variance amplification. We provide the first non-asymptotic convergence guarantees for TTSA in this setting under Markovian noise. Extensive experiments demonstrate the decisive advantage of this paradigm: compared to squared-residual and implicit-gradient baselines, our framework achieves a 2.6\% top-1 accuracy gain in SimCLR, 17$\times$ faster convergence in non-linear ODE control where baselines fail, significantly improved entropy stability in reinforcement learning, and an 11.1\% quality improvement in generative modeling.

Problem

• Variance Trap
• Root-Finding Bilevel Optimization
• Stochastic Root-Finding
• Hypergradient Estimation
• Implicit Jacobians

Innovation

• Root-Finding Bilevel Optimization
• Variance Trap
• Jacobian-Free
• Two-Time-Scale Stochastic Approximation
• Implicit Gradient