Implemented Jaccard distance (#2684)

Summary:
**Summary**
Implemented Jaccard distance requested in this issue: https://github.com/facebookresearch/faiss/issues/1299
**Test plan**
Run: make -C build test
Output: 100% tests passed, 0 tests failed out of 174

Pull Request resolved: https://github.com/facebookresearch/faiss/pull/2684

Reviewed By: mdouze

Differential Revision: D43398833

Pulled By: lvoursl

fbshipit-source-id: b38cf27a7858842efe967bcb1033977863716a76

contrib/exhaustive_search.py

@@ -20,7 +20,7 @@ def knn_ground_truth(xq, db_iterator, k, metric_type=faiss.METRIC_L2):
     LOG.info("knn_ground_truth queries size %s k=%d" % (xq.shape, k))
     t0 = time.time()
     nq, d = xq.shape
-    keep_max = metric_type == faiss.METRIC_INNER_PRODUCT
+    keep_max = faiss.is_similarity_metric(metric_type)
     rh = faiss.ResultHeap(nq, k, keep_max=keep_max)
 
     index = faiss.IndexFlat(d, metric_type)
@@ -265,7 +265,7 @@ def range_search_max_results(index, query_iterator, radius,
                      (totres, max_results))
             radius, totres = apply_maxres(
                 res_batches, min_results,
-                keep_max=index.metric_type == faiss.METRIC_INNER_PRODUCT
+                keep_max=faiss.is_similarity_metric(index.metric_type)
             )
         t2 = time.time()
         t_search += t1 - t0
@@ -281,7 +281,7 @@ def range_search_max_results(index, query_iterator, radius,
     if clip_to_min and totres > min_results:
         radius, totres = apply_maxres(
             res_batches, min_results,
-            keep_max=index.metric_type == faiss.METRIC_INNER_PRODUCT
+            keep_max=faiss.is_similarity_metric(index.metric_type)
         )
 
     nres = np.hstack([nres_i for nres_i, dis_i, ids_i in res_batches])

contrib/ivf_tools.py

@@ -138,7 +138,7 @@ class BigBatchSearcher:
         self.xq = xq
         self.index = index
         self.use_float16 = use_float16
-        keep_max = index.metric_type == faiss.METRIC_INNER_PRODUCT
+        keep_max = faiss.is_similarity_metric(index.metric_type)
         self.rh = faiss.ResultHeap(len(xq), k, keep_max=keep_max)
         self.t_accu = [0] * 5
         self.t_display = self.t0 = time.time()

faiss/Clustering.cpp

@@ -372,7 +372,7 @@ void Clustering::train_encoded(
     std::unique_ptr<float[]> dis(new float[nx]);
 
     // remember best iteration for redo
-    bool lower_is_better = index.metric_type != METRIC_INNER_PRODUCT;
+    bool lower_is_better = !is_similarity_metric(index.metric_type);
     float best_obj = lower_is_better ? HUGE_VALF : -HUGE_VALF;
     std::vector<ClusteringIterationStats> best_iteration_stats;
     std::vector<float> best_centroids;

faiss/IndexFlat.cpp

@@ -40,6 +40,10 @@ void IndexFlat::search(
     } else if (metric_type == METRIC_L2) {
         float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances};
         knn_L2sqr(x, get_xb(), d, n, ntotal, &res, nullptr, sel);
+    } else if (is_similarity_metric(metric_type)) {
+        float_minheap_array_t res = {size_t(n), size_t(k), labels, distances};
+        knn_extra_metrics(
+                x, get_xb(), d, n, ntotal, metric_type, metric_arg, &res);
     } else {
         FAISS_THROW_IF_NOT(!sel);
         float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances};

faiss/IndexHNSW.cpp

@@ -98,7 +98,7 @@ struct NegativeDistanceComputer : DistanceComputer {
 };
 
 DistanceComputer* storage_distance_computer(const Index* storage) {
-    if (storage->metric_type == METRIC_INNER_PRODUCT) {
+    if (is_similarity_metric(storage->metric_type)) {
         return new NegativeDistanceComputer(storage->get_distance_computer());
     } else {
         return storage->get_distance_computer();
@@ -346,7 +346,7 @@ void IndexHNSW::search(
         InterruptCallback::check();
     }
 
-    if (metric_type == METRIC_INNER_PRODUCT) {
+    if (is_similarity_metric(metric_type)) {
         // we need to revert the negated distances
         for (size_t i = 0; i < k * n; i++) {
             distances[i] = -distances[i];

faiss/IndexIVFAdditiveQuantizer.cpp

@@ -145,7 +145,7 @@ struct AQInvertedListScanner : InvertedListScanner {
             : ia(ia), aq(*ia.aq) {
         this->store_pairs = store_pairs;
         this->code_size = ia.code_size;
-        keep_max = ia.metric_type == METRIC_INNER_PRODUCT;
+        keep_max = is_similarity_metric(ia.metric_type);
         tmp.resize(ia.d);
     }
 

faiss/IndexNNDescent.cpp

@@ -88,7 +88,7 @@ struct NegativeDistanceComputer : DistanceComputer {
 };
 
 DistanceComputer* storage_distance_computer(const Index* storage) {
-    if (storage->metric_type == METRIC_INNER_PRODUCT) {
+    if (is_similarity_metric(storage->metric_type)) {
         return new NegativeDistanceComputer(storage->get_distance_computer());
     } else {
         return storage->get_distance_computer();

faiss/IndexNSG.cpp

@@ -112,7 +112,7 @@ void IndexNSG::search(
         InterruptCallback::check();
     }
 
-    if (metric_type == METRIC_INNER_PRODUCT) {
+    if (is_similarity_metric(metric_type)) {
         // we need to revert the negated distances
         for (size_t i = 0; i < k * n; i++) {
             distances[i] = -distances[i];

faiss/MetaIndexes.cpp

@@ -218,7 +218,7 @@ void IndexRandom::search(
         float dprev = 0;
         for (int j = 0; j < k; j++) {
             float step = rng.rand_float();
-            if (metric_type == METRIC_INNER_PRODUCT) {
+            if (is_similarity_metric(metric_type)) {
                 step = -step;
             }
             dprev += step;

faiss/MetricType.h

@@ -31,11 +31,20 @@ enum MetricType {
     METRIC_Canberra = 20,
     METRIC_BrayCurtis,
     METRIC_JensenShannon,
+    METRIC_Jaccard, ///< defined as: sum_i(min(a_i, b_i)) / sum_i(max(a_i, b_i))
+                    ///< where a_i, b_i > 0
 };
 
 /// all vector indices are this type
 using idx_t = int64_t;
 
+/// this function is used to distinguish between min and max indexes since
+/// we need to support similarity and dis-similarity metrics in a flexible way
+constexpr bool is_similarity_metric(MetricType metric_type) {
+    return ((metric_type == METRIC_INNER_PRODUCT) ||
+            (metric_type == METRIC_Jaccard));
+}
+
 } // namespace faiss
 
 #endif

faiss/gpu/impl/Distance.cuh

@@ -228,6 +228,13 @@ void allPairwiseDistanceOnDevice(
                     outDistances,
                     JensenShannonDistance(),
                     stream);
+        } else if (metric == faiss::MetricType::METRIC_Jaccard) {
+            runGeneralDistanceKernel(
+                    tVectorsDimInnermost,
+                    tQueriesDimInnermost,
+                    outDistances,
+                    JaccardSimilarity(),
+                    stream);
         } else {
             FAISS_THROW_FMT("unimplemented metric type %d", metric);
         }
@@ -375,6 +382,16 @@ void bfKnnOnDevice(
                     JensenShannonDistance(),
                     outDistances,
                     outIndices);
+        } else if (metric == faiss::MetricType::METRIC_Jaccard) {
+            runGeneralDistance(
+                    resources,
+                    stream,
+                    tVectorsDimInnermost,
+                    tQueriesDimInnermost,
+                    k,
+                    JaccardSimilarity(),
+                    outDistances,
+                    outIndices);
         } else {
             FAISS_THROW_FMT("unimplemented metric type %d", metric);
         }

faiss/gpu/impl/DistanceUtils.cuh

@@ -251,6 +251,35 @@ struct JensenShannonDistance {
     float dist;
 };
 
+struct JaccardSimilarity {
+    __host__ __device__ JaccardSimilarity() : numerator(0), denominator(0) {}
+
+    static constexpr bool kDirection = true; // maximize
+    static constexpr float kIdentityData = 0;
+    static constexpr float kMaxDistance = -std::numeric_limits<float>::max();
+
+    __host__ __device__ void handle(float a, float b) {
+        numerator += fmin(a, b);
+        denominator += fmax(a, b);
+    }
+
+    __host__ __device__ float reduce() {
+        return numerator / denominator;
+    }
+
+    __host__ __device__ void combine(const JaccardSimilarity& v) {
+        numerator += v.numerator;
+        denominator += v.denominator;
+    }
+
+    __host__ __device__ JaccardSimilarity zero() const {
+        return JaccardSimilarity();
+    }
+
+    float numerator;
+    float denominator;
+};
+
 template <typename T, bool InnerContig>
 Tensor<T, 2, InnerContig> sliceCentroids(
         Tensor<T, 2, InnerContig>& centroids,

faiss/gpu/test/TestGpuDistance.cu

@@ -37,7 +37,8 @@ void testTransposition(
     std::vector<float> vecs = randVecs(numVecs, dim);
     std::vector<float> queries = randVecs(numQuery, dim);
 
-    if (metric == faiss::MetricType::METRIC_JensenShannon) {
+    if ((metric == faiss::MetricType::METRIC_JensenShannon) ||
+        (metric == faiss::MetricType::METRIC_Jaccard)) {
         // make values positive
         for (auto& v : vecs) {
             v = std::abs(v);
@@ -192,6 +193,10 @@ TEST(TestGpuDistance, JensenShannon) {
     testTransposition(false, false, faiss::MetricType::METRIC_JensenShannon);
 }
 
+TEST(TestGpuDistance, Jaccard) {
+    testTransposition(false, false, faiss::MetricType::METRIC_Jaccard);
+}
+
 int main(int argc, char** argv) {
     testing::InitGoogleTest(&argc, argv);
 

faiss/gpu/test/test_gpu_basics.py

@@ -305,7 +305,8 @@ class TestAllPairwiseDistance(unittest.TestCase):
             faiss.METRIC_Linf,
             faiss.METRIC_Canberra,
             faiss.METRIC_BrayCurtis,
-            faiss.METRIC_JensenShannon
+            faiss.METRIC_JensenShannon,
+            faiss.METRIC_Jaccard
         ]
 
         for metric in metrics:
@@ -349,7 +350,7 @@ class TestAllPairwiseDistance(unittest.TestCase):
 
             # INNER_PRODUCT is in descending order, make sure it is the same
             # order
-            if metric == faiss.METRIC_INNER_PRODUCT:
+            if faiss.is_similarity_metric(metric):
                 ref_d = np.sort(ref_d, axis=1)
 
             print('f32', np.abs(ref_d - out_d).max())
@@ -382,7 +383,7 @@ class TestAllPairwiseDistance(unittest.TestCase):
 
             # INNER_PRODUCT is in descending order, make sure it is the same
             # order
-            if metric == faiss.METRIC_INNER_PRODUCT:
+            if faiss.is_similarity_metric(metric):
                 ref_d_f16 = np.sort(ref_d_f16, axis=1)
 
             print('f16', np.abs(ref_d_f16 - out_d_f16).max())

faiss/impl/NSG.cpp

@@ -57,7 +57,7 @@ struct NegativeDistanceComputer : DistanceComputer {
 } // namespace
 
 DistanceComputer* storage_distance_computer(const Index* storage) {
-    if (storage->metric_type == METRIC_INNER_PRODUCT) {
+    if (is_similarity_metric(storage->metric_type)) {
         return new NegativeDistanceComputer(storage->get_distance_computer());
     } else {
         return storage->get_distance_computer();

faiss/utils/extra_distances-inl.h

@@ -8,6 +8,7 @@
 /** In this file are the implementations of extra metrics beyond L2
  *  and inner product */
 
+#include <faiss/MetricType.h>
 #include <faiss/utils/distances.h>
 #include <type_traits>
 
@@ -17,12 +18,13 @@ template <MetricType mt>
 struct VectorDistance {
     size_t d;
     float metric_arg;
+    static constexpr bool is_similarity = is_similarity_metric(mt);
 
     inline float operator()(const float* x, const float* y) const;
 
     // heap template to use for this type of metric
     using C = typename std::conditional<
-            mt == METRIC_INNER_PRODUCT,
+            is_similarity_metric(mt),
             CMin<float, int64_t>,
             CMax<float, int64_t>>::type;
 };
@@ -114,4 +116,18 @@ inline float VectorDistance<METRIC_JensenShannon>::operator()(
     return 0.5 * accu;
 }
 
+template <>
+inline float VectorDistance<METRIC_Jaccard>::operator()(
+        const float* x,
+        const float* y) const {
+    // WARNING: this distance is defined only for positive input vectors.
+    // Providing vectors with negative values would lead to incorrect results.
+    float accu_num = 0, accu_den = 0;
+    for (size_t i = 0; i < d; i++) {
+        accu_num += fmin(x[i], y[i]);
+        accu_den += fmax(x[i], y[i]);
+    }
+    return accu_num / accu_den;
+}
+
 } // namespace faiss

faiss/utils/extra_distances.cpp

@@ -50,14 +50,14 @@ void pairwise_extra_distances_template(
     }
 }
 
-template <class VD>
+template <class VD, class C>
 void knn_extra_metrics_template(
         VD vd,
         const float* x,
         const float* y,
         size_t nx,
         size_t ny,
-        float_maxheap_array_t* res) {
+        HeapArray<C>* res) {
     size_t k = res->k;
     size_t d = vd.d;
     size_t check_period = InterruptCallback::get_period_hint(ny * d);
@@ -74,16 +74,21 @@ void knn_extra_metrics_template(
             float* simi = res->get_val(i);
             int64_t* idxi = res->get_ids(i);
 
-            maxheap_heapify(k, simi, idxi);
+            // maxheap_heapify(k, simi, idxi);
+            heap_heapify<C>(k, simi, idxi);
             for (j = 0; j < ny; j++) {
                 float disij = vd(x_i, y_j);
 
-                if (disij < simi[0]) {
-                    maxheap_replace_top(k, simi, idxi, disij, j);
+                // if (disij < simi[0]) {
+                if ((!vd.is_similarity && (disij < simi[0])) ||
+                    (vd.is_similarity && (disij > simi[0]))) {
+                    // maxheap_replace_top(k, simi, idxi, disij, j);
+                    heap_replace_top<C>(k, simi, idxi, disij, j);
                 }
                 y_j += d;
             }
-            maxheap_reorder(k, simi, idxi);
+            // maxheap_reorder(k, simi, idxi);
+            heap_reorder<C>(k, simi, idxi);
         }
         InterruptCallback::check();
     }
@@ -158,12 +163,14 @@ void pairwise_extra_distances(
         HANDLE_VAR(BrayCurtis);
         HANDLE_VAR(JensenShannon);
         HANDLE_VAR(Lp);
+        HANDLE_VAR(Jaccard);
 #undef HANDLE_VAR
         default:
             FAISS_THROW_MSG("metric type not implemented");
     }
 }
 
+template <class C>
 void knn_extra_metrics(
         const float* x,
         const float* y,
@@ -172,7 +179,7 @@ void knn_extra_metrics(
         size_t ny,
         MetricType mt,
         float metric_arg,
-        float_maxheap_array_t* res) {
+        HeapArray<C>* res) {
     switch (mt) {
 #define HANDLE_VAR(kw)                                            \
     case METRIC_##kw: {                                           \
@@ -187,12 +194,33 @@ void knn_extra_metrics(
         HANDLE_VAR(BrayCurtis);
         HANDLE_VAR(JensenShannon);
         HANDLE_VAR(Lp);
+        HANDLE_VAR(Jaccard);
 #undef HANDLE_VAR
         default:
             FAISS_THROW_MSG("metric type not implemented");
     }
 }
 
+template void knn_extra_metrics<CMax<float, int64_t>>(
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        MetricType mt,
+        float metric_arg,
+        HeapArray<CMax<float, int64_t>>* res);
+
+template void knn_extra_metrics<CMin<float, int64_t>>(
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        MetricType mt,
+        float metric_arg,
+        HeapArray<CMin<float, int64_t>>* res);
+
 FlatCodesDistanceComputer* get_extra_distance_computer(
         size_t d,
         MetricType mt,
@@ -213,6 +241,7 @@ FlatCodesDistanceComputer* get_extra_distance_computer(
         HANDLE_VAR(BrayCurtis);
         HANDLE_VAR(JensenShannon);
         HANDLE_VAR(Lp);
+        HANDLE_VAR(Jaccard);
 #undef HANDLE_VAR
         default:
             FAISS_THROW_MSG("metric type not implemented");

faiss/utils/extra_distances.h

@@ -33,6 +33,7 @@ void pairwise_extra_distances(
         int64_t ldb = -1,
         int64_t ldd = -1);
 
+template <class C>
 void knn_extra_metrics(
         const float* x,
         const float* y,
@@ -41,7 +42,7 @@ void knn_extra_metrics(
         size_t ny,
         MetricType mt,
         float metric_arg,
-        float_maxheap_array_t* res);
+        HeapArray<C>* res);
 
 /** get a DistanceComputer that refers to this type of distance and
  *  indexes a flat array of size nb */

tests/test_extra_distances.py

@@ -82,6 +82,18 @@ class TestExtraDistances(unittest.TestCase):
         self.run_simple_dis_test(scipy.spatial.distance.jensenshannon,
                                  faiss.METRIC_JensenShannon)
 
+    def test_jaccard(self):
+        xq, yb = self.make_example()
+        ref_dis = np.array([
+            [
+                (np.min([x, y], axis=0).sum() / np.max([x, y], axis=0).sum())
+                for y in yb
+            ]
+            for x in xq
+        ])
+        new_dis = faiss.pairwise_distances(xq, yb, faiss.METRIC_Jaccard)
+        self.assertTrue(np.allclose(ref_dis, new_dis))
+
 
 class TestKNN(unittest.TestCase):
     """ test that the knn search gives the same as distance matrix + argmin """