Implements a first version of projected gradient descent with decaying regularization.

Author: PierreQuinton

src/torchjd/aggregation/_dual_cone_utils.py

@@ -1,60 +1,30 @@
-from typing import Literal
-
-import numpy as np
 import torch
-from qpsolvers import solve_qp
 from torch import Tensor
 
 
-def project_weights(U: Tensor, G: Tensor, solver: Literal["quadprog"]) -> Tensor:
+def project_weights(U: Tensor, G: Tensor, max_iter: int = 200, eps: float = 1e-08) -> Tensor:
     """
     Computes the tensor of weights corresponding to the projection of the vectors in `U` onto the
     rows of a matrix whose Gramian is provided.
 
     :param U: The tensor of weights corresponding to the vectors to project, of shape `[..., m]`.
     :param G: The Gramian matrix of shape `[m, m]`. It must be symmetric and positive definite.
-    :param solver: The quadratic programming solver to use.
     :return: A tensor of projection weights with the same shape as `U`.
     """
 
-    G_ = _to_array(G)
-    U_ = _to_array(U)
-
-    W = np.apply_along_axis(lambda u: _project_weight_vector(u, G_, solver), axis=-1, arr=U_)
-
-    return torch.as_tensor(W, device=G.device, dtype=G.dtype)
-
-
-def _project_weight_vector(u: np.ndarray, G: np.ndarray, solver: Literal["quadprog"]) -> np.ndarray:
-    r"""
-    Computes the weights `w` of the projection of `J^T u` onto the dual cone of the rows of `J`,
-    given `G = J J^T` and `u`. In other words, this computes the `w` that satisfies
-    `\pi_J(J^T u) = J^T w`, with `\pi_J` defined in Equation 3 of [1].
-
-    By Proposition 1 of [1], this is equivalent to solving for `v` the following quadratic program:
-    minimize        v^T G v
-    subject to      u \preceq v
-
-    Reference:
-    [1] `Jacobian Descent For Multi-Objective Optimization <https://arxiv.org/pdf/2406.16232>`_.
-
-    :param u: The vector of weights `u` of shape `[m]` corresponding to the vector `J^T u` to
-        project.
-    :param G: The Gramian matrix of `J`, equal to `J J^T`, and of shape `[m, m]`. It must be
-        symmetric and positive definite.
-    :param solver: The quadratic programming solver to use.
-    """
-
-    m = G.shape[0]
-    w = solve_qp(G, np.zeros(m), -np.eye(m), -u, solver=solver)
-
-    if w is None:  # This may happen when G has large values.
-        raise ValueError("Failed to solve the quadratic programming problem.")
-
-    return w
-
-
-def _to_array(tensor: Tensor) -> np.ndarray:
-    """Transforms a tensor into a numpy array with float64 dtype."""
-
-    return tensor.cpu().detach().numpy().astype(np.float64)
+    shape = U.shape
+    m = shape[-1]
+    U_matrix = U.reshape([-1, m]).T
+    V = U_matrix.clone()
+
+    # torch.linalg.eigvals synchronizes G on the CPU.
+    lambda_max = torch.max(torch.linalg.eigvals(G).real).item()
+    for t in range(1, max_iter + 1):
+        step_size = 1.0 / (lambda_max + 1.0 / t)
+        V_new = torch.maximum(V - step_size * (G @ V), U_matrix)
+        gap = (V - V_new).norm()
+        if gap < eps:
+            print(t)
+            break
+        V = V_new
+    return V.T.reshape(shape)

src/torchjd/aggregation/dualproj.py

@@ -3,7 +3,7 @@
 from torch import Tensor
 
 from ._dual_cone_utils import project_weights
-from ._gramian_utils import compute_gramian, normalize, regularize
+from ._gramian_utils import compute_gramian, normalize
 from ._pref_vector_utils import pref_vector_to_str_suffix, pref_vector_to_weighting
 from .bases import _WeightedAggregator, _Weighting
 from .mean import _MeanWeighting
@@ -100,6 +100,6 @@ def __init__(
 
     def forward(self, matrix: Tensor) -> Tensor:
         u = self.weighting(matrix)
-        G = regularize(normalize(compute_gramian(matrix), self.norm_eps), self.reg_eps)
-        w = project_weights(u, G, self.solver)
+        G = normalize(compute_gramian(matrix), self.norm_eps)
+        w = project_weights(u, G)
         return w

src/torchjd/aggregation/upgrad.py

@@ -4,7 +4,7 @@
 from torch import Tensor
 
 from ._dual_cone_utils import project_weights
-from ._gramian_utils import compute_gramian, normalize, regularize
+from ._gramian_utils import compute_gramian, normalize
 from ._pref_vector_utils import pref_vector_to_str_suffix, pref_vector_to_weighting
 from .bases import _WeightedAggregator, _Weighting
 from .mean import _MeanWeighting
@@ -96,6 +96,6 @@ def __init__(
 
     def forward(self, matrix: Tensor) -> Tensor:
         U = torch.diag(self.weighting(matrix))
-        G = regularize(normalize(compute_gramian(matrix), self.norm_eps), self.reg_eps)
-        W = project_weights(U, G, self.solver)
+        G = normalize(compute_gramian(matrix), self.norm_eps)
+        W = project_weights(U, G)
         return torch.sum(W, dim=0)

tests/unit/aggregation/test_dual_cone_utils.py

@@ -1,9 +1,8 @@
-import numpy as np
 import torch
 from pytest import mark, raises
 from torch.testing import assert_close
 
-from torchjd.aggregation._dual_cone_utils import _project_weight_vector, project_weights
+from torchjd.aggregation._dual_cone_utils import project_weights
 
 
 @mark.parametrize("shape", [(5, 7), (9, 37), (2, 14), (32, 114), (50, 100)])
@@ -33,7 +32,7 @@ def test_solution_weights(shape: tuple[int, int]):
     G = J @ J.T
     u = torch.rand(shape[0])
 
-    w = project_weights(u, G, "quadprog")
+    w = project_weights(u, G)
     dual_gap = w - u
 
     # Dual feasibility
@@ -62,8 +61,8 @@ def test_scale_invariant(shape: tuple[int, int], scaling: float):
     G = J @ J.T
     u = torch.rand(shape[0])
 
-    w = project_weights(u, G, "quadprog")
-    w_scaled = project_weights(u, scaling * G, "quadprog")
+    w = project_weights(u, G)
+    w_scaled = project_weights(u, scaling * G)
 
     assert_close(w_scaled, w)
 
@@ -81,16 +80,16 @@ def test_tensorization_shape(shape: tuple[int, ...]):
 
     G = matrix @ matrix.T
 
-    W_tensor = project_weights(U_tensor, G, "quadprog")
-    W_matrix = project_weights(U_matrix, G, "quadprog")
+    W_tensor = project_weights(U_tensor, G)
+    W_matrix = project_weights(U_matrix, G)
 
     assert_close(W_matrix.reshape(shape), W_tensor)
 
 
-def test_project_weight_vector_failure():
-    """Tests that `_project_weight_vector` raises an error when the input G has too large values."""
+def test_project_weight_failure():
+    """Tests that `project_weight` raises an error when the input G has too large values."""
 
-    large_J = np.random.randn(10, 100) * 1e5
+    large_J = torch.randn(10, 100) * 1e5
     large_G = large_J @ large_J.T
     with raises(ValueError):
-        _project_weight_vector(np.ones(10), large_G, "quadprog")
+        project_weights(torch.ones(10), large_G)