Scalable One-Pass Optimisation of High-Dimensional Weight-Update Hyperparameters by Implicit Differentiation

1 INTRODUCTION . Many machine learning methods are governed by hyperparameters : quantities other than model parameters or weights which nonetheless influence training ( e.g . optimiser settings , dropout probabilities and dataset configurations ) . As suitable hyperparameter selection is crucial to system performance ( e.g . Kohavi & John ( 1995 ) ) , it is a pillar of efforts to automate machine learning ( Hutter et al. , 2018 , Chapter 1 ) , spawning several hyperparameter optimisation ( HPO ) algorithms ( e.g . Bergstra & Bengio ( 2012 ) ; Snoek et al . ( 2012 ; 2015 ) ; Falkner et al . ( 2018 ) ) . However , HPO is computationally intensive and random search is an unexpectedly strong ( but beatable ; Turner et al . ( 2021 ) ) baseline ; beyond random or grid searches , HPO is relatively underused in research ( Bouthillier & Varoquaux , 2020 ) . Recently , Lorraine et al . ( 2020 ) used gradient-based updates to adjust hyperparameters during training , displaying impressive optimisation performance and scalability to high-dimensional hyperparameters . Despite their computational efficiency ( since updates occur before final training performance is known ) , Lorraine et al. ’ s algorithm only applies to hyperparameters on which the loss function depends explicitly ( such as ` 2 regularisation ) , notably excluding optimiser hyperparameters . Our work extends Lorraine et al. ’ s algorithm to support arbitrary continuous inputs to a differentiable weight update formula , including learning rates and momentum factors . We demonstrate our algorithm handles a range of hyperparameter initialisations and datasets , improving test loss after a single training episode ( ‘ one pass ’ ) . Relaxing differentiation-through-optimisation ( Domke , 2012 ) and hypergradient descent ’ s ( Baydin et al. , 2018 ) exactness allows us to improve computational and memory efficiency . Our scalable one-pass method improves performance from arbitrary hyperparameter initialisations , and could be augmented with a further search over those initialisations if desired . 2 WEIGHT-UPDATE HYPERPARAMETER TUNING . In this section , we develop our method . Expanded derivations and a summary of differences from Lorraine et al . ( 2020 ) are given in Appendix C. exact hypergradients from the implicit function theorem ( IFT ) ( ) , with our method ’ s approximate hypergradients ( ) superimposed . We target optimal validation loss ( ) , adjusting weights w based on the training loss . Classical weight updates ( for fixed hyperparameters λ ) converge ( , ) to the best-response line w∗ ( λ ) ( ) ; the IFT gives hyperparameter updates ( ) leading to a minimum of validation loss along w∗ ( λ ) . Our approximate hyperparameter updates ( ) differ in magnitude from these exact updates , but still give useful guidance . 2.1 IMPLICIT FUNCTION THEOREM IN BILEVEL OPTIMISATION . Consider some model with learnable parameters w , training loss LT , optimisation hyperparameters λ and validation loss LV . We use w∗ , λ∗ to represent the optimal values of these quantities , found by solving the following bilevel optimisation problem : ( a ) λ∗ = arg min λ LV ( λ , w∗ ( λ ) ) , such that ( b ) w∗ ( λ ) = arg min w LT ( λ , w ) . ( 1 ) The optimal model parameters w∗ may vary with λ , making LV an implicit function of λ alone . We approach the outer optimisation ( 1a ) similarly to Lorraine et al . ( 2020 ) and Majumder et al . ( 2019 ) , using the hypergradient : the total derivative of LV with respect to the hyperparameters λ . One strategy for solving ( 1 ) is therefore to alternate between updating w for several steps using ∂LT∂w and updating λ using the hypergradient , as shown in Figure 1 ( by and ) . Carefully distinguishing the total differential dλ and the partial differential ∂λ , we have dLV dλ = ∂LV ∂λ + ∂LV ∂w∗ ∂w∗ ∂λ . ( 2 ) While derivatives of LV are easily computed for typical loss functions , the final derivative of the optimal model parameters ( ∂w ∗ ∂λ ) presents some difficulty . Letting square brackets indicate the evaluation of the interior at the subscripted values , we may rewrite ∂w ∗ ∂λ as follows : Theorem 1 ( Cauchy ’ s Implicit Function Theorem ( IFT ) ) Suppose for some λ′ and w′ that [ ∂LT ∂w ] λ′ , w′ = 0 . If ∂LT∂w is a continuously differentiable function with invertible Jacobian , then there exists a function w∗ ( λ ) over an open subset of hyperparameter space , such that λ′ lies in the open subset , [ ∂LT ∂w ] λ , w∗ ( λ ) = 0 and ∂w∗ ∂λ = − ( ∂2LT ∂w∂wT ) −1 ∂2LT ∂w∂λT . ( 3 ) w∗ ( λ ) is called the best response of w to λ ( Figure 1 ) . While ( 3 ) suggests a route to computing ∂w ∗ ∂λ , inverting a potentially high-dimensional Hessian in w is not computationally tractable . 2.2 APPROXIMATE BEST-RESPONSE DERIVATIVE . To develop and justify a computationally tractable approximation to ( 3 ) , we mirror the strategy of Lorraine et al . ( 2020 ) . Consider the broad class of weight optimisers with updates of the form wi ( λ ) = wi−1 ( λ ) − u ( λ , wi−1 ( λ ) ) ( 4 ) for some arbitrary differentiable function u , with i indexing each update iteration . We deviate here from the approach of Lorraine et al . ( 2020 ) by admitting general functions u ( λ , w ) , rather than assuming the particular choice uSGD = η ∂LT∂w ( see Appendix A.4 for more details ) . In particular , this allows λ to include optimiser hyperparameters . Differentiating ( 4 ) and unrolling the recursion gives [ ∂wi ∂λ ] λ′ = − ∑ 0≤j < i ( ∏ 0≤k < j [ I− ∂u ∂w ] λ′ , wi−1−k ) [ ∂u ∂λ ] λ′ , wi−1−j , ( 5 ) where j indexes our steps back through time from wi−1 to w0 , and all w depend on the current hyperparameters λ′ — see Appendix C.2 for the full derivation . Now , we follow Lorraine et al . and assume w0 , . . . , wi−1 to be equal to wi . With this , we simplify the product in ( 5 ) to a jth power by evaluating all derivatives at wi . This result is then used to approximate ∂w ∗ ∂λ by further assuming that wi ≈ w∗ . These two approximations lead to [ ∂wi ∂λ ] λ′ ≈ − ∑ 0≤j < i ( I− ∂u ∂w ) j ∂u ∂λ λ′ , wi ( λ′ ) ≈ [ ∂w∗ ∂λ ] λ′ . ( 6 ) We reinterpret i as a predefined look-back distance , trading off accuracy and computational efficiency . The combination of these approximations implies wi = w∗ for all i , which is initially inaccurate , but which we would expect to become more correct as training proceeds . In mitigation , we perform several weight updates prior to each hyperparameter update . This means derivatives in earlier terms of the series of ( 6 ) ( which are likely the largest , dominant terms ) are evaluated at weights closer to w∗ , therefore making the summation more accurate . In Section 4 , we show that the approximations described here result in an algorithm that is both practical and effective . Our approximate result ( 6 ) combines the general weight update of Majumder et al . ( 2019 ) with the overall approach and constant-weight assumption of Lorraine et al . ( 2020 ) . The latter empirically show that an approximation similar to ( 6 ) leads to a directionally-accurate approximate hypergradient ; we illustrate the approximate updates from our derivations in Figure 1 ( by ) . 2.3 CONVERGENCE TO BEST-RESPONSE DERIVATIVE . To justify the approximations in ( 6 ) , note that the central part of that equation is a truncated Neumann series . Taking the limit i→∞ , when such a limit exists , results in the closed form [ ∂w∗ ∂λ ] λ′ ≈ − [ ( ∂u ∂w ) −1 ∂u ∂λ ] λ′ , w∗ ( λ′ ) . ( 7 ) This is precisely the result of the IFT ( Theorem 1 ) applied to u instead of ∂LT∂w ; that is , substituting the simple SGD update uSGD ( λ , w ) = η ∂LT∂w into ( 7 ) recovers ( 3 ) exactly . Thus , under certain conditions , our approximation ( 6 ) converges to the true best-response Jacobian in the limit of infinitely long look-back windows . 2.4 HYPERPARAMETER UPDATES . Substituting ( 6 ) into ( 2 ) yields a tractable approximation for the hypergradient dLVdλ , with which we can update hyperparameters by gradient descent . Our implementation in Algorithm 1 ( Figure 4 ) closely parallels Lorraine et al. ’ s algorithm , invoking Jacobian-vector products ( Pearlmutter , 1994 ) during gradient computation for memory efficiency via the grad_outputs argument , which also provides the repeated multiplication for the jth power in ( 6 ) . Thus , we retain the O ( |w| + |λ| ) time and memory cost of Lorraine et al . ( 2020 ) , where | · | denotes cardinality . The core loop to compute the summation in ( 6 ) comes from an algorithm of Liao et al . ( 2018 ) . Note that Algorithm 1 approximates dLVdλ retrospectively by considering only the last weight update rather than any future weight updates . Unlike differentiation-through-optimisation ( Domke , 2012 ) , Algorithm 1 crucially estimates hypergradients without reference to old model parameters , thanks to the approximate-hypergradient construction of Lorraine et al . ( 2020 ) and ( 6 ) . We thus do not store network weights at multiple time steps , so gradient-based HPO becomes possible on previously-intractable large-scale problems . In essence , we develop an approximation to online hypergradient descent ( Baydin et al. , 2018 ) . Optimiser hyperparameters generally do not affect the optimal weights , suggesting their hypergradients should be zero . However , in practice , w∗ is better reinterpreted as the approximately optimal weights obtained after a finite training episode . These certainly depend on the optimiser hyperparameters , which govern the convergence of w , thus justifying our use of the bilevel framework . We emphasise training is not reset after each hyperparameter update — we simply continue training from where we left off , using the new hyperparameters . Consequently , Algorithm 1 avoids the time cost of multiple training restarts . While our locally greedy hyperparameter updates threaten a short-horizon bias ( Wu et al. , 2018 ) , we still realise practical improvements in our experiments . 2.5 REINTERPRETATION OF ITERATIVE OPTIMISATION . Originally , we stated the IFT ( 3 ) in terms of minima of LT ( zeros of ∂LT∂w ) , and substituting uSGD = η ∂LT∂w into ( 7 ) recovers this form of ( 3 ) . However , in general , ( 7 ) recovers the Theorem for zeros of u , which are not necessarily minima of the training loss . Despite this , our development can be compared to Lorraine et al . ( 2020 ) by expressing u as the derivative of an augmented ‘ pseudo-loss ’ function . Consider again the simple SGD update uSGD , which provides the weight update rule wi = wi−1 − η ∂LT∂w . By trivially defining a pseudo-loss L = ηLT , we may absorb η into a loss-like derivative , yielding wi = wi−1 − ∂L∂w . More generally , we may write L = ∫ u ( λ , w ) dw . Expressing the update in this form suggests a reinterpretation of the role of optimiser hyperparameters . Conventionally , our visualisation of gradient descent has the learning rate control the size of steps over some undulating landscape . Instead , we propose fixing a unit step size , with the ‘ learning rate ’ scaling the landscape underneath . Similarly , we suppose a ‘ momentum ’ could , at every point , locally squash the loss surface in the negative-gradient direction and stretch it in the positive-gradient direction . In aggregate , these transformations straighten out optimisation trajectories and bring local optima closer to the current point . While more complex hyperparameters lack a clear visualisation in this framework , it nevertheless allows a broader class of hyperparameters to ‘ directly alter the loss function ’ instead of remaining completely independent , circumventing the problem with optimiser hyperparameters noted by Lorraine et al . ( 2020 ) . Figure 2 and Appendix A.3 support this argument .

Summary: This paper tackles the hyperparameter optimization problem with a one-pass approach that alternatively optimizes over machine learning model parameters and hyperparameters. To optimize over a bilevel optimization problem, this paper generalizes Lorraine et al. (2019) by replacing the gradient update of parameters with a more general parametric form, thus allowing the hyperparameters to be updated alongside parameters by gradients coming from truncated Neumann series.