Address Valerian's review comments on ExcessMTL

Author: KhusPatel4450

docs/source/docs/aggregation/excess_mtl.rst

@@ -3,5 +3,8 @@
 ExcessMTL
 =========
 
+.. autoclass:: torchjd.aggregation.ExcessMTL
+    :members: __call__, reset
+
 .. autoclass:: torchjd.aggregation.ExcessMTLWeighting
     :members: __call__, reset

src/torchjd/aggregation/__init__.py

@@ -45,7 +45,7 @@
 from ._constant import Constant, ConstantWeighting
 from ._cr_mogm import CRMOGMWeighting
 from ._dualproj import DualProj, DualProjWeighting
-from ._excess_mtl import ExcessMTLWeighting
+from ._excess_mtl import ExcessMTL, ExcessMTLWeighting
 from ._fairgrad import FairGrad, FairGradWeighting
 from ._graddrop import GradDrop
 from ._gradvac import GradVac, GradVacWeighting
@@ -75,6 +75,7 @@
     "CRMOGMWeighting",
     "DualProj",
     "DualProjWeighting",
+    "ExcessMTL",
     "ExcessMTLWeighting",
     "FairGrad",
     "FairGradWeighting",

src/torchjd/aggregation/_excess_mtl.py

@@ -11,6 +11,7 @@
 from torchjd.aggregation._mixins import _NonDifferentiable
 from torchjd.linalg import Matrix
 
+from ._aggregator_bases import WeightedAggregator
 from ._weighting_bases import _MatrixWeighting
 
 
@@ -29,64 +30,37 @@ class ExcessMTLWeighting(_MatrixWeighting, Stateful, _NonDifferentiable):
         Must be positive.
     :param n_warmup_steps: Number of forward calls during which weights stay uniform
         (:math:`[1/m, \ldots, 1/m]`) and gradient statistics are collected. The baseline excess
-        risk is set to the average excess risk observed during warmup. When ``0`` (default), the
-        first call's excess risk is used as the baseline and weights are updated immediately
-        (matching the official implementation).
+        risk is then set to the average excess risk observed during warmup. When ``0``, the first
+        call's excess risk is used immediately as the baseline. The default ``1`` matches the
+        behavior of the official implementation and LibMTL. The paper (Appendix C.1) recommends
+        collecting statistics for 3 full epochs, i.e. ``n_warmup_steps = 3 * len(dataloader)``.
 
     .. warning::
         The state tensor :math:`S \in \mathbb{R}^{m \times n}` accumulates squared gradients
-        across **all** calls, where :math:`n` is the total number of model parameters. For large
+        across calls, where :math:`n` is the total number of model parameters. For large
         models this can be a significant memory cost. Call :meth:`reset` between experiments.
 
     .. note::
         The weight update is adapted from the `official implementation
         <https://github.com/uiuctml/ExcessMTL>`_ and `LibMTL
         <https://github.com/median-research-group/LibMTL/blob/main/LibMTL/weighting/ExcessMTL.py>`_.
-        The warmup strategy follows Appendix C.1 of the paper, which recommends collecting
-        gradient statistics for several epochs before beginning weight updates; set
-        ``n_warmup_steps`` accordingly (e.g. ``3 * len(dataloader)``).
-
-    .. admonition:: Example
-
-        .. testcode::
-
-            import torch
-            from torch.nn import Linear, MSELoss, ReLU, Sequential
-            from torch.optim import SGD
-
-            from torchjd import autojac
-            from torchjd.aggregation import ExcessMTLWeighting, WeightedAggregator
-            from torchjd.autojac import jac_to_grad
-
-            inputs = torch.randn(8, 5)
-            targets = torch.randn(8, 2)
-
-            model = Sequential(Linear(5, 4), ReLU(), Linear(4, 2))
-            optimizer = SGD(model.parameters())
-            criterion = MSELoss()
-            aggregator = WeightedAggregator(ExcessMTLWeighting())
-
-            outputs = model(inputs)
-            losses = [criterion(outputs[:, i], targets[:, i]) for i in range(2)]
-            autojac.backward(losses)
-            jac_to_grad(model.parameters(), aggregator)
-            optimizer.step()
-            optimizer.zero_grad()
+        Unlike those implementations, which initialize task weights to ``1``, we follow the paper
+        and initialize them to ``1/m`` so that they always lie on the probability simplex.
     """
 
     def __init__(
         self,
         robust_step_size: float = 1.0,
-        n_warmup_steps: int = 0,
+        n_warmup_steps: int = 1,
     ) -> None:
         super().__init__()
         self.robust_step_size = robust_step_size
         self.n_warmup_steps = n_warmup_steps
         self.register_buffer("_weights", None)
-        self.register_buffer("_grad_sum", None)
+        self.register_buffer("_sq_grad_sum", None)
         self.register_buffer("_initial_w", None)
         self.register_buffer("_warmup_w_sum", None)
-        self.register_buffer("_n_steps", torch.zeros((), dtype=torch.long))
+        self._n_steps: int = 0
         self._state_key: tuple[int, int, torch.dtype, torch.device] | None = None
 
     @property
@@ -118,31 +92,31 @@ def reset(self) -> None:
         warmup."""
 
         self._weights = None
-        self._grad_sum = None
+        self._sq_grad_sum = None
         self._initial_w = None
         self._warmup_w_sum = None
-        self._n_steps.zero_()
+        self._n_steps = 0
         self._state_key = None
 
     def forward(self, matrix: Matrix, /) -> Tensor:
         self._ensure_state(matrix)
 
+        sq_matrix = matrix.detach() ** 2
+
         # Accumulate squared gradients for AdaGrad-style diagonal Hessian (Equation 7)
-        grad_sum = cast(Tensor, self._grad_sum)
-        grad_sum = grad_sum + matrix.detach() ** 2
-        self._grad_sum = grad_sum
+        sq_grad_sum = cast(Tensor, self._sq_grad_sum)
+        sq_grad_sum += sq_matrix
 
         # Excess risk proxy: Ê_i ≈ g_i^T H_i^{-1} g_i (Equation 6)
-        h = torch.sqrt(grad_sum + 1e-7)
-        w = (matrix.detach() ** 2 / h).sum(dim=1)  # shape [m]
+        h = torch.sqrt(sq_grad_sum + 1e-7)
+        w = (sq_matrix / h).sum(dim=1)  # shape [m]
 
-        n_steps = int(self._n_steps.item())
-        self._n_steps = self._n_steps + 1
+        n_steps = self._n_steps
+        self._n_steps += 1
 
         # Warmup: collect excess risk stats but return uniform weights
         if n_steps < self._n_warmup_steps:
-            warmup_w_sum = self._warmup_w_sum
-            self._warmup_w_sum = w if warmup_w_sum is None else cast(Tensor, warmup_w_sum) + w
+            cast(Tensor, self._warmup_w_sum).add_(w)
             return cast(Tensor, self._weights)
 
         # Set baseline on the first non-warmup call
@@ -152,7 +126,7 @@ def forward(self, matrix: Matrix, /) -> Tensor:
                 self._initial_w = cast(Tensor, self._warmup_w_sum) / self._n_warmup_steps
                 w = w / (cast(Tensor, self._initial_w) + 1e-7)
             else:
-                # Official impl behaviour: first call's excess is the baseline; use w raw
+                # Official impl behavior: first call's excess is the baseline; use w raw
                 self._initial_w = w
         else:
             w = w / (cast(Tensor, self._initial_w) + 1e-7)
@@ -166,14 +140,14 @@ def forward(self, matrix: Matrix, /) -> Tensor:
 
     def _ensure_state(self, matrix: Matrix) -> None:
         key = (matrix.shape[0], matrix.shape[1], matrix.dtype, matrix.device)
-        if self._state_key == key and self._grad_sum is not None:
+        if self._state_key == key and self._sq_grad_sum is not None:
             return
         m, n = matrix.shape
-        self._grad_sum = matrix.new_zeros(m, n)
+        self._sq_grad_sum = matrix.new_zeros(m, n)
+        self._warmup_w_sum = matrix.new_zeros(m)
         self._weights = matrix.new_full((m,), 1.0 / m)
         self._initial_w = None
-        self._warmup_w_sum = None
-        self._n_steps.zero_()
+        self._n_steps = 0
         self._state_key = key
 
     def __repr__(self) -> str:
@@ -182,3 +156,58 @@ def __repr__(self) -> str:
             f"robust_step_size={self.robust_step_size!r}, "
             f"n_warmup_steps={self.n_warmup_steps!r})"
         )
+
+
+class ExcessMTL(WeightedAggregator, Stateful, _NonDifferentiable):
+    r"""
+    :class:`~torchjd.aggregation.WeightedAggregator` from `Robust Multi-Task Learning with Excess
+    Risks <https://proceedings.mlr.press/v235/he24n.html>`_ (ICML 2024).
+
+    At each call, task weights are updated via an exponentiated gradient step (Equation 9) driven
+    by per-task excess risk estimates. See :class:`~torchjd.aggregation.ExcessMTLWeighting` for
+    details on the algorithm and state management.
+
+    :param robust_step_size: Step size :math:`\eta_\alpha` for the exponentiated weight update.
+        Must be positive.
+    :param n_warmup_steps: Number of forward calls during which weights stay uniform
+        (:math:`[1/m, \ldots, 1/m]`) and gradient statistics are collected. When ``0``, the first
+        call's excess risk is used as the baseline immediately. Defaults to ``1``.
+    """
+
+    weighting: ExcessMTLWeighting
+
+    def __init__(
+        self,
+        robust_step_size: float = 1.0,
+        n_warmup_steps: int = 1,
+    ) -> None:
+        super().__init__(ExcessMTLWeighting(robust_step_size, n_warmup_steps))
+
+    @property
+    def robust_step_size(self) -> float:
+        return self.weighting.robust_step_size
+
+    @robust_step_size.setter
+    def robust_step_size(self, value: float) -> None:
+        self.weighting.robust_step_size = value
+
+    @property
+    def n_warmup_steps(self) -> int:
+        return self.weighting.n_warmup_steps
+
+    @n_warmup_steps.setter
+    def n_warmup_steps(self, value: int) -> None:
+        self.weighting.n_warmup_steps = value
+
+    def reset(self) -> None:
+        """Clears all state so the next forward starts from uniform weights and re-enters
+        warmup."""
+
+        self.weighting.reset()
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            f"robust_step_size={self.robust_step_size!r}, "
+            f"n_warmup_steps={self.n_warmup_steps!r})"
+        )

tests/unit/aggregation/test_excess_mtl.py

@@ -88,23 +88,6 @@ def test_weights_change_after_warmup() -> None:
     assert weights[0] > weights[1]
 
 
-def test_update_recurrence() -> None:
-    """Verify the first weight update manually (n_warmup_steps=0, LibMTL behaviour).
-
-    With J = [[2., 0.], [1., 0.]] and robust_step_size=1.0:
-      grad_sum = [[4., 0.], [1., 0.]]
-      h ≈ [[2., sqrt(eps)], [1., sqrt(eps)]]  (eps = 1e-7, negligible in float32 for nonzero entries)
-      w = [4/2 + 0, 1/1 + 0] = [2, 1]
-      initial_w = [2, 1]   (first call: save raw excess as baseline)
-      weights = [exp(2), exp(1)] / (exp(2) + exp(1))
-    """
-    J = tensor_([[2.0, 0.0], [1.0, 0.0]])
-    W = ExcessMTLWeighting(robust_step_size=1.0)
-    e2 = torch.exp(tensor_(2.0))
-    e1 = torch.exp(tensor_(1.0))
-    assert_close(W(J), tensor_([e2 / (e2 + e1), e1 / (e2 + e1)]))
-
-
 def test_two_consecutive_steps() -> None:
     """Verify warm-started carry-over across two calls.
 
@@ -119,7 +102,7 @@ def test_two_consecutive_steps() -> None:
     """
     J1 = tensor_([[2.0, 0.0], [1.0, 0.0]])
     J2 = tensor_([[1.0, 0.0], [2.0, 0.0]])
-    W = ExcessMTLWeighting(robust_step_size=1.0)
+    W = ExcessMTLWeighting(robust_step_size=1.0, n_warmup_steps=0)
 
     e2 = torch.exp(tensor_(2.0))
     e1 = torch.exp(tensor_(1.0))
@@ -140,17 +123,17 @@ def test_warmup_baseline_is_average() -> None:
 
     With n_warmup_steps=2 and J1=[[2,0],[1,0]], J2=[[1,0],[2,0]]:
 
-    Warmup call 1 — grad_sum_1 = J1**2 = [[4,0],[1,0]]:
+    Warmup call 1 — sq_grad_sum_1 = J1**2 = [[4,0],[1,0]]:
       h_1 ≈ [[2, sqrt(eps)], [1, sqrt(eps)]]
       w_1 = [4/2, 1/1] = [2, 1]
 
-    Warmup call 2 — grad_sum_2 = J1**2 + J2**2 = [[5,0],[5,0]]:
+    Warmup call 2 — sq_grad_sum_2 = J1**2 + J2**2 = [[5,0],[5,0]]:
       h_2 ≈ [[sqrt(5), sqrt(eps)], [sqrt(5), sqrt(eps)]]
       w_2 = [1/sqrt(5), 4/sqrt(5)]
 
     initial_w = (w_1 + w_2) / 2  (Appendix C.1 average)
 
-    Post-warmup call 3 with J3 = J1 — grad_sum_3 = [[9,0],[6,0]]:
+    Post-warmup call 3 with J3 = J1 — sq_grad_sum_3 = [[9,0],[6,0]]:
       h_3 ≈ [[3, sqrt(eps)], [sqrt(6), sqrt(eps)]]
       w_3 = [4/3, 1/sqrt(6)]
       w_norm = w_3 / (initial_w + 1e-7)
@@ -162,21 +145,21 @@ def test_warmup_baseline_is_average() -> None:
     J3 = tensor_([[2.0, 0.0], [1.0, 0.0]])
     W = ExcessMTLWeighting(n_warmup_steps=2, robust_step_size=1.0)
 
-    W(J1)  # warmup step 1 — grad_sum becomes J1**2
-    W(J2)  # warmup step 2 — grad_sum becomes J1**2 + J2**2
+    W(J1)  # warmup step 1 — sq_grad_sum becomes J1**2
+    W(J2)  # warmup step 2 — sq_grad_sum becomes J1**2 + J2**2
 
-    grad_sum_1 = J1**2
-    h_1 = torch.sqrt(grad_sum_1 + 1e-7)
+    sq_grad_sum_1 = J1**2
+    h_1 = torch.sqrt(sq_grad_sum_1 + 1e-7)
     w_1 = (J1**2 / h_1).sum(dim=1)
 
-    grad_sum_2 = grad_sum_1 + J2**2
-    h_2 = torch.sqrt(grad_sum_2 + 1e-7)
+    sq_grad_sum_2 = sq_grad_sum_1 + J2**2
+    h_2 = torch.sqrt(sq_grad_sum_2 + 1e-7)
     w_2 = (J2**2 / h_2).sum(dim=1)
 
     initial_w = (w_1 + w_2) / 2  # Appendix C.1 baseline
 
-    grad_sum_3 = grad_sum_2 + J3**2
-    h_3 = torch.sqrt(grad_sum_3 + 1e-7)
+    sq_grad_sum_3 = sq_grad_sum_2 + J3**2
+    h_3 = torch.sqrt(sq_grad_sum_3 + 1e-7)
     w_3 = (J3**2 / h_3).sum(dim=1)
     w_norm = w_3 / (initial_w + 1e-7)
     pre_norm = tensor_([0.5, 0.5]) * torch.exp(w_norm)