From b749eda548e0655f9c9ff3a42aab365a9cf3a838 Mon Sep 17 00:00:00 2001 From: AryanRG Date: Fri, 30 Oct 2026 22:53:05 +0530 Subject: [PATCH 1/3] fix(aggregation): Use squared distances for Krum scores --- CHANGELOG.md | 6 ++++ src/torchjd/aggregation/_krum.py | 10 +++--- tests/unit/aggregation/test_krum.py | 55 ++++++++++++++++++++++++++++- 3 files changed, 66 insertions(+), 5 deletions(-) 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..68073da1 100644 --- a/src/torchjd/aggregation/_krum.py +++ b/src/torchjd/aggregation/_krum.py @@ -29,12 +29,14 @@ def forward(self, gramian: PSDMatrix, /) -> Tensor: distances_squared = ( gradient_norms_squared.unsqueeze(0) + gradient_norms_squared.unsqueeze(1) - 2 * gramian ) - distances = torch.sqrt(distances_squared.clamp(min=0.0)) + distances_squared = distances_squared.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..6bc07bc1 100644 --- a/tests/unit/aggregation/test_krum.py +++ b/tests/unit/aggregation/test_krum.py @@ -2,11 +2,14 @@ from pytest import mark, raises from torch import Tensor +from torch.optim import SGD +from torch.testing import assert_close from utils.contexts import ExceptionContext -from utils.tensors import ones_, tensor_ +from utils.tensors import ones_, tensor_, zeros_ from torchjd.aggregation import Krum from torchjd.aggregation._krum import KrumWeighting +from torchjd.autojac import backward, jac_to_grad from ._asserts import assert_expected_structure from ._inputs import scaled_matrices_2_plus_rows, typical_matrices_2_plus_rows @@ -131,3 +134,53 @@ def test_negative_squared_distances_are_clamped() -> None: assert weights[:2].sum().item() == 1.0 assert weights[2:].abs().sum().item() == 0.0 + + +def test_squared_distance_scores_select_correct_row() -> None: + # Summing the two nearest squared distances gives scores [30, 44, 28, 38, 56]. + # Summing their square roots instead incorrectly selects row 0. + matrix = tensor_([[3, 0], [0, 7], [2, 5], [5, 0], [6, 7]]) + weights = KrumWeighting(n_byzantine=1)(matrix @ matrix.T) + + assert_close(weights, tensor_([0, 0, 1, 0, 0])) + assert_close(Krum(n_byzantine=1)(matrix), tensor_([2, 5])) + + +@mark.parametrize( + ["n_selected", "expected_weights"], + [ + (1, [1, 0, 0, 0, 0]), + (2, [0.5, 0.5, 0, 0, 0]), + (5, [0.2, 0.2, 0.2, 0.2, 0.2]), + ], +) +def test_multikrum_squared_distance_scores(n_selected: int, expected_weights: list[float]) -> None: + # Scores [21, 22, 26, 33, 43] rank rows 0 and 1 first. Unsquared distances + # rank rows 0 and 2 first, changing the Multi-Krum average. + matrix = tensor_([[3, 0], [3, 7], [2, 0], [5, 4], [0, 7]]) + weights = KrumWeighting(n_byzantine=1, n_selected=n_selected)(matrix @ matrix.T) + + assert_close(weights, tensor_(expected_weights)) + assert_close( + Krum(n_byzantine=1, n_selected=n_selected)(matrix), + tensor_(expected_weights) @ matrix, + ) + + +@mark.parametrize("optimize_gramian_computation", [False, True]) +def test_squared_distance_scores_in_optimizer_step(optimize_gramian_computation: bool) -> None: + param = zeros_(2, requires_grad=True) + optimizer = SGD([param], lr=0.1) + targets = -tensor_([[3, 0], [0, 7], [2, 5], [5, 0], [6, 7]]) + losses = 0.5 * (param - targets).square().sum(dim=1) + + backward(losses, inputs=[param]) + weights = jac_to_grad( + [param], + Krum(n_byzantine=1), + optimize_gramian_computation=optimize_gramian_computation, + ) + assert_close(weights, tensor_([0, 0, 1, 0, 0])) + assert_close(param.grad, tensor_([2, 5])) + optimizer.step() + assert_close(param, tensor_([-0.2, -0.5])) From e5e54d4ecbb1ce1c0b8842cebdbf9487e6a58a63 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Val=C3=A9rian=20Rey?= Date: Sat, 31 Oct 2026 14:38:36 +0100 Subject: [PATCH 2/3] test(aggregation): Simplify Krum tests and make clamp test meaningful Co-Authored-By: Claude Opus 5.5 --- tests/unit/aggregation/test_krum.py | 69 ++++++++--------------------- 1 file changed, 18 insertions(+), 51 deletions(-) diff --git a/tests/unit/aggregation/test_krum.py b/tests/unit/aggregation/test_krum.py index 6bc07bc1..02923f11 100644 --- a/tests/unit/aggregation/test_krum.py +++ b/tests/unit/aggregation/test_krum.py @@ -2,14 +2,12 @@ from pytest import mark, raises from torch import Tensor -from torch.optim import SGD from torch.testing import assert_close from utils.contexts import ExceptionContext -from utils.tensors import ones_, tensor_, zeros_ +from utils.tensors import ones_, tensor_ from torchjd.aggregation import Krum from torchjd.aggregation._krum import KrumWeighting -from torchjd.autojac import backward, jac_to_grad from ._asserts import assert_expected_structure from ._inputs import scaled_matrices_2_plus_rows, typical_matrices_2_plus_rows @@ -125,62 +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 - - -def test_squared_distance_scores_select_correct_row() -> None: - # Summing the two nearest squared distances gives scores [30, 44, 28, 38, 56]. - # Summing their square roots instead incorrectly selects row 0. - matrix = tensor_([[3, 0], [0, 7], [2, 5], [5, 0], [6, 7]]) - weights = KrumWeighting(n_byzantine=1)(matrix @ matrix.T) - - assert_close(weights, tensor_([0, 0, 1, 0, 0])) - assert_close(Krum(n_byzantine=1)(matrix), tensor_([2, 5])) + assert_close(weights, tensor_([0, 1, 0, 0, 0])) @mark.parametrize( - ["n_selected", "expected_weights"], + ["matrix", "n_selected", "expected_weights"], [ - (1, [1, 0, 0, 0, 0]), - (2, [0.5, 0.5, 0, 0, 0]), - (5, [0.2, 0.2, 0.2, 0.2, 0.2]), + # 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_multikrum_squared_distance_scores(n_selected: int, expected_weights: list[float]) -> None: - # Scores [21, 22, 26, 33, 43] rank rows 0 and 1 first. Unsquared distances - # rank rows 0 and 2 first, changing the Multi-Krum average. - matrix = tensor_([[3, 0], [3, 7], [2, 0], [5, 4], [0, 7]]) - weights = KrumWeighting(n_byzantine=1, n_selected=n_selected)(matrix @ matrix.T) +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)) - assert_close( - Krum(n_byzantine=1, n_selected=n_selected)(matrix), - tensor_(expected_weights) @ matrix, - ) - - -@mark.parametrize("optimize_gramian_computation", [False, True]) -def test_squared_distance_scores_in_optimizer_step(optimize_gramian_computation: bool) -> None: - param = zeros_(2, requires_grad=True) - optimizer = SGD([param], lr=0.1) - targets = -tensor_([[3, 0], [0, 7], [2, 5], [5, 0], [6, 7]]) - losses = 0.5 * (param - targets).square().sum(dim=1) - - backward(losses, inputs=[param]) - weights = jac_to_grad( - [param], - Krum(n_byzantine=1), - optimize_gramian_computation=optimize_gramian_computation, - ) - assert_close(weights, tensor_([0, 0, 1, 0, 0])) - assert_close(param.grad, tensor_([2, 5])) - optimizer.step() - assert_close(param, tensor_([-0.2, -0.5])) From 3996c84e750890539b61056de1f4481c17a3fe0e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Val=C3=A9rian=20Rey?= Date: Sat, 31 Oct 2026 14:38:37 +0100 Subject: [PATCH 3/3] refactor(aggregation): Chain clamp on Krum squared distances Co-Authored-By: Claude Opus 5.5 --- src/torchjd/aggregation/_krum.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/torchjd/aggregation/_krum.py b/src/torchjd/aggregation/_krum.py index 68073da1..03f138f4 100644 --- a/src/torchjd/aggregation/_krum.py +++ b/src/torchjd/aggregation/_krum.py @@ -28,8 +28,7 @@ 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_squared = distances_squared.clamp(min=0.0) + ).clamp(min=0.0) n_closest = gramian.shape[0] - self.n_byzantine - 2 smallest_distances_squared, _ = torch.topk(