Loss balancer

Description & Motivation

It's commonplace to combine many losses (l1 loss, feature matching loss, multiscale losses, etc.). However, when they are on different scales, it is fiddly to choose the weighting. And, the weighting might change over time.

A loss balancer would be a useful addition, using EMA.

The API could be this simple:

loss = self.balancer(loss1, loss2, ...)

(By default, each loss would have equal weight, but this could be parameterized in the balancer class.)

Pitch

Meta's encodec has a loss balancer and it makes it very easy for them to mix many different scaled losses.

This would probably be an easy win for many lightning users that mix losses.

Alternatives

There is available code for the balancer: https://github.com/facebookresearch/encodec/blob/main/encodec/balancer.py

However: a) the licensing is not permissive and b) it doesn't work in DDP.

Additional context

https://arxiv.org/pdf/2210.13438.pdf

cc @borda

Start by reading the referenced encodec/balancer.py implementation and the issue's proposed self.balancer(loss1, loss2, ...) API. Investigate how the balancer should support EMA weighting and DDP, then identify the appropriate Lightning entry point and tests. Done means a documented loss-balancing feature that combines differently scaled losses and works in distributed training.