feat: Add scale_mode options to AlignedMTL

Expose median and rmse scaling modes for the balance transformation to match
original behavior while keeping min as default.

Author: rkhosrowshahi

src/torchjd/aggregation/_aligned_mtl.py

@@ -25,6 +25,8 @@
 # SOFTWARE.
 
 
+from typing import Literal
+
 import torch
 from torch import Tensor
 
@@ -44,18 +46,29 @@ class AlignedMTL(GramianWeightedAggregator):
 
     :param pref_vector: The preference vector to use.  If not provided, defaults to
         :math:`\begin{bmatrix} \frac{1}{m} & \dots & \frac{1}{m} \end{bmatrix}^T \in \mathbb{R}^m`.
+    :param scale_mode: The scaling mode used to build the balance transformation. ``"min"`` uses
+        the smallest eigenvalue (default), ``"median"`` uses the median eigenvalue, and ``"rmse"``
+        uses the mean eigenvalue (as in the original implementation).
 
     .. note::
         This implementation was adapted from the `official implementation
         <https://github.com/SamsungLabs/MTL/tree/master/code/optim/aligned>`_.
     """
 
-    def __init__(self, pref_vector: Tensor | None = None):
+    def __init__(
+        self,
+        pref_vector: Tensor | None = None,
+        scale_mode: Literal["min", "median", "rmse"] = "min",
+    ):
         self._pref_vector = pref_vector
-        super().__init__(AlignedMTLWeighting(pref_vector))
+        self._scale_mode = scale_mode
+        super().__init__(AlignedMTLWeighting(pref_vector, scale_mode=scale_mode))
 
     def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(pref_vector={repr(self._pref_vector)})"
+        return (
+            f"{self.__class__.__name__}(pref_vector={repr(self._pref_vector)}, "
+            f"scale_mode={repr(self._scale_mode)})"
+        )
 
     def __str__(self) -> str:
         return f"AlignedMTL{pref_vector_to_str_suffix(self._pref_vector)}"
@@ -68,22 +81,32 @@ class AlignedMTLWeighting(Weighting[PSDMatrix]):
 
     :param pref_vector: The preference vector to use. If not provided, defaults to
         :math:`\begin{bmatrix} \frac{1}{m} & \dots & \frac{1}{m} \end{bmatrix}^T \in \mathbb{R}^m`.
+    :param scale_mode: The scaling mode used to build the balance transformation. ``"min"`` uses
+        the smallest eigenvalue (default), ``"median"`` uses the median eigenvalue, and ``"rmse"``
+        uses the mean eigenvalue (as in the original implementation).
     """
 
-    def __init__(self, pref_vector: Tensor | None = None):
+    def __init__(
+        self,
+        pref_vector: Tensor | None = None,
+        scale_mode: Literal["min", "median", "rmse"] = "min",
+    ):
         super().__init__()
         self._pref_vector = pref_vector
+        self._scale_mode = scale_mode
         self.weighting = pref_vector_to_weighting(pref_vector, default=MeanWeighting())
 
     def forward(self, gramian: PSDMatrix) -> Tensor:
         w = self.weighting(gramian)
-        B = self._compute_balance_transformation(gramian)
+        B = self._compute_balance_transformation(gramian, self._scale_mode)
         alpha = B @ w
 
         return alpha
 
     @staticmethod
-    def _compute_balance_transformation(M: Tensor) -> Tensor:
+    def _compute_balance_transformation(
+        M: Tensor, scale_mode: Literal["min", "median", "rmse"] = "min"
+    ) -> Tensor:
         lambda_, V = torch.linalg.eigh(M, UPLO="U")  # More modern equivalent to torch.symeig
         tol = torch.max(lambda_) * len(M) * torch.finfo().eps
         rank = sum(lambda_ > tol)
@@ -96,6 +119,17 @@ def _compute_balance_transformation(M: Tensor) -> Tensor:
         lambda_, V = lambda_[order][:rank], V[:, order][:, :rank]
 
         sigma_inv = torch.diag(1 / lambda_.sqrt())
-        lambda_R = lambda_[-1]
-        B = lambda_R.sqrt() * V @ sigma_inv @ V.T
+
+        if scale_mode == "min":
+            scale = lambda_[-1]
+        elif scale_mode == "median":
+            scale = torch.median(lambda_)
+        elif scale_mode == "rmse":
+            scale = lambda_.mean()
+        else:
+            raise ValueError(
+                f"Invalid scale_mode={scale_mode!r}. Expected 'min', 'median', or 'rmse'."
+            )
+
+        B = scale.sqrt() * V @ sigma_inv @ V.T
         return B

tests/unit/aggregation/test_aligned_mtl.py

@@ -23,9 +23,9 @@ def test_permutation_invariant(aggregator: AlignedMTL, matrix: Tensor):
 
 def test_representations():
     A = AlignedMTL(pref_vector=None)
-    assert repr(A) == "AlignedMTL(pref_vector=None)"
+    assert repr(A) == "AlignedMTL(pref_vector=None, scale_mode='min')"
     assert str(A) == "AlignedMTL"
 
     A = AlignedMTL(pref_vector=torch.tensor([1.0, 2.0, 3.0], device="cpu"))
-    assert repr(A) == "AlignedMTL(pref_vector=tensor([1., 2., 3.]))"
+    assert repr(A) == "AlignedMTL(pref_vector=tensor([1., 2., 3.]), scale_mode='min')"
     assert str(A) == "AlignedMTL([1., 2., 3.])"