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 is concerned with tuning continuous hyperparameters of a neural network setup (model, loss function and optimizer parameters) by way of a cheap approximation to the hypergradient. It is closely related to earlier work (Lorraine, 2019), which however overlooked that the idea can be generalized to optimizer parameters (e.g., learning rate, momentum), which is effectively done here. In a range of well-done experiments, the method is shown to very convincingly outperform (Lorraine, 2019), and in some cases even the "gold standard" approach to differentiate through the whole learning trajectory. Code is provided.