|
| 1 | +# Partly adapted from https://github.com/Cranial-XIX/FAMO — MIT License, Copyright (c) 2023 Bo Liu. |
| 2 | +# See NOTICES for the full license text. |
| 3 | +from collections.abc import Sequence |
| 4 | + |
| 5 | +import torch |
| 6 | +from torch import Tensor, nn |
| 7 | +from torch.nn.functional import softmax |
| 8 | +from torch.optim import Adam |
| 9 | + |
| 10 | +from torchjd._mixins import Stateful |
| 11 | + |
| 12 | +from ._scalarizer_base import Scalarizer |
| 13 | + |
| 14 | +_EPSILON = 1e-8 |
| 15 | + |
| 16 | + |
| 17 | +class FAMO(Scalarizer, Stateful): |
| 18 | + r""" |
| 19 | + :class:`~torchjd.Stateful` |
| 20 | + :class:`~torchjd.scalarization.Scalarizer` that combines the input tensor of values using Fast |
| 21 | + Adaptive Multitask Optimization (FAMO), proposed in `FAMO: Fast Adaptive Multitask Optimization |
| 22 | + <https://proceedings.neurips.cc/paper_files/paper/2023/file/b2fe1ee8d936ac08dd26f2ff58986c8f-Paper-Conference.pdf>`_. |
| 23 | +
|
| 24 | + FAMO decreases all task losses at an approximately equal rate while using only the loss values, |
| 25 | + so it never needs the per-task gradients. The values are combined as |
| 26 | +
|
| 27 | + .. math:: |
| 28 | + c \sum_i z_i \log(\ell_i - b_i + \epsilon), \qquad |
| 29 | + z = \mathrm{softmax}(w), \qquad |
| 30 | + c = \left( \sum_i \frac{z_i}{\ell_i - b_i + \epsilon} \right)^{-1} |
| 31 | +
|
| 32 | + where: |
| 33 | +
|
| 34 | + - :math:`\ell_i` is the :math:`i`-th value (typically the loss of task :math:`i`); |
| 35 | + - :math:`b_i` is the lower bound on the :math:`i`-th loss (the ``min_losses`` parameter, |
| 36 | + ``0`` by default); |
| 37 | + - :math:`w_i` is the task-weighting logit of task :math:`i`, learned internally by FAMO; |
| 38 | + - :math:`z = \mathrm{softmax}(w)` are the task weights; |
| 39 | + - :math:`c` is a normalization constant (treated as a constant in the backward pass) that makes |
| 40 | + the resulting update a convex combination of the task gradients; |
| 41 | + - :math:`\epsilon` is a small positive constant for numerical stability. |
| 42 | +
|
| 43 | + Backpropagating this scalarized loss gives FAMO's balanced update direction for the model. |
| 44 | +
|
| 45 | + The task-weighting logits :math:`w` are not learned through that backward pass. Instead, after |
| 46 | + the model has been updated, call :meth:`update` with the losses recomputed on the same batch. It |
| 47 | + measures how much each loss changed across the step, |
| 48 | +
|
| 49 | + .. math:: |
| 50 | + \delta_i = \log(\ell_i^{\text{before}} - b_i + \epsilon) |
| 51 | + - \log(\ell_i^{\text{after}} - b_i + \epsilon), |
| 52 | +
|
| 53 | + and takes an `Adam <https://docs.pytorch.org/docs/stable/generated/torch.optim.Adam.html>`_ step |
| 54 | + on :math:`w` in that direction. FAMO owns this ``Adam`` internally |
| 55 | + (configured by ``lr`` and ``weight_decay``), so you only call the scalarizer and then |
| 56 | + :meth:`update`; there is no second optimizer to manage. |
| 57 | +
|
| 58 | + :param shape: The shape of the values to scalarize, used to create one task-weighting logit per |
| 59 | + value. An ``int`` ``n`` is interpreted as the shape ``(n,)``. |
| 60 | + :param min_losses: The per-task lower bound :math:`b` subtracted from the values before the |
| 61 | + logarithm. If provided, it must have the shape given by ``shape``. If ``None``, zeros are |
| 62 | + used, in which case the values must be strictly positive. |
| 63 | + :param lr: Learning rate of the internal ``Adam`` that learns the task-weighting logits. Must be |
| 64 | + non-negative. The paper uses ``0.025``. |
| 65 | + :param weight_decay: Weight decay of the internal ``Adam``, i.e. the paper's regularization |
| 66 | + coefficient on the logits. Must be non-negative. Defaults to ``1e-3`` (as in the paper's |
| 67 | + Algorithm 2 and in LibMTL); the official implementation uses ``1e-5``. |
| 68 | +
|
| 69 | + The following example shows how to do one iteration of training of a model with FAMO. The losses |
| 70 | + are recomputed on the same batch after the model step so that :meth:`update` can adjust the |
| 71 | + weights. |
| 72 | +
|
| 73 | + >>> import torch |
| 74 | + >>> from torch.nn import Linear |
| 75 | + >>> |
| 76 | + >>> from torchjd.scalarization import FAMO |
| 77 | + >>> |
| 78 | + >>> model = Linear(3, 2) |
| 79 | + >>> scalarizer = FAMO(2) # Move to the right device with e.g. FAMO(2).to(device="cuda") |
| 80 | + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1) |
| 81 | + >>> |
| 82 | + >>> features = torch.randn(8, 3) |
| 83 | + >>> losses = model(features).pow(2).mean(dim=0) # One loss per output dimension. |
| 84 | + >>> loss = scalarizer(losses) |
| 85 | + >>> optimizer.zero_grad() |
| 86 | + >>> loss.backward() |
| 87 | + >>> optimizer.step() |
| 88 | + >>> |
| 89 | + >>> # Recompute the losses on the same batch, after the model update. |
| 90 | + >>> new_losses = model(features).pow(2).mean(dim=0) |
| 91 | + >>> scalarizer.update(new_losses) # Updates the task weights internally. |
| 92 | +
|
| 93 | + .. note:: |
| 94 | + FAMO takes the logarithm of :math:`\ell_i - b_i`, so each value must stay strictly above its |
| 95 | + lower bound :math:`b_i` (the paper assumes non-negative losses). With the default |
| 96 | + ``min_losses`` of zeros, this means the values must be strictly positive. This precondition |
| 97 | + is not enforced. |
| 98 | +
|
| 99 | + .. note:: |
| 100 | + This implementation was adapted from the `official implementation |
| 101 | + <https://github.com/Cranial-XIX/FAMO>`_. |
| 102 | + """ |
| 103 | + |
| 104 | + min_losses: Tensor |
| 105 | + |
| 106 | + def __init__( |
| 107 | + self, |
| 108 | + shape: int | Sequence[int], |
| 109 | + min_losses: Tensor | None = None, |
| 110 | + lr: float = 0.025, |
| 111 | + weight_decay: float = 1e-3, |
| 112 | + ) -> None: |
| 113 | + if lr < 0.0: |
| 114 | + raise ValueError(f"Parameter `lr` should be non-negative. Found `lr = {lr}`.") |
| 115 | + if weight_decay < 0.0: |
| 116 | + raise ValueError( |
| 117 | + f"Parameter `weight_decay` should be non-negative. Found `weight_decay = " |
| 118 | + f"{weight_decay}`." |
| 119 | + ) |
| 120 | + |
| 121 | + super().__init__() |
| 122 | + self._w = nn.Parameter(torch.zeros(shape)) |
| 123 | + |
| 124 | + if min_losses is None: |
| 125 | + min_losses = torch.zeros(self._w.shape) |
| 126 | + elif min_losses.shape != self._w.shape: |
| 127 | + raise ValueError( |
| 128 | + f"Parameter `min_losses` should have shape {tuple(self._w.shape)} (matching the " |
| 129 | + f"shape of the logits). Found `min_losses.shape = {tuple(min_losses.shape)}`." |
| 130 | + ) |
| 131 | + self.register_buffer("min_losses", min_losses) |
| 132 | + |
| 133 | + self.lr = lr |
| 134 | + self.weight_decay = weight_decay |
| 135 | + self._optimizer: Adam | None = None |
| 136 | + self._prev_losses: Tensor | None = None |
| 137 | + |
| 138 | + def forward(self, values: Tensor, /) -> Tensor: |
| 139 | + self._check_shape(values) |
| 140 | + |
| 141 | + self._prev_losses = values.detach().clone() |
| 142 | + |
| 143 | + weights = softmax(self._w.flatten(), dim=0).reshape(values.shape).detach() |
| 144 | + shifted = values - self.min_losses + _EPSILON |
| 145 | + normalizer = (weights / shifted).sum().detach() |
| 146 | + return ((weights / normalizer) * torch.log(shifted)).sum() |
| 147 | + |
| 148 | + def update(self, values: Tensor, /) -> None: |
| 149 | + """ |
| 150 | + Updates the task-weighting logits from the change in losses across the model update, by |
| 151 | + taking one step of the internal ``Adam``. Must be called after the scalarizer has been |
| 152 | + called on the batch's losses, with the losses recomputed on the same batch after the model |
| 153 | + step. |
| 154 | + """ |
| 155 | + |
| 156 | + if self._prev_losses is None: |
| 157 | + raise ValueError( |
| 158 | + "`update` must be called after the scalarizer is called on the losses." |
| 159 | + ) |
| 160 | + self._check_shape(values) |
| 161 | + |
| 162 | + before = self._prev_losses - self.min_losses + _EPSILON |
| 163 | + after = values.detach() - self.min_losses + _EPSILON |
| 164 | + delta = torch.log(before) - torch.log(after) |
| 165 | + |
| 166 | + with torch.enable_grad(): |
| 167 | + weights = softmax(self._w.flatten(), dim=0) |
| 168 | + grad = torch.autograd.grad(weights, self._w, grad_outputs=delta.flatten())[0] |
| 169 | + |
| 170 | + if self._optimizer is None: |
| 171 | + self._optimizer = Adam([self._w], lr=self.lr, weight_decay=self.weight_decay) |
| 172 | + self._w.grad = grad |
| 173 | + self._optimizer.step() |
| 174 | + # Clear the gradient so it cannot leak into a user optimizer that the logits were mistakenly |
| 175 | + # added to: FAMO is the only thing that should step them. |
| 176 | + self._w.grad = None |
| 177 | + |
| 178 | + def reset(self) -> None: |
| 179 | + with torch.no_grad(): |
| 180 | + self._w.zero_() |
| 181 | + self._optimizer = None |
| 182 | + self._prev_losses = None |
| 183 | + |
| 184 | + def _check_shape(self, values: Tensor) -> None: |
| 185 | + if values.shape != self._w.shape: |
| 186 | + raise ValueError( |
| 187 | + f"Parameter `values` should have shape {tuple(self._w.shape)} (matching the shape " |
| 188 | + f"of the logits). Found `values.shape = {tuple(values.shape)}`." |
| 189 | + ) |
| 190 | + |
| 191 | + def __repr__(self) -> str: |
| 192 | + return ( |
| 193 | + f"{self.__class__.__name__}(shape={tuple(self._w.shape)}, lr={self.lr}, " |
| 194 | + f"weight_decay={self.weight_decay})" |
| 195 | + ) |
0 commit comments