faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/Clustering.cpp

@@ -11,15 +11,14 @@
 #include <faiss/VectorTransform.h>
 #include <faiss/impl/AuxIndexStructures.h>
 
-#include <chrono>
 #include <cinttypes>
 #include <cmath>
 #include <cstdio>
 #include <cstring>
-
-#include <omp.h>
+#include <limits>
 
 #include <faiss/IndexFlat.h>
+#include <faiss/impl/ClusteringHelpers.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/kmeans1d.h>
 #include <faiss/utils/distances.h>
@@ -28,10 +27,10 @@
 
 namespace faiss {
 
-Clustering::Clustering(int d, int k) : d(d), k(k) {}
+Clustering::Clustering(int d_, int k_) : d(d_), k(k_) {}
 
-Clustering::Clustering(int d, int k, const ClusteringParameters& cp)
-        : ClusteringParameters(cp), d(d), k(k) {}
+Clustering::Clustering(int d_, int k_, const ClusteringParameters& cp)
+        : ClusteringParameters(cp), d(d_), k(k_) {}
 
 void Clustering::post_process_centroids() {
     if (spherical) {
@@ -58,213 +57,6 @@ void Clustering::train(
             weights);
 }
 
-namespace {
-
-uint64_t get_actual_rng_seed(const int seed) {
-    return (seed >= 0)
-            ? seed
-            : static_cast<uint64_t>(std::chrono::high_resolution_clock::now()
-                                            .time_since_epoch()
-                                            .count());
-}
-
-idx_t subsample_training_set(
-        const Clustering& clus,
-        idx_t nx,
-        const uint8_t* x,
-        size_t line_size,
-        const float* weights,
-        uint8_t** x_out,
-        float** weights_out) {
-    if (clus.verbose) {
-        printf("Sampling a subset of %zd / %" PRId64 " for training\n",
-               clus.k * clus.max_points_per_centroid,
-               nx);
-    }
-
-    const uint64_t actual_seed = get_actual_rng_seed(clus.seed);
-
-    std::vector<int> perm;
-    if (clus.use_faster_subsampling) {
-        // use subsampling with splitmix64 rng
-        SplitMix64RandomGenerator rng(actual_seed);
-
-        const idx_t new_nx = clus.k * clus.max_points_per_centroid;
-        perm.resize(new_nx);
-        for (idx_t i = 0; i < new_nx; i++) {
-            perm[i] = rng.rand_int(nx);
-        }
-    } else {
-        // use subsampling with a default std rng
-        perm.resize(nx);
-        rand_perm(perm.data(), nx, actual_seed);
-    }
-
-    nx = clus.k * clus.max_points_per_centroid;
-    uint8_t* x_new = new uint8_t[nx * line_size];
-    *x_out = x_new;
-
-    // might be worth omp-ing as well
-    for (idx_t i = 0; i < nx; i++) {
-        memcpy(x_new + i * line_size, x + perm[i] * line_size, line_size);
-    }
-    if (weights) {
-        float* weights_new = new float[nx];
-        for (idx_t i = 0; i < nx; i++) {
-            weights_new[i] = weights[perm[i]];
-        }
-        *weights_out = weights_new;
-    } else {
-        *weights_out = nullptr;
-    }
-    return nx;
-}
-
-/** compute centroids as (weighted) sum of training points
- *
- * @param x            training vectors, size n * code_size (from codec)
- * @param codec        how to decode the vectors (if NULL then cast to float*)
- * @param weights      per-training vector weight, size n (or NULL)
- * @param assign       nearest centroid for each training vector, size n
- * @param k_frozen     do not update the k_frozen first centroids
- * @param centroids    centroid vectors (output only), size k * d
- * @param hassign      histogram of assignments per centroid (size k),
- *                     should be 0 on input
- *
- */
-
-void compute_centroids(
-        size_t d,
-        size_t k,
-        size_t n,
-        size_t k_frozen,
-        const uint8_t* x,
-        const Index* codec,
-        const int64_t* assign,
-        const float* weights,
-        float* hassign,
-        float* centroids) {
-    k -= k_frozen;
-    centroids += k_frozen * d;
-
-    memset(centroids, 0, sizeof(*centroids) * d * k);
-
-    size_t line_size = codec ? codec->sa_code_size() : d * sizeof(float);
-
-#pragma omp parallel
-    {
-        int nt = omp_get_num_threads();
-        int rank = omp_get_thread_num();
-
-        // this thread is taking care of centroids c0:c1
-        size_t c0 = (k * rank) / nt;
-        size_t c1 = (k * (rank + 1)) / nt;
-        std::vector<float> decode_buffer(d);
-
-        for (size_t i = 0; i < n; i++) {
-            int64_t ci = assign[i];
-            assert(ci >= 0 && ci < k + k_frozen);
-            ci -= k_frozen;
-            if (ci >= c0 && ci < c1) {
-                float* c = centroids + ci * d;
-                const float* xi;
-                if (!codec) {
-                    xi = reinterpret_cast<const float*>(x + i * line_size);
-                } else {
-                    float* xif = decode_buffer.data();
-                    codec->sa_decode(1, x + i * line_size, xif);
-                    xi = xif;
-                }
-                if (weights) {
-                    float w = weights[i];
-                    hassign[ci] += w;
-                    for (size_t j = 0; j < d; j++) {
-                        c[j] += xi[j] * w;
-                    }
-                } else {
-                    hassign[ci] += 1.0;
-                    for (size_t j = 0; j < d; j++) {
-                        c[j] += xi[j];
-                    }
-                }
-            }
-        }
-    }
-
-#pragma omp parallel for
-    for (idx_t ci = 0; ci < k; ci++) {
-        if (hassign[ci] == 0) {
-            continue;
-        }
-        float norm = 1 / hassign[ci];
-        float* c = centroids + ci * d;
-        for (size_t j = 0; j < d; j++) {
-            c[j] *= norm;
-        }
-    }
-}
-
-// a bit above machine epsilon for float16
-#define EPS (1 / 1024.)
-
-/** Handle empty clusters by splitting larger ones.
- *
- * It works by slightly changing the centroids to make 2 clusters from
- * a single one. Takes the same arguments as compute_centroids.
- *
- * @return           nb of splitting operations (larger is worse)
- */
-int split_clusters(
-        size_t d,
-        size_t k,
-        size_t n,
-        size_t k_frozen,
-        float* hassign,
-        float* centroids) {
-    k -= k_frozen;
-    centroids += k_frozen * d;
-
-    /* Take care of void clusters */
-    size_t nsplit = 0;
-    RandomGenerator rng(1234);
-    for (size_t ci = 0; ci < k; ci++) {
-        if (hassign[ci] == 0) { /* need to redefine a centroid */
-            size_t cj;
-            for (cj = 0; true; cj = (cj + 1) % k) {
-                /* probability to pick this cluster for split */
-                float p = (hassign[cj] - 1.0) / (float)(n - k);
-                float r = rng.rand_float();
-                if (r < p) {
-                    break; /* found our cluster to be split */
-                }
-            }
-            memcpy(centroids + ci * d,
-                   centroids + cj * d,
-                   sizeof(*centroids) * d);
-
-            /* small symmetric perturbation */
-            for (size_t j = 0; j < d; j++) {
-                if (j % 2 == 0) {
-                    centroids[ci * d + j] *= 1 + EPS;
-                    centroids[cj * d + j] *= 1 - EPS;
-                } else {
-                    centroids[ci * d + j] *= 1 - EPS;
-                    centroids[cj * d + j] *= 1 + EPS;
-                }
-            }
-
-            /* assume even split of the cluster */
-            hassign[ci] = hassign[cj] / 2;
-            hassign[cj] -= hassign[ci];
-            nsplit++;
-        }
-    }
-
-    return nsplit;
-}
-
-} // namespace
-
 void Clustering::train_encoded(
         idx_t nx,
         const uint8_t* x_in,
@@ -272,7 +64,7 @@ void Clustering::train_encoded(
         Index& index,
         const float* weights) {
     FAISS_THROW_IF_NOT_FMT(
-            nx >= k,
+            nx >= static_cast<idx_t>(k),
             "Number of training points (%" PRId64
             ") should be at least "
             "as large as number of clusters (%zd)",
@@ -280,13 +72,13 @@ void Clustering::train_encoded(
             k);
 
     FAISS_THROW_IF_NOT_FMT(
-            (!codec || codec->d == d),
+            (!codec || static_cast<size_t>(codec->d) == d),
             "Codec dimension %d not the same as data dimension %d",
             int(codec->d),
             int(d));
 
     FAISS_THROW_IF_NOT_FMT(
-            index.d == d,
+            static_cast<size_t>(index.d) == d,
             "Index dimension %d not the same as data dimension %d",
             int(index.d),
             int(d));
@@ -309,16 +101,16 @@ void Clustering::train_encoded(
     std::unique_ptr<float[]> del3;
     size_t line_size = codec ? codec->sa_code_size() : sizeof(float) * d;
 
-    if (nx > k * max_points_per_centroid) {
+    if (static_cast<size_t>(nx) > k * max_points_per_centroid) {
         uint8_t* x_new;
         float* weights_new;
-        nx = subsample_training_set(
+        nx = detail::subsample_training_set(
                 *this, nx, x, line_size, weights, &x_new, &weights_new);
         del1.reset(x_new);
         x = x_new;
         del3.reset(weights_new);
         weights = weights_new;
-    } else if (nx < k * min_points_per_centroid) {
+    } else if (static_cast<size_t>(nx) < k * min_points_per_centroid) {
         fprintf(stderr,
                 "WARNING clustering %" PRId64
                 " points to %zd centroids: "
@@ -328,7 +120,7 @@ void Clustering::train_encoded(
                 idx_t(k) * min_points_per_centroid);
     }
 
-    if (nx == k) {
+    if (static_cast<size_t>(nx) == k) {
         // this is a corner case, just copy training set to clusters
         if (verbose) {
             printf("Number of training points (%" PRId64
@@ -397,7 +189,7 @@ void Clustering::train_encoded(
     t0 = getmillisecs();
 
     // initialize seed
-    const uint64_t actual_seed = get_actual_rng_seed(seed);
+    const uint64_t actual_seed = detail::get_actual_rng_seed(seed);
 
     // temporary buffer to decode vectors during the optimization
     std::vector<float> decode_buffer(codec ? d * decode_block_size : 0);
@@ -407,19 +199,52 @@ void Clustering::train_encoded(
             printf("Outer iteration %d / %d\n", redo, nredo);
         }
 
-        // initialize (remaining) centroids with random points from the dataset
+        // initialize centroids using the selected method
         centroids.resize(d * k);
-        std::vector<int> perm(nx);
 
-        rand_perm(perm.data(), nx, actual_seed + 1 + redo * 15486557L);
+        size_t k_to_init = k - n_input_centroids;
+        if (k_to_init > 0) {
+            // Fast path for RANDOM initialization - preserves exact original
+            // behavior
+            if (init_method == ClusteringInitMethod::RANDOM) {
+                std::vector<int> perm(nx);
+                rand_perm(perm.data(), nx, actual_seed + 1 + redo * 15486557L);
+                for (size_t i = 0; i < k_to_init; i++) {
+                    if (!codec) {
+                        memcpy(centroids.data() + (n_input_centroids + i) * d,
+                               x + perm[n_input_centroids + i] * line_size,
+                               line_size);
+                    } else {
+                        codec->sa_decode(
+                                1,
+                                x + perm[n_input_centroids + i] * line_size,
+                                centroids.data() + (n_input_centroids + i) * d);
+                    }
+                }
+            } else {
+                // For k-means++ and AFK-MC², we need all vectors decoded
+                const float* x_float = nullptr;
+                std::vector<float> x_decoded;
 
-        if (!codec) {
-            for (int i = n_input_centroids; i < k; i++) {
-                memcpy(&centroids[i * d], x + perm[i] * line_size, line_size);
-            }
-        } else {
-            for (int i = n_input_centroids; i < k; i++) {
-                codec->sa_decode(1, x + perm[i] * line_size, &centroids[i * d]);
+                if (!codec) {
+                    x_float = reinterpret_cast<const float*>(x);
+                } else {
+                    // Decode all vectors for initialization
+                    x_decoded.resize(nx * d);
+                    codec->sa_decode(nx, x, x_decoded.data());
+                    x_float = x_decoded.data();
+                }
+
+                ClusteringInitialization initializer(d, k_to_init);
+                initializer.method = init_method;
+                initializer.seed = actual_seed + 1 + redo * 15486557L;
+                initializer.afkmc2_chain_length = afkmc2_chain_length;
+                initializer.init_centroids(
+                        nx,
+                        x_float,
+                        centroids.data() + n_input_centroids * d,
+                        n_input_centroids,
+                        n_input_centroids > 0 ? centroids.data() : nullptr);
             }
         }
 
@@ -453,9 +278,10 @@ void Clustering::train_encoded(
             } else {
                 // search by blocks of decode_block_size vectors
                 size_t code_size = codec->sa_code_size();
-                for (size_t i0 = 0; i0 < nx; i0 += decode_block_size) {
+                for (size_t i0 = 0; i0 < static_cast<size_t>(nx);
+                     i0 += decode_block_size) {
                     size_t i1 = i0 + decode_block_size;
-                    if (i1 > nx) {
+                    if (i1 > static_cast<size_t>(nx)) {
                         i1 = nx;
                     }
                     codec->sa_decode(
@@ -474,7 +300,7 @@ void Clustering::train_encoded(
 
             // accumulate objective
             obj = 0;
-            for (int j = 0; j < nx; j++) {
+            for (idx_t j = 0; j < nx; j++) {
                 obj += dis[j];
             }
 
@@ -482,7 +308,7 @@ void Clustering::train_encoded(
             std::vector<float> hassign(k);
 
             size_t k_frozen = frozen_centroids ? n_input_centroids : 0;
-            compute_centroids(
+            detail::compute_centroids(
                     d,
                     k,
                     nx,
@@ -494,7 +320,7 @@ void Clustering::train_encoded(
                     hassign.data(),
                     centroids.data());
 
-            int nsplit = split_clusters(
+            int nsplit = detail::split_clusters(
                     d, k, nx, k_frozen, hassign.data(), centroids.data());
 
             // collect statistics
@@ -502,7 +328,7 @@ void Clustering::train_encoded(
                     obj,
                     (getmillisecs() - t0) / 1000.0,
                     t_search_tot / 1000,
-                    imbalance_factor(nx, k, assign.get()),
+                    imbalance_factor(nx, static_cast<int>(k), assign.get()),
                     nsplit};
             iteration_stats.push_back(stats);
 
@@ -529,6 +355,27 @@ void Clustering::train_encoded(
 
             index.add(k, centroids.data());
             InterruptCallback::check();
+
+            // Early stopping: if objective didn't change, we've converged.
+            // Safe to access iteration_stats[size - 2] because we push_back
+            // above, so size >= i + 1, and when i > 0 we have size >= 2.
+            if (i > 0) {
+                float prev_obj =
+                        iteration_stats[iteration_stats.size() - 2].obj;
+
+                double change = (prev_obj == 0)
+                        ? std::numeric_limits<double>::max()
+                        : std::abs(prev_obj - obj) / std::abs(prev_obj);
+
+                if (change >= 0 && change <= early_stop_threshold) {
+                    if (verbose) {
+                        printf("\n  Converged at iteration %d: "
+                               "objective did not change\n",
+                               i);
+                    }
+                    break;
+                }
+            }
         }
 
         if (verbose) {
@@ -555,19 +402,19 @@ void Clustering::train_encoded(
     }
 }
 
-Clustering1D::Clustering1D(int k) : Clustering(1, k) {}
+Clustering1D::Clustering1D(int k_) : Clustering(1, k_) {}
 
-Clustering1D::Clustering1D(int k, const ClusteringParameters& cp)
-        : Clustering(1, k, cp) {}
+Clustering1D::Clustering1D(int k_, const ClusteringParameters& cp)
+        : Clustering(1, k_, cp) {}
 
 void Clustering1D::train_exact(idx_t n, const float* x) {
     const float* xt = x;
 
     std::unique_ptr<uint8_t[]> del;
-    if (n > k * max_points_per_centroid) {
+    if (static_cast<size_t>(n) > k * max_points_per_centroid) {
         uint8_t* x_new;
         float* weights_new;
-        n = subsample_training_set(
+        n = detail::subsample_training_set(
                 *this,
                 n,
                 (uint8_t*)x,
@@ -592,7 +439,7 @@ float kmeans_clustering(
         size_t k,
         const float* x,
         float* centroids) {
-    Clustering clus(d, k);
+    Clustering clus(static_cast<int>(d), static_cast<int>(k));
     clus.verbose = d * n * k > (size_t(1) << 30);
     // display logs if > 1Gflop per iteration
     IndexFlatL2 index(d);
@@ -615,13 +462,14 @@ Index* ProgressiveDimIndexFactory::operator()(int dim) {
     return new IndexFlatL2(dim);
 }
 
-ProgressiveDimClustering::ProgressiveDimClustering(int d, int k) : d(d), k(k) {}
+ProgressiveDimClustering::ProgressiveDimClustering(int d_, int k_)
+        : d(d_), k(k_) {}
 
 ProgressiveDimClustering::ProgressiveDimClustering(
-        int d,
-        int k,
+        int d_,
+        int k_,
         const ProgressiveDimClusteringParameters& cp)
-        : ProgressiveDimClusteringParameters(cp), d(d), k(k) {}
+        : ProgressiveDimClusteringParameters(cp), d(d_), k(k_) {}
 
 namespace {
 
@@ -642,7 +490,7 @@ void ProgressiveDimClustering::train(
         ProgressiveDimIndexFactory& factory) {
     int d_prev = 0;
 
-    PCAMatrix pca(d, d);
+    PCAMatrix pca(static_cast<int>(d), static_cast<int>(d));
 
     std::vector<float> xbuf;
     if (apply_pca) {
@@ -667,7 +515,7 @@ void ProgressiveDimClustering::train(
         }
         std::unique_ptr<Index> clustering_index(factory(di));
 
-        Clustering clus(di, k, *this);
+        Clustering clus(di, static_cast<int>(k), *this);
         if (d_prev > 0) {
             // copy warm-start centroids (padded with 0s)
             clus.centroids.resize(k * di);

data/vendor/faiss/faiss/Clustering.h

@@ -10,6 +10,7 @@
 #ifndef FAISS_CLUSTERING_H
 #define FAISS_CLUSTERING_H
 #include <faiss/Index.h>
+#include <faiss/impl/ClusteringInitialization.h>
 
 #include <vector>
 
@@ -57,6 +58,23 @@ struct ClusteringParameters {
     /// Whether to use splitmix64-based random number generator for subsampling,
     /// which is faster, but may pick duplicate points.
     bool use_faster_subsampling = false;
+
+    /// Initialization method for centroids.
+    /// RANDOM: uniform random sampling (default, current behavior)
+    /// KMEANS_PLUS_PLUS: k-means++ (O(nkd), better quality)
+    /// AFK_MC2: Assumption-Free K-MC² (O(nd) + O(mk²d), fast approximation)
+    ClusteringInitMethod init_method = ClusteringInitMethod::RANDOM;
+
+    /// Chain length for AFK-MC² initialization.
+    /// Only used when init_method = AFK_MC2.
+    /// Longer chains give better approximation but are slower.
+    uint16_t afkmc2_chain_length = 50;
+
+    /// Early stop threshold, the range is [0, 1].
+    /// The value of 0 implies a default Faiss behavior,
+    /// so the training process stops only if an error
+    /// is unchanged from the previous iteration.
+    double early_stop_threshold = 0.0;
 };
 
 struct ClusteringIterationStats {