MTAdam: Automatic Balancing of Multiple Training Loss Terms

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 term can change as training progresses, especially for adversarial terms. In this work, we generalize the Adam optimization algorithm to handle multiple 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 arising 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 suboptimal initial loss weighting and to training outcomes that match conventional training with the prescribed hyperparameters of each method. MTAdam code can be found here 1 .

MetaBalance: Improving Multi-Task Recommendations via Adapting Gradient Magnitudes of Auxiliary Tasks

“…Gradient direction-based methods that are designed for adapting auxiliary tasks to improve the target task, which will be detailed in Section A.1.1, including: GradSimilarity [7], GradSurgery [34], OL-AUX [17]. Multi-Task balancing methods that treat all tasks equally, which will be detailed in Section A.1.2, including: Uncertainty [14], GradNorm [4], DWA [19], MTAdam [23] and MGDA [27]. And three simple baselines.…”

“…MTAdam [23] is an Adam-based optimizer that balances gradient magnitudes and then update parameters according to the rule of Adam [15]. Following MTAdam, we also directly manipulate the gradient magnitudes, instead of weighting task losses like Uncertainty [14], GradNorm [4] and DWA [19].…”

“…Multi-Task Balancing Methods. In multi-task learning, methods have been proposed to balance the joint learning of all tasks to avoid the situation where one or more tasks have a dominant influence on the network weight [4,14,19,23,27]. Although these methods have no special preference to the target task (as in our focus in this paper), we do discuss their connection to MetaBalance in Section A.1.2 (in Appendix) and experimentally compare with them in Section 5.…”

“…Finally, instead of using current magnitudes ∥G 𝑡 𝑡𝑎𝑟 ∥ and ∥G 𝑡 𝑎𝑢𝑥,𝑖 ∥ in Algorithm 1, following [23], we apply the moving average of magnitude of the corresponding gradient to take into account the variance among all gradient magnitudes over the training iterations:…”

Privacy-preserving and bandwidth-efficient federated learning

“…The regularization parameter 𝜆 is fixed to 5 for all experiments. The weights 𝛾 1 , 𝛾 2 , and 𝛾 3 of different losses in the objective of PROLIN (see Section 5.2.4) are adjusted dynamically during optimization using the technique in [40].…”