refactor(aggregation): Refactor weighting structure (#347)

* Move weighting base classes to _weighting_bases.py and aggregator base classes to aggregator_bases.py
* Add Matrix and PSDMatrix annotated types (not fully used yet)
* Make Weighting generic on the type of input stat it takes
* Add Composition
* Add _GramianWeightedAggregator as an alias for a _WeightedAggregator that composes its (gramian-based) weighting with the compute_gramian function
* Change gramian-based Weightings to not compute the gramian themselves
* Change gramian-based Aggregators to be _GramianWeightedAggregator
* Make some aggregators store some of their weighting's parameters to have access to them in __str__ and __repr__
* Adapt some weighting tests
---------
Co-authored-by: Pierre Quinton <pierre.quinton@gmail.com>

Author: ValerianRey

docs/source/docs/aggregation/bases.rst

@@ -3,7 +3,7 @@
 Aggregator (abstract)
 =====================
 
-.. automodule:: torchjd.aggregation.bases
+.. automodule:: torchjd.aggregation.aggregator_bases
     :members:
     :undoc-members:
     :show-inheritance:

src/torchjd/aggregation/__init__.py

@@ -1,8 +1,8 @@
 from ._utils.check_dependencies import (
     OptionalDepsNotInstalledError as _OptionalDepsNotInstalledError,
 )
+from .aggregator_bases import Aggregator
 from .aligned_mtl import AlignedMTL
-from .bases import Aggregator
 from .config import ConFIG
 from .constant import Constant
 from .dualproj import DualProj

src/torchjd/aggregation/_utils/pref_vector.py

@@ -1,12 +1,14 @@
 from torch import Tensor
 
-from torchjd.aggregation.bases import _Weighting
+from torchjd.aggregation._weighting_bases import Matrix, Weighting
 from torchjd.aggregation.constant import _ConstantWeighting
 
 from .str import vector_to_str
 
 
-def pref_vector_to_weighting(pref_vector: Tensor | None, default: _Weighting) -> _Weighting:
+def pref_vector_to_weighting(
+    pref_vector: Tensor | None, default: Weighting[Matrix]
+) -> Weighting[Matrix]:
     """
     Returns the weighting associated to a given preference vector, with a fallback to a default
     weighting if the preference vector is None.

src/torchjd/aggregation/_weighting_bases.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Annotated, Callable, Generic, TypeVar
+
+from torch import Tensor, nn
+
+_T = TypeVar("_T", contravariant=True)
+_FnInputT = TypeVar("_FnInputT")
+_FnOutputT = TypeVar("_FnOutputT")
+Matrix = Annotated[Tensor, "ndim=2"]
+PSDMatrix = Annotated[Matrix, "Positive semi-definite"]
+
+
+class Weighting(Generic[_T], nn.Module, ABC):
+    r"""
+    Abstract base class for all weighting methods. It has the role of extracting a vector of weights
+    of dimension :math:`m` from some statistic of a matrix of dimension :math:`m \times n`.
+    """
+
+    @abstractmethod
+    def forward(self, stat: _T) -> Tensor:
+        """Computes the vector of weights from the input stat."""
+
+    # Override to make type hints and documentation more specific
+    def __call__(self, stat: _T) -> Tensor:
+        """Computes the vector of weights from the input stat and applies all registered hooks."""
+
+        return super().__call__(stat)
+
+    def _compose(self, fn: Callable[[_FnInputT], _T]) -> Weighting[_FnInputT]:
+        return _Composition(self, fn)
+
+    __lshift__ = _compose
+
+
+class _Composition(Weighting[_T]):
+    """
+    Weighting that composes a Weighting with a function, so that the Weighting is applied to the
+    output of the function.
+    """
+
+    def __init__(self, weighting: Weighting[_FnOutputT], fn: Callable[[_T], _FnOutputT]):
+        super().__init__()
+        self.fn = fn
+        self.weighting = weighting
+
+    def forward(self, stat: _T) -> Tensor:
+        return self.weighting(self.fn(stat))

src/torchjd/aggregation/bases.py …/torchjd/aggregation/aggregator_bases.pysrc/torchjd/aggregation/bases.py renamed to src/torchjd/aggregation/aggregator_bases.py src/torchjd/aggregation/bases.py renamed to src/torchjd/aggregation/aggregator_bases.py

@@ -2,6 +2,9 @@
 
 from torch import Tensor, nn
 
+from ._utils.gramian import compute_gramian
+from ._weighting_bases import Matrix, PSDMatrix, Weighting
+
 
 class Aggregator(nn.Module, ABC):
     r"""
@@ -42,35 +45,15 @@ def __str__(self) -> str:
         return f"{self.__class__.__name__}"
 
 
-class _Weighting(nn.Module, ABC):
-    r"""
-    Abstract base class for all weighting methods. It has the role of extracting a vector of weights
-    of dimension :math:`m` from a matrix of dimension :math:`m \times n`.
-    """
-
-    def __init__(self):
-        super().__init__()
-
-    @abstractmethod
-    def forward(self, matrix: Tensor) -> Tensor:
-        """Computes the vector of weights from the input matrix."""
-
-    # Override to make type hints and documentation more specific
-    def __call__(self, matrix: Tensor) -> Tensor:
-        """Computes the vector of weights from the input matrix and applies all registered hooks."""
-
-        return super().__call__(matrix)
-
-
 class _WeightedAggregator(Aggregator):
     """
-    :class:`~torchjd.aggregation.bases.Aggregator` that combines the rows of the input matrix with
-    weights given by applying a :class:`~torchjd.aggregation.bases._Weighting` to the matrix.
+    Aggregator that combines the rows of the input jacobian matrix with weights given by applying a
+    Weighting to it.
 
     :param weighting: The object responsible for extracting the vector of weights from the matrix.
     """
 
-    def __init__(self, weighting: _Weighting):
+    def __init__(self, weighting: Weighting[Matrix]):
         super().__init__()
         self.weighting = weighting
 
@@ -91,3 +74,15 @@ def forward(self, matrix: Tensor) -> Tensor:
         weights = self.weighting(matrix)
         vector = self.combine(matrix, weights)
         return vector
+
+
+class _GramianWeightedAggregator(_WeightedAggregator):
+    """
+    WeightedAggregator that computes the gramian of the input jacobian matrix before applying a
+    Weighting to it.
+
+    :param weighting: The object responsible for extracting the vector of weights from the gramian.
+    """
+
+    def __init__(self, weighting: Weighting[PSDMatrix]):
+        super().__init__(weighting << compute_gramian)

src/torchjd/aggregation/aligned_mtl.py

@@ -28,15 +28,15 @@
 import torch
 from torch import Tensor
 
-from ._utils.gramian import compute_gramian
 from ._utils.pref_vector import pref_vector_to_str_suffix, pref_vector_to_weighting
-from .bases import _WeightedAggregator, _Weighting
+from ._weighting_bases import PSDMatrix, Weighting
+from .aggregator_bases import _GramianWeightedAggregator
 from .mean import _MeanWeighting
 
 
-class AlignedMTL(_WeightedAggregator):
+class AlignedMTL(_GramianWeightedAggregator):
     """
-    :class:`~torchjd.aggregation.bases.Aggregator` as defined in Algorithm 1 of
+    :class:`~torchjd.aggregation.aggregator_bases.Aggregator` as defined in Algorithm 1 of
     `Independent Component Alignment for Multi-Task Learning
     <https://openaccess.thecvf.com/content/CVPR2023/papers/Senushkin_Independent_Component_Alignment_for_Multi-Task_Learning_CVPR_2023_paper.pdf>`_.
 
@@ -65,7 +65,7 @@ def __init__(self, pref_vector: Tensor | None = None):
         weighting = pref_vector_to_weighting(pref_vector, default=_MeanWeighting())
         self._pref_vector = pref_vector
 
-        super().__init__(weighting=_AlignedMTLWrapper(weighting))
+        super().__init__(_AlignedMTLWrapper(weighting))
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(pref_vector={repr(self._pref_vector)})"
@@ -74,26 +74,24 @@ def __str__(self) -> str:
         return f"AlignedMTL{pref_vector_to_str_suffix(self._pref_vector)}"
 
 
-class _AlignedMTLWrapper(_Weighting):
+class _AlignedMTLWrapper(Weighting[PSDMatrix]):
     """
-    Wrapper of :class:`~torchjd.aggregation.bases._Weighting` that corrects the extracted
+    Wrapper of :class:`~torchjd.aggregation._weighting_bases.Weighting` that corrects the extracted
     weights with the balance transformation defined in Algorithm 1 of `Independent Component
     Alignment for Multi-Task Learning
     <https://openaccess.thecvf.com/content/CVPR2023/papers/Senushkin_Independent_Component_Alignment_for_Multi-Task_Learning_CVPR_2023_paper.pdf>`_.
 
-    :param weighting: The wrapped :class:`~torchjd.aggregation.bases._Weighting`
+    :param weighting: The wrapped :class:`~torchjd.aggregation._weighting_bases.Weighting`
         responsible for extracting weight vectors from the input matrices.
     """
 
-    def __init__(self, weighting: _Weighting):
+    def __init__(self, weighting: Weighting[PSDMatrix]):
         super().__init__()
         self.weighting = weighting
 
-    def forward(self, matrix: Tensor) -> Tensor:
-        w = self.weighting(matrix)
-
-        M = compute_gramian(matrix)
-        B = self._compute_balance_transformation(M)
+    def forward(self, gramian: Tensor) -> Tensor:
+        w = self.weighting(gramian)
+        B = self._compute_balance_transformation(gramian)
         alpha = B @ w
 
         return alpha

src/torchjd/aggregation/cagrad.py

@@ -1,4 +1,5 @@
 from ._utils.check_dependencies import check_dependencies_are_installed
+from ._weighting_bases import PSDMatrix, Weighting
 
 check_dependencies_are_installed(["cvxpy", "clarabel"])
 
@@ -7,14 +8,14 @@
 import torch
 from torch import Tensor
 
-from ._utils.gramian import compute_gramian, normalize
+from ._utils.gramian import normalize
 from ._utils.non_differentiable import raise_non_differentiable_error
-from .bases import _WeightedAggregator, _Weighting
+from .aggregator_bases import _GramianWeightedAggregator
 
 
-class CAGrad(_WeightedAggregator):
+class CAGrad(_GramianWeightedAggregator):
     """
-    :class:`~torchjd.aggregation.bases.Aggregator` as defined in Algorithm 1 of
+    :class:`~torchjd.aggregation.aggregator_bases.Aggregator` as defined in Algorithm 1 of
     `Conflict-Averse Gradient Descent for Multi-task Learning
     <https://arxiv.org/pdf/2110.14048.pdf>`_.
 
@@ -43,24 +44,24 @@ class CAGrad(_WeightedAggregator):
     """
 
     def __init__(self, c: float, norm_eps: float = 0.0001):
-        super().__init__(weighting=_CAGradWeighting(c=c, norm_eps=norm_eps))
+        super().__init__(_CAGradWeighting(c=c, norm_eps=norm_eps))
+        self._c = c
+        self._norm_eps = norm_eps
 
         # This prevents considering the computed weights as constant w.r.t. the matrix.
         self.register_full_backward_pre_hook(raise_non_differentiable_error)
 
     def __repr__(self) -> str:
-        return (
-            f"{self.__class__.__name__}(c={self.weighting.c}, norm_eps={self.weighting.norm_eps})"
-        )
+        return f"{self.__class__.__name__}(c={self._c}, norm_eps={self._norm_eps})"
 
     def __str__(self) -> str:
-        c_str = str(self.weighting.c).rstrip("0")
+        c_str = str(self._c).rstrip("0")
         return f"CAGrad{c_str}"
 
 
-class _CAGradWeighting(_Weighting):
+class _CAGradWeighting(Weighting[PSDMatrix]):
     """
-    :class:`~torchjd.aggregation.bases._Weighting` that extracts weights using the CAGrad
+    :class:`~torchjd.aggregation._weighting_bases.Weighting` that extracts weights using the CAGrad
     algorithm, as defined in algorithm 1 of `Conflict-Averse Gradient Descent for Multi-task
     Learning <https://arxiv.org/pdf/2110.14048.pdf>`_.
 
@@ -85,11 +86,7 @@ def __init__(self, c: float, norm_eps: float):
         self.c = c
         self.norm_eps = norm_eps
 
-    def forward(self, matrix: Tensor) -> Tensor:
-        gramian = compute_gramian(matrix)
-        return self._compute_from_gramian(gramian)
-
-    def _compute_from_gramian(self, gramian: Tensor) -> Tensor:
+    def forward(self, gramian: Tensor) -> Tensor:
         U, S, _ = torch.svd(normalize(gramian, self.norm_eps))
 
         reduced_matrix = U @ S.sqrt().diag()

src/torchjd/aggregation/config.py

@@ -30,14 +30,15 @@
 
 from ._utils.non_differentiable import raise_non_differentiable_error
 from ._utils.pref_vector import pref_vector_to_str_suffix, pref_vector_to_weighting
-from .bases import Aggregator
+from .aggregator_bases import Aggregator
 from .sum import _SumWeighting
 
 
 class ConFIG(Aggregator):
     """
-    :class:`~torchjd.aggregation.bases.Aggregator` as defined in Equation 2 of `ConFIG: Towards
-    Conflict-free Training of Physics Informed Neural Networks <https://arxiv.org/pdf/2408.11104>`_.
+    :class:`~torchjd.aggregation.aggregator_bases.Aggregator` as defined in Equation 2 of `ConFIG:
+    Towards Conflict-free Training of Physics Informed Neural Networks
+    <https://arxiv.org/pdf/2408.11104>`_.
 
     :param pref_vector: The preference vector used to weight the rows. If not provided, defaults to
         equal weights of 1.

src/torchjd/aggregation/constant.py

@@ -1,13 +1,14 @@
 from torch import Tensor
 
 from ._utils.str import vector_to_str
-from .bases import _WeightedAggregator, _Weighting
+from ._weighting_bases import Matrix, Weighting
+from .aggregator_bases import _WeightedAggregator
 
 
 class Constant(_WeightedAggregator):
     """
-    :class:`~torchjd.aggregation.bases.Aggregator` that makes a linear combination of the rows of
-    the provided matrix, with constant, pre-determined weights.
+    :class:`~torchjd.aggregation.aggregator_bases.Aggregator` that makes a linear combination of the
+    rows of the provided matrix, with constant, pre-determined weights.
 
     :param weights: The weights associated to the rows of the input matrices.
 
@@ -37,9 +38,9 @@ def __str__(self) -> str:
         return f"{self.__class__.__name__}([{weights_str}])"
 
 
-class _ConstantWeighting(_Weighting):
+class _ConstantWeighting(Weighting[Matrix]):
     """
-    :class:`~torchjd.aggregation.bases._Weighting` that returns constant, pre-determined
+    :class:`~torchjd.aggregation._weighting_bases.Weighting` that returns constant, pre-determined
     weights.
 
     :param weights: The weights associated to the rows of the input matrices.