Source-linked AI summary
A Two-Timescale Framework for Bilevel Optimization: Complexity Analysis and Application to Actor-Critic
Mingyi Hong, Hoi-To Wai, Zhaoran Wang, Zhuoran Yang
TL;DR
The paper addresses the difficulty of stochastic bilevel optimization when inner and outer updates must proceed together. It proposes and analyzes a single-loop two-timescale stochastic approximation method, obtaining convergence guarantees and a natural actor-critic application with O(K^-1/4) convergence in expected discounted-reward gap.
Problem
Stochastic bilevel optimization is difficult because outer updates depend on inner solutions, especially when inner and outer gradients are revealed sequentially.
Method
The paper uses a single-loop TTSA algorithm that simultaneously updates the inner solution with a larger step size and the projected outer solution with a smaller step size.
Results
O(K^-1/4) convergence in expected discounted-reward gap is obtained for a two-timescale natural actor-critic policy optimization algorithm with linear function approximation.
Takeaways & Limitations
The TTSA framework provides convergence-rate results for bilevel problems with strongly convex, convex, or non-convex outer objectives and applies to actor-critic reinforcement learning.
Takeaways & Limitations
The convexity of the outer objective must be verified case by case because no general sufficient verification conditions are known.
Abstract
from arXiv · showhide
This paper analyzes a two-timescale stochastic algorithm framework for bilevel optimization. Bilevel optimization is a class of problems which exhibit a two-level structure, and its goal is to minimize an outer objective function with variables which are constrained to be the optimal solution to an (inner) optimization problem. We consider the case when the inner problem is unconstrained and strongly convex, while the outer problem is constrained and has a smooth objective function. We propose a two-timescale stochastic approximation (TTSA) algorithm for tackling such a bilevel problem. In the algorithm, a stochastic gradient update with a larger step size is used for the inner problem, while a projected stochastic gradient update with a smaller step size is used for the outer problem. We analyze the convergence rates for the TTSA algorithm under various settings: when the outer problem is strongly convex (resp.~weakly convex), the TTSA algorithm finds an $\mathcal{O}(K^{-2/3})$-optimal (resp.~$\mathcal{O}(K^{-2/5})$-stationary) solution, where $K$ is the total iteration number. As an application, we show that a two-timescale natural actor-critic proximal policy optimization algorithm can be viewed as a special case of our TTSA framework. Importantly, the natural actor-critic algorithm is shown to converge at a rate of $\mathcal{O}(K^{-1/4})$ in terms of the gap in expected discounted reward compared to a global optimal policy.
1 Introduction
Bilevel optimization couples an outer problem to an inner solution, making simultaneous stochastic optimization difficult. The paper develops a single-loop TTSA method and analyzes its convergence, including an actor-critic application.
- Problem setting: Bilevel optimization minimizes an outer objective whose variables depend on the solution of an inner optimization problem.The framework allows stochastic objectives and has applications including reinforcement learning, hyperparameter optimization, and game theory.
- Problem setting: Simultaneous optimization is difficult because gradient-based outer updates require sufficiently accurate solutions of the inner problem.Existing double-loop methods repeatedly solve the inner problem or use large batches for accurate outer-gradient estimates.
- Method: The paper studies a single-loop TTSA algorithm for unconstrained strongly convex inner problems and smooth outer objectives.Inner and outer solutions are updated simultaneously using cheap stochastic-gradient estimates.
- Method: TTSA uses a larger step size for the inner update and a smaller step size for the projected outer update.The inner and outer updates therefore operate on faster and slower timescales, respectively.
- Theory: The convergence analysis couples inner and outer optimality quantities through one-step inequalities, with the timescale distinction playing a crucial role.The framework covers strongly convex, convex, and weakly convex outer objectives.
- Application: O(K^-1/4) convergence to an optimal policy is obtained for a two-timescale natural actor-critic policy optimization algorithm with linear function approximation.The rate is reported as comparable to state-of-the-art results.
2 Two-Timescale Stochastic Approximation Algorithm for (1)
The paper develops a single-loop TTSA method for bilevel problems with strongly convex unconstrained inner objectives and constrained smooth outer objectives. It uses faster inner updates and slower projected outer updates, with gradient surrogates and convergence analysis under stated regularity assumptions.
- Problem class: The problem class assumes a twice continuously differentiable inner function that is strongly convex in its unconstrained inner variable.The inner gradient is Lipschitz continuous in that variable, with strong-convexity modulus µg > 0.
- TTSA framework: The outer update is a projected stochastic-gradient recursion, while the inner update uses a stochastic gradient recursion for the inner objective.The projection is onto the convex feasible set X, and the stochastic estimates target the corresponding gradients.
- TTSA framework: The TTSA algorithm updates inner and outer solutions simultaneously using larger inner and smaller outer step sizes.The timescale condition αk/βk → 0 keeps the outer update slower while the inner iterate tracks y⋆(x).
- Gradient estimation: The method constructs a surrogate outer gradient by replacing y⋆(x) with the current inner iterate and uses it in TTSA.The surrogate equals the outer gradient when evaluated at y⋆(x), and the construction relies on the inner problem being strongly convex and unconstrained.
- Gradient estimation: The surrogate-gradient estimator may be biased, but its bias is required to decay polynomially to zero, with variance and bias bounds established for the proposed construction.The estimator samples gradient, Jacobian, and Hessian information from distributions designed to provide unbiased component estimates.
- Convergence analysis: The convergence analysis couples outer and inner optimality quantities through one-step inequalities tailored to strongly convex, convex, and weakly convex outer objectives.The paper also uses Lipschitz properties of the solution map y⋆(x) and the gradient surrogate; these imply weak convexity of the composite outer objective.
- Applications: The framework covers applications including hyperparameter optimization, Stackelberg games, and model-agnostic meta-learning.In MAML, the inner problem adapts model parameters to a training set, while the outer problem selects shared weights using validation performance.
- Applications: Whether the composite outer objective is convex or strongly convex generally must be verified case by case.The paper states that it does not provide general sufficient conditions for establishing this property.
3 Main Results
This section analyzes TTSA convergence through coupled optimality and tracking errors under strongly convex and weakly convex outer objectives. The resulting rates and oracle complexities depend on diminishing or constant step sizes, with the inner update operating on the faster timescale.
- Convergence analysis: TTSA bounds outer optimality and inner tracking errors together because the single-loop updates couple the two problems.The tracking error measures distance from the inner optimum corresponding to the current outer iterate.
- Step-size effects: Diminishing step sizes make the transient and fluctuation terms vanish asymptotically, whereas non-vanishing step sizes may leave a nonzero limiting bound.The bound contains a transient term that decays sub-geometrically.
- Strongly convex outer objective: O(k^-2/3) convergence holds for both the optimality gap and tracking error when the outer objective is strongly convex.The result is stated for step sizes satisfying the theorem’s conditions.
- Strongly convex outer objective: O(log(1/ϵ)/ϵ^3/2) stochastic-oracle calls suffice for an ϵ-optimal solution in the strongly convex case.The oracle calls cover stochastic gradient, Hessian, and Jacobian estimates for both outer and inner functions.
- Convex outer objective: O(ϵ^-4 log(1/ϵ)) stochastic-oracle calls find an ϵ-optimal solution for convex objectives over bounded constraint sets.The constraint-set diameter appears in the convergence constants but is omitted from the big-O notation.
4 Application to Reinforcement Learning
The paper formulates policy optimization with linear function approximation as bilevel optimization: policy evaluation is the inner problem, while policy improvement is the outer problem. TT-NAC applies the two-timescale framework, with convergence established under stated MDP and regularity assumptions.
- Policy optimization: Policy optimization maximizes expected discounted reward, with the action-value function defined by a Bellman equation.The paper writes Qπ as the solution of Q(s, a) = r(s, a) + γ(PπQ)(s, a).
- Bilevel formulation: Linear approximation represents the action-value function as Qθ(s, a) := φ(s, a)⊤θ, yielding a bilevel formulation.The outer variable is the policy, while the inner variable parameterizes policy evaluation.
- TT-NAC algorithm: The inner critic minimizes mean squared Bellman error using TD learning and a faster-timescale update.The critic uses consecutive state-action transitions sampled under the current policy.
- TT-NAC algorithm: The outer actor improves the policy through a KL-based mirror-descent or proximal-policy-optimization step.The update uses a gradient surrogate and step size α.
- TT-NAC algorithm: The resulting recursions define the two-timescale natural actor-critic algorithm, which can be implemented with low memory through a running critic sum.The policy update is connected to the natural policy gradient method.
- Convergence analysis: O(K^-1/4) is the reported TT-NAC convergence rate for finding an optimal policy in objective value.The rate is comparable to another TT-NAC analysis and is derived from the general TTSA framework.
5 Numerical Experiments
The numerical study evaluates TTSA on FashionMNIST data hyper-cleaning with corrupted labels and compares it with BSA, stocBiO, and HOAG. TTSA outperforms the comparison methods in the reported experiment.
- Experimental setup: TTSA is compared with BSA, stocBiO using different batch sizes, and HOAG on the data hyper-cleaning task.The comparison includes both stochastic and deterministic bilevel optimization methods.
- Evaluation: The plots report training loss and testing accuracy against the number of gradients evaluated at corruption rate p = 0.4.The figure uses total gradient evaluations as the computational axis.
- Experimental setup: The experiment trains logistic regression with cross-entropy loss on FashionMNIST using 50k training, 10k validation, and 10k testing samples.Training labels are corrupted independently with probability 0.4, and the outer problem is non-convex while the lower-level problem is strongly convex.
- Results: TTSA outperforms BSA, stocBiO, and HOAG when performance is compared against total outer samples accessed.HOAG requires full-batch gradients, while stocBiO relies on large batches and is reported to converge relatively slowly.
6 Conclusion
The paper develops efficient TTSA algorithms for bilevel problems with unconstrained, strongly convex inner objectives and analyzes them across several outer-objective settings. It also specializes the framework to TT-NAC and obtains a convergence rate comparable to existing literature.
- Conclusion: The paper develops efficient two-timescale stochastic approximation algorithms for bilevel problems with unconstrained, strongly convex inner problems.This is the stated scope of the proposed framework.
- Conclusion: The analysis establishes TTSA convergence rates when the outer objective is strongly convex, convex, or non-convex.The conclusion summarizes convergence across these three outer-objective settings.
- Conclusion: The theory is customized to a two-timescale actor-critic proximal policy optimization algorithm in reinforcement learning.The resulting rate is reported as comparable to existing literature.
A Omitted Proofs of Theorem 1
This section defines auxiliary quantities used to simplify notation in the omitted proofs.
- Omitted proofs: The omitted proofs introduce shorthand quantities for arbitrary natural numbers n and m.The definitions are provided for notational brevity.
A.1 Proof of Lemma 3
The proof derives conditional and total-expectation recursions for the inner iterate, using unbiasedness, strong convexity, and the inner optimum condition. It then bounds the resulting optimality gap under specified step-size conditions.
- Recursion: The proof expands yk+1 and conditions on Fk to derive a recursion for the inner iterate.The analysis uses conditional expectation before taking total expectation.
- Inner-problem properties: Unbiasedness of hk and strong convexity of g support the inner-iterate bound.The proof also invokes the condition ∇yg(xk, y⋆(xk)) = 0 for the optimal inner solution.
- Sequence bounds: The resulting recursion is solved using bounds on the step-size sequence and auxiliary sequence inequalities.The proof applies Lemmas 10 and 11 with parameters tied to µg, µℓ, αi, and βi.
B.1 Proof of Lemma 5
The proof of Lemma 5 combines inequalities over iterations and substitutes the resulting bounds to establish both inequalities in (32).
- Conclusion: Substituting the summed bound into the companion relations yields the second and first inequalities in (32), respectively.The derivation assumes c0 > 0 and d0 > 0 for the corresponding simplifications.
B.2 Proof of Lemma 6
The proof develops descent estimates for the outer objective and controls inner and outer errors through projection, smoothness, strong convexity, and auxiliary inequalities. It also states the reinforcement-learning assumptions used for the actor-critic application.
- Outer descent: The proof defines OPTk as the expected outer-objective gap and derives a descent estimate for it.The estimate uses Lipschitz continuity, the step-size restriction µgβ < 1, and strong convexity of the outer objective.
- Update control: The optimality condition of the projected outer step yields bounds that are combined with Young’s inequality and projection properties.The argument controls terms involving the outer update and the inner-gradient approximation.
- RL assumptions: The actor-critic application assumes bounded rewards and bounded linear features representing each policy’s action-value function.The reward bound is |r(s, a)| ≤ r, while Qπ is linear in φ with ∥φ(s, a)∥2 ≤ 1.
D Proof of Theorem 3
Theorem 3 analyzes the actor-critic specialization by controlling policy, critic, and objective errors under the stated assumptions. The proof concludes that the averaged objective gap is O(Kmax^-1/4).
- Policy update: The proof measures policy progress using the ℓ1 distance and divergence terms under the optimal policy’s visitation measure.It introduces policy norms and applies a three-point inequality to the proximal update.
- Lipschitz control: Lemma 8 supplies Lipschitz bounds for Qπ and θ⋆(π) under bounded rewards, stationary-distribution, feature, and concentrability assumptions.The proof derives these bounds using the Bellman equation, the performance difference lemma, Hölder’s inequality, and feature covariance conditioning.
- Critic error: The critic update is analyzed through conditional expectations of hk and a semigradient of the MSBE function.The analysis also bounds the critic error using the second-order moment of the stochastic update.
- Objective gap: The policy objective analysis combines the performance difference lemma with the critic approximation and proximal-update inequalities.The resulting recursion is summed from k = 0 to Kmax −1 and divided by Kmax.
- Rate: O(Kmax^-1/4) bounds the right-hand side after applying Cauchy-Schwarz, completing the convergence proof.The rate is stated for the averaged objective-gap bound obtained in the theorem proof.
F.3 Proof of Lemma 1
The proof bounds products of independent random matrices and uses these estimates to control the variance of the stochastic quantity h_k. It derives norm and variance bounds under contraction and moment conditions.
- Variance estimate: For h_k, the proof decomposes the variance into terms and bounds the first term using σ2 and the stochastic-gradient second-moment estimate.The estimate E[∥∇yf(x, y; ξ(1))∥2] ≤ σ2 is stated explicitly.
- Application to Lemma 1: The product-matrix estimate is invoked after verifying that the relevant random matrices satisfy Lemma 12’s assumptions under the paper’s stated assumptions.This connects the abstract matrix-product lemma to the quantities used in the proof of Lemma 1.
- Application to Lemma 1: The proof bounds the expected squared norm of Hyy by d1 Lg(µ2g + σ2gxy) and separately bounds ∥E[Hyy]∥ by 1/µg.Together, these estimates are used to obtain a variance estimate.
- Random-matrix product bound: Lemma 12 bounds the growth of recursively defined products Zi = YiZi−1 when the random matrices are independent, symmetric, and satisfy contraction and variance conditions.The key condition is (1 − µ)^2 + σ^2 < 1.
- Random-matrix product bound: The proof applies spectral–Schatten norm equivalence and independence to obtain a recursive second-moment bound for Zt.The recursion contracts by the factor (1 − µ)^2 + σ^2, with ∥I∥2^2 = d supplying the initial dimension factor.