[Paper Review] Optimizing Millions of Hyperparameters by Implicit Differentiation

Jonathan Lorraine, Paul Vicol+1 more|arXiv (Cornell University)|Nov 6, 2019

Machine Learning and Data Classification65 references43 citations

TL;DR

The paper presents a gradient-based hyperparameter optimization method that uses the implicit function theorem with a Neumann-series inverse Hessian to scale HO to millions of hyperparameters in large neural networks.

ABSTRACT

We propose an algorithm for inexpensive gradient-based hyperparameter optimization that combines the implicit function theorem (IFT) with efficient inverse Hessian approximations. We present results about the relationship between the IFT and differentiating through optimization, motivating our algorithm. We use the proposed approach to train modern network architectures with millions of weights and millions of hyper-parameters. For example, we learn a data-augmentation network - where every weight is a hyperparameter tuned for validation performance - outputting augmented training examples. Jointly tuning weights and hyperparameters with our approach is only a few times more costly in memory and compute than standard training.

Motivation & Objective

• Motivate and formalize hyperparameter optimization (HO) as nested optimization where hyperparameters influence the validation loss through the best-response of model weights.
• Introduce a stable, memory-efficient inverse Hessian approximation to enable differentiating through optimization at scale.
• Show scalability of the IFT-based HO to large architectures and high-dimensional hyperparameters.
• Demonstrate practical HO applications such as per-parameter regularization, data distillation, and learned data augmentation.

Proposed method

• Formulate HO as a nested optimization problem with a training loss L_T and a validation loss L_V; the hyperparameters lambda are optimized to minimize L_V* where w*(lambda) minimizes L_T.
• Decompose the hypergradient into a direct term and an indirect term that involves the best-response Jacobian ∂w*/∂lambda; emphasize the indirect term is typically the bottleneck.
• Use the implicit function theorem to express ∂w*/∂lambda via the inverse training Hessian ∂^2 L_T/∂w∂w^T and the mixed partials ∂^2 L_T/∂w∂lambda^T.
• Propose a Neumann-series based inverse-Hessian approximation that enables constant-memory computation and is tied to unrolled differentiation; provide a vector-Jacobian product framework for efficiency.
• Present three algorithms: (i) gradient-based HO with implicit differentiation; (ii) an approximate inverse-HVP routine; (iii) a Neumann-series-based hypergradient computation that trades iterations i for accuracy.

Experimental results

Research questions

• RQ1 Can the implicit function theorem be used to compute hypergradients efficiently for high-dimensional hyperparameters?
• RQ2 Does a stable, constant-memory inverse-Hessian approximation enable practical HO on modern large neural networks?
• RQ3 What is the relationship between unrolled differentiation and IFT in the context of HO, and how can they be combined effectively?
• RQ4 How does joint optimization of weights and hyperparameters impact data augmentation, dataset distillation, and per-parameter regularization?
• RQ5 What are the memory and compute trade-offs when scaling HO to millions of hyperparameters?

Key findings

• The proposed Neumann-series inverse-Hessian approximation achieves stable, constant-memory hypergradient computation.
• IFT-based hypergradients converge to the true solution as the unrolled approximation becomes more accurate, establishing a link between unrolled optimization and IFT.
• The approach scales to large architectures like AlexNet and LSTM language models with millions of weights and hyperparameters.
• The method enables applications such as per-parameter regularization, data-distillation (one distilled image per class), and learned data augmentation.
• Empirical results indicate standard train/validation splits may be insufficient when optimizing many hyperparameters, motivating re-training on joint training-validation data.