Multi-task learning: weight selection for combining loss functions

almog: <p>Hi,</p>
<p>I have two tasks in my model- regression and classification (2 heads). I’m using both MSE and CE loss respectively.<br>
As for now, I am combining the losses linearly:     combined_loss = mse_loss+ce_loss,<br>
and then doing: combined_loss.backward()</p>
<p>The main problem is that the scaling of the 2 losses is really different, and the MSE’a range is bigger than the CE’s range. The MSE can be between 60-140 (depends on the dataset) while the CE is between 0.2-0.6. Therefore the CE doesn’t really effect on the combined loss.</p>
<p>How can I scale the 2 losses in the most automated way? without doing grid search on different hyper parameters for weights?</p>
<p>Thanks,<br>
Almog</p>

KFrank: <p>Hi Almog!</p>
<aside class="quote no-group" data-username="almog" data-post="1" data-topic="152332" data-full="true">
<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/a/9fc29f/48.png" class="avatar"> almog:</div>
<blockquote>
<p>How can I scale the 2 losses in the most automated way?</p>
</blockquote>
</aside>
<p>There is a proposed scheme for training the relative weights of the per-task<br>
losses when training a multi-task model.  (I think I saw this discussed in a<br>
previous thread on this forum, but I couldn’t find it.)  I haven’t ever tried it,<br>
but it looks sensible to me, and I imagine that it would work.</p>
<p>Here is a pytorch implementation and the reference it is based on:</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>

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

almog: <p>Hi Frank,<br>
Thank you!</p>
<p>I have tried this solution, I am not sure if it improved the performance more than a linear combination of the two losses. Looks like the model performance remained similar.</p>
<p>I would appreciate more solutions to try please<br>
Thanks</p>