LABO : Towards Learning Optimal Label Regularization via Bi-level Optimization (2023)

Abstract: Regularization techniques are crucial to improving the generalization performance and training efficiency of deep neural networks. Many deep learning algorithms rely on weight decay, dropout, batch/layer normalization to converge faster and generalize. Label Smoothing (LS) is another simple, versatile and efficient regularization which can be applied to various supervised classification tasks. Conventional LS, however, regardless of the training instance assumes that each non-target class is equally likely. In this work, we present a general framework for training with label regularization, which includes conventional LS but can also model instance-specific variants. Based on this formulation, we propose an efficient way of learning LAbel regularization by devising a Bi-level Optimization (LABO) problem. We derive a deterministic and interpretable solution of the inner loop as the optimal label smoothing without the need to store the parameters or the output of a trained model. Finally, we conduct extensive experiments and demonstrate our LABO consistently yields improvement over conventional label regularization on various fields, including seven machine translation and three image classification tasks across various neural network architectures while maintaining training efficiency.

## introduction deep neural networks (dnns) form the backbone of current state-of-the-art algorithms in various fields including natural language processing @xcite , computer vision @xcite and speech recognition @xcite . however, heavily overparameterized models may incur overfitting and suffer from poor generalizations @xcite . to address the issue, many regularization techniques have been developed in the literature: weight decay which constrains the optimization space @xcite , batch or layer normalization which speeds up the training of feed-forward nns @xcite , and dropout which implicitly approximates the effect of averaging the predictions of all sparse subnetworks networks @xcite . label smoothing (ls) is another simple regularization technique; it is widely applied to many applications including image classification @xcite and token-level sequence generation @xcite for enhancing the generalization, without suffering additional computational costs. it encourages a model to treat each non-target class as equally likely for classification by using a uniform distribution to smooth one-hot labels @xcite . although combining the uniform distribution with the original one-hot label is beneficial for regularization, conventional ls does not take into account the true relations between different label categories. more specifically, for token-level generation, uniformly allocating the probability mass on non-target words disregards the semantic relationship between non-target words and the context. on the other hand, the probability of the target is predefined and unchanged. however, the distribution of natural language exhibits remarkable variations in the per-token perplexity @xcite , which encourages us to adapt corresponding target probabilities for different contexts. one of the instance-dependent techniques of learning the relation between different target categories is knowledge distillation (kd) @xcite , which is a popular technique of transfer learning utilizing knowledge from related tasks or teacher models @xcite @xcite . it is widely applied for model compression and ensembling across applications ranging from computer vision @xcite @xcite to natural language processing @xcite . however, kd requires training a separate model as a teacher for every new task. besides, it either introduces an extra inference pass to get the teacher's prediction for each instance during the training or requires saving the teacher's prediction for all training samples to avoid the extra forward pass of teacher models. this greatly increases the time and space complexity in practice. especially for token-level generation tasks, e.g. machine translation, to save all the output probabilities of teacher models costs o(n lv ) space, where n is the number of sequences, l is the averaged length of all sequences and v is the vocabulary size. besides the empirical success of kd, it is unclear how student networks benefit from these smoothed labels. a series of investigations have looked at regularization and have demonstrated that the success of both kd and label smoothing is due to a similar regularization effect of smoothed targets @xcite . based on this finding and the low training overhead of ls, there is a significant interest, in the community, in algorithms that can enhance conventional ls. @xcite demonstrates the importance of an instance specific ls regularization. they demonstrate better performance compared to ls, but use a trained model to infer prior knowledge on the label space and thereby sacrifice some of the efficiency of ls. in this work, we first revisit the conventional ls and generalize it to an instance-dependent label regularization framework with a constraint on overconfidence. within this framework, we demonstrate that both ls and kd can be interpreted as instances of a smoothing distribution with a confidence penalty. finally, we propose to learn the optimal smoothing function along with the model training, by devising a bi-level optimization problem. we solve the inner problem by giving a deterministic and interpretable solution and applying gradient-based optimization to solve the outer optimization. our contributions can be summarized as follows: we will provide a brief overview of existing label regularization methods. label smoothing. ls is a regularization technique to improve the generalization performance of neural networks by preventing the model from predicting the training examples overconfidently. it smoothes a one-hot target label with the uniform distribution u (•) = 1 k , where k is the number of classes. as a result, ls training is equivalent to training with a smoothed label pu , where: and k is the ground-truth class. we note from eq. 1 that a higher α can lead to a smoother label distribution and a lower α to a peakier label. confidence penalty. another important technique similar to label smoothing is the confidence penalty (cp) @xcite , in which a regularization term h(p θ (x)) is introduced into the objective function to directly punish the overconfidence of predictions: ## methodology in this section, we interpret our optimal label regularization from the bi-level optimization perspective. first, we provide a close look at the conventional uniform label regularization and show the generalized framework bridges the objectives of conventional ls and kd methods. then, we introduce a closed-form solution to our optimal instance-dependent label smoothing and describe an online implementation under this formulation. ## experiments we evaluate the performance of our proposed labo method on both machine translation and computer vision. for machine translation, we evaluate on the iwslt'14 @xcite , iwlst'17 @xcite and wmt'14 @xcite datasets using transformer-base models and for image classification on cifar-10, cifar-100 @xcite and ima-genet2012 @xcite using resnet-based update θ t+1 = θ t -η∇ θ r θ ( p * ls , s); 9: end for models of various sizes (parameters). all experiments were performed on one or more nvidia tesla (v100) gpus. ## discussion one of important problems of neural language models is the large discrepancy of predicted probabilities between tokens with low and high frequency, in other words, the long-tailed phenomena in the neural language models @xcite . in this section, we study the impact of our method on the prediction of tokens with different frequencies. we first computed the averaged frequency of tokens in every source sentence @xmath0. this frequency score (fs) is defined in @xcite : ## related work there is a great body of works inspired by the kd technique. some of them focus on boosting performance with explicitly regularized smoothed targets. zhang and sabuncu (2020) interpret studentteacher training as an amortized maximum aposteriori estimation and derive an equivalence between self-distillation and instance-dependent label smoothing. this analysis helped them to devise a regularization scheme, beta smoothing. however, they still use an extra model to infer a prior distribution on their smoothing technique during the training. @xcite introduce an additional regularization to penalize the predictive output between different samples of the same class. there are several works discussing the empirical impact of kd and giving different practical suggestions. @xcite conducted extensive experiments to explore the effect of label regularization to the generalization ability of compact pre-trained language models. other works propose to learn class-wise label smoothing or progressive refining of smoothed la-bels. @xcite propose to capture the relation of classes by introducing decoupled labels table which increases space complexity by o(k × k). the concurrent work @xcite utilizes the model trained at i-th epoch as the teacher to regularize the training at (i + 1)-th epoch along with annealing techniques. however, contrary to our work, it either requires a separate model as teacher, or stores the labels generated at the i-th epoch for training the next epoch, whereby increasing space complexity by o(n × k). @xcite investigate the relationship between generalized entropy regularization and label smoothing and provide empirical results on two text generation tasks. @xcite propose a masked label smoothing to remove the conflict by reallocating the smoothed probabilities based on the difference among languages. @xcite propose to learn confidence for nmt by jointly training a connet model to estimate the confidence of the prediction. we develop labo motivated by a principled starting point: to generalize the smoothing distribution to a general form with neither time nor space complexity increase during the training and inference. ## conclusion our aim in this work is to fill the accuracy gap between label smoothing and knowledge distillation techniques while maintaining the training efficiency of ls regularization. we proposed learning an instance-dependent label smoothing regularization simultaneously with training our model on the target. we began by generalizing the classical ls method and introduced our objective function by substituting the uniform distribution with a general, instance-dependent, discrete distribution. within this formulation, we explained the relationship between the ls and kd. then, using a bi-level optimization approach, we obtained an approximation for the optimal smoothing function. we conducted extensive experiments to compare our model with conventional ls, kd and various state-of-the-art label regularization methods on popular mt and v benchmarks and showed the effectiveness and efficiency of our technique. finally, we analyze the impact of our methods on the prediction of neural machine translation systems under different averaged token frequency settings and show our methods can greatly reduce the discrepancy between predicted probabilities of different tokens. in practice, apart from general regularization techniques like dropout and weight decay, many advanced techniques are designed for specific tasks like fixnorm @xcite , cut-mix @xcite and switchout @xcite . we leave the future work to combine these methods with our technique. moreover, we plan to explore the practical applications of our method for large-scale model training. one specific application is to improve the pre-training of large language models and vision transformers.