Meta-learning to Improve Pre-training

1 Introduction . A popular and important learning paradigm for neural networks is pre-training ( PT ) followed by finetuning ( FT ) , an approach commonly used in transfer learning [ 13 , 59 , 19 , 27 , 52 , 11 , 37 , 74 , 35 , 28 ] , and semi-supervised learning [ 9 , 8 , 24 ] . This paradigm has led to performance improvements in many domains , including computer vision [ 13 , 59 , 19 , 37 , 74 , 35 ] , natural language processing [ 27 , 52 , 11 , 40 , 34 ] , graph structured prediction [ 28 ] , and clinical machine learning [ 45 , 46 , 2 , 48 ] , and is especially helpful in settings where downstream tasks have limited training data . The PT & FT paradigm introduces high-dimensional , complex PT hyperparameters , such as parameterized data augmentation policies used in contrastive representation learning [ 8 , 22 ] or the use of task , class , or instance weighting variables in multi-task PT to avoid negative transfer [ 70 ] . These hyperparameters can significantly affect the quality of pre-trained models [ 8 ] , and thus finding techniques to set their values optimally is an important area of research . Choosing optimal PT hyperparameter values is challenging , and existing methods do not work well . Simple approaches such as random or grid search are inefficient since evaluating a hyperparameter setting requires performing the full , two-stage PT & FT optimization , which may be prohibitively computationally expensive . Gradient-free approaches , such as Bayesian optimization or evolutionary algorithms [ 33 , 61 , 47 ] , are also limited in how well they scale to this setting . Gradient-based 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . approaches [ 44 , 41 , 43 , 42 ] can be used online to jointly learn hyperparameters and model parameters and can scale to millions of hyperparameters [ 42 ] , but typically deal with a standard single-stage learning problem ( e.g. , normal supervised learning ) and are therefore not directly applicable to the two-stage PT & FT learning problem . In this work , we address this gap and propose a method for high-dimensional PT hyperparameter optimization . We first formalize a variant of the PT & FT paradigm , which we call meta-parameterized pre-training ( Figure 1 ) , where meta-parameters refer to arbitrary PT hyperparameters or parameterizable architectural choices that can be optimized to improve the learned representations.1 We outline a meta-learning problem characterizing the optimal meta-parameters propose a gradient-based method to learn meta-parameters . Our contributions are : • We formalize meta-parameterized pre-training , a variant of the pre-training and fine-tuning ( PT & FT ) paradigm where PT is augmented to incorporate meta-parameters : arbitrary structures that can be optimized to improve learned representations . • We propose a scalable gradient-based algorithm to learn meta-parameters using a novel method to obtain meta-parameter gradients through the two-stage PT & FT process . Our gradient estimator composes a constant-memory implicit differentiation approximation for the longer PT stage and exact backpropagation through training for the shorter FT stage . • We show that our algorithm recovers optimal meta-parameters in toy experiments on synthetic data . • In two real-world experimental domains , we demonstrate our algorithm improves performance . Firstly , on a multitask PT benchmark over biological graph-structured data [ 28 ] , using our method to optimize meta-parameters representing task weights improves performance by up to 3.9 % AUROC . Secondly , for semi-supervised learning using SimCLR [ 8 ] over electrocardiography data , using our algorithm to optimize meta-parameters representing the weights of a data augmentation neural network improves performance by up to 1.9 % AUROC . 2 Problem Setup and Preliminaries . In this section , we define the meta-parameterized pre-training meta-learning problem , and compare it to traditional fine-tuning and pre-training . A full glossary of notation is in Appendix B , Table 3 . Notation . Let the subscript • be a placeholder for either PT ( pre-training ) or FT ( fine-tuning ) , X ⊆ Rd be our input domain , Y• and Ŷ• be the true and predicted output spaces for some model respectively , and Θ , Ψ• , Φ be spaces of parameters for models . We will use f• : X ; ( Θ , Ψ• ) → Ŷ• to refer to a parametric model , with the semicolon separating the input space from the parameter spaces . We then define f• = f ( head ) • ◦ f ( feat ) , such that f ( feat ) ( · ; θ ∈ Θ ) is a feature extractor that is transferable across learning stages ( e.g. , pre-training to fine-tuning ) , and f ( head ) • ( · ; ψ ∈ Ψ• ) is a stage-specific head that is not transferable . Given a data distribution x• , y• ∼ D• , parametric model f• , and loss function L• : Ŷ• × Y• → R , we will also define for convenience a corresponding expected loss L• : Θ , Ψ• → R via L• ( θ , ψ• ; D• ) = ED• [ L• ( f• ( x• ; θ , ψ• ) , y• ) ] . We also adopt the convention that the output of the argmin operator is any arbitrary minimum , rather than the set of possible minima , to avoid complications in notation . 2.1 Problem Formulation . Supervised Learning ( Fig . 1A ) . In a fully-supervised setting ( our fine-tuning domain ) , we are given a data distribution DFT , model f , and loss LFT . Using a learning algorithm AlgFT ( e.g. , SGD ) that takes as input initial parameters θ ( 0 ) FT , ψ ( 0 ) FT , our goal is to approximate the LFT-optimal parameters : θ∗FT , ψ ∗ FT = AlgFT ( θ ( 0 ) FT , ψ ( 0 ) FT ; DFT ) ≈ argminθ∈Θ , ψ∈ΨFT LFT ( θ , ψ ; DFT ) Pre-training ( Fig . 1B ) . For tasks where data is scarce , we can additionally incorporate a pretraining step and approximate the optimal initial parameters for FT ( i.e. , the final pre-trained weights are used as initialization weights of the FT stage ) , again via an optimization algorithm AlgPT : θ∗PT = AlgPT ( θ ( 0 ) PT , ψ ( 0 ) PT ; DPT ) ≈ argminθ∈Θ LFT ( AlgFT ( θ , ψ ( 0 ) FT ; DFT ) ; DFT ) . 2 1We use the term meta-parameter since these structures do not directly affect inference of the final model after FT , but instead inform the process of learning this model ( by modulating the PT process ) . 2Note that we discard the PT head ψ∗PT here as only the PT feature extractor θ ∗ PT is transferred . Figure ( 1 ) Meta-Parameterized Pre-Training . A paradigm where meta-parameters — rich , potentially high dimensional structures that generalize PT hyperparameters — are incorporated in PT to improve the learned representations . Meta-parameters are optimized in a meta-PT phase , using data from FT task ( s ) in a meta-FT dataset . The FT and meta-FT datasets are ( potentially overlapping ) samples from the FT data distribution . Meta-Parameterized PT ( Fig . 1C ) . In Meta-Parameterized PT , we recognize that , in addition to taking as input the PT parameters θ , AlgPT is itself parameterized by a set of meta-parameters φ ∈ Φ : arbitrary , potentially high dimensional quantities that inform the structure of the algorithm directly . These could represent weighting strategies , data augmentation policies , or sampling processes . The optimal meta-parameters φ ( opt ) are the solution to the following meta-PT optimization problem : φ ( opt ) = argmin φ∈Φ LFT ( AlgFT ( AlgPT ( θ ( 0 ) PT , ψ ( 0 ) PT ; DPT , φ ) , ψ ( 0 ) FT ; DFT ) ; DFT ) . 2.2 Example : Multitask Meta-Parameterized Pre-Training . To make our notation concrete , here we instantiate our setup for a multitask pre-training problem . Problem : Suppose we have a multitask classification dataset , ( X × Y ) N such that Y = Y1 × · · · × YK consists of labels for K distinct tasks . Of this full set of tasks , we are interested only in a subset of M tasks , S = { t1 , . . . , tM } ⊆ { 1 , . . . , K } . Supervised FT : Under supervised FT alone , we can directly average a cross-entropy loss LCE over only the tasks in S , LFT ( ŷ , y ) = 1M ∑M j=1 LCE ( ŷ ( tj ) , y ( tj ) ) , and then solve this problem via SGD . PT : If we assume that S is a random subset of the full set of tasks , we can introduce a PT stage over all tasks : LPT ( ŷ , y ) = 1K ∑K i=1 LCE ( ŷ ( i ) , y ( i ) ) , followed by FT on S alone . As S is a random subset , leveraging all tasks for PT is well motivated and may improve performance . Meta-Parameterized PT : In the case where T is not a random subset , the PT strategy described above is no longer well-motivated . However , using meta-parameterized PT , we can still effectively pre-train by introducing the meta-parameters that weight the tasks φ = [ φ1 . . . φK ] and modulate the loss function LPT : LPT ( ŷ , y ; φ ) = ∑K i=1 φiLCE ( ŷ ( i ) , yi ) . With optimal meta-parameters φ ( opt ) , the PT stage will leverage only that subset of tasks that best informs the final FT performance . This setting mirrors our real-world experiment in Section 5 . 3 Methods : Optimizing Meta-Parameters for Two-Stage Training . We now introduce our gradient-based algorithm to optimize meta-parameters . We first describe how to efficiently approximate meta-parameter gradients through the two-stage PT and FT optimization . We then present our algorithm , and outline practical considerations when using it . 3.1 Efficient Computation of Meta-Parameter Gradients . We begin by defining : g ( φ ; θ ( 0 ) PT , ψ ( 0 ) PT , ψ ( 0 ) FT ) = LFT ( AlgFT ( Parameter θPT︷ ︸︸ ︷ AlgPT ( θ ( 0 ) PT , ψ ( 0 ) PT ; DPT , φ ) , ψ ( 0 ) FT ; DFT ) ︸ ︷︷ ︸ Parameters θFT , ψFT ; DFT ) , ( 1 ) so that φ ( opt ) = argminφ∈Φ g ( φ ) . We also define two best-response values : θ∗PT ( φ ) = AlgPT ( θ ( 0 ) PT , ψ ( 0 ) PT ; DPT , φ ) , θ∗FT ( φ ) , ψ ∗ FT ( φ ) = AlgFT ( θ ∗ PT ( φ ) , ψ ( 0 ) FT ; DFT ) . We do not explicitly include the dependence of the best responses on the initialization values for notational convenience . With these defined , we now consider the desired gradient term , ∂g∂φ . Under our definitions , the direct partial derivatives ∂LFT∂φ and ∂AlgFT ∂φ are zero , so ∂g ∂φ reduces to a simple expression of the chain rule : ∂g ∂φ ∣∣∣∣ φ′ = ∂LFT ∂ [ θFT , ψFT ] ∣∣∣∣ θ∗FT ( φ ′ ) , ψ∗FT ( φ ′ ) ︸ ︷︷ ︸ FT Loss Gradient × FT Best Response Jacobian︷ ︸︸ ︷ ∂AlgFT ∂θPT ∣∣∣∣ θ∗PT ( φ ′ ) × ∂AlgPT ∂φ ∣∣∣∣ φ′︸ ︷︷ ︸ PT Best Response Jacobian . ( 2 ) The FT Loss Gradient term on the RHS of ( 2 ) is easily computed using backpropagation . Computing the other two terms is more involved , and we detail each below , beginning with the PT best response Jacobian . The full algorithm with both gradient estimation terms is provided in Algorithm 1 . PT Best Response Jacobian ∂AlgPT∂φ . Using recent work in hyperparameter optimization with implicit differentiation [ 42 ] , we re-express this term using the implicit function theorem ( IFT ) . If we assume that θ∗PT ( φ ) = AlgPT ( θ ( 0 ) PT ; DPT , φ ) is a good approximation of argminθ∈Θ LPT ( θ ; DPT , φ ) ( i.e. , the PT model converges to LPT-optimal parameters ) , then under certain smoothness and regularity assumptions on the PT parameters and meta-parameters , the IFT allows us to re-express ∂AlgPT∂φ as : ∂AlgPT ∂φ ∣∣∣∣ φ′ = − [ ∂2LPT ∂θPT ∂θ > PT ] −1 × ∂ 2LPT ∂θPT ∂φ > ∣∣∣∣ θ∗PT ( φ ′ ) , φ′ , ( 3 ) which is the product of the inverse Hessian and a matrix of mixed partial derivatives . Following [ 42 ] , the inverse can be efficiently approximated using a truncated Neumann series . FT Best Response Jacobian ∂AlgFT∂θPT . First , note that without additional constraints on AlgFT , the FT best response Jacobian may be zero . This is because LFT has no functional dependence on the variable θPT and , if we assume the convergence point θ∗FT is stable ( as we did for the PT best response Jacobian ) , this implies that the gradient of θ∗FT with respect to θPT would be zero . To enable effective learning , we must therefore either ( 1 ) impose restrictions on AlgFT to ensure there is a dependence between the initialization point and the final loss value ( e.g. , proximal regularization [ 55 ] ) or ( 2 ) leverage methods that do not differentiate through AlgFT through convergence , as at non-converged points we will still observe nonzero LFT-gradients [ 29 , 51 ] . Given that the FT phase often involves shorter optimization horizons than PT , we take approach 2 here , and iteratively update θFT for K steps . We first initialize the FT head ψ ( 0 ) FT and then compute : θ ( 0 ) FT = copy ( θ ∗ PT ) ( init with PT solution , implicitly performing stop gradient ) θ ( k ) FT , ψ ( k ) FT = [ θ ( k−1 ) FT , ψ ( k−1 ) FT ] − ηFT ∂LFT ∂ [ θFT , ψFT ] ∣∣∣∣ θ ( k−1 ) FT , ψ ( k−1 ) FT k = 1 , . . . , K θ∗FT , ψ ∗ FT ≈ θ ( K ) FT , ψ ( K ) FT , ( 4 ) and compute the gradient ∂AlgFT∂θPT ∣∣∣ θ∗PT ( φ ′ ) by differentiating through this optimization.3 We can also choose to freeze the feature extractor parameters θFT and update only the head parameters ψFT during truncated FT , and use this to obtain meta-parameter gradients . This resembles linear evaluation , where a linear classifier is trained on top of fixed , pre-trained feature extractors [ 50 , 3 , 63 ] . Together , these two approximations allow for efficient computation of meta-parameter gradients . 3While Equation 4 uses standard gradient descent , we could use other differentiable optimizers ( e.g. , Adam ) . Algorithm 1 Gradient-based algorithm to learn meta-parameters . Notation defined in Appendix B , Table 3 . Vector-Jacobian products ( VJPs ) can be efficiently computed by standard autodifferentiation . 1 : Initialize PT parameters θ ( init ) PT , ψ ( init ) PT , ψ ( 0 ) FT and meta-parameters φ ( 0 ) 2 : for n = 1 , . . . , N iterations do 3 : Initialize θ ( 0 ) PT = θ ( init ) PT and ψ ( 0 ) PT = ψ ( init ) PT . 4 : for p = 1 , . . . , P PT iterations do 5 : [ θ ( p ) PT , ψ ( p ) PT ] = [ θ ( p−1 ) PT , ψ ( p−1 ) PT ] − ηPT ∂LPT ∂ [ θPT , ψPT ] ∣∣∣∣ θ ( p−1 ) PT , ψ ( p−1 ) PT 6 : end for 7 : Initialize FT encoder with PT solution : θ ( 0 ) FT = copy ( θ ( P ) PT ) . 8 : Approximate θ∗FT , ψ ∗ FT using Eq . 4 . 9 : Compute g1 = ∂LFT ∂ [ θFT , ψFT ] ∣∣∣∣ θ∗FT , ψ ∗ FT 10 : Compute VJP g2 = g1 ∂AlgFT ∂θPT ∣∣∣ θ ( P ) PT , ψ ( 0 ) FT using the unrolled learning step from line 8 . 11 : Approximate VJP ∂g∂φ ∣∣∣ φ ( n−1 ) = g2 ∂AlgPT ∂φ ∣∣∣ φ ( n−1 ) using the IFT ( Eq . 3 ) . 12 : φ ( n ) = φ ( n−1 ) − ηV ∂g∂φ ∣∣∣ φ ( n−1 ) 13 : Update PT initialization by setting : θ ( init ) PT = θ ( P ) PT and ψ ( init ) PT = ψ ( P ) PT . 14 : end for

Summary: This paper propose to solve a very important problem, meta-learning hyperparameters in pre-training stage, followed by fine-tuning a target task. Considering that the hyperparameters can be high-dimensional, the authors consider gradient-based hyperparameter optimization (HO) methods such as IFT based methods and unrolled differentiation. The main difficulty comes from the huge computational cost of dealing with long PT and FT trajectory and repeating such meta-optimization steps. Authors thus simply reduce the number of gradient steps used for each PT and FT stage. Yet, the experimental results demonstrate the effectiveness of the proposed methods over some of the simple baselines.