6. Custom Multi-Objective Loss with Configurable Weights

Problem

RLHF-style training pipelines rarely optimize a single loss. A typical step combines a task loss, a KL penalty against a reference policy, and maybe an auxiliary regularizer, each with its own weight that gets tuned independently.

The combining function itself is simple arithmetic, but it is easy to break the one property that actually matters: the combined scalar must still support .backward() flowing gradients back into every term that produced it, with none of the pieces accidentally detached along the way.

Source: src/6_weighted_multi_objective_loss.py

def combined_loss(losses: dict[str, torch.Tensor], weights: dict[str, float]) -> torch.Tensor: ...

>>> losses = {"task": torch.tensor(2.0), "kl": torch.tensor(0.5)}
>>> weights = {"task": 1.0, "kl": 0.1}
>>> combined_loss(losses, weights).item()
2.049999952316284

Step-by-Step Approach

1. Initialize an accumulator to None, not to a plain 0.0 or a fresh torch.tensor(0.0) — starting from the first real term keeps dtype/device consistent with whatever tensors were actually passed in.
2. Iterate the losses dict; for each name, multiply the loss tensor by its matching weight from the weights dict.
3. Never call .item(), float(...), or .detach() on any term along the way — any of those breaks the autograd graph for that term.
4. Accumulate by tensor addition (total + term), which preserves the graph for every summed term simultaneously.
5. Handle the empty-dict edge case explicitly: return a plain scalar zero tensor rather than crashing or returning None.
6. Sanity check: build two toy leaf tensors with requires_grad=True, combine losses derived from each, call .backward() once on the result, and confirm both leaves received the expected gradient.

The key insight is that "weighted sum" is trivial arithmetic — the actual skill being tested is not breaking the autograd graph while doing it, which is exactly the kind of subtle bug that silently produces a model that trains but never actually receives gradient from one of its loss terms.

Reference solution

import torch


def combined_loss(losses: dict, weights: dict) -> torch.Tensor:
    # start from None, not 0.0 or torch.tensor(0.0): a plain python/tensor
    # zero would still work numerically here, but starting from the first
    # real term keeps dtype/device inference correct and avoids an
    # unnecessary extra add op for the common non-empty case
    total = None
    for name, loss in losses.items():
        term = weights[name] * loss  # weight * tensor keeps the graph attached
        total = term if total is None else total + term
    if total is None:
        return torch.tensor(0.0)
    return total

Key Functions & Tricks

• weight * loss — multiplying a Python float by a tensor keeps the result a tensor with the graph intact; no manual wrapping needed.
• total = term if total is None else total + term — the standard "fold with a None seed" idiom, avoiding a spurious extra addition against a zero of the wrong dtype/device.
• torch.tensor(0.0) — the explicit fallback for the empty-input edge case, so the function always returns a tensor, never None.
• result.backward() then reading leaf.grad — the standard way to verify a combining function didn't silently detach one of its inputs.

How to Recognize This Pattern