MTAdam: Automatic Balancing of Multiple Training Loss Terms

Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10713–10729 November 7–11, 2021. c⃝2021 Association for Computational Linguistics 10713 MTAdam: Automatic Balancing of Multiple Training Loss Terms Itzik Malkiel Tel Aviv University itzik.malkiel@gmail.com Lior Wolf Tel Aviv University wolf@cs.tau.ac.il Abstract When training neural models, it is common to combine multiple loss terms. The balancing of these terms requires considerable human effort and is computationally demanding. Moreover, the optimal trade-off between the loss terms can change as training progresses, e.g., for ad- versarial terms. In this work, we generalize the Adam optimization algorithm to handle multi- ple loss terms. The guiding principle is that for every layer, the gradient magnitude of the terms should be balanced. To this end, the Multi-Term Adam (MTAdam) computes the derivative of each loss term separately, infers the first and second moments per parameter and loss term, and calculates a first moment for the magnitude per layer of the gradients aris- ing from each loss. This magnitude is used to continuously balance the gradients across all layers, in a manner that both varies from one layer to the next and dynamically changes over time. Our results show that training with the new method leads to fast recovery from sub- optimal initial loss weighting and to training outcomes that match or improve conventional training with the prescribed hyperparameters of each method. 1 Introduction In both supervised and unsupervised learning, adding loss terms often leads to improved perfor- mance. However, as more loss terms are added, the space of possible balancing weights increases expo- nentially, and more resources need to be allocated to identify good configurations that would justify the added terms. Another common challenge is that at different stages of the training process, the opti- mal balance may change (Mescheder et al., 2017). It is, therefore, necessary to have the balancing terms update dynamically during training, which further increases the hyperparameter space. In this work, we introduce Multi-Term Adam (MTAdam), an optimization algorithm for multi-term loss functions. MTAdam extends Adam (Kingma and Ba, 2014) and allows an ef- fective training of an unweighted multi-term loss objective. Thus, MTAdam can streamline the com- putationally demanding task of hyperparameter search, which is required for effectively weight- ing multi-term loss objectives. At every training iteration, a dynamic weight is assigned to each of the loss terms, based on the magnitude of the gradient that each term entails. The weights are assigned in a way that balances be- tween these gradients and equates their magnitude. This, however, would be an ineffective balancing method, without two crucial components: (i) the balancing needs to occur independently for each layer of the neural network and separately over the losses, since the relative contributions of the losses vary, depending on the layer, and (ii) the update step needs to take into account the maxi- mal variance among all losses, to support sufficient explorations in places of the parameter space, in which one of the losses becomes more sensitive. The main focus …

Content truncated to 500 words