From 0fd8c9c4c1a09e0883bd57cebb12267f429a15ef Mon Sep 17 00:00:00 2001 From: ppraneth Date: Sat, 13 Jun 2026 16:16:26 +0530 Subject: [PATCH 1/4] add FAMO --- CHANGELOG.md | 9 ++ docs/source/docs/scalarization/famo.rst | 7 + docs/source/docs/scalarization/index.rst | 1 + src/torchjd/scalarization/__init__.py | 2 + src/torchjd/scalarization/_famo.py | 157 ++++++++++++++++++++ tests/unit/scalarization/test_famo.py | 179 +++++++++++++++++++++++ 6 files changed, 355 insertions(+) create mode 100644 docs/source/docs/scalarization/famo.rst create mode 100644 src/torchjd/scalarization/_famo.py create mode 100644 tests/unit/scalarization/test_famo.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a8b4cba..cea6e64e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,15 @@ 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. Its task-weighting logits are an `nn.Parameter`; after the model + step, call its `update()` method with the losses recomputed on the same batch to set their + gradient, then step your own optimizer. + ## [0.14.0] - 2026-06-10 ### Added 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..558954a8 --- /dev/null +++ b/src/torchjd/scalarization/_famo.py @@ -0,0 +1,157 @@ +from collections.abc import Sequence + +import torch +from torch import Tensor, nn +from torch.nn.functional import softmax + +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 learnable task-weighting logit of task :math:`i`, stored as an + ``nn.Parameter``; + - :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 sets the gradient of :math:`w` to the matching softmax vector-Jacobian product. Your own + optimizer then steps :math:`w`. To match the paper, use ``Adam`` with a weight decay equal to + the paper's regularization coefficient. + + :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. + + The following example shows how to train a model with FAMO. The model and the task-weighting + logits use two separate optimizers, and 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) + >>> weight_optimizer = torch.optim.Adam( + ... scalarizer.parameters(), lr=0.025, weight_decay=0.001 + ... ) + >>> + >>> 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) # Sets the gradient of the task-weighting logits. + >>> weight_optimizer.step() + + .. 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. + """ + + min_losses: Tensor + + def __init__(self, shape: int | Sequence[int], min_losses: Tensor | None = None) -> None: + 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._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: + """ + Sets the gradient of the task-weighting logits from the change in losses across the model + update, overwriting any existing gradient. 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. The logits are not stepped here; call your own optimizer's ``step`` afterwards. + """ + + if self._prev_losses is None: + raise ValueError("`update` must be called after `forward`, which stores 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] + self.w.grad = grad + + def reset(self) -> None: + with torch.no_grad(): + self.w.zero_() + 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)})" diff --git a/tests/unit/scalarization/test_famo.py b/tests/unit/scalarization/test_famo.py new file mode 100644 index 00000000..99834422 --- /dev/null +++ b/tests/unit/scalarization/test_famo.py @@ -0,0 +1,179 @@ +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 set 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_sets_finite_logit_grad() -> None: + famo = _famo(2) + famo(tensor_([1.0, 2.0])) + famo.update(tensor_([0.5, 1.5])) + assert famo.w.grad is not None + assert famo.w.grad.isfinite().all() + + +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 logit gradient is zero. + 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.grad, zeros_((2,))) + + +def test_is_trainable_via_update() -> None: + # The full two-call protocol: forward, then update with the recomputed losses, then step the + # logit optimizer. The logits should move away from their initial zeros. + famo = _famo(2) + optimizer = torch.optim.Adam(famo.parameters(), lr=0.1) + famo(tensor_([2.0, 1.0])) + famo.update(tensor_([1.0, 1.0])) + optimizer.step() + assert not torch.equal(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) + + +def test_reset() -> None: + famo = _famo(2) + with torch.no_grad(): + famo.w.add_(1.0) + famo(tensor_([1.0, 2.0])) + famo.reset() + torch.testing.assert_close(famo.w.detach(), zeros_((2,))) + assert famo._prev_losses is None + + +def test_logits_are_a_learnable_parameter() -> None: + assert len(list(FAMO(2).parameters())) == 1 + + +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,))" + assert repr(FAMO((2, 3))) == "FAMO(shape=(2, 3))" + assert str(FAMO(3)) == "FAMO" From 978161249af7e05dab91f8fd8574986c8db2b71e Mon Sep 17 00:00:00 2001 From: ppraneth Date: Mon, 15 Jun 2026 18:51:59 +0530 Subject: [PATCH 2/4] minor fixes --- CHANGELOG.md | 5 +- NOTICES | 28 ++++++++ src/torchjd/scalarization/_famo.py | 96 ++++++++++++++++++--------- tests/unit/scalarization/test_famo.py | 62 +++++++++-------- 4 files changed, 130 insertions(+), 61 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cea6e64e..eaf30c8b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,9 +13,8 @@ changelog does not include internal changes that do not affect the user. - 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. Its task-weighting logits are an `nn.Parameter`; after the model - step, call its `update()` method with the losses recomputed on the same batch to set their - gradient, then step your own optimizer. + 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 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/src/torchjd/scalarization/_famo.py b/src/torchjd/scalarization/_famo.py index 558954a8..9e214961 100644 --- a/src/torchjd/scalarization/_famo.py +++ b/src/torchjd/scalarization/_famo.py @@ -1,8 +1,11 @@ +# 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 @@ -31,8 +34,7 @@ class FAMO(Scalarizer, Stateful): - :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 learnable task-weighting logit of task :math:`i`, stored as an - ``nn.Parameter``; + - :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; @@ -48,19 +50,23 @@ class FAMO(Scalarizer, Stateful): \delta_i = \log(\ell_i^{\text{before}} - b_i + \epsilon) - \log(\ell_i^{\text{after}} - b_i + \epsilon), - and sets the gradient of :math:`w` to the matching softmax vector-Jacobian product. Your own - optimizer then steps :math:`w`. To match the paper, use ``Adam`` with a weight decay equal to - the paper's regularization coefficient. + 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 train a model with FAMO. The model and the task-weighting - logits use two separate optimizers, and the losses are recomputed on the same batch after the - model step so that :meth:`update` can adjust the weights. + The following example shows how to train 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 @@ -70,9 +76,6 @@ class FAMO(Scalarizer, Stateful): >>> 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) - >>> weight_optimizer = torch.optim.Adam( - ... scalarizer.parameters(), lr=0.025, weight_decay=0.001 - ... ) >>> >>> features = torch.randn(8, 3) >>> losses = model(features).pow(2).mean(dim=0) # One loss per output dimension. @@ -83,31 +86,51 @@ class FAMO(Scalarizer, Stateful): >>> >>> # Recompute the losses on the same batch, after the model update. >>> new_losses = model(features).pow(2).mean(dim=0) - >>> scalarizer.update(new_losses) # Sets the gradient of the task-weighting logits. - >>> weight_optimizer.step() + >>> 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) -> None: + 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)) + 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: + 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"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: @@ -115,21 +138,23 @@ def forward(self, values: Tensor, /) -> Tensor: self._prev_losses = values.detach().clone() - weights = softmax(self.w.flatten(), dim=0).reshape(values.shape).detach() + 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: """ - Sets the gradient of the task-weighting logits from the change in losses across the model - update, overwriting any existing gradient. Must be called after the scalarizer has been + 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. The logits are not stepped here; call your own optimizer's ``step`` afterwards. + step. """ if self._prev_losses is None: - raise ValueError("`update` must be called after `forward`, which stores the losses.") + 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 @@ -137,21 +162,32 @@ def update(self, values: Tensor, /) -> None: 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] - self.w.grad = 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._w.zero_() + self._optimizer = None self._prev_losses = None def _check_shape(self, values: Tensor) -> None: - if values.shape != self.w.shape: + if values.shape != self._w.shape: raise ValueError( - f"Parameter `values` should have shape {tuple(self.w.shape)} (matching the shape " + 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)})" + 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 index 99834422..75b021c8 100644 --- a/tests/unit/scalarization/test_famo.py +++ b/tests/unit/scalarization/test_famo.py @@ -44,7 +44,7 @@ def test_value_gradient_matches_formula() -> None: def test_int_shape_matches_tuple_shape() -> None: values = tensor_([1.0, 2.0, 4.0]) - assert FAMO(3).w.shape == (3,) + assert FAMO(3)._w.shape == (3,) torch.testing.assert_close(_famo(3)(values), _famo((3,))(values)) @@ -60,41 +60,39 @@ def test_grad_flow(values: Tensor) -> None: 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 set by `update` instead). + # 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 + assert famo._w.grad is None -def test_update_sets_finite_logit_grad() -> 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.grad is not None - assert famo.w.grad.isfinite().all() + 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 logit gradient is zero. + # 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.grad, zeros_((2,))) - - -def test_is_trainable_via_update() -> None: - # The full two-call protocol: forward, then update with the recomputed losses, then step the - # logit optimizer. The logits should move away from their initial zeros. - famo = _famo(2) - optimizer = torch.optim.Adam(famo.parameters(), lr=0.1) - famo(tensor_([2.0, 1.0])) - famo.update(tensor_([1.0, 1.0])) - optimizer.step() - assert not torch.equal(famo.w.detach(), zeros_((2,))) + torch.testing.assert_close(famo._w.detach(), zeros_((2,))) def test_update_before_forward_raises() -> None: @@ -152,18 +150,26 @@ def test_min_losses_shifts_values() -> None: 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) - with torch.no_grad(): - famo.w.add_(1.0) famo(tensor_([1.0, 2.0])) + famo.update(tensor_([0.5, 1.5])) famo.reset() - torch.testing.assert_close(famo.w.detach(), zeros_((2,))) + torch.testing.assert_close(famo._w.detach(), zeros_((2,))) assert famo._prev_losses is None - - -def test_logits_are_a_learnable_parameter() -> None: - assert len(list(FAMO(2).parameters())) == 1 + assert famo._optimizer is None def test_nan_propagates_for_value_below_bound() -> None: @@ -174,6 +180,6 @@ def test_nan_propagates_for_value_below_bound() -> None: def test_representations() -> None: - assert repr(FAMO(3)) == "FAMO(shape=(3,))" - assert repr(FAMO((2, 3))) == "FAMO(shape=(2, 3))" + 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" From 0fe401182c30d5fc8d8e4e5f71b8c29c7e6a71f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Val=C3=A9rian=20Rey?= Date: Mon, 15 Jun 2026 16:10:16 +0200 Subject: [PATCH 3/4] Add Adam link --- src/torchjd/scalarization/_famo.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/torchjd/scalarization/_famo.py b/src/torchjd/scalarization/_famo.py index 9e214961..2170ec49 100644 --- a/src/torchjd/scalarization/_famo.py +++ b/src/torchjd/scalarization/_famo.py @@ -50,7 +50,8 @@ class FAMO(Scalarizer, Stateful): \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 + 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. From ec0402b530901e4c02a35c4f25bb8211bafbe5e2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Val=C3=A9rian=20Rey?= Date: Mon, 15 Jun 2026 16:12:00 +0200 Subject: [PATCH 4/4] Make explicit that the example is about 1 step of FAMO --- src/torchjd/scalarization/_famo.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/torchjd/scalarization/_famo.py b/src/torchjd/scalarization/_famo.py index 2170ec49..7fca1b41 100644 --- a/src/torchjd/scalarization/_famo.py +++ b/src/torchjd/scalarization/_famo.py @@ -66,8 +66,9 @@ class FAMO(Scalarizer, Stateful): 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 train a model with FAMO. The losses are recomputed on the - same batch after the model step so that :meth:`update` can adjust the weights. + 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