diff --git a/python_bindings/bindings.cpp b/python_bindings/bindings.cpp index 9d2c1082..d16d4488 100644 --- a/python_bindings/bindings.cpp +++ b/python_bindings/bindings.cpp @@ -1,3 +1,5 @@ +#include +#include #include #include #include @@ -244,9 +246,25 @@ class Index { float norm = 0.0f; for (int i = 0; i < dim; i++) norm += data[i] * data[i]; - norm = 1.0f / (sqrtf(norm) + 1e-30f); + + // Bound relative error from underflow across all dimensions. + const float min_norm = std::numeric_limits::min() * dim; + if (norm > min_norm && std::isfinite(norm)) { + norm = 1.0f / sqrtf(norm); + for (int i = 0; i < dim; i++) + norm_array[i] = data[i] * norm; + return; + } + + // Recompute tiny or overflowing norms in double; keep zero vectors zero. + double norm_double = 0.0; + for (int i = 0; i < dim; i++) { + double value = data[i]; + norm_double += value * value; + } + norm_double = norm_double == 0.0 ? 0.0 : 1.0 / sqrt(norm_double); for (int i = 0; i < dim; i++) - norm_array[i] = data[i] * norm; + norm_array[i] = (float)(data[i] * norm_double); } @@ -795,9 +813,25 @@ class BFIndex { float norm = 0.0f; for (int i = 0; i < dim; i++) norm += data[i] * data[i]; - norm = 1.0f / (sqrtf(norm) + 1e-30f); + + // Bound relative error from underflow across all dimensions. + const float min_norm = std::numeric_limits::min() * dim; + if (norm > min_norm && std::isfinite(norm)) { + norm = 1.0f / sqrtf(norm); + for (int i = 0; i < dim; i++) + norm_array[i] = data[i] * norm; + return; + } + + // Recompute tiny or overflowing norms in double; keep zero vectors zero. + double norm_double = 0.0; + for (int i = 0; i < dim; i++) { + double value = data[i]; + norm_double += value * value; + } + norm_double = norm_double == 0.0 ? 0.0 : 1.0 / sqrt(norm_double); for (int i = 0; i < dim; i++) - norm_array[i] = data[i] * norm; + norm_array[i] = (float)(data[i] * norm_double); } diff --git a/tests/python/bindings_test_spaces.py b/tests/python/bindings_test_spaces.py index 901cadcc..6d84f710 100644 --- a/tests/python/bindings_test_spaces.py +++ b/tests/python/bindings_test_spaces.py @@ -37,3 +37,132 @@ def testRandomSelf(self): diff=np.mean(np.abs(distances-expected_distances)) self.assertAlmostEqual(diff, 0, delta=1e-3) + + +class CosineNormalizationTestCase(unittest.TestCase): + def _check_distances(self, index_type, data, queries, expected): + index = index_type(space='cosine', dim=data.shape[1]) + if index_type is hnswlib.Index: + index.init_index(max_elements=len(data), random_seed=100) + index.set_ef(10) + else: + index.init_index(max_elements=len(data)) + index.set_num_threads(1) + index.add_items(data, np.arange(len(data))) + labels, distances = index.knn_query( + queries, k=len(data), num_threads=1) + + # Every point is requested; compare distances by ID, including ties. + np.testing.assert_array_equal( + np.sort(labels, axis=1), + np.broadcast_to(np.arange(len(data)), labels.shape)) + reference_by_id = expected[np.arange(len(queries))[:, None], labels] + np.testing.assert_allclose( + distances, reference_by_id, rtol=0, atol=2e-6) + + def test_small_nonzero_vectors(self): + # These literal distances follow from identical/opposite directions. + expected = np.array([[0.0, 2.0], [2.0, 0.0]]) + for index_type in (hnswlib.Index, hnswlib.BFIndex): + for dim in (3, 8, 16): + for magnitude in (1.0, 1e-12, 1e-25): + with self.subTest(index=index_type.__name__, dim=dim, + magnitude=magnitude): + data = np.zeros((2, dim), dtype=np.float32) + data[:, 0] = [magnitude, -magnitude] + self._check_distances( + index_type, data, data, expected) + + def test_cosine_distances_across_stored_and_query_scales(self): + smallest = np.nextafter(np.float32(0), np.float32(1)) + scales = ( + (1.0, 1.0), + (1e-25, 1.0), + (1.0, 1e-25), + (1e-25, 1e-25), + (1e-40, 1.0), + (1.0, 1e-40), + (1e-40, 1e-40), + (smallest, smallest), + (1e30, 1.0), + (1.0, 1e30), + (1e30, 1e30), + ) + for index_type in (hnswlib.Index, hnswlib.BFIndex): + for dim in (3, 8, 16): + directions = np.zeros((4, dim), dtype=np.float32) + directions[:, :2] = [[1, 0], [3, 4], [-1, 0], [0, 1]] + for stored_scale, query_scale in scales: + with self.subTest(index=index_type.__name__, dim=dim, + stored_scale=stored_scale, + query_scale=query_scale): + data = directions * np.float32(stored_scale) + queries = directions[:2] * np.float32(query_scale) + + # Independent cosine formula on the actual float32 + # inputs, promoted before squaring or dot products. + wide_data = data.astype(np.float64) + wide_queries = queries.astype(np.float64) + dot = wide_queries @ wide_data.T + lengths = ( + np.linalg.norm(wide_queries, axis=1)[:, None] + * np.linalg.norm(wide_data, axis=1)[None, :]) + expected = 1.0 - dot / lengths + self._check_distances( + index_type, data, queries, expected) + + def test_cosine_distances_near_float_norm_limits(self): + limits = np.finfo(np.float32) + root_min = np.float32(np.sqrt(float(limits.tiny))) + for index_type in (hnswlib.Index, hnswlib.BFIndex): + for dim in (3, 8, 16, 128): + root_max = np.float32(np.sqrt(float(limits.max) / dim)) + scales = ( + np.float32(1e-22), + np.float32(1e-21), + root_min / np.float32(2), + np.nextafter(root_min, np.float32(0)), + root_min, + np.nextafter(root_min, np.float32(np.inf)), + np.nextafter(root_max, np.float32(0)), + root_max, + np.nextafter(root_max, np.float32(np.inf)), + np.float32(limits.max / np.float32(4)), + ) + # Dense directions exercise rounding across many components. + directions = np.ones((4, dim), dtype=np.float32) + directions[1, ::2] = 3 + directions[1, 1::2] = 4 + directions[2] = -1 + directions[3, 1::2] = 0 + for scale in scales: + for stored_scale, query_scale in ( + (scale, 1.0), (1.0, scale), (scale, scale)): + with self.subTest(index=index_type.__name__, dim=dim, + stored_scale=stored_scale, + query_scale=query_scale): + data = directions * np.float32(stored_scale) + queries = directions[:2] * np.float32(query_scale) + wide_data = data.astype(np.float64) + wide_queries = queries.astype(np.float64) + dot = wide_queries @ wide_data.T + lengths = ( + np.linalg.norm(wide_queries, axis=1)[:, None] + * np.linalg.norm(wide_data, axis=1)[None, :]) + expected = 1.0 - dot / lengths + self._check_distances( + index_type, data, queries, expected) + + def test_zero_vector_compatibility(self): + # Existing convention: zero has dot product zero with every vector. + expected = np.array([ + [1.0, 1.0, 1.0], + [1.0, 0.0, 2.0], + [1.0, 2.0, 0.0], + ]) + for index_type in (hnswlib.Index, hnswlib.BFIndex): + data = np.zeros((3, 8), dtype=np.float32) + data[1, 0] = 1.0 + data[2, 0] = -1.0 + with self.subTest(index=index_type.__name__): + self._check_distances(index_type, data, data, expected)