Improved Regularization and Robustness for Fine-tuning in Neural Networks

1 Introduction . Learning from limited labeled data is a fundamental problem in many real-world applications ( Ratner et al. , 2016 , 2017 ) . A common approach to address this problem is fine-tuning a large model that has been pre-trained on publicly available labeled data ( He et al. , 2019 ) . Since fine-tuning is typically applied to a target task with limited labels , this algorithm is prone to overfitting or “ memorization ” issues ( Tan et al. , 2018 ) . These issues worsen when the target task contains noisy labels ( Zhang et al. , 2016 ) . In this paper , we analyze regularization methods for fine-tuning from both theoretical and empirical perspectives . Based on the analysis , we propose a regularized self-labeling approach that improves the generalization and robustness properties of fine-tuning . Previous works ( Li et al. , 2018a , b ) have proposed regularization methods to constrain the distance between a fine-tuned model and the pre-trained model in the Euclidean norm . Li et al . ( 2020 ) provides extensive study to show that the performance of fine-tuning and the benefit of adding regularization depend on the hyperparameter choices . Salman et al . ( 2020 ) empirically find that performing adversarial training during the pre-training phase helps learn pre-trained models that transfer better to downstream tasks . The work of Gouk et al . ( 2021 ) generalizes the above ideas to various norm choices and finds that projected gradient descent methods perform well for implementing distance-based regularization . Additionally , they derive generalization bounds for fine-tuning using Rademacher complexity . These works focus on settings where there is no label noise in the target data set . When 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . label noise is present , for example , due to applying weak supervision techniques ( Ratner et al. , 2016 ) , an important question is to design methods that are robust to such noise . The problem of learning from noisy labels has a rich history of study in supervised learning ( Natarajan et al. , 2013 ) . In contrast , little is known in the transfer learning setting . These considerations motivate us to analyze the generalization and robustness properties of fine-tuning . In Section 4.1 , we begin by conducting a PAC-Bayesian analysis of regularized fine-tuning . This is inspired by recent works that have found PAC-Bayesian analysis correlates with empirical performance better than Rademacher complexity ( Jiang et al. , 2020 ) . We identify two critical measures for analyzing the generalization performance of fine-tuning . The first measure is the ` 2 norm of the distance between the pre-trained model ( initialization ) and the fine-tuned model . The second measure is the perturbed loss of the fine-tuned model , i.e . its loss after the model weights get perturbed by random noise . First , we observe that the fine-tuned weights remain closed to the pre-trained model . Moreover , the top layers travel much further away from the pre-trained model than the bottom layers . Second , we find that fine-tuning from a pre-trained model implies better noise stability than training from a randomly initialized model . In Section 4.2 , we evaluate regularized fine-tuning for target tasks with noisy labels . We find that fine-tuning is prone to “ memorizing the noisy labels ” , and regularization helps alleviate such memorization behavior . Moreover , we observe that the neural network has not yet overfitted to the noisy labels during the early phase of fine-tuning . Thus , its prediction could be used to relabel the noisy labels . We propose an algorithm that incorporates layerwise regularization and self-labeling for improved regularization and robustness based on our results . Figure 1 illustrates the two components . First , we encode layer-wise distance constraints to regularize the model weights at different levels . Compared to ( vanilla ) fine-tuning , our algorithm reduces the gap between the training and test accuracy , thus alleviating overfitting . Second , we add a self-labeling mechanism that corrects and reweights “ noisy labels ” based on the neural network ’ s predictions . Figure 1 shows that our algorithm effectively hinders the model from learning the incorrect labels by relabeling them to correct ones . In Section 5 , we evaluate our proposed algorithm for both transfer learning and few-shot classification tasks with image and text data sets . First , using ResNet-101 ( He et al. , 2016 ) as the pre-trained model , our algorithm outperforms previous fine-tuning methods on seven image classification tasks by 1.76 % on average and 3.56 % when their labels are noisy . Second , we find qualitatively similar results for applying our approach to medical image classification tasks ( ChestX-ray14 ( Wang et al. , 2017 ; Rajpurkar et al. , 2017 ) ) and vision transformers ( Dosovitskiy et al. , 2020 ) . Finally , we extend our approach to few-shot learning and sentence classification . For these related but different tasks and data modalities , we find an improvement of 0.75 % and 0.46 % over previous methods , respectively . In summary , our contributions are threefold . First , we provide a PAC-Bayesian analysis of regularized fine-tuning . Our result implies empirical measures that explain the generalization performance of regularized fine-tuning . Second , we present a regularized self-labeling approach to enhance the generalization and robustness properties of fine-tuning . Third , we validate our approach on an extensive collection of classification tasks and pre-trained model architectures . 2 Related work . Fine-tuning is widely used in multi-task and transfer learning , meta-learning , and few-shot learning . Previous works ( Li et al. , 2018b , a ) find that injecting ` 2 regularization helps improve the performance of fine-tuning . Li et al . ( 2018b ) propose a ` 2 distance regularization method that penalizes the ` 2 distance between the fine-tuned weights and the pre-trained weights . Li et al . ( 2018a ) penalize the distance between the feature maps as opposed to the layer weights . Chen et al . ( 2019 ) design a regularization method that suppresses the spectral components ( of the feature maps ) with small singular values to avoid negative transfer . Gouk et al . ( 2021 ) instead encode distance constraints in constrained minimization and use projected gradient descent to ensure the weights are close to the pre-trained model . Salman et al . ( 2020 ) show that fine-tuning from adversarially robust pre-trained models outperforms fine-tuning from ( standard ) pre-trained models . The robustness of learning algorithms in the presence of label noise has been extensively studied in supervised learning ( Natarajan et al. , 2013 ) . Three broad ideas for designing robust algorithms include defining novel losses , identifying noisy labels , and regularization methods . Zhang and Sabuncu ( 2018 ) design the robust Generalized Cross Entropy ( GCE ) loss which is a mixture of Cross Entropy ( CE ) and mean absolute error . The Symmetric Cross Entropy ( SCE ) ( Wang et al. , 2019 ) loss combines reverse cross-entropy with the CE loss . Ma et al . ( 2020 ) proposes APL , which normalizes the loss to be robust to noisy labels and combines active and passive loss functions . Thulasidasan et al . ( 2019 ) propose DAC , which identifies and suppresses the signals of noisy samples by abstention-based training . Liu et al . ( 2020 ) introduce ELR , an early learning regularization approach to mitigate label “ memorization ” . Huang et al . ( 2020 ) propose a self-adaptive training method that corrects noisy labels and reweights training data to suppress erroneous signals . These works primarily concern the supervised learning setting . To the best of our knowledge , fine-tuning algorithms under label noise are relatively under-explored in transfer learning . Our approach draws inspiration from the semi-supervised learning literature , which has evaluated pseudo-labeling and self-training approaches given a limited amount of labeled data , and a large amount of unlabeled data ( Zou et al. , 2019 ; Xu et al. , 2021 ; Tai et al. , 2021 ) . A very recent work considered a sharpness-aware approach and evaluated its performance for fine-tuning from noisy labels ( Foret et al. , 2021 ) . It would be interesting to see if combining their approach with our ideas could lead to better results . From a theoretical perspective , the seminal work of Ben-David et al . ( 2010 ) considers a setting where labeled data from many source tasks and a target task is available . They show that minimizing a weighted combination of the source and target empirical risks leads to the best result . In the supervised setting , Arora et al . ( 2019 ) provide data-dependent generalization bounds for neural networks . Nagarajan and Kolter ( 2018 ) and Wei and Ma ( 2019a , b ) provide improved generalization bounds that depend only polynomially on the depth of the neural network characterized by a margin condition . Recent work has found that the PAC-Bayes theory provides generalization measures that correlate with empirical generalization performance better than other alternatives ( Jiang et al. , 2020 ) . We defer a more extensive review of PAC-Bayesian generalization theory to Section A . 3 Preliminaries . Problem setup . We begin by formally introducing the setup . Suppose we would like to solve a target task . We have a training data set of size n ( t ) . Let ( x ( t ) 1 , y ( t ) 1 ) , . . . , ( x ( t ) n ( t ) , y ( t ) n ( t ) ) be the feature vectors and the labels of the training data set . We assume that every x ( t ) i lies in a d-dimensional space denoted by X ⇢ Rd . Following standard terminologies in statistical learning , we assume all data are drawn from some unknown distribution supported on X ⇥ Y , where Y ✓ R is the label space . Denote the underlying data distribution as P ( t ) . For our result in Section 4.1 , we consider feedforward neural networks , and our results can also be extended to convolutional neural networks ( e.g. , using ideas from Long and Sedghi ( 2020 ) ) . Consider an L layer neural network . For each layer i from 1 to L , let i be the activation function and let Wi be the weight matrix at layer i . Given an input z to the i-th layer , the output is then denoted as i ( z ) = i ( Wiz ) . Thus , the final output of the network is given by fW ( x ) = L L 1 · · · 1 ( x ) , for any input x 2 X , where we use W = [ W1 , . . . , WL ] to include all the parameters of the network for ease of notation . The prediction error of fW is measured by a loss function ` ( · ) that is both convex and 1-Lipschitz . L ( t ) ( fW ) = E ( x , y ) ⇠P ( t ) [ ` ( fW ( x ) , y ) ] . ( 1 ) Suppose we have access to a pre-trained source model Ŵ ( s ) . Using Ŵ ( s ) as an initialization , we then fine-tune the layers Ŵ ( s ) 1 , . . . , Ŵ ( s ) L to predict the target labels . Applications . We consider seven image classification data sets described in Table 1 . We use ResNets pre-trained on ImageNet as the initialization Ŵ ( s ) ( Russakovsky et al. , 2015 ; He et al. , 2016 ) . We perform fine-tuning using the pre-trained network on the above data sets . See Section 5.1 for further description of the training procedure . Modeling label noise . In many settings , fine-tuning is applied to a target task whose labels may contain noise ; for example , if the target labels are created using weak supervision techniques ( Ratner et al. , 2019 ; Saab et al. , 2021 ) . To capture such settings , we denote a noisy data set as ( x ( t ) 1 , ỹ ( t ) 1 ) , . . . , ( x ( t ) n ( t ) , ỹ ( t ) n ( t ) ) , where ỹ ( t ) i is a noisy version of y ( t ) i . We consider two types of label noise : independent noise and correlated noise . We say that the label noise is independent if it is independent of the input feature vector Pr ( ỹ ( t ) i = k|y ( t ) i = j , x ( t ) i ) = Pr ( ỹ ( t ) i = k|y ( t ) i = j ) = ⌘j , k for some fixed ⌘j , k between 0 and 1 . On the other hand , we say that the label noise is correlated if it depends on the input feature vector ( i.e . the above equation does not hold ) .

Summary: This paper studies the aspects of fine-tuning a model (pre-trained on a separate dataset) in order to learn a new task. It has the following contribution:  (i) Generalization properties of fine-tuning explained by the PAC-Bayes generalization bound that depends on two terms: (a) distance between each layer from the pre-trained model and (b) noise stability of the pre-trained model.  (ii) Inspired by the generalization analysis, the paper presents an algorithm to perform fine-tuning. This algorithm includes (a) regularization term that computes the distance between the current parameters and the pre-trained model, and (b) iteratively corrects mislabeled examples where model has high confidence and down-weighting less confident examples.  (iii) Proposed algorithm is evaluated on a suite of benchmark tasks including transfer learning and few shot classification tasks.