Parameter-Efficient Transfer Learning with Diff Pruning

1 INTRODUCTION . Task-specific finetuning of pretrained deep networks has become the dominant paradigm in contemporary NLP , achieving state-of-the-art results across a suite of natural language understanding tasks ( Devlin et al. , 2019 ; Liu et al. , 2019c ; Yang et al. , 2019 ; Lan et al. , 2020 ) . While straightforward and empirically effective , this approach is difficult to scale to multi-task , memory-constrained settings ( e.g . for on-device applications ) , as it requires shipping and storing a full set of model parameters for each task . Inasmuch as these models are learning generalizable , task-agnostic language representations through self-supervised pretraining , finetuning the entire model for each task is an especially inefficient use of model parameters . A popular approach to parameter-efficiency is to learn sparse models for each task where a subset of the final model parameters are exactly zero ( Gordon et al. , 2020 ; Sajjad et al. , 2020 ; Zhao et al. , 2020 ; Sanh et al. , 2020 ) . Such approaches often face a steep sparsity/performance tradeoff , and a substantial portion of nonzero parameters ( e.g . 10 % -30 % ) are still typically required to match the performance of the dense counterparts . An alternative is to use multi-task learning or feature-based transfer for more parameter-efficient transfer learning with pretrained models ( Liu et al. , 2019b ; Clark et al. , 2019 ; Stickland & Murray , 2019 ; Reimers & Gurevych , 2019 ; Feng et al. , 2020 ) . These methods learn only a small number of additional parameters ( e.g . a linear layer ) on top of a shared model . However , multi-task learning generally requires access to all tasks during training to prevent catastrophic forgetting ( French , 1999 ) , while feature-based transfer learning ( e.g . based on taskagnostic sentence representations ) is typically outperformed by full finetuning ( Howard & Ruder , 2018 ) . Adapters ( Rebuffi et al. , 2018 ) have recently emerged as a promising approach to parameterefficient transfer learning within the pretrain-finetune paradigm ( Houlsby et al. , 2019 ; Pfeiffer et al. , 2020a ; b ; c ) . Adapter layers are smaller , task-specific modules that are inserted between layers of a pretrained model , which remains fixed and is shared across tasks . These approaches do not require access to all tasks during training making them attractive in settings where one hopes to obtain and share performant models as new tasks arrive in stream . Houlsby et al . ( 2019 ) find that adapter layers trained on BERT can match the performance of fully finetuned BERT on the GLUE benchmark ( Wang et al. , 2019a ) while only requiring 3.6 % additional parameters ( on average ) per task . In this work , we consider a similar setting as adapters but propose a new diff pruning approach with the goal of even more parameter-efficient transfer learning . Diff pruning views finetuning as learning a task-specific difference vector that is applied on top of the pretrained parameter vector , which remains fixed and is shared across different tasks . In order to learn this vector , we reparameterize the task-specific model parameters as θtask = θpretrained +δtask , where the pretrained parameter vector θpretrained is fixed and the task-specific diff vector δtask is finetuned . The diff vector is regularized with a differentiable approximation to the L0-norm penalty ( Louizos et al. , 2018 ) to encourage sparsity . This approach can become parameter-efficient as the number of tasks increases as it only requires storing the nonzero positions and weights of the diff vector for each task . The cost of storing the shared pretrained model remains constant and is amortized across multiple tasks . On the GLUE benchmark ( Wang et al. , 2019a ) , diff pruning can match the performance of the fully finetuned BERT baselines while finetuning only 0.5 % of the pretrained parameters per task , making it a potential alternative to adapters for parameter-efficient transfer learning . 2 BACKGROUND : TRANSFER LEARNING FOR NLP . The field of NLP has recently seen remarkable progress through transfer learning with a pretrainand-finetune paradigm , which initializes a subset of the model parameters for all tasks from a pretrained model and then finetunes on a task specific objective . Pretraining objectives include context prediction ( Mikolov et al. , 2013 ) , autoencoding ( Dai & Le , 2015 ) , machine translation ( McCann et al. , 2017 ) , and more recently , variants of language modeling ( Peters et al. , 2018 ; Radford et al. , 2018 ; Devlin et al. , 2019 ) objectives . Here we consider applying transfer learning to multiple tasks . We consider a setting with a potentially unknown set of tasks , where each τ ∈ T has an associated training set { x ( n ) τ , y ( n ) τ } Nn=1.1 For all tasks , the goal is to produce ( possibly tied ) model parameters θτ to minimize the empirical risk , min θτ 1 N N∑ n=1 L ( f ( x ( n ) τ ; θτ ) , y ( n ) τ ) + λR ( θτ ) where f ( · ; θ ) is a parameterized function over the input ( e.g . a neural network ) , L ( · , · ) is a loss function ( e.g . cross-entropy ) , and R ( · ) is an optional regularizer with hyperparameter λ . This multi-task setting can use the pretrain-then-finetune approach by simply learning independent parameters for each task ; however the large size of pretrained models makes this approach exceedingly parameter inefficient . For example , widely-adopted models such as BERTBASE and BERTLARGE have 110M and 340M parameters respectively , while their contemporaries such as T5 ( Raffel et al. , 2020 ) , Megatron-LM ( Shoeybi et al. , 2019 ) , and Turing-NLG ( Rajbhandari et al. , 2019 ) have parameter counts in the billions . Storing the fully finetuned models becomes difficult even for a moderate number of tasks.2 A classic approach to tackling this parameterinefficiency ( Caruana , 1997 ) is to train a single shared model ( along with a task-specific output layer ) against multiple tasks through joint training . However , the usual formulation of multi-task learning requires the set of tasks T to be known in advance in order to prevent catastrophic forgetting ( French , 1999 ) ,3 making it unsuitable for applications in which the set of tasks is unknown ( e.g . when tasks arrive in stream ) . 3 DIFF PRUNING . Diff pruning formulates task-specific finetuning as learning a diff vector δτ that is added to the pretrained model parameters θpretrained . We first reparameterize the task-specific model parameters , θτ = θpretrained + δτ , 1Therefore our setup is different from the classic multitask setting which usually assumes that set of tasks is known 2An intriguing line of work suggests that large-scale language models can be used without finetuning for a variety of tasks if given the appropriate context ( Radford et al. , 2019 ; Brown et al. , 2020 ) . While interesting , these models generally underperform task-specific models and require billions of parameters , though recent work suggests that they can be made substantially smaller ( Schick & Schutze , 2020 ) . 3However , work on continual learning mitigates these issues to an extent ( Shin et al. , 2017 ; Lopez-Paz & Ranzato , 2017 ; Lee et al. , 2017 ; Kirkpatrick et al. , 2017 ; Parisi et al. , 2018 ) . which results in the following empirical risk minimization problem , min δτ 1 N N∑ n=1 L ( f ( x ( n ) τ ; θpretrained + δτ ) , y ( n ) τ ) + λR ( θpretrained + δτ ) . This trivial reparameterization is equivalent to the original formulation . Its benefit comes in the multi-task setting where the cost of storing the pretrained parameters θpretrained is amortized across tasks , and the only marginal cost for new tasks is the diff vector . If we can regularize δτ to be sparse such that ‖δτ‖0 ‖θpretrained‖0 , then this approach can become more parameter-efficient as the number of tasks increases . We can specify this goal with an L0-norm penalty on the diff vector , R ( θpretrained + δτ ) = ‖δτ‖0 = d∑ i=1 1 { δτ , i 6= 0 } . 3.1 DIFFERENTIABLE APPROXIMATION TO THE L0-NORM This regularizer is difficult to directly optimize as it is non-differentiable . In order to approximate this L0 objective , we follow the standard approach for gradient-based learning with L0 sparsity using a relaxed mask vector ( Louizos et al. , 2018 ) . This approach involves relaxing a binary vector into continuous space , and then multiplying it with a dense weight vector to determine how much of the weight vector is applied during training . After training , the mask is deterministic and a large portion of the diff vector is true zero . To apply this method we first decompose δτ into a binary mask vector multiplied with a dense vector , δτ = zτ wτ , zτ ∈ { 0 , 1 } d , wτ ∈ Rd We can now instead optimize an expectation with respect to zτ , whose distribution p ( zτ ; ατ ) is initially Bernoulli with parameters ατ , min ατ , wτ Ezτ∼p ( zτ ; ατ ) [ 1 N N∑ n=1 L ( f ( x ( n ) τ ; θpretrained + zτ wτ , ) , y ( n ) τ ) + λ‖δτ‖0 ] . This objective is still difficult in practice due to zτ ’ s being discrete ( which requires the score function gradient estimator ) , but the expectation provides some guidance for empirically effective relaxations . We follow prior work ( Louizos et al. , 2018 ; Wang et al. , 2019b ) and relax zτ into continuous space [ 0 , 1 ] d with a stretched Hard-Concrete distribution ( Jang et al. , 2017 ; Maddison et al. , 2017 ) , which allows for the use of pathwise gradient estimators . Specifically , zτ is now defined to be a deterministic and ( sub ) differentiable function of a sample u from a uniform distribution , u ∼ U ( 0,1 ) , sτ = σ ( logu− log ( 1− u ) + ατ ) , s̄τ = sτ × ( r − l ) + l , zτ = min ( 1 , max ( 0 , s̄τ ) ) . Here l < 0 and r > 1 are two constants used to stretch sτ into the interval ( l , r ) d before it is clamped to [ 0 , 1 ] d with the min ( 1 , max ( 0 , · ) ) operation . In this case we have a differentiable closed-form expression for the expected L0-norm , E [ ‖δτ‖0 ] = d∑ i=1 E [ 1 { zτ , i > 0 } ] = d∑ i=1 σ ( ατ , i − log −l r ) . Thus the final optimization problem is given by , min ατ , wτ Eu∼U [ 0,1 ] [ 1 N N∑ n=1 L ( f ( x ( n ) τ ; θpretrained + zτ wτ , ) , y ( n ) τ ) ] + λ d∑ i=1 σ ( ατ , i − log −l r ) , and we can now utilize pathwise gradient estimators to optimize the first term with respect to ατ since the expectation no longer depends on it.4 After training we obtain the final diff vector δτ by sampling u once to obtain zτ ( which is not necessarily a binary vector but has a significant number of dimensions equal to exactly zero due to the clamping function ) , then setting δτ = zτ wτ .5 4To reduce notation clutter we subsume the parameters of the task-specific output layer , which is not pretrained , into θpretrained . We do not apply the L0-norm penalty on these parameters during training . 5We found sampling once to work as well as more complicated alternatives ( e.g . based on multiple samples ) . 3.2 L0-BALL PROJECTION WITH MAGNITUDE PRUNING FOR SPARSITY CONTROL Differentiable L0 regularization provides a strong way to achieve high sparsity rate . However , it would be ideal to have more fine-grained control into the exact sparsity rate in the diff vector , especially considering applications which require specific parameter budgets . As λ is just the Lagrangian multiplier for the constraint E [ ‖δτ‖0 ] < η for some η , this could be achieved in principle by searching over different values of λ . However we found it more efficient and empirically effective to achieve an exact sparsity rate by simply projecting onto the L0-ball after training . Specifically we use magnitude pruning on the diff vector δτ and target a sparsity rate t % by only keeping the top t % × d values in δτ .6 Note that unlike standard magnitude pruning , this is based on the magnitude of the diff vector values and not the model parameters . As is usual in magnitude pruning , we found it important to further finetune δτ with the nonzero masks fixed to maintain good performance ( Han et al. , 2016 ) . Since this type of parameter-efficiency through projection onto the L0-ball can be applied without adaptive diff pruning,7 such an approach will serve as one of our baselines in the empirical study .

Summary: This paper proposes diff pruning, an alternative paradigm for parameter-efficient transfer learning of pre-trained models. Similar to adapters, diff pruning leaves the body of the pre-trained model unchanged. Rather than inserting additional task-specific parameters into the pre-trained model, diff pruning adds reparameterizes the parameters of the transferred model $\theta_\tau$ by adding a diff vector $\delta_\tau$ to them: $\theta_\tau = \theta_{\text{pretrained}} + \delta_\tau$. Parameter efficiency is achieved by regularizing $\theta_\tau$ to be sparse. The authors achieve this by using a relaxed mask vector to approximate the $L_0$ norm. They also propose a way to control for a specific sparsity rate via projection onto the $L_0$ ball after training and to enforce group sparsity that takes the model's structure into account. The approach is evaluated on the GLUE benchmark where it achieves competitive performance to full fine-tuning a BERT Large model and adapters while being more parameter-efficient than both of them.