MTA dam: Automatic Balancing of Multiple Training Loss Terms (2021)

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 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 or improve conventional training with the prescribed hyperparameters of each method.

## introduction in both supervised and unsupervised learning, adding loss terms often leads to improved performance. however, as more loss terms are added, the space of possible balancing weights increases exponentially, 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 optimal balance may change @xcite . 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 @xcite and allows an effective training of an unweighted multi-term loss objective. thus, mtadam can streamline the computationally demanding task of hyperparameter search, which is required for effectively weighting 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 between 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 maximal 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 of our experiments is in the domains of natural language understanding and conditional adversarial image generation, in which the leading methods often utilize multiple loss terms. our results show that mtadam is able to recover from suboptimal starting points, in which the weight parameters are set inappropriately, while adam and other baseline methods cannot. adam @xcite combines both principles and employs the first and second moments of the gradients of each learned parameter and applies them during the backward step. adam has become a dominant optimizer, that is applied across many applications, and, in particular, it is the de-facto standard in the field of adversarial training. it is known for improving convergence to work well with the default values of its own hyperparameters. multiple methods have been suggested for selecting hyperparameters @xcite @xcite . hyperband @xcite performs hyperparameter search as an infinite-armed bandit problem utilizing a predefined amount of resources, while searching for the best configuration of hyperparameters that maximizes the given success criterion. it can provide an order-of-magnitude speedup compared to bayesian optimization @xcite @xcite . in our experiments, we employ multiple methods from natural language understanding and the field of image generation. the methods vary in the type of supervision employed or the task being solved: (i) bert @xcite is a pretrained transformer-based @xcite language model, that applies fine-tuning on various language understanding tasks using supervision. (ii) recobert @xcite is a text-similarity model that fine-tunes bert on a given catalog, by jointly optimizing a standard masked language modeling, along with a unique title-description model. (iii) pix2pix @xcite generates an image in domain b based on an input image in domain a, after observing matching pairs during training. (iv) cyclegan @xcite performs the same task, while training in an unsupervised manner on unmatched images from the two domains. (v) srgan @xcite generates high-resolution images from lowresolution ones and is trained in a supervised way. ## motivation: multi-loss dynamics to examine the challenges that arise when not scaling the losses correctly, we consider a simple regression problem that involves the l1 or the l2 loss. specifically, we fit two coefficients to learn the polynomial f (x) = ax + bx 2 . to this end, we randomly sample batches of @xmath0, and train a network to minimize the loss, i.e, the network performs a dot product between (x, x 2 ) and the vector of parameters @xmath1, and the objective compares between (x, x 2 )w and f (x): for l1 |(x, x 2 )w -f (x)| and for l2 (x, x 2 )w -f (x) 2 2 . we track and analyze the loss and its gradients on the two parameters. we run three adam experiments with the values @xmath2.52, @xmath3.29 (picked at random, results are very similar for other values). each time we train the model with a different objective: l1, l2, and l1 where the gradients are scaled with the magnitudes of the l2 gradients. the latter utilizes the l1 loss and scales the gradient magnitudes of the coefficients to match the magnitudes created by the l2 loss, i.e. we apply we further explore this in a scenario with multiple loss terms. it is often the case that new loss terms are added during training. however, one would like to avoid tuning each one separately. as a computationally efficient sample, we consider supervised image superresolution using two out of the three loss terms of @xcite : mse, a gan which verifies that the result is in the target domain, and a perceptual loss. we consider the case in which either the gan loss or the perceptual loss are weighted equally, in comparison to the mse (a weight that is larger by two orders of magnitude than the optimal) and observe, in fig. 2 that this leads to suboptimal results. however, using the gradient magnitude of a balanced term fixes this and allows the method to train well. thus, this motivating experiment demonstrates that we can balance a loss term by using the gradient magnitude of another term. experiments with the three terms together are shown in sec. 5. ## method the adam algorithm optimized one stochastic objective function f t (θ) over the set of parameters θ, where t is an index of the current mini-batch of samples. in contrast, mtadam optimizes a set of such terms f 1 t (θ), . . . , f i t (θ). while adam's task is to minimize the expected value e t [f t (θ)] w.r.t. the parameters θ, mtadam minimizes a weighted average of the i terms. the weights of these mixtures are all positive, but otherwise unknown. the guiding principle for the determination of the weights at each iteration t is that the moving average of the magnitude of the gradient of each term is equal across terms. this magnitude is evaluated and bal-anced at every layer of the network. in adam, two moments are continuously updated, using a moving average scheme: m t is the first moment of the gradient ∇ θ f t and v t is the second moment. both are vectors of the same size of θ. the moving averages are computed using the mixing coefficients β 1 and β 2 for the two moments. mtadam records such moments for each term @xmath4...i separately. in addition, it uses a mixing coefficient β 3 in order to maintain the moving average of the gradient magnitude per each layer ℓ, which is denoted by n i l,t . adam borrows from the sgd with momentum method (nesterov) and updates the vector of parameters based on the weighted first moment of the gradient. in mtadam, the first moment is computed based on a weighted gradient, in which the parameters of each layer ℓ for every term i are weighted such that their magnitude is normalized by the factor n i ℓ,t . this way, across all layers, and at every time point, the i terms contribute equally to the gradient step. the adam algorithm is depicted in the left side of alg. 1 and mtadam on the right. in line 1, mtadam initializes i pairs of first and second moment vectors. this is similar to adam, except for initializing a pair of moments for each loss term. in line 2, and different from adam, mtadam initializes i first moments for the magnitude of the gradients, per layer. in line 3, both mtadam and adam iterate over the stochastic mini-batches, performing t training steps. in line 4, mtadam iterates over the loss terms. ## experiments we compare the results of mtadam with five baselines: (1-3) adam, rmsprop, and sgd with momentum, applied with suboptimal weightings. (4) hyperband @xcite applied to perform a hyperparameter search for the weighting parameters (lambdas). ( 5 ) adam optimizer applied with the prescribed weighting. we note that baseline (4) benefits from running training multiple times and that baseline (5) employs the weights proposed for each method after a development process that is likely to have included a hyper-parameter search, in which multiple runs were evaluated by the developers of each method. for each optimization method, we employ the default parameters in pytorch. for all adam and mtadam experiments, β 1 = 0.9 and β 2 = 0.999. for mtadam β 3 = 0.9. ## conclusions mtadam is shown to be a widely applicable optimizer, which can dynamically balance multiple loss terms in an effective way. it is a general algorithm, which can find its usage in additional types of tasks that require the optimization of multiple terms, such as domain adaptation and some forms of self-supervised learning. our code is attached as supplementary. mtadam is implemented as a generic pytorch optimizer and applying it is almost as simple as applying adam. score on the validation set. in the maps dataset, a value of 59.61 was selected. fig. 8 compares the convergence of three cycle-gan models: unbalanced mtadam, unbalanced adam, and balanced adam. all models are applied on the horse2zebra dataset @xcite . as can be seen, mtadam exhibits a competitive convergence to the adam experiment applied with balanced weighting, which is much better than the performance of adam on the unbalanced weights. fig. 9 presents a few representative samples from the cyclegan @xcite experiments, employed to transfer horse images to zebras. as can be seen, the unbalanced adam training completely fails to generate zebras images and collapses to the identity mapping. this can be attributed to the domination of cyclic loss in the unbalanced settings, which dictates the convergence, leaveing the other loss terms ineffective. on the other hand, the unbalanced mtadam model was able to successfully generate zebra images, in a quality that is similar or better to the balanced adam model (which uses the prescribed weights). in particular, in the first row, we can see that cyclegan-unbalanced-mtadam was able to outperform the cyclegan-balanced-adam experiment, as the latter fails to generate a zebra image for this particular sample, while mtadam was able to generate a fairly good quality image. fig. 10 depicts a few samples from the same models described above, this time employed to generate horse images from zebra images. as can be seen, the unbalanced adam fails again to generate horse images, collapses to the identity mapping, and this time also introduces visual artifacts in a few images (see the green artifact in the bottom row). in addition, mtadam yields images of the same quality as the balanced adam experiment. fig. 11 presents visual results for the srgan @xcite experiments. the images were taken from set14 @xcite . all models were trained to generate high-resolution images from low-resolution images, with a factor of 4x upscaling. as can be seen, the srgan model that employs unbalanced weights and adam optimizer yields images with visual artifacts, and low fidelity. this can be attributed to the domination of the gan and perceptual loss terms over the pixel-wise term. in contrast, the srgan model that employs our mtadam optimizer with unbalanced weights yields images of the same quality as the srgan that utilizes the prescribed weights @xcite and adam. (a) (b) (c) (d) (e) figure 13: cyclegan results for mapping (a) real maps to fake images in the aerial photos domain for (b) adam with suboptimal weights, (c) mtadam with suboptimal weights, and (d) adam with the prescribed weights. (e) the ground truth.