Balance between 2 MSE losses

username: ays
cooked: <p>I have an MLP and trying to predict 2 values from the MLP (so i’m doing a multi task learning). these values have two losses of different scales (one is in the tens and the other is in 100,000s). so I cant simply add the losses together before back propagating. I’m currently trying to learn a weight so I can use it to balance the losses but I dont know how do go about that. Below is a brief summary of the code</p>
<pre><code class="lang-auto">

class MLP(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(input_dim, 100)
        self.out1 = nn.Linear(100, output_dim - 1)
        self.out2 = nn.Linear(100, 1)
        self.relu = nn.ReLU()

        #create a learnable parameter that is used to weight the two losses
        self.weight = nn.Parameter(torch.ones(1, requires_grad=True))

        

    def forward(self, x):
        out1 = self.out1(self.relu(self.fc1(x)))
        out2 = self.out2(self.relu(self.fc1(x)))
        return out1, out2, self.weight

#in train method (assume we have the target values and data)
predicted_one, predicted_two, weight = model(data)

loss_one = nn.Functional.mse_loss(predicted_one, target_one)
loss_two = nn.Functional.mse_loss(predicted_two, target_two)

total_loss = loss_one + weight * loss_two


</code></pre>
<p>I dont think this is how to go about it because the the weight doesn’t learn. Please Kindly assist.</p>

username: KFrank
cooked: <p>Hi Ays!</p>
<aside class="quote no-group" data-username="ays" data-post="1" data-topic="185315" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ays/48/56534_2.png" class="avatar"> ays:</div>
<blockquote>
<p>I’m currently trying to learn a weight so I can use it to balance the losses but I dont know how do go about that.</p>
</blockquote>
</aside>
<p>If it were me, I would just treat <code>weight</code> as a hyperparameter to be tuned<br>
by hand.</p>
<p>There is, however, a proposed <code>MultiTaskLoss</code> (that I have never used) that<br>
seems to be a theoretically-sound approach to automatically tuning the<br>
relative weight between two losses.</p>
<p>Here is an implementation:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/ywatanabe1989/custom_losses_pytorch/blob/master/multi_task_loss.py">
  <header class="source">

      <a href="https://github.com/ywatanabe1989/custom_losses_pytorch/blob/master/multi_task_loss.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/ywatanabe1989/custom_losses_pytorch/blob/master/multi_task_loss.py" target="_blank" rel="noopener nofollow ugc">ywatanabe1989/custom_losses_pytorch/blob/master/multi_task_loss.py</a></h4>


      <pre><code class="lang-py">import torch

class MultiTaskLoss(torch.nn.Module):
  '''https://arxiv.org/abs/1705.07115'''
  def __init__(self, is_regression, reduction='none'):
    super(MultiTaskLoss, self).__init__()
    self.is_regression = is_regression
    self.n_tasks = len(is_regression)
    self.log_vars = torch.nn.Parameter(torch.zeros(self.n_tasks))
    self.reduction = reduction

  def forward(self, losses):
    dtype = losses.dtype
    device = losses.device
    stds = (torch.exp(self.log_vars)**(1/2)).to(device).to(dtype)
    self.is_regression = self.is_regression.to(device).to(dtype)
    coeffs = 1 / ( (self.is_regression+1)*(stds**2) )
    multi_task_losses = coeffs*losses + torch.log(stds)

    if self.reduction == 'sum':
</code></pre>



  This file has been truncated. <a href="https://github.com/ywatanabe1989/custom_losses_pytorch/blob/master/multi_task_loss.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>And here is the paper about the scheme that is referenced in the code:</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://arxiv.org/abs/1705.07115">
  <header class="source">
      <img src="https://static.arxiv.org/static/browse/0.3.4/images/icons/favicon-32x32.png" class="site-icon" width="" height="">

      <a href="https://arxiv.org/abs/1705.07115" target="_blank" rel="noopener nofollow ugc">arXiv.org</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/402;"><img src="https://static.arxiv.org/static/browse/0.3.4/images/arxiv-logo-fb.png" class="thumbnail" width="" height=""></div>

<h3><a href="https://arxiv.org/abs/1705.07115" target="_blank" rel="noopener nofollow ugc">Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and...</a></h3>

  <p>Numerous deep learning applications benefit from multi-task learning with multiple regression and classification objectives. In this paper we make the observation that the performance of such systems is strongly dependent on the relative weighting...</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Best.</p>
<p>K. Frank</p>

username: ays
cooked: <aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="185315">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p><code>https://arxiv.org/abs/1705.07115</code></p>
</blockquote>
</aside>
<p>Hi Frank. Thanks a lot for your response. I tried using the Multi-task Learning approach but im getting a loss as in the graph below.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/3/63175de3898dd71b44196545e147c3051d9fb207.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/63175de3898dd71b44196545e147c3051d9fb207" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/3/63175de3898dd71b44196545e147c3051d9fb207_2_690x240.jpeg" alt="image" data-base62-sha1="e8BkCyOTBpYLP0K9kO5eSmE46jR" width="690" height="240" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/3/63175de3898dd71b44196545e147c3051d9fb207_2_690x240.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/6/3/63175de3898dd71b44196545e147c3051d9fb207_2_1035x360.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/3/63175de3898dd71b44196545e147c3051d9fb207_2_1380x480.jpeg 2x" data-dominant-color="FAFBFB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1878×654 90.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: InnovArul
cooked: <p>Silly question:</p>
<p>Did you try normalizing the target values?</p>
<p>(say with -1 to 1, maybe <code>tanh()</code> final layer + MSE loss / L1 loss, and during inference, rescale the values accordingly)</p>

username: ays
cooked: <p>Hi Arul,<br>
Thanks for the suggestion. So my first loss has target values ranging between 0 and 1, my second loss has target values ranging between 0 and 4202496. So do you mean I should normalise the second target variable to values between 0 and 1? if yes, does that mean I wouldn’t need a weight in the loss because it would be of the same scale?</p>
<p>NB: sorry I’m a bit new to all these so its still trying to get a hang of it.</p>

username: InnovArul
cooked: <p>No worries. We all are learning:)</p>
<p>Yes, I meant to keep both the target variables on the same scale. (either [0,1] or [-1,1]).<br>
The method proposed by Kendall et. al’s paper is from a different perspective to learn to weigh the losses based on the uncertainties (if the variance of a particular loss is very high, its contribution to the gradients will be low and vice versa). I would still keep it (or you can experiment without it too).</p>
<p>Please let me know how it goes.</p>

username: J_Johnson
cooked: <p>Here is a simple function you can apply to normalize 2 or more rewards in order to balance various objectives:</p>
<pre><code class="lang-auto">import torch

def norm_rewards(raw_rewards, alphas, freq_weights, max_clip, scale = 1.):
    raw_rewards = torch.clip(raw_rewards, max = max_clip)
    return torch.sum(alphas*freq_weights*raw_rewards/max_clip)*scale

raw_rewards = torch.tensor([0.5, 2_000_000.]) # these are the dynamic raw rewards from each objective
alphas = torch.tensor([2., 1.]) #this should be priority of each reward; if all equal, can use torch.ones_like(raw_rewards)
freq_weights = torch.tensor([4., 1.]) #this should be the inverse of the frequency the reward occurs, i.e. 25% of the time would be 4. and 100% of the time would be 1.
max_clip = torch.tensor([1., 4_202_496.]) #this should be the max expected value of each reward and can also be used to clamp reward values for stability
scale = 1. #optional scale value which scales the overall reward

print(norm_rewards(raw_rewards, alphas, freq_weights, max_clip, scale))
</code></pre>
<p>Just to clarify the difference between the freq_weights and the alphas. The alphas are what you do to set priority of a given objective. I’ll give an example from a thought experiment recently in an alignment discussion.</p>
<p>Suppose you’re training a drone to strike a target. But once in every four missions, your operator sends an abort command, which hard codes the drone to stop the attack. In the simulation, the drone learns that by taking out the operator first, it can obtain a higher reward.</p>
<p>But now we also provide a second reward for the drone model for every time it follows the extra instructions from the operator, when given. So you might, in this case, set the alpha for that reward to some value higher than 1, to ensure the model prioritizes following instructions over hitting the target.</p>
<p>However, because that only occurs once every four missions, we also need to factor in this reduced frequency. If we do not, then the model may still simply ignore the command because the cumulative reward for hitting the target is greater. And that is what the freq_weights do. They balance the cumulative reward across instances.</p>
<p>Let me know if you have any questions and if it works out. Cheers</p>

username: ays
cooked: <p>Hi <a class="mention" href="/u/innovarul">@InnovArul</a> and <a class="mention" href="/u/j_johnson">@J_Johnson</a>,</p>
<p>Thanks both for your response. I scaled my second target value between 0 and 1 using minmaxscaler and it works okay. I ended up not using a learned parameter to weigh my loss cos it seemed as if the parameter was not learning - and everything was okay. Currently I’m facing an issue with my loss, it seems as if my model is not learning, I have a loss graph as shown below. please kindly assist.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/1/c136880543cda4ae5ad2ff645d5414a0824754ce.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/c136880543cda4ae5ad2ff645d5414a0824754ce" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/1/c136880543cda4ae5ad2ff645d5414a0824754ce_2_690x245.jpeg" alt="image" data-base62-sha1="rzeZ2VDhb0L33oFDFE7dgO8SbPU" width="690" height="245" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/1/c136880543cda4ae5ad2ff645d5414a0824754ce_2_690x245.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/c/1/c136880543cda4ae5ad2ff645d5414a0824754ce_2_1035x367.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/c/1/c136880543cda4ae5ad2ff645d5414a0824754ce_2_1380x490.jpeg 2x" data-dominant-color="FCFCFD"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1906×678 67.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: J_Johnson
cooked: <p>It looks to me like it’s learning just fine. Loss starts out high and gets lower. What is the anticipated behavior and how is the accuracy? And how big is your dataset?</p>
<p>Things you might try to bring the validation loss lower are:</p>
<ol>
<li>Adjusting hyperparameters, such as learning rate, dropout, etc.;</li>
<li>Trying different model architectures, including skip connections;</li>
<li>Increase the number of hidden dims;</li>
<li>Normalizing all input values to be between 0 and 1;</li>
</ol>
<p>Regarding the learnable weight, this can be applied as the alphas in the previous function provided. But set it as <code>torch.ones(2, requires_grad=True)</code> for 2 rewards/losses.</p>
<p>By the way, you might want to separate your losses for printing metrics, as that may help as a sanity check to ensure both losses are improving.</p>

username: ays
cooked: <p>The training loss flattens out after 1 epoch.</p>