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: The paper describes a gradient-based hyperparameter optimization for differentiable pre-training hyperparameters in a pre-training and fine-tuning setup. The proposed approach approximately differentiates through the two-stage learning process. The authors show the benefit of their approach in two experiments, where they outperform multiple baselines.