feat(svs): LeanVec OOD support (facebookresearch#4773)

Summary:
This PR will add LeanVec OOD (out-of-distribution) support.

Pull Request resolved: facebookresearch#4773

Reviewed By: alibeklfc

Differential Revision: D94607022

Pulled By: mnorris11

fbshipit-source-id: df1b023919b6a2753fdddcd7d5490475f8630d36

Author: ahuber21

faiss/CMakeLists.txt

@@ -460,7 +460,7 @@ if(FAISS_OPT_LEVEL STREQUAL "dd")
 endif()
 
 if(FAISS_ENABLE_SVS)
-  find_package(svs_runtime REQUIRED)
+  find_package(svs_runtime 0.2.0 REQUIRED)
 
   target_link_libraries(faiss PUBLIC svs::svs_runtime)
   target_link_libraries(faiss_avx2 PUBLIC svs::svs_runtime)

faiss/Index.cpp

@@ -24,6 +24,14 @@ void Index::train(idx_t /*n*/, const float* /*x*/) {
     // does nothing by default
 }
 
+void Index::train(
+        idx_t /*n*/,
+        const float* /*x*/,
+        idx_t /*n_train_q*/,
+        const float* /*xq_train*/) {
+    // does nothing by default
+}
+
 void Index::range_search(
         idx_t,
         const float*,

faiss/Index.h

@@ -131,6 +131,20 @@ struct Index {
      */
     virtual void train(idx_t n, const float* x);
 
+    /** Perfrom training on a representative set of vectors and a representative
+     * set of queries
+     *
+     * @param n         nb of training vectors
+     * @param x         training vectors, size n * d
+     * @param n_train_q nb of training queries
+     * @param xq_train  training queries, size n_train_q * d
+     */
+    virtual void train(
+            idx_t n,
+            const float* x,
+            idx_t n_train_q,
+            const float* xq_train);
+
     virtual void train_ex(idx_t n, const void* x, NumericType numeric_type) {
         if (numeric_type == NumericType::Float32) {
             train(n, static_cast<const float*>(x));

faiss/python/class_wrappers.py

@@ -332,21 +332,57 @@ def replacement_assign(self, x, k, labels=None):
         self.assign_c(n, swig_ptr(x), swig_ptr(labels), k)
         return labels
 
-    def replacement_train(self, x, numeric_type = faiss.Float32):
+    def replacement_train(
+        self, x, *, numeric_type=faiss.Float32, xq_train=None
+    ):
         """Trains the index on a representative set of vectors.
         The index must be trained before vectors can be added to it.
+        Optionally accepts numeric_type to specify the type of
+        input vectors.
+        Optionally accepts a set of training query vectors for
+        out-of-distribution training.
 
         Parameters
         ----------
         x : array_like
-            Query vectors, shape (n, d) where d is appropriate for the index.
+            Query vectors, shape (n, d) where d is appropriate
+            for the index. `dtype` must be float32.
+        numeric_type : type
+            Numeric type of the input vectors.
+        xq_train : array_like, optional
+            Training query vectors, shape (n_train_q, d) where
+            d is appropriate for the index.
             `dtype` must be float32.
         """
+        # Prepare training data
         n, d = x.shape
         assert d == self.d
         x = np.ascontiguousarray(x, dtype=_numeric_to_str(numeric_type))
+
+        # Prepare training queries if provided
+        n_train_q, train_q = 0, None
+        if xq_train is not None:
+            if numeric_type != faiss.Float32:
+                raise TypeError(
+                    "xq_train is only supported for numeric_type faiss.Float32"
+                )
+            n_train_q, d_train = xq_train.shape
+            assert d_train == self.d
+            train_q = swig_ptr(
+                np.ascontiguousarray(
+                    xq_train,
+                    dtype=_numeric_to_str(numeric_type),
+                )
+            )
+
+        # Dispatch to train_c / train_ex
         if numeric_type == faiss.Float32:
-            self.train_c(n, swig_ptr(x))
+            if train_q is not None:
+                self.train_c(
+                    n, swig_ptr(x), n_train_q, train_q
+                )
+            else:
+                self.train_c(n, swig_ptr(x))
         else:
             self.train_ex(n, swig_ptr(x), numeric_type)
 

faiss/svs/IndexSVSFaissUtils.h

@@ -45,18 +45,6 @@
 // create svs_runtime as alias for svs::runtime::FAISS_SVS_RUNTIME_VERSION
 SVS_RUNTIME_CREATE_API_ALIAS(svs_runtime, FAISS_SVS_RUNTIME_VERSION);
 
-// SVS forward declarations
-namespace svs {
-namespace runtime {
-inline namespace v0 {
-struct FlatIndex;
-struct VamanaIndex;
-struct DynamicVamanaIndex;
-struct LeanVecTrainingData;
-} // namespace v0
-} // namespace runtime
-} // namespace svs
-
 namespace faiss {
 
 inline svs_runtime::MetricType to_svs_metric(faiss::MetricType metric) {
@@ -119,8 +107,8 @@ struct InputBufferConverter {
     std::vector<T> buffer;
 };
 
-// Specialization for reinterpret cast when types are integral and have the same
-// size
+// Specialization for reinterpret cast when types are integral and have
+// the same size
 template <typename T, typename U>
 struct InputBufferConverter<
         T,
@@ -178,8 +166,8 @@ struct OutputBufferConverter {
     std::vector<T> buffer;
 };
 
-// Specialization for reinterpret cast when types are integral and have the same
-// size
+// Specialization for reinterpret cast when types are integral and have
+// the same size
 template <typename T, typename U>
 struct OutputBufferConverter<
         T,

faiss/svs/IndexSVSFlat.h

@@ -26,6 +26,8 @@
 #include <faiss/Index.h>
 #include <faiss/svs/IndexSVSFaissUtils.h>
 
+#include <svs/runtime/flat_index.h>
+
 #include <iostream>
 
 namespace faiss {

faiss/svs/IndexSVSVamana.h

@@ -27,6 +27,7 @@
 #include <faiss/svs/IndexSVSFaissUtils.h>
 
 #include <svs/runtime/api_defs.h>
+#include <svs/runtime/dynamic_vamana_index.h>
 
 #include <iostream>
 

faiss/svs/IndexSVSVamanaLeanVec.cpp

@@ -66,6 +66,14 @@ void IndexSVSVamanaLeanVec::add(idx_t n, const float* x) {
 }
 
 void IndexSVSVamanaLeanVec::train(idx_t n, const float* x) {
+    train(n, x, 0, nullptr);
+}
+
+void IndexSVSVamanaLeanVec::train(
+        idx_t n,
+        const float* x,
+        idx_t n_train_q,
+        const float* queries) {
     FAISS_THROW_IF_MSG(
             training_data || impl, "Index already trained or contains data.");
 
@@ -74,7 +82,7 @@ void IndexSVSVamanaLeanVec::train(idx_t n, const float* x) {
             "LVQ/LeanVec support not available on this platform or build");
 
     auto status = svs_runtime::LeanVecTrainingData::build(
-            &training_data, d, n, x, leanvec_d);
+            &training_data, d, n, x, n_train_q, queries, leanvec_d);
     if (!status.ok()) {
         FAISS_THROW_MSG(status.message());
     }

faiss/svs/IndexSVSVamanaLeanVec.h

@@ -41,8 +41,17 @@ struct IndexSVSVamanaLeanVec : IndexSVSVamana {
 
     void add(idx_t n, const float* x) override;
 
+    /* Default train assumes in-distribution data */
     void train(idx_t n, const float* x) override;
 
+    /* Generic train with out-of-distribution parameters.
+     * Out-of-distribution (OOD) means database vectors and queries _can_ be
+     * sampled from different distributions (e.g., cross-modal). More details in
+     * the original publication, arXiv:2312.16335.
+     */
+    void train(idx_t n, const float* x, idx_t n_train_q, const float* xq_train)
+            override;
+
     void serialize_training_data(std::ostream& out) const;
     void deserialize_training_data(std::istream& in);
 

tests/test_svs_py.py

@@ -580,5 +580,41 @@ def _create_instance(self):
         return idx
 
 
+@unittest.skipIf(_SKIP_SVS_LL, _SKIP_SVS_LL_REASON)
+class TestSVSLeanVecOOD(unittest.TestCase):
+    """Test out-of-distribution training for LeanVec SVS indices"""
+
+    def setUp(self):
+        self.d = 256
+        self.idx = faiss.IndexSVSVamanaLeanVec(
+            self.d, 64, faiss.METRIC_INNER_PRODUCT, 64, faiss.SVS_LeanVec4x8
+        )
+        self.idx.alpha = 0.95
+
+        self.x = np.random.rand(1000, self.d).astype("float32")
+        self.tq = np.random.rand(1000, self.d).astype("float32")
+
+    def test_svs_leanvec_ood_training(self):
+        self.assertIsNone(self.idx.training_data)
+        self.idx.train(self.x, xq_train=self.tq)
+        self.assertIsNotNone(self.idx.training_data)
+
+    def test_svs_leanvec_ood_training_smaller(self):
+        self.idx.train(self.x, xq_train=self.tq[:500])
+
+    def test_svs_leanvec_ood_training_wrong_dim(self):
+        wrong_dim = np.random.rand(1000, self.d + 1).astype("float32")
+        with self.assertRaises(AssertionError):
+            self.idx.train(self.x, xq_train=wrong_dim)
+
+    def test_svs_leanvec_ood_training_wrong_type(self):
+        with self.assertRaises(TypeError):
+            self.idx.train(
+                self.x,
+                xq_train=self.tq,
+                numeric_type=faiss.Float16,
+            )
+
+
 if __name__ == '__main__':
     unittest.main()