From 199f18c91549da13fdc0d8fa505788046ee24923 Mon Sep 17 00:00:00 2001 From: Amit T Subhash Date: Thu, 23 Jul 2026 19:37:44 +0530 Subject: [PATCH] Chunk OrthogonalProcrustesAlignment to cut peak memory (#307) fit() materialized the full (n_label, n_ref) distance matrix and its row-wise argsort before selecting top_k references, so peak memory was O(n_label * n_ref) regardless of subsample. Search top_k in blocks of label rows: each block is scored against all references, so the per-row selection is bit-identical to the unchunked path (ties included), while peak memory drops to O(block * n_ref). On a 15000x15000 case this is 5.4 GB -> 0.6 GB. The block size is internal, derived from a module-level constant, so there is no public API change. Addresses onford's report; regression tests compare the chunked path against an independent full-matrix reference across shapes, top_k (including top_k == n_ref), 1D/2D labels and ties, and assert the search really runs in bounded blocks. --- cebra/data/helper.py | 25 +++++++-- tests/test_data_helper.py | 104 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 125 insertions(+), 4 deletions(-) diff --git a/cebra/data/helper.py b/cebra/data/helper.py index 8582edae..39496a23 100644 --- a/cebra/data/helper.py +++ b/cebra/data/helper.py @@ -82,6 +82,15 @@ def _require_numpy_array(array: Union[npt.NDArray, torch.Tensor]): return array +# Target block size for the ``top_k`` search: the rows of ``label`` are +# processed in blocks, so only a ``(chunk_size, n_ref)`` distance matrix and +# its argsort indices are held at once. 10 million elements is roughly 80 MB +# per array in float64/intp. This is a target rather than a hard cap: a block +# is always at least one row, so the working set never goes below +# ``(1, n_ref)``. +_PROCRUSTES_MAX_DISTANCE_ELEMENTS = 10_000_000 + + class OrthogonalProcrustesAlignment: """Aligns two dataset by solving the orthogonal Procrustes problem. @@ -219,10 +228,18 @@ def fit( f"got ref_data:{data.shape[0]} samples and ref_labels:{label.shape[0]} samples." ) - distance = self._distance(label, ref_label) - - # keep indexes of the {self.top_k} labels the closest to the reference labels - target_idx = np.argsort(distance, axis=1)[:, :self.top_k] + # Search the top_k closest reference labels for each label, chunking + # over the rows of `label` so the full (n_label, n_ref) distance matrix + # is never materialized. Each row is scored against all references, so + # the per-row selection is identical to the unchunked computation. + n_label, n_ref = label.shape[0], ref_label.shape[0] + chunk_size = max(1, _PROCRUSTES_MAX_DISTANCE_ELEMENTS // n_ref) + target_idx = np.empty((n_label, self.top_k), dtype=np.intp) + for start in range(0, n_label, chunk_size): + stop = start + chunk_size + distance = self._distance(label[start:stop], ref_label) + target_idx[start:stop] = np.argsort(distance, + axis=1)[:, :self.top_k] if data.shape[1] != ref_data.shape[1]: raise ValueError( diff --git a/tests/test_data_helper.py b/tests/test_data_helper.py index a3d1e774..5361bcc4 100644 --- a/tests/test_data_helper.py +++ b/tests/test_data_helper.py @@ -359,3 +359,107 @@ def test_ensembling_performances(n_models=3): score = decoder.score(valid_embedding, valid_label) assert all(score_i < score for score_i in scores) + + +@pytest.mark.parametrize("n_label, n_ref, dim, top_k", [ + (50, 40, 4, 5), + (37, 60, 3, 1), + (100, 20, 2, 10), + (23, 7, 3, 7), +]) +@pytest.mark.parametrize("label_dim", [1, 2]) +def test_orthogonal_alignment_chunking_matches_full(n_label, n_ref, dim, top_k, + label_dim, monkeypatch): + """Chunked top_k search reproduces the full-matrix path (issue #307).""" + import scipy.linalg + rng = np.random.RandomState(0) + ref_data = rng.uniform(0, 1, (n_ref, dim)) + data = rng.uniform(0, 1, (n_label, dim)) + # distinct (non-tie) labels so the reference ordering is unambiguous + ref_label = (rng.permutation(n_ref).astype(float) + + rng.uniform(0, 1e-3, n_ref))[:, None] + label = (rng.permutation(n_label).astype(float) + + rng.uniform(0, 1e-3, n_label))[:, None] + if label_dim == 1: + ref_label, label = ref_label.ravel(), label.ravel() + + model = cebra_data_helper.OrthogonalProcrustesAlignment(top_k=top_k) + # independent reference: full distance matrix, then top_k per row + distance = model._distance(label.reshape(n_label, -1), + ref_label.reshape(n_ref, -1)) + full_idx = np.argsort(distance, axis=1)[:, :top_k] + X = data[:, None].repeat(top_k, axis=1).reshape(-1, dim) + Y = ref_data[full_idx].reshape(-1, dim) + ref_transform, _ = scipy.linalg.orthogonal_procrustes(X, Y) + + # force many small chunks: 3 rows of `label` at a time + monkeypatch.setattr(cebra_data_helper, "_PROCRUSTES_MAX_DISTANCE_ELEMENTS", + 3 * n_ref) + chunked = cebra_data_helper.OrthogonalProcrustesAlignment(top_k=top_k) + chunked.fit(ref_data, data, ref_label, label) + np.testing.assert_allclose(chunked._transform, + ref_transform, + rtol=1e-10, + atol=1e-10) + + # force a single chunk: must be bit-identical to the chunked run + monkeypatch.setattr(cebra_data_helper, "_PROCRUSTES_MAX_DISTANCE_ELEMENTS", + n_label * n_ref + 1) + single = cebra_data_helper.OrthogonalProcrustesAlignment(top_k=top_k) + single.fit(ref_data, data, ref_label, label) + np.testing.assert_array_equal(single._transform, chunked._transform) + np.testing.assert_array_equal(single.transform(data), + chunked.transform(data)) + + +def test_orthogonal_alignment_chunking_bit_identical_with_ties(monkeypatch): + """Row chunking is bit-identical, even with tied distances (#307).""" + rng = np.random.RandomState(1) + n_label, n_ref, dim, top_k = 60, 30, 3, 5 + ref_data = rng.uniform(0, 1, (n_ref, dim)) + data = rng.uniform(0, 1, (n_label, dim)) + # integer labels with many duplicates -> tied distances + ref_label = rng.randint(0, 5, size=(n_ref, 1)).astype(float) + label = rng.randint(0, 5, size=(n_label, 1)).astype(float) + + monkeypatch.setattr(cebra_data_helper, "_PROCRUSTES_MAX_DISTANCE_ELEMENTS", + n_label * n_ref + 1) + single = cebra_data_helper.OrthogonalProcrustesAlignment(top_k=top_k) + single.fit(ref_data, data, ref_label, label) + + monkeypatch.setattr(cebra_data_helper, "_PROCRUSTES_MAX_DISTANCE_ELEMENTS", + 4 * n_ref) + chunked = cebra_data_helper.OrthogonalProcrustesAlignment(top_k=top_k) + chunked.fit(ref_data, data, ref_label, label) + np.testing.assert_array_equal(single._transform, chunked._transform) + + +def test_orthogonal_alignment_chunking_bounds_block_size(monkeypatch): + """The top_k search runs in bounded blocks, not one full matrix (#307).""" + rng = np.random.RandomState(2) + n_label, n_ref, dim, top_k = 50, 10, 3, 4 + ref_data = rng.uniform(0, 1, (n_ref, dim)) + data = rng.uniform(0, 1, (n_label, dim)) + ref_label = rng.uniform(0, 1, (n_ref, 1)) + label = rng.uniform(0, 1, (n_label, 1)) + + rows_per_call = [] + original = cebra_data_helper.OrthogonalProcrustesAlignment._distance + + def spy(self, label_i, label_j): + rows_per_call.append(label_i.shape[0]) + return original(self, label_i, label_j) + + # 3 rows of `label` per block + monkeypatch.setattr(cebra_data_helper, "_PROCRUSTES_MAX_DISTANCE_ELEMENTS", + 3 * n_ref) + monkeypatch.setattr(cebra_data_helper.OrthogonalProcrustesAlignment, + "_distance", spy) + + model = cebra_data_helper.OrthogonalProcrustesAlignment(top_k=top_k) + model.fit(ref_data, data, ref_label, label) + + # the full (n_label, n_ref) distance matrix is never materialized + assert max(rows_per_call) <= 3 + assert len(rows_per_call) == (n_label + 2) // 3 + assert sum(rows_per_call) == n_label