JingsongLi commented on code in PR #82: URL: https://github.com/apache/paimon-vector-index/pull/82#discussion_r3869324179
########## core/src/projected_assign.rs: ########## @@ -0,0 +1,1195 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Exact coarse assignment through a low-dimensional projection. +//! +//! `IVFPQIndex::add` must find the nearest coarse centroid of every row. The +//! exact scan is one `n × nlist × d` GEMM, which dominates `add` once PQ +//! encoding is vectorized. This module keeps the result exact but does most of +//! the work in a `d'`-dimensional PCA subspace of the centroids: +//! +//! 1. project the row and centroids with a contractive PCA projection `P` +//! (`d' × d`); +//! 2. for every centroid compute a conservative lower bound on the true +//! squared distance: the projected distance plus the reverse triangle +//! inequality on the components outside the subspace +//! (`|x-c|² >= |P₀(x-c)|² + (|P₀⊥x| - |P₀⊥c|)²`), with every f32 +//! projection, GEMM and f64 rounding error accounted for explicitly; +//! 3. evaluate exact distances in ascending bound order and stop as soon as +//! the next bound cannot beat the best exact distance (branch-and-bound). +//! +//! Step 2 is a valid lower bound for any contractive `P`, so the assignment +//! never depends on how well the PCA converged; `P` only decides how many +//! exact evaluations step 3 needs. Rows without low-dimensional structure +//! degrade to checking every centroid, i.e. the exact scan plus one small +//! GEMM. At `train`, candidate widths are therefore timed on a sample of +//! training vectors, including projection, bound construction, candidate +//! collection, and exact checks. Automatic mode keeps the fastest projection +//! only when it is clearly faster than an exact scan of the same sample. +//! Without a sample, forced mode takes the widest candidate and automatic +//! mode leaves the exact scan in place. +//! +//! The projection is a training artifact: it is derived from the centroids, +//! shared by `IVFPQIndex::from_trained`, and never serialized. + +use crate::blas::{dgemm_a_bt, sgemm_a_bt}; +use crate::distance::{fvec_l2sqr, fvec_l2sqr_four}; +use nalgebra::{DMatrix, SymmetricEigen}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use rayon::prelude::*; +use std::time::{Duration, Instant}; + +/// Automatic mode only times widths with `d' * MAX_AUTO_DP_DIVISOR <= d`. +const MAX_AUTO_DP_DIVISOR: usize = 3; +/// Calibration samples below this count as absent. +const MIN_CALIBRATION_ROWS: usize = 256; +/// Rows of the training data used to calibrate `d'`. +pub(crate) const CALIBRATION_ROWS: usize = 2048; +/// With a calibration sample, automatic mode keeps the projection only when +/// its measured time is below this fraction of the exact scan. +const MAX_AUTO_COST_FRACTION: f64 = 0.85; +/// Block subspace iterations; the bound stays valid regardless of convergence. +const SUBSPACE_ITERATIONS: usize = 8; +const SUBSPACE_SEED: u64 = 0x7a5e_c7ed; +/// Smallest projection worth a GEMM; also the rounding unit for `d'`. +const MIN_DP: usize = 8; +/// Rows per parallel block; the projected score matrix is `rows × nlist`. +const MAX_BLOCK_ROWS: usize = 1024; +/// The projection is re-orthonormalized in f64 and scaled by +/// `1 - ORTHO_MARGIN`, so its Gram matrix is `(1 - ORTHO_MARGIN)² · I` up to +/// rounding. `orthonormal_f64` verifies (Gershgorin, including the f64 +/// rounding of the Gram entries) that every eigenvalue of `P Pᵀ` lies within +/// `GRAM_MARGIN` of that diagonal, and declines the projection otherwise, so +/// the constant singular-value bounds below hold for every stored +/// projection. Gram-Schmidt typically lands within ~1e-15. +const ORTHO_MARGIN: f64 = 1e-9; +const GRAM_MARGIN: f64 = 1e-9; +/// Bounds on the squared singular values of the stored projection (see +/// `ORTHO_MARGIN`): `|P v|² / SIGMA_MAX_SQ <= |P₀ v|² <= |P v|² / SIGMA_MIN_SQ` +/// for the orthonormal basis `P₀` of its row space. +const SIGMA_MAX_SQ: f64 = 1.0; +const SIGMA_MIN_SQ: f64 = (1.0 - ORTHO_MARGIN) * (1.0 - ORTHO_MARGIN) * (1.0 - GRAM_MARGIN); +/// Candidates evaluated per branch-and-bound stage before re-pruning. +const STAGE: usize = 32; + +/// How `IVFPQIndex::add` chooses between the exact centroid scan and the +/// projected branch-and-bound. Both produce the exact nearest centroid. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub enum ProjectedAssignment { + /// Use the projection when the centroids are compressible enough for it + /// to pay off; otherwise scan. + #[default] + Auto, + /// Always build and use the projection (still exact). + Enabled, + /// Always scan. + Disabled, +} + +/// Contractive projection of the coarse centroids plus the per-centroid data +/// the lower bound needs. +#[derive(Clone, Debug)] +pub struct CoarseProjection { + d: usize, + dp: usize, + explained_variance: f64, + /// Centroid mean (f64), subtracted before projecting rows and centroids. + /// Translation cancels in the projected term but not in the residual + /// term, whose norms must not be dominated by a common offset. + mean: Vec<f64>, + /// `dp × d`, scaled so its operator norm does not exceed one. Stored in + /// f64 (exactly the f32 values used for the singular-value bounds) so the + /// row projection carries only f64 rounding. + proj: Vec<f64>, + /// `nlist × dp`: projected centroids. + cents_p: Vec<f32>, + /// `|cents_p[c]|²`. + cents_p_norms: Vec<f64>, + /// `|cents_p[c]|`, cached to keep square roots out of assignment's hot loop. + cents_p_norms_sqrt: Vec<f64>, + /// Upper bound on the f32 GEMM error in each projected centroid. + cents_p_errors: Vec<f64>, + /// Relative error factors of the `rows × nlist` f32 GEMM and of the f64 + /// arithmetic in the bound, fixed by `dp`. + gemm_error_factor: f64, + f64_error_factor: f64, + /// Interval for `|P₀⊥ c|`, the norm of each centroid's component outside + /// the projected subspace, including every rounding error above. + cents_res_lo: Vec<f64>, + cents_res_hi: Vec<f64>, +} + +impl CoarseProjection { + /// Fit a projection to `nlist` centroids of dimension `d`. + /// + /// `calibration` holds `calibration_rows` sample vectors in the centroid + /// space (typically training vectors). Candidate widths up to `d / 3` + /// (`d / 2` when forced) are timed on that sample and the fastest wins; + /// automatic mode keeps the projection only when it also beats the exact + /// scan on the sample by `MAX_AUTO_COST_FRACTION`. Without a sample, + /// forced mode takes the widest width and automatic mode builds nothing. + /// + /// Returns `None` when the projection is not worth it (`force == false`) + /// or cannot be built (too few centroids, zero variance). + pub(crate) fn train( + cents: &[f32], + nlist: usize, + d: usize, + force: bool, + calibration: &[f32], + calibration_rows: usize, + ) -> Option<Self> { + if nlist == 0 || d == 0 { + return None; + } + // Largest d' auto mode may pick, rounded down so the gate below holds; + // forced mode searches up to d / 2 (beyond that the projected GEMM + // costs about as much as the scan it is meant to replace). + let auto_limit = (d / MAX_AUTO_DP_DIVISOR) / MIN_DP * MIN_DP; + let block = if force { + (d / 2).max(MIN_DP).min(d) + } else { + auto_limit + }; + if !force && (block < MIN_DP || nlist < 2 * block) { + return None; + } + let block = block.min(nlist); + + let mean = column_mean(cents, nlist, d); + let centered: Vec<f32> = cents + .iter() + .enumerate() + .map(|(i, v)| v - mean[i % d] as f32) + .collect(); + let total_variance = centered + .iter() + .map(|v| (*v as f64) * (*v as f64)) + .sum::<f64>() + / nlist as f64; + if total_variance.is_nan() || total_variance <= 0.0 { + return None; + } + Review Comment: [P2] Skip PCA when Auto has no calibration sample `set_quantizer_centroids` is now the public path for manually loading a trained coarse quantizer, and a newly constructed index has the default Auto mode but no `calibration_sample`. This function only discovers that Auto cannot select a projection after `top_subspace` has already fitted the full PCA. In a release/8-thread probe with `nlist=4096,d=768`, the default setter spent 423.8 ms here and then returned `None`; the same setter with projected assignment disabled took about 1 µs. Please clamp/validate `calibration_rows` and return early for `!force && calibration_rows < MIN_CALIBRATION_ROWS` before computing the mean, centered matrix, or PCA, and cover the default manual-load path. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
