From afc5949a825ed2161cadbaa9a1e30a34f3da90ff Mon Sep 17 00:00:00 2001 From: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> Date: Tue, 6 Oct 2026 00:31:34 +0100 Subject: [PATCH] Rank associative memories by normalized cosine similarity Signed-off-by: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> --- .../basic_associative_memory.py | 13 +- .../cosine_retrieval_test.py | 123 ++++++++++++++++++ 2 files changed, 135 insertions(+), 1 deletion(-) create mode 100644 concordia/associative_memory/cosine_retrieval_test.py diff --git a/concordia/associative_memory/basic_associative_memory.py b/concordia/associative_memory/basic_associative_memory.py index 4566bd84f..b627a1a41 100644 --- a/concordia/associative_memory/basic_associative_memory.py +++ b/concordia/associative_memory/basic_associative_memory.py @@ -133,6 +133,8 @@ def get_data_frame(self) -> pd.DataFrame: def _get_top_k_cosine(self, x: np.ndarray, k: int): """Returns the top k most cosine similar rows to an input vector x. + Zero-norm query or memory embeddings have similarity zero. + Args: x: The input vector. k: The number of rows to return. @@ -140,10 +142,19 @@ def _get_top_k_cosine(self, x: np.ndarray, k: int): Returns: Rows, sorted by cosine similarity in descending order. """ + x = np.asarray(x) + x_norm = np.linalg.norm(x) + x = x / x_norm if x_norm > 0 else np.zeros_like(x) + + def cosine_similarity(y): + y = np.asarray(y) + y_norm = np.linalg.norm(y) + return np.dot(x, y / y_norm) if y_norm > 0 else 0.0 + with self._memory_bank_lock: self._flush_pending() cosine_similarities = self._memory_bank['embedding'].apply( - lambda y: np.dot(x, y) + cosine_similarity ) # Sort the cosine similarities in descending order. diff --git a/concordia/associative_memory/cosine_retrieval_test.py b/concordia/associative_memory/cosine_retrieval_test.py new file mode 100644 index 000000000..f601354d5 --- /dev/null +++ b/concordia/associative_memory/cosine_retrieval_test.py @@ -0,0 +1,123 @@ +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Memory retrieval ranks directions independently of embedding magnitude.""" + +from absl.testing import absltest +from absl.testing import parameterized +from concordia.associative_memory import basic_associative_memory +import numpy as np + + +def make_bank(vectors): + bank = basic_associative_memory.AssociativeMemoryBank( + lambda text: np.asarray(vectors[text], dtype=float) + ) + bank.extend(key for key in vectors if key != 'query') + return bank + + +class CosineRetrievalTest(parameterized.TestCase): + + def test_exact_match_beats_a_larger_off_axis_embedding(self): + bank = make_bank( + {'query': [1.0, 0.0], 'exact': [1.0, 0.0], 'large': [100.0, 100.0]} + ) + self.assertEqual(bank.retrieve_associative('query', k=1), ['exact']) + + @parameterized.parameters( + (1.0, 1.0, 1.0), (100.0, 0.01, 20.0), (0.1, 30.0, 0.001) + ) + def test_independent_positive_rescaling_does_not_change_ranking( + self, a, b, c + ): + vectors = { + 'query': np.array([1.0, 2.0]) * a, + 'near': np.array([1.0, 2.0]) * b, + 'partial': np.array([2.0, 1.0]) * c, + 'opposite': np.array([-1.0, -2.0]) * 1000, + } + self.assertEqual( + make_bank(vectors).retrieve_associative('query', k=3), + ['near', 'partial', 'opposite'], + ) + + def test_random_embeddings_match_explicit_normalized_dot_products(self): + rng = np.random.default_rng(3) + query = rng.normal(size=7) + values = rng.normal(size=(20, 7)) * np.geomspace(0.001, 1000, 20)[:, None] + vectors = {'query': query} | {str(i): x for i, x in enumerate(values)} + scores = (values / np.linalg.norm(values, axis=1, keepdims=True)) @ ( + query / np.linalg.norm(query) + ) + expected = [str(i) for i in np.argsort(-scores)[:8]] + self.assertEqual( + make_bank(vectors).retrieve_associative('query', k=8), expected + ) + + def test_restore_uses_stored_embeddings_without_mutating_them(self): + vectors = { + 'query': [1.0, 0.0], + 'exact': [1.0, 0.0], + 'large': [100.0, 100.0], + } + original = make_bank(vectors) + state = original.get_state() + restored = basic_associative_memory.AssociativeMemoryBank( + lambda text: np.asarray(vectors[text]) + ) + restored.set_state(state) + before = restored.get_state() + self.assertEqual(restored.retrieve_associative('query'), ['exact']) + self.assertEqual(restored.get_state(), before) + self.assertEqual(restored.retrieve_recent(2), ['exact', 'large']) + + def test_zero_vectors_have_zero_similarity_without_warnings(self): + bank = make_bank({ + 'query': [1.0, 0.0], + 'opposite': [-1.0, 0.0], + 'zero': [0.0, 0.0], + 'same': [2.0, 0.0], + }) + with np.errstate(all='raise'): + self.assertEqual( + bank.retrieve_associative('query', k=3), ['same', 'zero', 'opposite'] + ) + zero_query = make_bank( + {'query': [0.0, 0.0], 'a': [1.0, 0.0], 'b': [3.0, 2.0]} + ) + with np.errstate(all='raise'): + self.assertCountEqual( + zero_query.retrieve_associative('query', k=2), ['a', 'b'] + ) + + def test_empty_bank_and_unit_vectors_keep_existing_behavior(self): + vectors = { + 'query': [1.0, 0.0], + 'north': [1.0, 0.0], + 'east': [0.0, 1.0], + 'south': [-1.0, 0.0], + } + bank = basic_associative_memory.AssociativeMemoryBank( + lambda text: np.asarray(vectors[text]) + ) + self.assertEqual(bank.retrieve_associative('query'), []) + bank.extend(['east', 'south', 'north']) + self.assertEqual( + bank.retrieve_associative('query', k=4), ['north', 'east', 'south'] + ) + + +if __name__ == '__main__': + absltest.main()