refactor(gradvac): base on GramianWeightedAggregator with GradVacWeighting

GradVac only needs gradient norms and dot products, which are fully
determined by the Gramian. This makes GradVac compatible with the autogram path.

- Remove grouping parameters (group_type, encoder, shared_params) from GradVac
- Export GradVacWeighting publicly

Author: rkhosrowshahi

docs/source/docs/aggregation/gradvac.rst

@@ -7,3 +7,8 @@ GradVac
     :members:
     :undoc-members:
     :exclude-members: forward
+
+.. autoclass:: torchjd.aggregation.GradVacWeighting
+    :members:
+    :undoc-members:
+    :exclude-members: forward

src/torchjd/aggregation/__init__.py

@@ -66,7 +66,7 @@
 from ._dualproj import DualProj, DualProjWeighting
 from ._flattening import Flattening
 from ._graddrop import GradDrop
-from ._gradvac import GradVac
+from ._gradvac import GradVac, GradVacWeighting
 from ._imtl_g import IMTLG, IMTLGWeighting
 from ._krum import Krum, KrumWeighting
 from ._mean import Mean, MeanWeighting
@@ -94,6 +94,7 @@
     "GeneralizedWeighting",
     "GradDrop",
     "GradVac",
+    "GradVacWeighting",
     "IMTLG",
     "IMTLGWeighting",
     "Krum",

src/torchjd/aggregation/_gradvac.py

@@ -1,158 +1,124 @@
 from __future__ import annotations
 
-from collections.abc import Iterable
-from typing import Literal, cast
+from typing import cast
 
 import torch
-import torch.nn as nn
 from torch import Tensor
 
-from torchjd._linalg import Matrix
+from torchjd._linalg import PSDMatrix
 
-from ._aggregator_bases import Aggregator
+from ._aggregator_bases import GramianWeightedAggregator
 from ._utils.non_differentiable import raise_non_differentiable_error
+from ._weighting_bases import Weighting
 
 
-def _all_layer_group_sizes(encoder: nn.Module) -> tuple[int, ...]:
-    """
-    Block sizes per leaf submodule with parameters, matching the ``all_layer`` grouping: iterate
-    ``encoder.modules()`` and append the total number of elements in each module that has no child
-    submodules and registers at least one parameter.
-    """
-
-    return tuple(
-        sum(w.numel() for w in module.parameters())
-        for module in encoder.modules()
-        if len(list(module.children())) == 0 and next(module.parameters(), None) is not None
-    )
-
-
-def _all_matrix_group_sizes(shared_params: Iterable[Tensor]) -> tuple[int, ...]:
-    """One block per tensor in ``shared_params`` order (``all_matrix`` / shared-parameter layout)."""
-
-    return tuple(p.numel() for p in shared_params)
-
-
-class GradVac(Aggregator):
+class GradVac(GramianWeightedAggregator):
     r"""
-    :class:`~torchjd.aggregation._aggregator_bases.Aggregator` implementing Gradient Vaccine
-    (GradVac) from `Gradient Vaccine: Investigating and Improving Multi-task Optimization in
-    Massively Multilingual Models (ICLR 2021 Spotlight)
+    :class:`~torchjd.aggregation._aggregator_bases.Aggregator` implementing the aggregation step of
+    Gradient Vaccine (GradVac) from `Gradient Vaccine: Investigating and Improving Multi-task
+    Optimization in Massively Multilingual Models (ICLR 2021 Spotlight)
     <https://openreview.net/forum?id=F1vEjWK-lH_>`_.
 
-    The input matrix is a Jacobian :math:`J \in \mathbb{R}^{m \times n}` whose rows are per-task
-    gradients. For each task :math:`i` and each parameter block :math:`k`, the order in which other
-    tasks :math:`j` are visited is drawn at random (independently for each :math:`k`); for each pair
-    :math:`(i, j)` on block :math:`k`, the cosine correlation :math:`\phi_{ijk}` between the
+    For each task :math:`i`, the order in which other tasks :math:`j` are visited is drawn at
+    random. For each pair :math:`(i, j)`, the cosine similarity :math:`\phi_{ij}` between the
     (possibly already modified) gradient of task :math:`i` and the original gradient of task
-    :math:`j` on that block is compared to an EMA target :math:`\hat{\phi}_{ijk}`. When
-    :math:`\phi_{ijk} < \hat{\phi}_{ijk}`, a closed-form correction adds a scaled copy of
-    :math:`g_j` to the block of :math:`g_i^{(\mathrm{PC})}`. The EMA is then updated with
-    :math:`\hat{\phi}_{ijk} \leftarrow (1-\beta)\hat{\phi}_{ijk} + \beta \phi_{ijk}`. The aggregated
+    :math:`j` is compared to an EMA target :math:`\hat{\phi}_{ij}`. When
+    :math:`\phi_{ij} < \hat{\phi}_{ij}`, a closed-form correction adds a scaled copy of
+    :math:`g_j` to :math:`g_i^{(\mathrm{PC})}`. The EMA is then updated with
+    :math:`\hat{\phi}_{ij} \leftarrow (1-\beta)\hat{\phi}_{ij} + \beta \phi_{ij}`. The aggregated
     vector is the sum of the modified rows.
 
     This aggregator is stateful: it keeps :math:`\hat{\phi}` across calls. Use :meth:`reset` when
-    the number of tasks, parameter dimension, grouping, device, or dtype changes.
-
-    **Parameter granularity** is selected by ``group_type`` (default ``"whole_model"``). It defines
-    how each task gradient row is partitioned into blocks :math:`k` so that cosines and EMA targets
-    :math:`\hat{\phi}_{ijk}` are computed **per block** rather than only globally:
-
-    * ``"whole_model"``: the full row of length :math:`n` is a single block. Cosine similarity is
-      taken between entire task gradients. Do not pass ``encoder`` or ``shared_params``.
-    * ``"all_layer"``: one block per leaf ``nn.Module`` under ``encoder`` that holds parameters
-      (same rule as iterating ``encoder.modules()`` and selecting leaves with parameters). Pass
-      ``encoder``; ``shared_params`` must be omitted.
-    * ``"all_matrix"``: one block per tensor in ``shared_params``, in iteration order. That order
-      must match how Jacobian columns are laid out for those shared parameters. Pass
-      ``shared_params``; ``encoder`` must be omitted.
-
-    :param beta: EMA decay for :math:`\hat{\phi}` (paper default ``0.5``). You may read or assign the
-        :attr:`beta` attribute between steps to tune the EMA update.
-    :param group_type: Granularity of parameter grouping; see **Parameter granularity** above.
-    :param encoder: Module whose subtree defines ``all_layer`` blocks when
-        ``group_type == "all_layer"``.
-    :param shared_params: Iterable of parameter tensors defining ``all_matrix`` block sizes and
-        order when ``group_type == "all_matrix"``. It is materialized once at construction.
+    the number of tasks or dtype changes.
+
+    :param beta: EMA decay for :math:`\hat{\phi}` (paper default ``0.5``). You may read or assign
+        the :attr:`beta` attribute between steps to tune the EMA update.
     :param eps: Small non-negative constant added to denominators when computing cosines and the
         vaccine weight (default ``1e-8``); set to ``0`` to omit this stabilization. You may read or
         assign the :attr:`eps` attribute between steps to tune numerical behavior.
 
     .. note::
-        GradVac is not compatible with autogram: it needs full Jacobian rows and per-block inner
-        products, not only a Gram matrix. Only the autojac path is supported.
+        For each task :math:`i`, the order of other tasks :math:`j` is shuffled independently
+        using the global PyTorch RNG (``torch.randperm``). Seed it with ``torch.manual_seed`` if
+        you need reproducibility.
 
     .. note::
-        For each task :math:`i` and block :math:`k`, the order of other tasks :math:`j` is shuffled
-        independently using the global PyTorch RNG (``torch.randperm``). Seed it with
-        ``torch.manual_seed`` if you need reproducibility.
+        To apply GradVac with per-layer or per-parameter-group granularity, first aggregate the
+        Jacobian into groups, apply GradVac per group, and sum the results. See the grouping usage
+        example for details.
+    """
+
+    def __init__(self, beta: float = 0.5, eps: float = 1e-8) -> None:
+        weighting = GradVacWeighting(beta=beta, eps=eps)
+        super().__init__(weighting)
+        self._gradvac_weighting = weighting
+        self.register_full_backward_pre_hook(raise_non_differentiable_error)
+
+    @property
+    def beta(self) -> float:
+        """EMA decay coefficient for :math:`\\hat{\\phi}` (paper default ``0.5``)."""
+
+        return self._gradvac_weighting.beta
+
+    @beta.setter
+    def beta(self, value: float) -> None:
+        self._gradvac_weighting.beta = value
+
+    @property
+    def eps(self) -> float:
+        """Small non-negative constant added to denominators for numerical stability."""
+
+        return self._gradvac_weighting.eps
+
+    @eps.setter
+    def eps(self, value: float) -> None:
+        self._gradvac_weighting.eps = value
+
+    def reset(self) -> None:
+        """Clears EMA state so the next forward starts from zero targets."""
+
+        self._gradvac_weighting.reset()
+
+    def __repr__(self) -> str:
+        return f"GradVac(beta={self.beta!r}, eps={self.eps!r})"
+
+
+class GradVacWeighting(Weighting[PSDMatrix]):
+    r"""
+    :class:`~torchjd.aggregation._weighting_bases.Weighting` giving the weights of
+    :class:`~torchjd.aggregation.GradVac`.
+
+    All required quantities (gradient norms, cosine similarities, and their updates after the
+    vaccine correction) are derived purely from the Gramian, without needing the full Jacobian.
+    If :math:`g_i^{(\mathrm{PC})} = \sum_k c_{ik} g_k`, then:
+
+    .. math::
+
+        \|g_i^{(\mathrm{PC})}\|^2 = \mathbf{c}_i G \mathbf{c}_i^\top,\qquad
+        g_i^{(\mathrm{PC})} \cdot g_j = \mathbf{c}_i G_{:,j}
+
+    where :math:`G` is the Gramian.  The correction :math:`g_i^{(\mathrm{PC})} \mathrel{+}= w
+    g_j` then becomes :math:`c_{ij} \mathrel{+}= w`, and the updated dot products follow
+    immediately.
+
+    This weighting is stateful: it keeps :math:`\hat{\phi}` across calls. Use :meth:`reset` when
+    the number of tasks or dtype changes.
+
+    :param beta: EMA decay for :math:`\hat{\phi}` (paper default ``0.5``).
+    :param eps: Small non-negative constant added to denominators (default ``1e-8``).
     """
 
-    def __init__(
-        self,
-        beta: float = 0.5,
-        group_type: Literal["whole_model", "all_layer", "all_matrix"] = "whole_model",
-        encoder: nn.Module | None = None,
-        shared_params: Iterable[Tensor] | None = None,
-        eps: float = 1e-8,
-    ) -> None:
+    def __init__(self, beta: float = 0.5, eps: float = 1e-8) -> None:
         super().__init__()
         if not (0.0 <= beta <= 1.0):
             raise ValueError(f"Parameter `beta` must be in [0, 1]. Found beta={beta!r}.")
-        params_tuple: tuple[Tensor, ...] = ()
-        fixed_block_sizes: tuple[int, ...] | None
-        if group_type == "whole_model":
-            if encoder is not None:
-                raise ValueError(
-                    'Parameter `encoder` must be None when `group_type == "whole_model"`.'
-                )
-            if shared_params is not None:
-                raise ValueError(
-                    'Parameter `shared_params` must be None when `group_type == "whole_model"`.'
-                )
-            fixed_block_sizes = None
-        elif group_type == "all_layer":
-            if encoder is None:
-                raise ValueError(
-                    'Parameter `encoder` is required when `group_type == "all_layer"`.'
-                )
-            if shared_params is not None:
-                raise ValueError(
-                    'Parameter `shared_params` must be None when `group_type == "all_layer"`.'
-                )
-            fixed_block_sizes = _all_layer_group_sizes(encoder)
-            if sum(fixed_block_sizes) == 0:
-                raise ValueError("Parameter `encoder` has no parameters in any leaf module.")
-        else:
-            if shared_params is None:
-                raise ValueError(
-                    'Parameter `shared_params` is required when `group_type == "all_matrix"`.'
-                )
-            if encoder is not None:
-                raise ValueError(
-                    'Parameter `encoder` must be None when `group_type == "all_matrix"`.'
-                )
-            params_tuple = tuple(shared_params)
-            if len(params_tuple) == 0:
-                raise ValueError(
-                    'Parameter `shared_params` must be non-empty when `group_type == "all_matrix"`.'
-                )
-            fixed_block_sizes = _all_matrix_group_sizes(params_tuple)
-
         if eps < 0.0:
             raise ValueError(f"Parameter `eps` must be non-negative. Found eps={eps!r}.")
 
         self._beta = beta
-        self._group_type = group_type
-        self._encoder = encoder
-        self._shared_params_len = len(params_tuple)
-        self._fixed_block_sizes = fixed_block_sizes
         self._eps = eps
-
         self._phi_t: Tensor | None = None
-        self._state_key: tuple[int, int, tuple[int, ...], torch.device, torch.dtype] | None = None
-
-        self.register_full_backward_pre_hook(raise_non_differentiable_error)
+        self._state_key: tuple[int, torch.dtype] | None = None
 
     @property
     def beta(self) -> float:
@@ -184,82 +150,56 @@ def reset(self) -> None:
         self._phi_t = None
         self._state_key = None
 
-    def __repr__(self) -> str:
-        enc = "None" if self._encoder is None else f"{self._encoder.__class__.__name__}(...)"
-        sp = "None" if self._group_type != "all_matrix" else f"n_params={self._shared_params_len}"
-        return (
-            f"{self.__class__.__name__}(beta={self._beta!r}, group_type={self._group_type!r}, "
-            f"encoder={enc}, shared_params={sp}, eps={self._eps!r})"
-        )
-
-    def _resolve_segment_sizes(self, n: int) -> tuple[int, ...]:
-        if self._group_type == "whole_model":
-            return (n,)
-        sizes = cast(tuple[int, ...], self._fixed_block_sizes)
-        if sum(sizes) != n:
-            raise ValueError(
-                "The Jacobian width `n` must equal the sum of block sizes implied by "
-                f"`encoder` or `shared_params` for this `group_type`. Found n={n}, "
-                f"sum(block_sizes)={sum(sizes)}.",
-            )
-        return sizes
-
-    def _ensure_state(
-        self,
-        m: int,
-        n: int,
-        sizes: tuple[int, ...],
-        device: torch.device,
-        dtype: torch.dtype,
-    ) -> None:
-        key = (m, n, sizes, device, dtype)
-        num_groups = len(sizes)
-        if self._state_key != key or self._phi_t is None:
-            self._phi_t = torch.zeros(m, m, num_groups, device=device, dtype=dtype)
-            self._state_key = key
+    def forward(self, gramian: PSDMatrix, /) -> Tensor:
+        device = gramian.device
+        dtype = gramian.dtype
+        cpu = torch.device("cpu")
 
-    def forward(self, matrix: Matrix, /) -> Tensor:
-        grads = matrix
-        m, n = grads.shape
-        if m == 0 or n == 0:
-            return torch.zeros(n, dtype=grads.dtype, device=grads.device)
+        G = cast(PSDMatrix, gramian.to(device=cpu))
+        m = G.shape[0]
 
-        sizes = self._resolve_segment_sizes(n)
-        device = grads.device
-        dtype = grads.dtype
-        self._ensure_state(m, n, sizes, device, dtype)
+        self._ensure_state(m, dtype)
         phi_t = cast(Tensor, self._phi_t)
-        beta = self.beta
-        eps = self.eps
 
-        pc_grads = grads.clone()
-        offsets = [0]
-        for s in sizes:
-            offsets.append(offsets[-1] + s)
+        beta = self._beta
+        eps = self._eps
+
+        # C[i, :] holds coefficients such that g_i^PC = sum_k C[i,k] * g_k (original gradients).
+        # Initially each modified gradient equals the original, so C = I.
+        C = torch.eye(m, device=cpu, dtype=dtype)
 
         for i in range(m):
+            # Dot products of g_i^PC with every original g_j, shape (m,).
+            cG = C[i] @ G
+
             others = [j for j in range(m) if j != i]
-            for k in range(len(sizes)):
-                perm = torch.randperm(len(others))
-                shuffled_js = [others[idx] for idx in perm.tolist()]
-                beg, end = offsets[k], offsets[k + 1]
-                for j in shuffled_js:
-                    slice_i = pc_grads[i, beg:end]
-                    slice_j = grads[j, beg:end]
-
-                    norm_i = slice_i.norm()
-                    norm_j = slice_j.norm()
-                    denom = norm_i * norm_j + eps
-                    phi_ijk = slice_i.dot(slice_j) / denom
-
-                    phi_hat = phi_t[i, j, k]
-                    if phi_ijk < phi_hat:
-                        sqrt_1_phi2 = (1.0 - phi_ijk * phi_ijk).clamp(min=0.0).sqrt()
-                        sqrt_1_hat2 = (1.0 - phi_hat * phi_hat).clamp(min=0.0).sqrt()
-                        denom_w = norm_j * sqrt_1_hat2 + eps
-                        w = norm_i * (phi_hat * sqrt_1_phi2 - phi_ijk * sqrt_1_hat2) / denom_w
-                        pc_grads[i, beg:end] = slice_i + slice_j * w
-
-                    phi_t[i, j, k] = (1.0 - beta) * phi_hat + beta * phi_ijk
-
-        return pc_grads.sum(dim=0)
+            perm = torch.randperm(len(others))
+            shuffled_js = [others[idx] for idx in perm.tolist()]
+
+            for j in shuffled_js:
+                dot_ij = cG[j]
+                norm_i_sq = (cG * C[i]).sum()
+                norm_i = norm_i_sq.clamp(min=0.0).sqrt()
+                norm_j = G[j, j].clamp(min=0.0).sqrt()
+                denom = norm_i * norm_j + eps
+                phi_ijk = dot_ij / denom
+
+                phi_hat = phi_t[i, j]
+                if phi_ijk < phi_hat:
+                    sqrt_1_phi2 = (1.0 - phi_ijk * phi_ijk).clamp(min=0.0).sqrt()
+                    sqrt_1_hat2 = (1.0 - phi_hat * phi_hat).clamp(min=0.0).sqrt()
+                    denom_w = norm_j * sqrt_1_hat2 + eps
+                    w = norm_i * (phi_hat * sqrt_1_phi2 - phi_ijk * sqrt_1_hat2) / denom_w
+                    C[i, j] = C[i, j] + w
+                    cG = cG + w * G[j]
+
+                phi_t[i, j] = (1.0 - beta) * phi_hat + beta * phi_ijk
+
+        weights = C.sum(dim=0)
+        return weights.to(device)
+
+    def _ensure_state(self, m: int, dtype: torch.dtype) -> None:
+        key = (m, dtype)
+        if self._state_key != key or self._phi_t is None:
+            self._phi_t = torch.zeros(m, m, dtype=dtype)
+            self._state_key = key