diff --git a/CHANGELOG.md b/CHANGELOG.md index dfb2616e..d04177a3 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/torchjd/aggregation/_krum.py b/src/torchjd/aggregation/_krum.py index d12d758b..03f138f4 100644 --- a/src/torchjd/aggregation/_krum.py +++ b/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]) diff --git a/tests/unit/aggregation/test_krum.py b/tests/unit/aggregation/test_krum.py index 54e775b9..02923f11 100644 --- a/tests/unit/aggregation/test_krum.py +++ b/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))