Title: YOU ONLY TRAIN ONCE: LOSS-CONDITIONAL TRAINING OF DEEP NETWORKS

Abstract: In many machine learning problems, loss functions are weighted sums of several terms. A typical approach to dealing with these is to train multiple separate models with different selections of weights and then either choose the best one according to some criterion or keep multiple models if it is desirable to maintain a diverse set of solutions. This is inefficient both at training and at inference time. We propose a method that allows replacing multiple models trained on one loss function each by a single model trained on a distribution of losses. At test time a model trained this way can be conditioned to generate outputs corresponding to any loss from the training distribution of losses. We demonstrate this approach on three tasks with parameterized losses: β-VAE, learned image compression, and fast style transfer. The code will be released at www.github.

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 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 (Ballé 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 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 ∼ P x,y with x ∈ X ⊂ R d X , y ∈ Y ⊂ R d Y , and a loss function L(•, •) : Y × Y → R, we aim to learn a model F : X → Y, with parameters θ, such that its predictions ŷ = F (x, θ) minimize the expected value of the loss L(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 λ ∈ Λ ⊂ R d λ . The most common case of such a family of losses is a weighted sum of several loss terms:

L(•, •, λ) = i λ i L i (•, •).

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 λ:

EQUATION

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:

EQUATION

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) E x,y∼Px,y L(y, F (x), λ) (2) min G∈C(X×Λ) E λ∼P λ E x,y∼Px,y L(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.

3.1. PRACTICAL DETAILS:

In this work we focus on training deep networks with methods based on stochastic gradient descent.

We therefore approximate the expectations in equation 2 using Monte Carlo estimates: we estimate the expectation over the data distribution by averaging over a training set D = {(x i , y i )} N i=1 , and the expectation over the loss parameters by sampling from the loss parameter distribution:

EQUATION

In our experiments, we sample a new loss parameter for each training data point every time it is encountered during training.

The distribution of the loss parameters to be used for training is typically not given and has to be chosen. Clearly, the support of the training distribution has to include the range of parameter values we are interested in covering, but this still leaves a lot of freedom in the specific choice of the distribution. In our experiments we use the log-uniform distribution, which is commonly used for sampling the parameters of machine learning models (Bergstra & Bengio, 2012) .

In our experiments the model F is a convolutional network. To condition the network on the loss parameters, we use Feature-wise Linear Modulation (FiLM) (Perez et al., 2018) . First, we select the layers of the network to be conditioned (can be all layers or a subset). Next, we condition each of the layers on the given weight parameters λ. Assume the layer outputs a feature map f of dimensions W × H × C, with W and H corresponding to the spatial dimensions and C to the channels. We feed the parameter vector λ to two multi-layer perceptrons (MLPs) M σ and M µ to generate two vectors, σ and µ, of dimensionality C each. We then multiply the feature map channel-wise by σ and add µ to get the transformed feature map f :

EQUATION

4. EXPERIMENTS:

We evaluate the proposed method both quantitatively and qualitatively on three problems with multiterm loss functions: β-VAE, learned image compression, and fast style transfer. All details regarding the architectures and the training, as well as additional results, are provided in the appendix.

4.1. β-VARIATIONAL AUTOENCODERS:

Our first problem is the loss of the β-variational autoencoder (Higgins et al., 2017 )

EQUATION

where q φ (z | x) is the amortized approximate posterior implemented by an encoder network with parameters φ, p θ (x | z) is a stochastic decoder network with parameters θ, and p(z) is the prior distribution, in our case a standard multivariate Gaussian. The loss has a single positive parameter β trading off between the reconstruction quality and the divergence between the approximate posterior and the prior.

We consider two settings: the CIFAR-10 dataset (Krizhevsky, 2009) with Gaussian outputs, and the Shapes3D dataset (Burgess & Kim, 2018) with Bernoulli outputs. In both cases we use convolutional encoder-decoder networks with Gaussian latents. We condition after each downsampling or upsampling operation. Both for YOTO and fixed-weight models, we train a set of networks with varying capacities, by proportionally changing the width of all layers. This is done with the goal In most cases, the proposed method performs close to models trained independently for each loss weight value. This is especially the case for high capacity models.

of measuring how does the capacity of a YOTO model affect its performance relative to a set of fixed-weight models. We train each network 3 times and report means and standard deviations over these three runs.

Quantitative results are shown in Figure 2 . For each dataset, we provide two plots: the frontier of the two loss components plotted against each other ((a), (c)), and the value of the full loss plotted against the value of the parameter β ((b), (d)). The first plot allows to assess how well the proposed model fits the two loss terms, while the second one shows the actual optimization objective -the weighted sum of the losses. To better highlight the differences between the methods, when plotting the full loss, we normalize it by subtracting the loss of one of the models from all other models.

On both datasets, we observe several common tendencies. First, for small-capacity models the proposed method somewhat underperforms relative to the fixed-weight models, especially on CIFAR-10. This is not surprising, since it is difficult for a small model to cover all loss variants. Second, both the YOTO models and the fixed-weight ones improve substantially with increasing model capacity, so a set of fixed-weight models of a certain width is consistently matched or outperformed by a twice wider YOTO model (this is additionally highlighted in Figure 4 that plots the average full loss versus the width of the network). Third, the larger the network capacity, the closer the performance of fixed-models and the YOTO model. In particular, on both datasets for the widest networks the two frontiers become very close. This is quite intuitive and matches the theoretical result that in the limit of infinite model capacity the proposed method should be as powerful as training separate fixed-weight models.

Qualitative results are shown in Figure 3 Figure 4 : Loss averaged over all loss weight values as a function of the model capacity. We compare the proposed method to a set of fixed-weight models on two datasets: CIFAR-10 (a) and Shapes3D (b). We vary the capacity of the network by multiplying the width of the network by a factor. The performance of the proposed method overall matches that of fixed-based models with a roughly 1.5x wider model. Moreover, the performance of the proposed method gets closer to that of the fixed-weight models when the capacity of the model is increased.

independently trained model. An interesting property of the YOTO model is that sampling with the same noise vector for different beta values generates similar images at different quality levels.

Additional results, including a comparison to two simple baselines (a single fixed-weight model and a baseline based on interpolation), an ablation study of the proposed method, and qualitative results on CIFAR-10, are shown in the appendix.

4.2. LEARNED IMAGE COMPRESSION:

We build on the learned compression model of Ballé et al. (2018) . The model is similar to a variational autoencoder, but with a few modifications that allow for the quantization of the latent representation and its use as a compressed image encoding. Moreover, it includes a "hyperprior" learned by a second variational autoencoder on top of the latents of the first one. The loss function has a single parameter -the "rate-distortion weight", trading off compression rate versus reconstruction quality. We refer the reader to (Ballé et al., 2018) for details. The architecture of the model includes four networks: a convolutional encoder and hyper-encoder, as well as a transposed convolutional decoder and a hyper-decoder. The networks are relatively shallow, with three convolutional layers each. We apply conditioning to all layers of the model. Same as for VAEs, we train three models for each setting to estimate the stability of training. Further details are provided in the appendix.

We evaluate the compression models on two datasets: Kodak (Kodak, 1993) and Tecnick (Asuni & Giachetti, 2014) . The quantitative results on Tecnick are shown in Figure 5 , while the results on Kodak are deferred to the appendix. We show rate-distortion plots with standard PSNR on pixel values as a measure of compression quality and bits per pixel as a measure of compression rate. We plot three models per method, but since the training is stable, there difference between them is almost not visible. The gap between separate fixed-weight models ("Fixed weight") and a single joint model of the same size ("YOTO") is fairly small, but consistent across all compression rates and is especially pronounced in the higher quality regime. However, making the model larger ("YOTO width x1.33") recovers a large fraction of the performance, in particular the wider model matches the performance of a smaller fixed-weight model in the high quality regime. However, in the highcompression regime there is still a small gap from the fixed-weight models.

We hypothesize that the reason for relatively worse performance of YOTO on the compression task compared to β-VAEs is that the networks used by Ballé et al. (2018) are relatively shallow. They therefore lack the representational power to efficiently make use of the conditioning input. We expect that one could devise combinations of architectures and conditioning methods for the compression task which would lead to a smaller gap between the fixed-weight and YOTO models; however, this additional tuning is beyond the scope of this paper.

4.3. STYLE TRANSFER:

The final problem we consider is style transfer (Gatys et al., 2016) -given a content image x c and a style image x s , the goal is to synthesize a stylized image y that contains the same content as x c , but in a style that resembles that of x s . This is typically cast as an optimization problem -we want to find an image y, that when passed through a pre-trained neural network has activations similar to those of x c , but whose aggregate feature statistics resemble those of x s . There is a lot of freedom in deciding which layers of the network to use and how to weigh them. These parameters are not easy to tune, since we are optimizing for visual appearance, which cannot be easily quantified. Hence, a single model capable of representing all loss variants would be of great benefit. These experiments are similar to those presented by Babaeizadeh & Ghiasi (2018) , but with more focus on a quantitative evaluation, same as in other applications reported above.

Published as a conference paper at ICLR 2020 We build on the work of Ghiasi et al. (2017) , who propose training a single deep network that can stylize images in arbitrary styles. To this end, the content image is fed to the stylization network itself, while the style image is fed to a different "style prediction" network that extracts a feature embedding. This embedding is then used to modulate the stylization network via conditional instance normalization. We refer the reader to (Ghiasi et al., 2017) for a detailed explanation.

Following Ghiasi et al. (2017) , we extract the activations from the VGG network (Simonyan & Zisserman, 2015) , and define a total of six losses: four style losses L s,1 (y, x s ), . . . , L s,4 (y, x s ), one content loss L c (y, x c ), and a total variation loss L tv (y) encouraging smoothness between the neighbouring pixels of y. The final loss is then a linear weighting of these, which we will denote as

L(y) = n i=1 λ s,i L s,i (y, x s ) + λ c L c (x, c) + λ tv L tv (y).

We train a single YOTO model over five of these parameters (while keeping the TV loss fixed to 10 4 ), with log-uniform distributions. To this end, in addition to the style embeddings we condition the stylization network on the vector of loss parameters λ. We sample the content images form ImageNet (Deng et al., 2009) and use 14 pointillism paintings as the style images. Further details are provided in the appendix.

We start by reporting quantitative results. As the problem has 5 loss parameters, it is not possible to visualize the frontier with respect to all losses. Instead, we take two visualization approaches. First, we show the full loss achieved by different methods on a set of loss parameter values selected as follows: we start with the default parameter settings of Ghiasi et al. (2017) λ s,1 = λ s,2 = λ s,3 = 10 -3 , λ s,4 = 10 -2 , λ c = 2, and then either increase or decrease one of these parameter by a factor of 10 while keeping the rest fixed. This adds up to a total of 10 parameter values: 2 variants per each of the 5 parameters. To better understand the variability of the learned models due to the noisy training process we further trained each model four times. We provide the results in Figure 6a . As a baseline, for each parameter configuration we report the full loss value for the model trained with the default parameters ("Single model baseline"). The YOTO model trained on all parameters closely matches the performance of fixed-weight models trained for each of the parameter configurations.

Next, we focus on just one aspect of the loss surface: the trade-off between the content loss and the sum of style losses. In Figure 6b we plot the frontier of these two values corresponding to different methods. Here we add a YOTO "specialist" model to the comparison, which has been trained with only 1 parameter, the content loss λ c . Moreover, we report a baseline method that simply interpolates the images produced by two models trained with the extreme values of the content weight. Both YOTO models are close to the fixed-wight models, with the frontier of the "specialist" model closely tracing the fixed models. The interpolation baseline performs very poorly.

Finally, we show select qualitative results in Figure 7 , where we compare fixed weight models to a YOTO model trained on all five parameters. We vary the content weight λ c and the third style weight λ s,3 , which contribute most to the full loss. We can see that YOTO captures the effect of changing each of these weights and obtains results similar to those of the separately trained models.

5. CONCLUSION:

We have presented an approach to training a single deep model to fit a parametric family of loss functions. This way, instead of training multiple models corresponding to different loss variants, we can train a single model that provides solutions for all loss variants. We demonstrated the successful application of the method on three applications: β-VAE, learned image compression, and fast style transfer. These initial results are promising, but several challenges remain. First, it would be interesting to devise provably effective strategies of selecting the loss parameter distribution, especially for losses with multiple parameters. Second, further experiments with various architectures and conditioning methods can lead to an even smaller discrepancy between the separate models and our YOTO model. Third, the method could be used for many other applications, within or beyond the domain of image generation. Fourth, one might use YOTO-trained models to initialize usual fixed-loss training or, vice-versa, use a pre-trained fixed-loss model to initialize the training of YOTO. Finally, we believe that this work can also open a new family of theoretical questions centered around the analysis of problems that have to optimize a continuum of loss functions. We see all these directions as exciting avenues for future research. with a single layer of dimension 256 and ReLU activations, on which two affine maps are applied to compute the scale and shift.

B ADDITIONAL RESULTS

B.1 BETA-VAE:

Figure 8 shows qualitative results for β-VAE on CIFAR-10. Same as for Shapes3D, both reconstructions and samples from the YOTO model are qualitatively very similar to those generated with models trained separately for each weight.

We compare the proposed method to two baselines. The first one is a naive baseline -training a VAE with a fixed β, and computing the loss value of this fixed model for all other β values ("Single model baseline"). We select the fixed β so that it minimizes the average validation loss over all β values. The second baseline interpolates between two models trained with fixed weights. Namely, we train two beta-VAEs for the extreme values of beta, and interpolate both the latents and the reconstructions produced by these models. We then measure the KL and reconstruction losses for these interpolated values. Note that this interpolation baseline is not a VAE any more: the latents and the reconstrutcions are averaged independently, so the interpolated reconstruction cannot be generated from the interpolated latent. Therefore, it is not clear how would one sample from such a model, and we do not report these results in the main paper to avoid confusion.

The results are shown in Figure 9 . Surprisingly, on Shapes3D the baseline performs very close to the fixed-model and YOTO frontiers, and sometimes even outperforms them. In contrast, on CIFAR-10 the baseline is much worse than the other methods. We believe the reason is that while averaging the logits for Bernoulli outputs on Shapes3D is quite meaningful, averaging the standard deviations of Gaussian outputs on CIFAR-10 strongly affects Gaussian log-likelihood.

Timing. We timed the training of a fixed-weight β-VAE and a YOTO model. To this end, we trained both models on a single CPU core over 200 mini-batch iterations. We found the YOTO model to be only 8% slower than its fixed-weight counterpart: 4.64 vs 5.04 training steps per second with the "width x2" network architecture. 

B.2 ABLATION STUDY:

We now ablate different design choices made in our implementation of YOTO to check their importance. We perform this analysis on the CIFAR-10 dataset, with the "width x4" architecture. The design decisions we study are: the weight distribution used during training (log-uniform or uniform), the method of sampling the loss weights (per sample or per mini-batch), the method of conditioning the network on the weights (FiLM, concatenating the conditioning vectors with the feature maps, or feeding the conditioning vectors to the network as an additional input), pre-processing of the conditioning inputs (taking their logarithm or not).

The results are presented in Figure 11 . Clearly, the choice of the parameter distribution during training is the most important choice, and log-uniform distribution performs much better than uniform.

Other choices are of less importance, but more advanced conditioning schemes have an advantage over simply feeding the loss parameters as one of the network inputs.

B.3 STYLE TRANSFER:

Figure 12 shows additional qualitative results on style transfer for different content and style images.

C THEORY:

We start by proving Proposition 1 from the main paper and then comment on the cases when the conditions of the proposition are satisfied.

Proof of Proposition 1. Denote the values minimized in optimization problems (1) and ( 2) by V 1 (•) and V 2 (•) respectively. From the definition of F * it follows that V 1 (G(•, λ)) V 1 (F * (•, λ)) for any λ ∈ Λ and G ∈ C(X × Λ). Therefore, V 2 (G) V 2 (F * ), so F * minimizes (2). Thus, for any G * that minimizes (2):

EQUATION

Moreover, V 1 (G * (•, λ)) -V 1 (F * (•, λ)) 0 for all λ ∈ Λ. Expectation of a non-negative function is 0 only if it equals 0 almost surely, which means that V 1 (G * (•, λ)) = V 1 (F * (•, λ)) a.s., q.e.d.

The most restrictive condition of the proposition is the continuity of the function F * (•, •) as a function of λ. Indeed, generally arg max is a multi-valued mapping that can be shown to be hemicontinuious using Berge's Maximum theorem (Berge, 1963) under weak technical conditions. However, there does not necessarily exist a continuous selection of this multi-valued mapping. Its existence can be proven under very restrictive conditions, such as the uniqueness of the minimizer, which clearly do not hold for optimization problems encountered in deep learning. While empirically the existence of a continuous F * seems natural, it remains to be seen if theoretical guarantees can be provided for realistically complex scenarios.

Another direction for extending Proposition 1 is explicitly incorporating function approximation.

With function approximation, it would likely not be possible to prove that minimizers of the two problems coincide almost surely, but rather that they can be made arbitrarily close with a probability arbitrarily close to 1. Overall, we believe our work raises a range of interesting theoretical questions to be addressed in future research.

:



Figures and Tables:

Figure 1: Illustration of the proposed method. At training time (left), for each training data point loss parameters λ are sampled from a distribution P λ , and the model is conditioned on these. Then, at test time (right), the model can be conditioned on any desired loss parameters to achieve behavior corresponding to the loss function with these parameters.

x0.5 YOTO, width x1.0 YOTO, width x4.0 Fixed weights, width x0.5 Fixed weights, width x1.0 Fixed weights, width x4.

Figure 2: Quantitative β-VAE results on CIFAR-10 (a,b), and Shapes3D (c,d) for models of varying capacity (width). In most cases, the proposed method performs close to models trained independently for each loss weight value. This is especially the case for high capacity models.

. Both for reconstruction and sampling, for each value of the weight β that we show YOTO produces results closely resembling those of the corresponding Image reconstructions using the trained β-VAE models. Samples from the trained β-VAE models.

Figure 3: Qualitative β-VAE results on the Shapes3D dataset. For each loss weight β, the reconstruction and sampling results of YOTO are very similar to those of the separately trained models.

Figure 5: Quantitative compression results on the Tecnick and Kodak datasets. A basic YOTO model under-performs compared to a set of fixed-weight models, but a wider network nearly matches the performance of the fixed-weight models, especially in the high-quality regime.

= 0.2 c =20 s, 1 = 10 2 s, 1 = 10 4 s, 2 = 10 2 s, 2 = 10 4 s, 3 = 10 2 s, 3 = 10 4 s, 4 = 10 1 s, 4Performance of YOTO when we unilaterally vary each weight. We show means and standard deviations over four runs.

Sum of all style losses plotted against the content loss.

Figure6: Quantitative results on style transfer. In both cases YOTO performance approaches or matches that of separate models trained per loss parameter vector.

Figure 8: Qualitative VAE results on the CIFAR-10 dataset. Both for reconstruction and sampling, the YOTO results are qualitatively very similar to those of separate models trained for each value of the loss weight.

Figure 9: Quantitative comparison of the proposed method with baselines for beta-VAE on CIFAR-10 (a,b), and Shapes3D (c,d). The proposed method outperforms the "single model" baseline by a large model. The "interpolation" baseline is reported for completeness, but it is not a fair baseline, since it does not correspond to a VAE.

Figure 10: Quantitative results of compression with the MS-SSIM metric on Tecnick (a) and Kodak (b).

Figure 11: Ablation study of the CIFAR-10 β-VAE. The distribution of the loss weight used at training time has the largest impact on the performance. Moreover, conditioning the model by providing the loss weights as inputs to the model underperforms compared to more advanced conditioning methods. Other variants all perform close to the full method.

Figure 12: Additional qualitative comparison of image stylization models on images from the validation set of ImageNet. The first row shows the results of YOTO trained on all parameters, the second row -of models trained independently per configuration. In both cases parameter value increases from left ot right. YOTO results are very similar to those of independently trained models.

Peer review (iclr2020):
Summary ------------- The authors propose a methodology to train a single Deep Neural Network (DNN) that minimize a parametrized (weighted) sum of different losses. Since the model is itself conditioned by the weights, it allows to train a single model for all weights values, instead of retraining when the weights change. Experiments suggest that this methodology does not degrade the performances to much w.r.t. retraining on every weight changes. The proposed conditioning of the layers is done via a reparametrization (FiLM, Perez et al. 2018) of the weights with a scale and a bias where and are MLP. This allows to condition the layer on , while keeping the number of parameters low. Novelty ---------- The idea of integrating a family of loss functions with model conditioning has also been proposed by Brault et al. [1], in the context of multi-task kernel learning. Hence I believe this work should be acknowledged. As the product of kernels is the tensor product of their feature maps, it would suggest to condition the network*s layers by taking the tensor product of the weights with respect to an MLP on . This could be applied on each layers or simply on the last Fully Connected layer. Note however that it would drastically increase the number of parameters and hence not be a viable solution (or maybe with some channel pooling?). References: [1] Infinite Task Learning in RKHSs; Brault, Romain and Lambert, Alex and Szabo, Zoltan and Sangnier, Maxime and d*Alche-Buc, Florence; Proceedings of Machine Learning Research; 2019. Quality ---------- The paper is self content and well written. Experiments are well detailed and seems to be reproducible. I would be a great addition to release the code in a public repository (with a link in the paper or appendices) if the paper is accepted. I would also suggest the following experiments: * An experiment showing the time penalty induce by training the loss conditional model. The authors claims that training multiple separate models is inefficient compared to their proposed method. While it seems obvious, it deserve an experiment as one of the claim. * The authors propose to sample one per SGD iteration. However it ma be useful to sample more of them. Especially when the set of is large (high dimensional) * Possibly use a pre-trained model and only tune the , MLPs Overall my decision is weak accept, the paper lacks of novelty and the experiments could be more extensive.