🎓 A1Learner
🔍

PyTorch Combining Losses

A combined loss adds two or more loss terms into one number. A weight in front of each term sets how much that term counts.

Two terms

The fit term scores the predictions. The size term scores the weights themselves.

import torch

pred = torch.tensor([2.0, 4.0])
target = torch.tensor([3.0, 3.0])
weights = torch.tensor([1.0, -2.0])
fit = ((pred - target) ** 2).mean()
size = (weights ** 2).mean()
print(f"fit {fit.item():.1f} size {size.item():.1f}")

Weight and add

One multiplication and one addition join the terms into a single loss.

import torch

fit = torch.tensor(1.0)
size = torch.tensor(2.5)
total = fit + 0.5 * size
print(total)

The weight steers

The same two terms give different totals under different weights. The weight is a steering dial.

import torch

fit = torch.tensor(1.0)
size = torch.tensor(2.5)
print(fit + 0.0 * size)
print(fit + 2.0 * size)

Example

import torch

torch.manual_seed(0)

# Each row has five features. The target mixes all five.
X = torch.randn(60, 5)
true_w = torch.tensor([2.0, -1.0, 0.5, 0.0, 3.0])
y = (X @ true_w + 0.1 * torch.randn(60)).reshape(-1, 1)


def train(size_weight):
    torch.manual_seed(1)
    model = torch.nn.Linear(5, 1)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
    for step in range(300):
        fit = torch.mean((model(X) - y) ** 2)
        size = torch.mean(model.weight ** 2)
        loss = fit + size_weight * size
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    fit = torch.mean((model(X) - y) ** 2).item()
    size = torch.mean(model.weight ** 2).item()
    return fit, size


fit, size = train(0.0)
print(f"Fit only: fit term {fit:.4f}  size term {size:.4f}")
fit, size = train(0.5)
print(f"Fit + 0.5 * size: fit term {fit:.4f}  size term {size:.4f}")

Key points

  • Term: one part of a combined loss, scored on its own
  • Weight: a number multiplied by a term to set its strength
  • Use weights to steer training when two goals pull in different directions
  • A bigger weight makes its term matter more, and the other terms matter less
  • The combined loss is still one number, so the training loop stays unchanged
  • Adding weighted terms turns two goals into one score, and the weights set the balance

In this example, the target mixes five features with true weights that include 2.0 and 3.0. The first model trains on the fit term alone. Its fit term falls to 0.0164, and its size term ends at 2.8425. The second model trains on fit plus 0.5 times size, starting from the same seed. Its size term ends lower, at 2.2765. Its fit term ends higher, at 0.1497. The weight bought smaller weights and paid with a slightly worse fit. That trade is the steering the weights provide.

Note

A size penalty on the weights is the same idea as weight decay in an optimizer

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A combined loss is fit + 0.5 * size. The size weight rises to 2.0 and training runs again. What should happen?