Add cosine similarity support for faiss engine #2376

VijayanB · 2025-01-09T07:37:53Z

Description

FAISS engine doesn't support cosine similarity natively. However we can use inner product to achieve the same, because, when vectors are normalized then inner product will be same as cosine similarity. Hence, before ingestion, normalize the input vector, and add it to faiss index with type as inner product, and, before search, normalize query vector if space type is cosine and engine is faiss.

Since we will be storing normalized vector in segments, we don't have to normalize whenever segments are merged. This will keep force merge time and search at competitive, provided we will face additional latency during indexing (one time where we normalize). To avoid this additional latency, customers can normalize their data set and create inner product.

This also adds support to radial search, for both max distance and min score.

Related Issues

#2242

Check List

New functionality includes testing.
New functionality has been documented.
API changes companion pull request created.
Commits are signed per the DCO using --signoff.
Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

VijayanB · 2025-01-09T07:49:56Z

Adding additional unit and integration test for radial search. Will mark it as ready once i add those tests

FAISS engine doesn't support cosine similarity natively. However we can use inner product to achieve the same, because, when vectors are normalized then inner product will be same as cosine similarity. Hence, before ingestion and perform search, normalize the input vector and add it to faiss index with type as inner product. Since we will be storing normalized vector in segments, to get actual vectors, source can be used. By saving as normalized vector, we don't have to normalize whenever segments are merged. This will keep force merge time and search at competitive, provided we will face additional latency during indexing (one time where we normalize). We also support radial search for cosine similarity. Signed-off-by: Vijayan Balasubramanian <[email protected]>

jmazanec15

Thanks @VijayanB - completed a first pass review

jmazanec15 · 2025-01-10T16:36:51Z

src/main/java/org/opensearch/knn/index/SpaceType.java

@@ -77,6 +77,11 @@ public float scoreTranslation(float rawScore) {
            return Math.max((2.0F - rawScore) / 2.0F, 0.0F);
        }

+        @Override
+        public float scoreToDistanceTranslation(float score) {
+            return score;


Confused - Why is this correct?

jmazanec15 · 2025-01-10T16:37:46Z

src/main/java/org/opensearch/knn/index/engine/AbstractKNNMethod.java

+    /**
+     * Gets the compatible space type for the given space type parameter.
+     * The subclass can override this method and returns the appropriate space type that
+     * is compatible with the library.


Can you add context on why this is needed via an example?

jmazanec15 · 2025-01-10T16:38:02Z

src/main/java/org/opensearch/knn/index/engine/AbstractKNNMethod.java

+     *         space type if it's already compatible
+     * @see SpaceType
+     */
+


nit: extra line

jmazanec15 · 2025-01-10T16:39:19Z

src/main/java/org/opensearch/knn/index/engine/AbstractKNNMethod.java

+     * @see SpaceType
+     */
+
+    protected SpaceType getCompatibleSpaceType(SpaceType spaceType) {


Can we call something like convertUserToMethodSpaceType? The naming here doesnt seem to give it enough description.

jmazanec15 · 2025-01-10T16:39:31Z

src/main/java/org/opensearch/knn/index/engine/KNNLibraryIndexingContext.java

@@ -47,4 +48,6 @@ public interface KNNLibraryIndexingContext {
     * @return Get the per dimension processor
     */
    PerDimensionProcessor getPerDimensionProcessor();
+
+    VectorTransformer getVectorTransformer();


jmazanec15 · 2025-01-10T16:41:23Z

src/main/java/org/opensearch/knn/index/mapper/ModelFieldMapper.java

+        KNNMethodConfigContext knnMethodConfigContext = getKNNMethodConfigContextFromModelMetadata(modelMetadata);
+        // Need to handle BWC case
+        if (knnMethodContext == null || knnMethodConfigContext == null) {
+            vectorTransformer = VectorTransformerFactory.getVectorTransformer(modelMetadata.getKnnEngine(), modelMetadata.getSpaceType());


Can we add log here?

jmazanec15 · 2025-01-10T16:42:31Z

src/main/java/org/opensearch/knn/index/mapper/VectorTransformer.java

+public interface VectorTransformer {
+
+    /**
+     * Transforms a float vector into a new vector of the same type.


Will this be in place or always a copy?

IMO, Transformer should generally return a new value rather than modifying the original value it receives, to preserve immutability.

jmazanec15 · 2025-01-10T16:46:34Z

src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java

-    private float[] getVectorForCreatingQueryRequest(VectorDataType vectorDataType, KNNEngine knnEngine) {
+    private float[] getVectorForCreatingQueryRequest(VectorDataType vectorDataType, KNNEngine knnEngine, SpaceType spaceType) {
+
+        // Cosine similarity is supported as Inner product by FAISS by normalizing input vector, hence, we have to normalize


Can we move this check out of this class? This class is already very crowded and I want to avoid adding more checks around engines. Instead, could we investigate either adding it to https://github.com/opensearch-project/k-NN/blob/main/src/main/java/org/opensearch/knn/index/engine/KNNLibrarySearchContext.java and/or adding a method in KNNVectorFieldType (https://github.com/opensearch-project/k-NN/blob/main/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java#L85) that says "should normalize query" or, better yet, transformQuery?

jmazanec15 · 2025-01-10T16:46:58Z

src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java

@@ -681,6 +681,8 @@ protected void validatePreparse() {
     */
    protected abstract PerDimensionProcessor getPerDimensionProcessor();

+    protected abstract VectorTransformer getVectorTransformer();


Can you add java doc for this one? Also, can you update javadoc for getVectorValidator() to say that it is validated before any transform calls?

jmazanec15 · 2025-01-10T16:50:00Z

src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java

@@ -700,7 +703,8 @@ protected void parseCreateField(ParseContext context, int dimension, VectorDataT
            }
            final float[] array = floatsArrayOptional.get();
            getVectorValidator().validateVector(array);
-            context.doc().addAll(getFieldsForFloatVector(array));
+            final float[] transformedArray = getVectorTransformer().transform(array);


This will be called before the per-dimension processor too? What should contract around these 2 be? Im wondering if we even need the per-dimension or if we can wrap that in this new full vector transform.

Signed-off-by: Vijayan Balasubramanian <[email protected]>

VijayanB force-pushed the faiss-cosine branch 3 times, most recently from d6f16a1 to 6671e64 Compare January 9, 2025 07:43

VijayanB force-pushed the faiss-cosine branch 5 times, most recently from ecfffb4 to eee45d5 Compare January 10, 2025 02:59

VijayanB force-pushed the faiss-cosine branch from eee45d5 to 4658bee Compare January 10, 2025 02:59

VijayanB marked this pull request as ready for review January 10, 2025 04:15

VijayanB requested review from heemin32, navneet1v, vamshin, jmazanec15, naveentatikonda, junqiu-lei, martin-gaievski, ryanbogan, luyuncheng, shatejas and 0ctopus13prime as code owners January 10, 2025 04:15

jmazanec15 reviewed Jan 10, 2025

View reviewed changes

Fixed code review comments

67c8eac

Signed-off-by: Vijayan Balasubramanian <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add cosine similarity support for faiss engine #2376

Add cosine similarity support for faiss engine #2376

VijayanB commented Jan 9, 2025 •

edited

Loading

VijayanB commented Jan 9, 2025 •

edited

Loading

jmazanec15 left a comment

jmazanec15 Jan 10, 2025

jmazanec15 Jan 10, 2025

VijayanB Jan 10, 2025

jmazanec15 Jan 10, 2025

VijayanB Jan 10, 2025

jmazanec15 Jan 10, 2025

VijayanB Jan 10, 2025

jmazanec15 Jan 10, 2025

VijayanB Jan 10, 2025

jmazanec15 Jan 10, 2025

jmazanec15 Jan 10, 2025

VijayanB Jan 10, 2025

jmazanec15 Jan 10, 2025

jmazanec15 Jan 10, 2025

VijayanB Jan 10, 2025

jmazanec15 Jan 10, 2025

Add cosine similarity support for faiss engine #2376

Are you sure you want to change the base?

Add cosine similarity support for faiss engine #2376

Conversation

VijayanB commented Jan 9, 2025 • edited Loading

Description

Related Issues

Check List

VijayanB commented Jan 9, 2025 • edited Loading

jmazanec15 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

VijayanB commented Jan 9, 2025 •

edited

Loading

VijayanB commented Jan 9, 2025 •

edited

Loading