Source-linked AI summary
Optimizing Millions of Hyperparameters by Implicit Differentiation
Jonathan Lorraine, Paul Vicol, David Duvenaud
TL;DR
High-dimensional HO is difficult because validation performance depends on how optimized weights respond to hyperparameters, and exact Hessian inversion is infeasible for modern networks. The paper combines the IFT with efficient inverse-Hessian approximations, showing that IFT corresponds to the limit of differentiating through optimization. The approach scales to modern architectures and millions of hyperparameters, while requiring only a few times the memory and compute of standard training.
Problem
High-dimensional HO requires accounting for how optimized weights change with hyperparameters, but exact Hessian inversion is infeasible for modern neural networks.
Method
The method uses the IFT to formulate hypergradients and approximates the inverse Hessian through a stable, constant-memory approximation connected to unrolled differentiation.
Results
The approach scales IFT-based HO to modern neural architectures, including settings with millions of weights and hyperparameters.
Takeaways & Limitations
The method supports fitting per-parameter regularization, distilled data, and learned-from-scratch augmentation through gradient-based HO.
Takeaways & Limitations
The approach requires continuous hyperparameters or continuous relaxations and is not straightforwardly applicable to optimizer hyperparameters.
Abstract
from arXiv · showhide
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.
1 Introduction
High-dimensional hyperparameter optimization becomes difficult because gradient-based methods must differentiate through optimized weights, while conventional methods break down as dimensionality grows. The paper proposes IFT-based approximations that scale to modern networks and explores fitting hyperparameters nearly as easily as weights.
- Hyperparameter optimization can investigate per-feature data distillation, per-loss-term weighting, and per-input augmentation when hyperparameters scale with model parameters.
- Random search works for low-dimensional HO, while Bayesian optimization often breaks down above 100 hyperparameters.
- Gradient-based HO for high-dimensional problems requires approximating how optimized weights change with hyperparameters.
- The proposed inverse-Hessian approximation scales to large neural networks, is more stable than conjugate gradient, and uses constant memory.
- When fitting many hyperparameters, large validation partitions and retraining with optimized hyperparameters are important for strong test performance.
- The paper contributes a stable constant-memory inverse approximation, establishes the IFT–unrolled-differentiation connection, and scales HO to large architectures and many hyperparameters.
2 Overview of Proposed Algorithm
The algorithm treats HO as nested optimization, uses the IFT to estimate the best-response Jacobian, and approximates the required inverse Hessian with a tractable Neumann-series procedure. Its hypergradient computation uses vector-Jacobian products rather than explicitly instantiating large matrices.
- The objective is to choose hyperparameters that minimize validation loss after training the neural-network weights.
- The hypergradient combines direct hyperparameter effects with an indirect term accounting for how optimized weights vary with hyperparameters.
- The IFT estimates how optimal weights change with hyperparameters by differentiating the fixed-point condition of the training gradient field.
- Exact inversion of an m × m Hessian costs O(m^3), so the method approximates the inverse using a Neumann series.
- Differentiating through i optimization steps around locally optimal weights is equivalent to using the first i Neumann-series terms.
- Vector-Jacobian products implement the approximation without materializing matrices, supporting efficient hypergradient computation.
3 Related Work
Prior work frames hypergradient computation through implicit differentiation or unrolled optimization, with scalability and memory costs remaining central concerns.
- Implicit differentiation applies the IFT to nested optimization, regularization, recurrent networks, and efficient cross-validation.
- Unrolled differentiation computes hypergradients by differentiating through gradient-based inner optimization and has been applied to hyperparameter optimization for neural networks.
- Increasing the number of unrolled steps raises memory and computational costs, while truncation can introduce bias.
- The related-work comparison covers hypergradient approximations and the memory costs of gradient-based hyperparameter-optimization methods.
- Table 1 organizes methods as approximate IFT inverses or differentiation through optimization around an evaluation point.
- Table 2 contrasts scaling with unrolled iterations I, bottleneck size K, and weight and hyperparameter sizes P and H.
4 Method
The method combines the IFT with unrolled optimization to approximate the indirect hypergradient using a constant-memory inverse-Hessian scheme, under stated differentiability and contractivity conditions.
- 4 Method: Hyperparameter optimization is a nested problem whose indirect gradient accounts for how optimized weights vary with hyperparameters.
- 4.1 Hyperparameter Opt. is Pure-Response: In hyperparameter optimization, the direct gradient can be identically zero, making the problem pure-response and requiring approximation of the indirect gradient.
- 4.2 Unrolled Optimization and the IFT: The fixed points of the unrolled-SGD derivative recurrence recover the IFT solution under suitable initialization and contractivity conditions.
- 4.2 Unrolled Optimization and the IFT: The exposition assumes α = 1 and presents the recurrence and its derivative to establish convergence to the IFT under conditions.
- 4.2 Unrolled Optimization and the IFT: The proposed approximation uses an index i to trade computation against error bounds and avoids instantiating matrices through vector-Jacobian products.
- 4.2 Unrolled Optimization and the IFT: The method requires constant memory and is described as more stable than conjugate gradient and other unrolled differentiation methods.
- 4.3 Scope and Limitations: The IFT assumptions include differentiable validation loss, twice-differentiable training loss with an invertible Hessian, and a differentiable best-response.
- 4.3 Scope and Limitations: The approach requires continuous hyperparameters or continuous relaxations and is not straightforwardly applicable to optimizer hyperparameters.
5 Experiments
Experiments evaluate inverse-Hessian approximations and demonstrate high-dimensional hyperparameter optimization across overfitting, dataset distillation, learned augmentation, and language modeling. The approach supports many hyperparameters, but validation-data allocation becomes important as dimensionality grows.
- 5.1 Approximate Inversion Algorithms: The Neumann approximation improves cosine similarity over CG with enough HO steps, while CG achieves smaller ℓ2 distance.These comparisons use approximate versus true hypergradients on small-scale problems where the exact inverse is feasible.
- 5.1 Approximate Inversion Algorithms: The true inverse Hessian is diagonally dominant, while additional Neumann terms produce structure closer to the true inverse.This supports identity-based approximations while motivating richer Neumann approximations.
- 5.2 Overfitting Validation Data: More than 50 000 000 hyperparameters allow Algorithm 1 to perfectly classify a small validation set, reaching 100 % training and validation accuracy but lower test accuracy.The experiment tunes a separate weight-decay hyperparameter for each neural-network parameter.
- 5.3 Dataset Distillation: Dataset distillation scales to 7 840, 30 720, and 300 720 hyperparameters for MNIST, CIFAR-10, and CIFAR-100, respectively, producing recognizable digits and reasonable color averages.Each synthetic image pixel is a hyperparameter, and one distilled image is learned per class.
- 5.4 Learned Data Augmentation: The learned augmentation network uses 6 659 hyperparameters and improves validation and test accuracy by 2-3%, with smaller variance across random restarts.A stochastic U-net augmentation network is tuned through validation loss while training ResNet18 on CIFAR-10.
- 5.6 Effects of Many Hyperparameters: For many hyperparameters, larger validation partitions are preferable, and retraining on all data is critical after fitting hyperparameters.For few hyperparameters, the standard 10% validation split performs similarly with or without retraining.
6 Conclusion
The paper concludes that its gradient-based hyperparameter optimization algorithm scales to modern neural networks with high-dimensional hyperparameters by combining the IFT with approximate Hessian inversion. It presents this as a basis for practical nested optimization and for fitting hyperparameters nearly as easily as weights.
- 6 Conclusion: The algorithm formulates the hypergradient with the IFT and approximately inverts the training-loss Hessian using its relationship with unrolled differentiation.This avoids directly inverting the large Hessians arising in modern neural networks.
- 6 Conclusion: The paper proposes a stable inverse-Hessian approximation with constant memory cost.The contribution targets scalable hypergradient computation for large neural networks.
- 6 Conclusion: The paper shows that the IFT is the limit of differentiating through optimization.This relationship motivates the proposed approximation strategy.
- 6 Conclusion: The authors identify practical nested optimization with structured Hessians, including GANs and other multi-agent games, as a possible application area.This is presented as a path rather than an experimentally established result.
C Implicit Function Theorem
The implicit function theorem states that, under an invertibility condition on the training-loss Hessian, the optimized weights locally form a continuously differentiable function of the hyperparameters. It also provides their partial derivatives through a matrix product.
- C Implicit Function Theorem: If the relevant training-loss derivative with respect to weights is invertible, a locally defined continuously differentiable best-response function w∗ exists.The theorem applies on an open neighborhood around the selected hyperparameter and weight values.
- C Implicit Function Theorem: The theorem gives the partial derivatives of w∗ through a matrix product.This derivative supplies the sensitivity of optimized weights to hyperparameters needed for hypergradient computation.
- C Implicit Function Theorem: The theorem is stated for mappings between hyperparameter space Λ and weight space W, with λ and w as their respective variables.The notation identifies Λ = R^m and W = R^n in the theorem’s general setup.
D Proofs
The proofs analyze when differentiating through unrolled SGD connects to the implicit-function solution, culminating in the Neumann-SGD theorem under a contractivity condition.
- Lemma 1 considers the fixed point w∞ satisfying ∂LT/∂w = 0 at λ,w∞(λ).
- Lemma 2 provides the recurrence result used in the Neumann-SGD proof.
- The Neumann-SGD theorem assumes the SGD recurrence starts at w0 = w∗(λ).
- The theorem additionally requires I + ∂2LT/∂w∂wT to be contractive.
- The derivation invokes a contractive Neumann series to approximate the inverse-Hessian-related expression.
E Experiments
The experiments use PyTorch on NVIDIA TITAN Xp GPUs, with Adam for network weights and RMSprop for hyperparameters in CNN experiments.
- The computational framework is PyTorch, and all experiments run on NVIDIA TITAN Xp GPUs.
- CNN network weights use Adam with learning rate 1e-4, while hyperparameters use RMSprop with learning rate 1e-2.
E.1 Overfitting a Small Validation Set
The experiments show that the algorithm can overfit small validation datasets across several model classes, achieving perfect training and validation accuracy while testing accuracy is substantially lower.
- The evaluation uses 50 training inputs and 50 validation inputs for MNIST and CIFAR-10, alongside a standard testing partition.
- Models include logistic regression, a one-layer fully connected network, LeNet, AlexNet, and ResNet44.
- 100% training and validation accuracy is achieved in all tested examples, while testing accuracy is significantly lower.
- Figure 10 plots separate training, validation, and testing error lines for the overfitting experiment.
E.2 Dataset Distillation
The experiments cover dataset distillation, stochastic image augmentation, and large-scale hyperparameter tuning for an LSTM with millions of parameters.
- E.2 Dataset Distillation: The dataset-distillation setup uses the entire MNIST validation dataset and 300 CIFAR validation data points.
- E.2 Dataset Distillation: The augmentation network is a U-Net that transforms images and receives a concatenated random-noise channel, producing four-channel inputs.
- E.2 Dataset Distillation: The LSTM implementation uses two layers, 650 hidden units per layer, and 650-dimensional word embeddings.
- E.2 Dataset Distillation: The LSTM experiment tunes separate weight-decay hyperparameters for each of 13 280 400 LSTM weights.
- E.2 Dataset Distillation: Large-scale training uses SGD for network weights and Adam with learning rate 0.001 for hyperparameters.
- E.2 Dataset Distillation: Figure 11 presents the complete dataset distillation for CIFAR-100 and references Figure 6.