Variational Information Bottleneck for Effective Low-Resource Fine-Tuning

1 INTRODUCTION . Transfer learning has emerged as the de facto standard technique in natural language processing ( NLP ) , where large-scale language models are pretrained on an immense amount of text to learn a general-purpose representation , which is then transferred to the target domain with fine-tuning on target task data . This method has exhibited state-of-the-art results on a wide range of NLP benchmarks ( Devlin et al. , 2019 ; Liu et al. , 2019 ; Radford et al. , 2019 ) . However , such pretrained models have a huge number of parameters , potentially making fine-tuning susceptible to overfitting . In particular , the task-universal nature of large-scale pretrained sentence representations means that much of the information in these representations is irrelevant to a given target task . If the amount of target task data is small , it can be hard for fine-tuning to distinguish relevant from irrelevant information , leading to overfitting on statistically spurious correlations between the irrelevant information and target labels . Learning low-resource tasks is an important topic in NLP ( Cherry et al. , 2019 ) because annotating more data can be very costly and time-consuming , and because in several tasks access to data is limited . In this paper , we propose to use the Information Bottleneck ( IB ) principle ( Tishby et al. , 1999 ) to address this problem of overfitting . More specifically , we propose a fine-tuning method that uses Variational Information Bottleneck ( VIB ; Alemi et al . 2017 ) to improve transfer learning in low-resource scenarios . VIB addresses the problem of overfitting by adding a regularization term to the training loss that directly suppresses irrelevant information . As illustrated in Figure 1 , the VIB component maps the sentence embedding from the pretrained model to a latent representation z , which is the only input to the taskspecific classifier . The information that is represented in z is chosen based on the IB principle , namely that all the information about the input that is represented in z should be necessary for the task . In particular , VIB directly tries to remove the irrelevant information , making it easier for the task classifier to avoid overfitting when trained on a small amount of data . We find that in low-resource scenarios , using VIB to suppress irrelevant features in pretrained sentence representations substantially improves accuracy on the target task . ∗Supported by the Viterbi Fellowship in the Center for Computer Engineering at the Technion . Removing unnecessary information from the sentence representation also implies removing redundant information . VIB tries to find the most concise representation which can still solve the task , so even if a feature is useful alone , it may be removed if it isn ’ t useful when added to other features because it is redundant . We hypothesize that this provides a useful inductive bias for some tasks , resulting in better generalization to out-of-domain data . In particular , it has recently been demonstrated that annotation biases and artifacts in several natural language understanding benchmarks ( Kaushik & Lipton , 2018 ; Gururangan et al. , 2018 ; Poliak et al. , 2018 ; Schuster et al. , 2019 ) allow models to exploit superficial shortcuts during training to perform surprisingly well without learning the underlying task . However , models that rely on such superficial features do not generalize well to out-of-domain datasets , which do not share the same shortcuts ( Belinkov et al. , 2019a ) . We investigate whether using VIB to suppress redundant features in pretrained sentence embeddings has the effect of removing these superficial shortcuts and keeping the deep semantic features that are truly useful for learning the underlying task . We find that using VIB does reduce the model ’ s dependence on shortcut features and substantially improves generalization to out-of-domain datasets . We evaluate the effectiveness of our method on fine-tuning BERT ( Devlin et al. , 2019 ) , which we call the VIBERT model ( Variational Information Bottleneck for Effective Low-Resource Fine-Tuning ) . On seven different datasets for text classification , natural language inference , similarity , and paraphrase tasks , VIBERT shows greater robustness to overfitting than conventional fine-tuning and other regularization techniques , improving accuracies on low-resource datasets . Moreover , on NLI datasets , VIBERT shows robustness to dataset biases , obtaining substantially better generalization to out-of-domain NLI datasets . Further analysis demonstrates that VIB regularization results in less biased representations . Our approach is highly effective and simple to implement , involving a small additional MLP classifier on top of the sentence embeddings . It is model agnostic and end-to-end trainable . In summary , we make the following contributions : 1 ) Proposing VIB for low-resource fine-tuning of large pretrained language models . 2 ) Showing empirically that VIB reduces overfitting , resulting in substantially improved accuracies on seven low-resource benchmark datasets against conventional fine-tuning and prior regularization techniques . 3 ) Showing empirically that training with VIB is more robust to dataset biases in NLI , resulting in significantly improved generalization to out-of-domain NLI datasets . To facilitate future work , we will release our code . 2 FINE-TUNING IN LOW-RESOURCE SETTINGS . The standard fine-tuning paradigm starts with a large-scale pretrained model such as BERT , adds a task-specific output component which uses the pretrained model ’ s sentence representation , and trains this model end-to-end on the task data , fine-tuning the parameters of the pretrained model . As depicted in Figure 1 , we propose to add a VIB component that controls the flow of information from the representations of the pretrained model to the output component . The goal is to address overfitting in resource-limited scenarios by removing irrelevant and redundant information from the pretrained representation . Problem Formulation We consider a general multi-class classification problem with a low-resource dataset D= { xi , yi } Ni=1 consisting of inputs xi∈X , and labels yi∈Y . We assume we are also given a large-scale pretrained encoder fϕ ( . ) parameterized by ϕ that computes sentence embeddings for the input xi . Our goal is to fine-tune fϕ ( . ) onD to maximize generalization . Information Bottleneck To specifically optimize for the removal of irrelevant and redundant information from the input representations , we adopt the Information Bottleneck principle . The objective of IB is to find a maximally compressed representation Z of the input representationX ( compression loss ) that maximally preserves information about the output Y ( prediction loss ) ,1 by minimizing : LIB = βI ( X , Z ) ︸ ︷︷ ︸ Compression Loss − I ( Z , Y ) ︸ ︷︷ ︸ Prediction Loss , ( 1 ) where β≥0 controls the balance between compression and prediction , and I ( . , . ) is the mutual information . Variational Information Bottleneck Alemi et al . ( 2017 ) derive an efficient variational estimate of ( 1 ) : LVIB =β E x [ KL [ pθ ( z|x ) , r ( z ) ] ] + E z∼pθ ( z|x ) [ −logqφ ( y|z ) ] , ( 2 ) where qφ ( y|z ) is a parametric approximation of p ( y|z ) , r ( z ) is an estimate of the prior probability p ( z ) of z , and pθ ( z|x ) is an estimate of the posterior probability of z . During training , the compressed sentence representation z is sampled from the distribution pθ ( z|x ) , meaning that a specific pattern of noise is added to the input of the output classifier qφ ( y|z ) . Increasing this noise decreases the information conveyed by z . In this way , the VIB module can block the output classifier qφ ( y|z ) from learning to use specific information . At test time , the expected value of z is used for predicting labels with qφ ( y|z ) . We refer to the dimensionality of z as K , which specifies the bottleneck size . Note that there is an interaction between decreasingK and increasing the compression by increasing β ( Shamir et al. , 2010 ; Harremoës & Tishby , 2007 ) . K and β are hyper-parameters ( Alemi et al. , 2017 ) . We consider parametric Gaussian distributions for prior r ( z ) and pθ ( z|x ) to allow an analytic computation for their Kullback-Leibler divergence,2 namely r ( z ) = N ( z|µ0 , Σ0 ) and pθ ( z|x ) = N ( z|µ ( x ) , Σ ( x ) ) , where µ and µ0 are K−dimensional mean vectors , and Σ and Σ0 are diagonal covariance matrices . We use the reparameterization trick ( Kingma & Welling , 2013 ) to estimate the gradients , namely z=µ ( x ) +Σ ( x ) , where ∼N ( 0 , I ) . To compute the compressed sentence representations pθ ( z|x ) , as shown in Figure 1 , we first feed sentence embeddings fϕ ( x ) through a shallow MLP . It is then followed by two linear layers , each withK hidden units to compute µ ( x ) and Σ ( x ) ( after a softplus transform to ensure non-negativity ) . We also use another linear layer to approximate qφ ( y|z ) . 3 EXPERIMENTS . Datasets We evaluate the performance on seven different benchmarks for multiple tasks , in particular text classification , natural language inference , similarity , and paraphrase detection . For NLI , we experiment with two well-known NLI benchmarks , namely SNLI ( Bowman et al. , 2015 ) and MNLI ( Williams et al. , 2018 ) . For text classification , we evaluate on two sentiment analysis datasets , namely IMDB ( Maas et al. , 2011 ) and Yelp2013 ( YELP ) ( Zhang et al. , 2015 ) . We additionally evaluate on three low-resource datasets in the GLUE benchmark ( Wang et al. , 2019 ) :3 paraphrase detection using MRPC ( Dolan & Brockett , 2005 ) , semantic textual similarity using STS-B ( Cer et al. , 2017 ) , and textual entailment using RTE ( Dagan et al. , 2006 ) . For the GLUE benchmark , SNLI , and Yelp , we evaluate on the standard validation and test splits . For MNLI , since the test sets are not available , we tune on the matched dev set and evaluate on the mismatched dev set ( MNLI-M ) or vice versa . See Appendix A for datasets statistics and Appendix B for hyper-parameters of all methods . Base Model We use the BERTBase ( 12 layers , 110M parameters ) and BERTLarge ( 24 layers , 340M parameters ) uncased ( Devlin et al. , 2019 ) implementation of Wolf et al . ( 2019 ) as our base models,4 known to work well for these tasks . We use the default hyper-parameters of BERT , i.e. , we use a sequence length of 128 , with batch size 32 . We use the stable variant of the Adam optimizer ( Zhang et al. , 2021 ; Mosbach et al. , 2021 ) with the default learning rate of 2e−5 through all experiments . We do not use warm-up or weight decay . Baselines We compare against prior regularization techniques , including previous state-of-the-art , Mixout : 1In this work , Z , X , and Y are random variables , and z , x and y are instances of these random variables . 2KL ( N ( µ0 , Σ0 ) ‖N ( µ1 , Σ1 ) ) = 12 ( tr ( Σ −1 1 Σ0 ) + ( µ1−µ0 ) TΣ −1 1 ( µ1−µ0 ) −K+log ( det ( Σ1 ) det ( Σ0 ) ) ) . 3We did not evaluate on WNLI and CoLA due to the irregularities in these datasets and the reported instability during the fine-tuning https : //gluebenchmark.com/faq . 4To have a controlled comparison , all results are computed with this PyTorch implementation , which might slightly differ from the TensorFlow variant ( Devlin et al. , 2019 ) . • Dropout ( Srivastava et al. , 2014 ) , a widely used stochastic regularization techniques used in multiple large-scale language models ( Devlin et al. , 2019 ; Yang et al. , 2019 ; Vaswani et al. , 2017 ) to mitigate overfitting . Following Devlin et al . ( 2019 ) , we apply dropout on all layers of BERT . • Mixout ( Lee et al. , 2019 ) is a stochastic regularization technique inspired by Dropout with the goal of preventing catastrophic forgetting during fine-tuning . Mixout regularizes the learning to minimize the deviation of a fine-tuned model from the pretrained initialization . It replaces the model parameters with the corresponding value from the pretrained model with probability p. • Weight Decay ( WD ) is a common regularization technique to improve generalization ( Krogh & Hertz , 1992 ) . It regularizes the large weights w by adding a penalization term λ2‖w‖ to the loss , where λ is a hyperparameter specifying the strength of regularization . Chelba & Acero ( 2004 ) and Daumé III ( 2007 ) adapt WD for fine-tuning of the pretrained models , and propose to replace this regularization term with λ‖w−w0‖ , where w0 are the weights of the pretrained models . Recently , Lee et al . ( 2019 ) demonstrated that the latter formulation of WD works better for finetuning of BERT than conventional WD and can improve generalization on small training sets .

Summary: The paper proposes a method to avoid overfitting while finetuning the large pretrained models for downstream tasks on small scale datasets. It has been shown that many SOTA models usually overfit w.r.t. spurious correlations in the data and as a result fail miserably when tested for generalization on the out of domain datasets. The proposed method tries to maximally filter out task-irrelevant information in the feature vectors by minimizing the mutual information between the original features and the bottleneck features while simultaneously optimizing for performance. Experiments on several datasets show improved performance on both in-domain and out-of-domain datasets.