fix(aggregation): Use squared distances for Krum scores

CHANGELOG.md

@@ -8,6 +8,12 @@ changelog does not include internal changes that do not affect the user.
 
 ## [Unreleased]
 
+### Fixed
+
+- Fixed `Krum` and `KrumWeighting` ranking rows by sums of Euclidean distances instead of
+  squared Euclidean distances, as required by Krum and Multi-Krum. This could select a different
+  gradient or a different set of gradients to average.
+
 ## [0.18.0] - 2026-10-08
 
 ### Added

src/torchjd/aggregation/_krum.py

@@ -28,13 +28,14 @@ def forward(self, gramian: PSDMatrix, /) -> Tensor:
         gradient_norms_squared = torch.diagonal(gramian)
         distances_squared = (
             gradient_norms_squared.unsqueeze(0) + gradient_norms_squared.unsqueeze(1) - 2 * gramian
-        )
-        distances = torch.sqrt(distances_squared.clamp(min=0.0))
+        ).clamp(min=0.0)
 
         n_closest = gramian.shape[0] - self.n_byzantine - 2
-        smallest_distances, _ = torch.topk(distances, k=n_closest + 1, largest=False)
-        smallest_distances_excluding_self = smallest_distances[:, 1:]
-        scores = smallest_distances_excluding_self.sum(dim=1)
+        smallest_distances_squared, _ = torch.topk(
+            distances_squared, k=n_closest + 1, largest=False
+        )
+        smallest_distances_squared_excluding_self = smallest_distances_squared[:, 1:]
+        scores = smallest_distances_squared_excluding_self.sum(dim=1)
 
         _, selected_indices = torch.topk(scores, k=self.n_selected, largest=False)
         one_hot_selected_indices = F.one_hot(selected_indices, num_classes=gramian.shape[0])

tests/unit/aggregation/test_krum.py

@@ -2,6 +2,7 @@
 
 from pytest import mark, raises
 from torch import Tensor
+from torch.testing import assert_close
 from utils.contexts import ExceptionContext
 from utils.tensors import ones_, tensor_
 
@@ -122,12 +123,31 @@ def test_weighting_n_selected_setter_rejects_non_positive() -> None:
 
 
 def test_negative_squared_distances_are_clamped() -> None:
-    x = tensor_([1.0, 1.0, 2.0, 2.1, 11.0])
+    # Perturbing the gramian makes the squared distances from row 0 to rows 1 and 2 equal to -3.
+    # Clamped scores are [1, 0, 8, 1, 1], but unclamped scores would be [-2, 0, 8, 1, 1].
+    x = tensor_([3, 2, 4, 2, 2])
     gramian = x.unsqueeze(1) * x.unsqueeze(0)
-    gramian[0, 1] = gramian[0, 1] + 1e-6
-    gramian[1, 0] = gramian[1, 0] + 1e-6
+    gramian[0, 1:3] += 2
+    gramian[1:3, 0] += 2
 
-    weights = KrumWeighting(n_byzantine=1)(gramian)
+    weights = KrumWeighting(n_byzantine=0)(gramian)
 
-    assert weights[:2].sum().item() == 1.0
-    assert weights[2:].abs().sum().item() == 0.0
+    assert_close(weights, tensor_([0, 1, 0, 0, 0]))
+
+
+@mark.parametrize(
+    ["matrix", "n_selected", "expected_weights"],
+    [
+        # Scores are [30, 44, 28, 38, 56]. Unsquared distances would select row 0.
+        ([[3, 0], [0, 7], [2, 5], [5, 0], [6, 7]], 1, [0, 0, 1, 0, 0]),
+        # Scores are [21, 22, 26, 33, 43]. Unsquared distances would select rows 0 and 2.
+        ([[3, 0], [3, 7], [2, 0], [5, 4], [0, 7]], 2, [0.5, 0.5, 0, 0, 0]),
+    ],
+)
+def test_squared_distance_scores(
+    matrix: list[list[float]], n_selected: int, expected_weights: list[float]
+) -> None:
+    matrix_ = tensor_(matrix)
+    weights = KrumWeighting(n_byzantine=1, n_selected=n_selected)(matrix_ @ matrix_.T)
+
+    assert_close(weights, tensor_(expected_weights))