diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a8b4cba..eaf30c8b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,14 @@ changelog does not include internal changes that do not affect the user. ## [Unreleased] +### Added + +- Added `FAMO` (Fast Adaptive Multitask Optimization) from [FAMO: Fast Adaptive Multitask + Optimization](https://proceedings.neurips.cc/paper_files/paper/2023/file/b2fe1ee8d936ac08dd26f2ff58986c8f-Paper-Conference.pdf) + (NeurIPS 2023), a stateful `Scalarizer` that decreases all task losses at an approximately equal + rate using only the loss values. It learns the task weights internally; after the model step, + call its `update()` method with the losses recomputed on the same batch to adjust them. + ## [0.14.0] - 2026-06-10 ### Added diff --git a/NOTICES b/NOTICES index b7dbf617..098695c3 100644 --- a/NOTICES +++ b/NOTICES @@ -168,3 +168,31 @@ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +------------------------------------------------------------------------------- + +Project: FAMO +Source: https://github.com/Cranial-XIX/FAMO +Used in: src/torchjd/scalarization/_famo.py + +MIT License + +Copyright (c) 2023 Bo Liu + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/docs/source/docs/scalarization/famo.rst b/docs/source/docs/scalarization/famo.rst new file mode 100644 index 00000000..b6c36c37 --- /dev/null +++ b/docs/source/docs/scalarization/famo.rst @@ -0,0 +1,7 @@ +:hide-toc: + +FAMO +==== + +.. autoclass:: torchjd.scalarization.FAMO + :members: __call__, update, reset diff --git a/docs/source/docs/scalarization/index.rst b/docs/source/docs/scalarization/index.rst index fff5d797..d38708c0 100644 --- a/docs/source/docs/scalarization/index.rst +++ b/docs/source/docs/scalarization/index.rst @@ -16,6 +16,7 @@ Abstract base class constant.rst dwa.rst + famo.rst geometric_mean.rst imtl_l.rst mean.rst diff --git a/src/torchjd/scalarization/__init__.py b/src/torchjd/scalarization/__init__.py index 6f596edd..f1d22029 100644 --- a/src/torchjd/scalarization/__init__.py +++ b/src/torchjd/scalarization/__init__.py @@ -21,6 +21,7 @@ from ._constant import Constant from ._dwa import DWA +from ._famo import FAMO from ._geometric_mean import GeometricMean from ._imtl_l import IMTLL from ._mean import Mean @@ -33,6 +34,7 @@ __all__ = [ "Constant", "DWA", + "FAMO", "GeometricMean", "IMTLL", "Mean", diff --git a/src/torchjd/scalarization/_famo.py b/src/torchjd/scalarization/_famo.py new file mode 100644 index 00000000..7fca1b41 --- /dev/null +++ b/src/torchjd/scalarization/_famo.py @@ -0,0 +1,195 @@ +# Partly adapted from https://github.com/Cranial-XIX/FAMO — MIT License, Copyright (c) 2023 Bo Liu. +# See NOTICES for the full license text. +from collections.abc import Sequence + +import torch +from torch import Tensor, nn +from torch.nn.functional import softmax +from torch.optim import Adam + +from torchjd._mixins import Stateful + +from ._scalarizer_base import Scalarizer + +_EPSILON = 1e-8 + + +class FAMO(Scalarizer, Stateful): + r""" + :class:`~torchjd.Stateful` + :class:`~torchjd.scalarization.Scalarizer` that combines the input tensor of values using Fast + Adaptive Multitask Optimization (FAMO), proposed in `FAMO: Fast Adaptive Multitask Optimization + `_. + + FAMO decreases all task losses at an approximately equal rate while using only the loss values, + so it never needs the per-task gradients. The values are combined as + + .. math:: + c \sum_i z_i \log(\ell_i - b_i + \epsilon), \qquad + z = \mathrm{softmax}(w), \qquad + c = \left( \sum_i \frac{z_i}{\ell_i - b_i + \epsilon} \right)^{-1} + + where: + + - :math:`\ell_i` is the :math:`i`-th value (typically the loss of task :math:`i`); + - :math:`b_i` is the lower bound on the :math:`i`-th loss (the ``min_losses`` parameter, + ``0`` by default); + - :math:`w_i` is the task-weighting logit of task :math:`i`, learned internally by FAMO; + - :math:`z = \mathrm{softmax}(w)` are the task weights; + - :math:`c` is a normalization constant (treated as a constant in the backward pass) that makes + the resulting update a convex combination of the task gradients; + - :math:`\epsilon` is a small positive constant for numerical stability. + + Backpropagating this scalarized loss gives FAMO's balanced update direction for the model. + + The task-weighting logits :math:`w` are not learned through that backward pass. Instead, after + the model has been updated, call :meth:`update` with the losses recomputed on the same batch. It + measures how much each loss changed across the step, + + .. math:: + \delta_i = \log(\ell_i^{\text{before}} - b_i + \epsilon) + - \log(\ell_i^{\text{after}} - b_i + \epsilon), + + and takes an `Adam `_ step + on :math:`w` in that direction. FAMO owns this ``Adam`` internally + (configured by ``lr`` and ``weight_decay``), so you only call the scalarizer and then + :meth:`update`; there is no second optimizer to manage. + + :param shape: The shape of the values to scalarize, used to create one task-weighting logit per + value. An ``int`` ``n`` is interpreted as the shape ``(n,)``. + :param min_losses: The per-task lower bound :math:`b` subtracted from the values before the + logarithm. If provided, it must have the shape given by ``shape``. If ``None``, zeros are + used, in which case the values must be strictly positive. + :param lr: Learning rate of the internal ``Adam`` that learns the task-weighting logits. Must be + non-negative. The paper uses ``0.025``. + :param weight_decay: Weight decay of the internal ``Adam``, i.e. the paper's regularization + coefficient on the logits. Must be non-negative. Defaults to ``1e-3`` (as in the paper's + Algorithm 2 and in LibMTL); the official implementation uses ``1e-5``. + + The following example shows how to do one iteration of training of a model with FAMO. The losses + are recomputed on the same batch after the model step so that :meth:`update` can adjust the + weights. + + >>> import torch + >>> from torch.nn import Linear + >>> + >>> from torchjd.scalarization import FAMO + >>> + >>> model = Linear(3, 2) + >>> scalarizer = FAMO(2) # Move to the right device with e.g. FAMO(2).to(device="cuda") + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + >>> + >>> features = torch.randn(8, 3) + >>> losses = model(features).pow(2).mean(dim=0) # One loss per output dimension. + >>> loss = scalarizer(losses) + >>> optimizer.zero_grad() + >>> loss.backward() + >>> optimizer.step() + >>> + >>> # Recompute the losses on the same batch, after the model update. + >>> new_losses = model(features).pow(2).mean(dim=0) + >>> scalarizer.update(new_losses) # Updates the task weights internally. + + .. note:: + FAMO takes the logarithm of :math:`\ell_i - b_i`, so each value must stay strictly above its + lower bound :math:`b_i` (the paper assumes non-negative losses). With the default + ``min_losses`` of zeros, this means the values must be strictly positive. This precondition + is not enforced. + + .. note:: + This implementation was adapted from the `official implementation + `_. + """ + + min_losses: Tensor + + def __init__( + self, + shape: int | Sequence[int], + min_losses: Tensor | None = None, + lr: float = 0.025, + weight_decay: float = 1e-3, + ) -> None: + if lr < 0.0: + raise ValueError(f"Parameter `lr` should be non-negative. Found `lr = {lr}`.") + if weight_decay < 0.0: + raise ValueError( + f"Parameter `weight_decay` should be non-negative. Found `weight_decay = " + f"{weight_decay}`." + ) + + super().__init__() + self._w = nn.Parameter(torch.zeros(shape)) + + if min_losses is None: + min_losses = torch.zeros(self._w.shape) + elif min_losses.shape != self._w.shape: + raise ValueError( + f"Parameter `min_losses` should have shape {tuple(self._w.shape)} (matching the " + f"shape of the logits). Found `min_losses.shape = {tuple(min_losses.shape)}`." + ) + self.register_buffer("min_losses", min_losses) + + self.lr = lr + self.weight_decay = weight_decay + self._optimizer: Adam | None = None + self._prev_losses: Tensor | None = None + + def forward(self, values: Tensor, /) -> Tensor: + self._check_shape(values) + + self._prev_losses = values.detach().clone() + + weights = softmax(self._w.flatten(), dim=0).reshape(values.shape).detach() + shifted = values - self.min_losses + _EPSILON + normalizer = (weights / shifted).sum().detach() + return ((weights / normalizer) * torch.log(shifted)).sum() + + def update(self, values: Tensor, /) -> None: + """ + Updates the task-weighting logits from the change in losses across the model update, by + taking one step of the internal ``Adam``. Must be called after the scalarizer has been + called on the batch's losses, with the losses recomputed on the same batch after the model + step. + """ + + if self._prev_losses is None: + raise ValueError( + "`update` must be called after the scalarizer is called on the losses." + ) + self._check_shape(values) + + before = self._prev_losses - self.min_losses + _EPSILON + after = values.detach() - self.min_losses + _EPSILON + delta = torch.log(before) - torch.log(after) + + with torch.enable_grad(): + weights = softmax(self._w.flatten(), dim=0) + grad = torch.autograd.grad(weights, self._w, grad_outputs=delta.flatten())[0] + + if self._optimizer is None: + self._optimizer = Adam([self._w], lr=self.lr, weight_decay=self.weight_decay) + self._w.grad = grad + self._optimizer.step() + # Clear the gradient so it cannot leak into a user optimizer that the logits were mistakenly + # added to: FAMO is the only thing that should step them. + self._w.grad = None + + def reset(self) -> None: + with torch.no_grad(): + self._w.zero_() + self._optimizer = None + self._prev_losses = None + + def _check_shape(self, values: Tensor) -> None: + if values.shape != self._w.shape: + raise ValueError( + f"Parameter `values` should have shape {tuple(self._w.shape)} (matching the shape " + f"of the logits). Found `values.shape = {tuple(values.shape)}`." + ) + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(shape={tuple(self._w.shape)}, lr={self.lr}, " + f"weight_decay={self.weight_decay})" + ) diff --git a/tests/unit/scalarization/test_famo.py b/tests/unit/scalarization/test_famo.py new file mode 100644 index 00000000..75b021c8 --- /dev/null +++ b/tests/unit/scalarization/test_famo.py @@ -0,0 +1,185 @@ +from contextlib import nullcontext as does_not_raise + +import torch +from pytest import mark, raises +from settings import DEVICE, DTYPE +from torch import Tensor +from utils.contexts import ExceptionContext +from utils.tensors import ones_, rand_, tensor_, zeros_ + +from torchjd.scalarization import FAMO + +from ._asserts import assert_grad_flow, assert_returns_scalar +from ._inputs import shapes + +# FAMO takes the log of the values, so they must be strictly positive. +positive_inputs = [rand_(shape) + 1.0 for shape in shapes] + + +def _famo(shape: int | tuple[int, ...]) -> FAMO: + """Builds a `FAMO` whose logits and lower bounds live on the test device and dtype.""" + return FAMO(shape).to(device=DEVICE, dtype=DTYPE) + + +def test_value() -> None: + # With logits initialized to 0, the task weights are uniform. The result is the normalized, + # weighted sum of the log-values. + values = tensor_([1.0, 2.0]) + z = tensor_([0.5, 0.5]) + shifted = values + 1e-8 + c = 1.0 / (z / shifted).sum() + expected = (c * z * torch.log(shifted)).sum() + torch.testing.assert_close(_famo(2)(values), expected) + + +def test_value_gradient_matches_formula() -> None: + # The gradient w.r.t. the values is FAMO's balanced update direction: c * z / (values - b). + values = tensor_([1.0, 2.0]).requires_grad_() + _famo(2)(values).backward() + z = tensor_([0.5, 0.5]) + shifted = values.detach() + 1e-8 + expected_grad = (1.0 / (z / shifted).sum()) * z / shifted + torch.testing.assert_close(values.grad, expected_grad) + + +def test_int_shape_matches_tuple_shape() -> None: + values = tensor_([1.0, 2.0, 4.0]) + assert FAMO(3)._w.shape == (3,) + torch.testing.assert_close(_famo(3)(values), _famo((3,))(values)) + + +@mark.parametrize("values", positive_inputs) +def test_expected_structure(values: Tensor) -> None: + assert_returns_scalar(_famo(tuple(values.shape)), values) + + +@mark.parametrize("values", positive_inputs) +def test_grad_flow(values: Tensor) -> None: + assert_grad_flow(_famo(tuple(values.shape)), values) + + +def test_forward_does_not_write_logit_grad() -> None: + # The weights are detached inside `forward`, so backward populates the values' gradient but + # never the logits' gradient (those are updated by `update` instead). + famo = _famo(2) + values = tensor_([1.0, 2.0]).requires_grad_() + famo(values).backward() + assert values.grad is not None + assert famo._w.grad is None + + +def test_update_steps_the_logits() -> None: + famo = _famo(2) + famo(tensor_([1.0, 2.0])) + famo.update(tensor_([0.5, 1.5])) + assert famo._w.detach().isfinite().all() + assert not torch.equal(famo._w.detach(), zeros_((2,))) + + +def test_update_clears_logit_grad() -> None: + # After stepping its own optimizer, FAMO clears the logit gradient so it cannot leak into a user + # optimizer the logits were mistakenly added to. + famo = _famo(2) + famo(tensor_([1.0, 2.0])) + famo.update(tensor_([0.5, 1.5])) + assert famo._w.grad is None + + +def test_update_uses_last_forward_losses() -> None: + # `update` compares the losses from the most recent `forward` against the ones it receives. When + # they are equal, the change is zero, so the logits do not move. + famo = _famo(2) + famo(tensor_([5.0, 5.0])) + famo(tensor_([1.0, 4.0])) + famo.update(tensor_([1.0, 4.0])) + torch.testing.assert_close(famo._w.detach(), zeros_((2,))) + + +def test_update_before_forward_raises() -> None: + with raises(ValueError): + _famo(2).update(tensor_([1.0, 2.0])) + + +@mark.parametrize( + ["param_shape", "values_shape", "expectation"], + [ + ((5,), (5,), does_not_raise()), + ((3, 4), (3, 4), does_not_raise()), + ((), (), does_not_raise()), + ((5,), (4,), raises(ValueError)), + ((5,), (5, 1), raises(ValueError)), + ((3, 4), (4, 3), raises(ValueError)), + ], +) +def test_forward_shape_check( + param_shape: tuple[int, ...], + values_shape: tuple[int, ...], + expectation: ExceptionContext, +) -> None: + scalarizer = _famo(param_shape) + values = ones_(values_shape) + with expectation: + _ = scalarizer(values) + + +def test_update_shape_check() -> None: + famo = _famo(2) + famo(tensor_([1.0, 2.0])) + with raises(ValueError): + famo.update(tensor_([1.0, 2.0, 3.0])) + + +def test_min_losses_defaults_to_zeros() -> None: + torch.testing.assert_close(FAMO(2).min_losses, torch.zeros(2)) + + +def test_min_losses_wrong_shape_raises() -> None: + with raises(ValueError): + FAMO(2, min_losses=tensor_([0.0, 0.0, 0.0])) + + +def test_min_losses_shifts_values() -> None: + # With a lower bound, the log is taken on values - min_losses. + values = tensor_([2.0, 3.0]) + bound = tensor_([1.0, 1.0]) + famo = FAMO(2, min_losses=bound).to(device=DEVICE, dtype=DTYPE) + z = tensor_([0.5, 0.5]) + shifted = values - bound + 1e-8 + c = 1.0 / (z / shifted).sum() + expected = (c * z * torch.log(shifted)).sum() + torch.testing.assert_close(famo(values), expected) + + +@mark.parametrize("lr", [-1.0]) +def test_negative_lr_raises(lr: float) -> None: + with raises(ValueError): + FAMO(2, lr=lr) + + +@mark.parametrize("weight_decay", [-1.0]) +def test_negative_weight_decay_raises(weight_decay: float) -> None: + with raises(ValueError): + FAMO(2, weight_decay=weight_decay) + + +def test_reset() -> None: + famo = _famo(2) + famo(tensor_([1.0, 2.0])) + famo.update(tensor_([0.5, 1.5])) + famo.reset() + torch.testing.assert_close(famo._w.detach(), zeros_((2,))) + assert famo._prev_losses is None + assert famo._optimizer is None + + +def test_nan_propagates_for_value_below_bound() -> None: + # log(values - min_losses) is undefined when a value is not above its bound; the nan must + # propagate rather than being silently clamped. + out = _famo(2)(tensor_([-1.0, 2.0])) + assert out.isnan() + + +def test_representations() -> None: + assert repr(FAMO(3)) == "FAMO(shape=(3,), lr=0.025, weight_decay=0.001)" + assert repr(FAMO((2, 3))) == "FAMO(shape=(2, 3), lr=0.025, weight_decay=0.001)" + assert str(FAMO(3)) == "FAMO"