A Stochastic Approach to Bi-Level Optimization for Hyperparameter Optimization and Meta Learning

TL;DR

This paper introduces a stochastic reformulation of bi-level optimization problems in deep learning, utilizing SGLD for sampling and a novel approximation to improve scalability and robustness in hyperparameter and meta learning tasks.

Contribution

It presents a new stochastic perspective on bi-level optimization, enabling scalable, robust solutions for large models and diverse meta learning applications.

Findings

Achieves promising results on various meta learning benchmarks.

Scales to learning 87 million hyperparameters in Vision Transformers.

Provides more stable and reliable solutions compared to existing methods.

Abstract

We tackle the general differentiable meta learning problem that is ubiquitous in modern deep learning, including hyperparameter optimization, loss function learning, few-shot learning, invariance learning and more. These problems are often formalized as Bi-Level optimizations (BLO). We introduce a novel perspective by turning a given BLO problem into a stochastic optimization, where the inner loss function becomes a smooth probability distribution, and the outer loss becomes an expected loss over the inner distribution. To solve this stochastic optimization, we adopt Stochastic Gradient Langevin Dynamics (SGLD) MCMC to sample inner distribution, and propose a recurrent algorithm to compute the MC-estimated hypergradient. Our derivation is similar to forward-mode differentiation, but we introduce a new first-order approximation that makes it feasible for large models without needing to store huge Jacobian matrices. The main benefits are two-fold: i) Our stochastic formulation takes into account uncertainty, which makes the method robust to suboptimal inner optimization or non-unique multiple inner minima due to overparametrization; ii) Compared to existing methods that often exhibit unstable behavior and hyperparameter sensitivity in practice, our method leads to considerably more reliable solutions. We demonstrate that the new approach achieves promising results on diverse meta learning problems and easily scales to learning 87M hyperparameters in the case of Vision Transformers.