You Only Train Once: Loss-Conditional Training of Deep Networks

1 INTRODUCTION . When designing a model for a machine learning problem at hand , a practitioner often thinks of multiple qualities that the model should have . For example , in addition to requiring a good prediction performance , it may be desirable to have a classifier that is simple , so that it would better generalize to unseen data . Similarly , when training a lossy image compression model , one optimizes both the size of the compressed images and their quality . In such scenarios , multiple loss functions have to be minimized simultaneously , each of them modelling a different facet of the considered problem . Typically , it is not possible to simultaneously optimize all these losses — either due to limited model capacity , or because some of the losses may be fundamentally in conflict ( for instance , image quality and the compression ratio in image compression ) . One thus has to decide how to balance the losses during optimization . Such multi-objective problems are most commonly scalarized by linearly combining the losses , with weights defining the trade-off between the loss terms . However , the exact values of the chosen weights can strongly affect the performance of the model and their tuning can be cumbersome . It might be also unpredictable how the weights affect the final values of the individual losses ; moreover , the actual goal may be to optimize a down-stream metric that is not being explicitly minimized during training ( for instance , classification accuracy ) . Furthermore , in many tasks there may not be a single optimal model , but it is instead important to make different trade-offs at test time . For instance , in image compression , to achieve a fine-grained coverage of rate-distortion trade-offs , one would first have to train at least ten models with different parameter settings , and then at inference time store them all in memory and select the appropriate one depending on the required compression level . This is inefficient both at training and at inference time . However , intuitively , the models trained for different loss variants are related and could share a large fraction of computation . Is it possible to improve the efficiency by making use of this model redundancy ? In this paper we propose a simple and broadly applicable approach that efficiently deals with multiterm loss functions and , more generally , arbitrarily parameterized loss functions . Instead of training one model for each parameter combination of the loss , we train one model that covers the whole space of different loss weightings . We achieve this by ( i ) training the model on a distribution of losses instead of a single loss , and ( ii ) conditioning the model outputs on the parameters of the The code will be released at www.github.com/google-research/google-research/yoto . loss function . This way , at inference time the conditioning vector can be varied , allowing us to traverse the space of models corresponding to different loss functions . We dub the method You Only Train Once ( YOTO ) . The conceptual simplicity of the approach makes it easily applicable to many problem domains , with only minimal changes to existing codebases . We experimentally showcase the efficacy of the proposed optimization scheme on image synthesis problems , where multi-loss problems are especially widespread . We first show that ( β- ) variational autoencoders ( Kingma & Welling , 2014 ; Higgins et al. , 2017 ) can be trained for a large range of β parameters at the cost of training a single fixed-β model . Then , we train a single deep image compression model ( Ballé et al. , 2018 ) that can adjust the rate-distortion trade-off on the fly at test time . Finally , we apply YOTO to style transfer ( Gatys et al. , 2016 ; Ghiasi et al. , 2017 ) , where it is provides a clear benefit given that there is no clear single optimum and that the effect of the various losses on the stylized image is difficult to characterize . 2 RELATED WORK . Losses with multiple terms are most commonly associated with multi-task learning ( Caruana , 1997 ; Kokkinos , 2017 ; Zamir et al. , 2018 ) , where the goal is to train a single model that simultaneously solves several learning tasks . We refer the reader to Ruder ( 2017 ) for a recent review . Most related to our approach are methods that condition the network architecture on the task at hand ( Rebuffi et al. , 2017 ; Mallya et al. , 2018 ) , which an be seen as a special case of our method with “ one hot ” loss weight vectors . While ideas from our work could also be applied to the multi-task learning scenario , here we choose to focus on problems where it is desirable to learn a diverse set of models , rather than a single well-performing model . Recently , Brault et al . ( 2019 ) studied simultaneous optimization of multiple loss functions in the context of multi-task kernel learning . Our work is similar in spirit , but differs drastically both technically and in the application domains . Conditional normalization schemes have recently been used for a range of supervised learning tasks , including Conditional Batch Normalization ( CBN ) for visual question answering ( de Vries et al. , 2017 ) , Conditional Instance Normalization and Adaptive Instance Normalization ( AdaIN ) for fast style transfer ( Dumoulin et al. , 2017 ; Huang & Belongie , 2017 ; Ghiasi et al. , 2017 ) , Feature-wise Linear Modulation ( FiLM ) for visual reasoning ( Perez et al. , 2018 ) , and a variety of conditioning schemes for generative models ( Karras et al. , 2019 ; Brock et al. , 2019 ) . A comprehensive overview of conditioning methods is provided by Dumoulin et al . ( 2018 ) . All these methods condition a convolutional network on some side information by linearly transforming the activations with coefficients computed from the side information . The side information typically informs the model about the prediction target : for instance , it can be the text of a question in the question answering setting or the style image in style transfer . The only exception we are aware of is the recent work of Babaeizadeh & Ghiasi ( 2018 ) , which conditions a style transfer network on the parameters of the loss function , same as in our method . The approach is technically very similar , but we formulate it as a generally applicable technique and evaluate on several tasks both qualitatively and quantitatively . Conditional models are also often used in sensorimotor control to train a single model capable of addressing a family of tasks , such as throwing darts at different targets ( da Silva et al. , 2012 ) , block stacking and ball hitting ( Deisenroth et al. , 2014 ) , navigating towards different locations ( Schaul et al. , 2015 ) , moving objects with a robotic arm ( Andrychowicz et al. , 2017 ) , driving in the desired direction ( Codevilla et al. , 2018 ) or trading off the reward components ( Dosovitskiy & Koltun , 2017 ) . A single model , conditioned on the task parameters , is trained to simultaneously maximize the rewards for a distribution of tasks . While technically similar to our setting , the conceptual difference is that between the tasks the training targets are changing ( either rewards or expert actions used as labels for imitation learning ) , rather than the loss function itself , as in our case . Another recent line of related work explores the use hypernetworks — networks that generate the weights of other networks — for hyperparameter optimization . Lorraine & Duvenaud ( 2018 ) propose to parametrize the weights of a network by a hypernetwork , which is conditioned on the values of the hyperparameters . The model can then be trained on a distribution of hyperparameters , and the hyperparameters can be optimized via gradient descent on the validation data . MacKay et al . ( 2019 ) improve the scalability of the method by proposing efficient approximations to hypernetworks . These methods are technically similar to our approach , but they focus on hyperparameter optimization , while we are more interested in obtaining the full set of solutions corresponding to the different loss coefficients . Recently , Shoshan et al . ( 2019 ) and Wang et al . ( 2019 ) addressed the problem of interpolating between predictions of two or more image generation models . Both methods start by training several separate models and then propose different ways of interpolating between them : either by using specialized adaptor blocks in the networks , or by directly interpolating the network weights . We , in contrast , train only a single model by directly optimizing the performance over a distribution of losses . This is conceptually simpler , can be easily integrated in existing models , and requires training a single model instead of several . 3 METHOD . We consider the following learning scenario : given a training distribution of pairs x , y ∼ Px , y with x ∈ X ⊂ RdX , y ∈ Y ⊂ RdY , and a loss function L ( · , · ) : Y × Y → R , we aim to learn a model F : X → Y , with parameters θ , such that its predictions ŷ = F ( x , θ ) minimize the expected value of the loss L ( y , ŷ ) over the dataset . This formulation is very broad and covers a variety of machine learning problems , including ( self- ) supervised learning learning , and some types of generative models . We focus on this setting for a more streamlined presentation , although the proposed method can be applied to models that go beyond this formulation . Instead of a single fixed loss function L ( · , · ) , assume that we are interested in a family of losses L ( · , · , λ ) , parameterized by a vector λ ∈ Λ ⊂ Rdλ . The most common case of such a family of losses is a weighted sum of several loss terms : L ( · , · , λ ) = ∑ i λi Li ( · , · ) . However , other parameters , such as , for instance , the power p in ` p losses , might also be included in the vector λ . In practice , one then typically minimizes the loss independently for each choice of the parameter vector λ : θ∗λ = arg min θ Ex , y∼Px , yL ( y , F ( x , θ ) , λ ) . ( 1 ) Instead of fixing λ , we propose to solve an optimization problem where the parameters λ are sampled from a distribution Pλ . Hence , during training time the model observes many losses , and can learn to utilize the relationships between them to optimize all of them simultaneously . The distribution Pλ can be seen as a prior over the losses , and we discuss its choice in the next subsection . At inference time , the joint model can be conditioned on an arbitrary desired parameter value λ , yielding the corresponding predictions . We thus solve the following optimization problem : θ∗ = arg min θ Eλ∼PλEx , y∼Px , y L ( y , Fc ( x , θ , λ ) , λ ) . ( 2 ) In the limit of infinite model capacity , this new proposed optimization problem is equivalent to the original independent optimization scheme , in the sense that the minimal values for both problems coincide . Further discussion , as well as the proof , are provided in the appendix . Proposition 1 . Consider two optimization problems : ( 1 ) min F∈C ( X ) Ex , y∼Px , yL ( y , F ( x ) , λ ) ( 2 ) min G∈C ( X×Λ ) Eλ∼PλEx , y∼Px , yL ( y , G ( x , λ ) , λ ) , where C ( S ) denotes the space of continuous functions on S. Assume there exists a continuous function F ∗ ( x , λ ) such that for every λ ∈ Λ the function F ∗ ( · , λ ) solves the problem ( 1 ) . Then , if G∗ is a solution of the problem ( 2 ) , G∗ ( · , λ ) also minimizes the problem ( 1 ) almost surely w.r.t . Pλ . This theoretical result shows that the method is fundamentally as powerful as the standard perparameter training , while only requiring a single model . However , in practice the infinite capacity assumption does not hold , so implementation details become important , which we describe next .

Summary: The problem tackled by the paper is related to the sensitivity of deep learning models to hyperparameters. While most of the hyperparameters correspond to the choice of architecture and optimization scheme, some influence the loss function. This paper assumes that the loss function consists of multiple weighted terms and proposes the method of finding the optimal neural network for each set of parameters by only training it once.