ivfadc.py

1"""2IVFADC：倒排文件 + 残差乘积量化（第18章 §18.4.3）34正文的思路：PQ 把距离计算压到 m 次查表，但还是要扫全库（O(N)）。5IVFADC 再加一层粗聚类：67    离线：粗 K-Means 得到 n_list 个中心；每个向量 y 归到最近的 c_i，8          对它做残差编码 r = y − c_i（而不是直接编码 y），9          倒排表 i 里存 (cluster id, PQ code)。10    在线：1) 算查询到所有粗中心的距离；11          2) 取最近的 n_probe 个倒排表；12          3) 对每个表算残差查询 x − c_i 的 ADC 查找表；13          4) 查表算近似距离，合并返回 Top-k。1415距离恒等式（正文的推导）：16    ‖x − (c_i + q(y − c_i))‖ = ‖(x − c_i) − q(y − c_i)‖17所以「减去粗中心」只是给查询与残差同时平移一个偏置，不改变算法结构。1819本脚本：201) 在 20000 条 64 维数据上建 IVFADC，核对倒排表规模；212) 扫 n_probe，画 Recall@10 / 候选数 / 查询耗时——正文的「调 n_probe 换22   Recall 与 Latency」；233) 对比「残差 PQ」与「直接 PQ 向量」在同一 n_probe 下的召回，验证残差编码24   确实降低了量化误差；254) 测一下 n_list 的影响。2627运行：python3 ch04_pq/ivfadc.py28"""2930from __future__ import annotations3132import pathlib33import sys34import time3536import numpy as np3738sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))39sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))4041from common import (  # noqa: E40242    knn_exact, make_queries, make_vectors, recall_at_k, sq_dists, table,43)44from product_quantization import PQ, kmeans  # noqa: E4024546SEED = 2026474849class IVFADC:50    """倒排 + 残差 PQ。residual=False 时退化成「倒排 + 直接 PQ」，用于对照。"""5152    def __init__(self, n_list: int = 64, m: int = 8, K: int = 256,53                 residual: bool = True, seed: int = SEED) -> None:54        self.n_list = n_list55        self.m = m56        self.K = K57        self.residual = residual58        self.seed = seed5960    def train(self, X: np.ndarray) -> None:61        self.centroids = kmeans(X, self.n_list, seed=self.seed)62        # sq_dists(base, queries) 返回 [len(queries), len(base)]，这里 base 是粗中心63        assign = np.argmin(sq_dists(self.centroids, X), axis=1)   # [N]64        self.assign = assign65        self.pq = PQ(m=self.m, K=self.K, seed=self.seed)66        if self.residual:67            R = X - self.centroids[assign]        # 残差：粗聚类已消掉全局偏移68        else:69            R = X70        self.pq.train(R)71        codes = self.pq.encode(R)72        # 倒排表：列表 i 里是分配给中心 i 的向量编号73        self.postings = [np.where(assign == i)[0] for i in range(self.n_list)]74        self.codes = codes7576    def search(self, q: np.ndarray, k: int = 10, n_probe: int = 8):77        """返回 (Top-k 编号, 本次扫描的候选数)。"""78        d_c = sq_dists(self.centroids, q[None, :])[0]79        probes = np.argpartition(d_c, n_probe - 1)[:n_probe]80        probes = probes[np.argsort(d_c[probes])]      # 近的倒排表先查8182        cand_list, dist_list = [], []83        for i in probes:84            ids = self.postings[i]85            if len(ids) == 0:86                continue87            shift = self.centroids[i] if self.residual else np.zeros_like(q)88            L = self.pq.adc_lut(q - shift)            # 查询侧也减去同一个偏置89            cand_list.append(ids)90            dist_list.append(self.pq.adc_dists(L, self.codes[ids]))9192        cand = np.concatenate(cand_list)93        dist = np.concatenate(dist_list)94        kk = min(k, len(cand))95        part = np.argpartition(dist, kk - 1)[:kk]96        part = part[np.argsort(dist[part])]97        return cand[part], len(cand)9899100def main() -> None:101    print("IVFADC：倒排文件 + 残差乘积量化（§18.4.3）\n")102    rng = np.random.default_rng(SEED)103104    N, D = 20000, 64105    X, _ = make_vectors(n=N, dim=D, n_clusters=80, spread=0.45)106    Q = make_queries(X, n_queries=100)107    k = 10108    truth, _ = knn_exact(X, Q, k)109110    # ---------------- 1) 建索引 ----------------111    n_list, m, K = 256, 16, 256112    idx = IVFADC(n_list=n_list, m=m, K=K, residual=True, seed=SEED)113    t0 = time.perf_counter()114    idx.train(X)115    t_train = time.perf_counter() - t0116117    sizes = np.array([len(p) for p in idx.postings])118    print(f"① 建索引（N={N:,}, D={D}, n_list={n_list}, m={m}, K={K}）：")119    rows = [120        ["粗聚类中心", f"{n_list}", f"{idx.centroids.nbytes / 1e3:.0f} KB",121         "k-means，O(n_list·N·D)"],122        ["PQ 码本", f"m×K×D* = {m}×{K}×{D // m}",123         f"{idx.pq.nbytes() / 1e3:.0f} KB", "训练在残差上"],124        ["倒排编码", f"{N:,} × {m} B", f"{N * m / 1e3:.0f} KB",125         "cluster id 省掉：同表内隐式已知"],126    ]127    table(["结构", "规模", "大小", "说明"], rows)128    print(f"\n  训练总耗时 {t_train:.1f} s。倒排表大小："129          f"min {sizes.min()} / 中位 {int(np.median(sizes))} / max {sizes.max()}")130    print("  分布不均是因为合成数据的簇中心是随机高斯，粗聚类无法完全均分；")131    print("  工业实现会用 balanced k-means 或在建表后拆分超长表。\n")132133    # ---------------- 2) n_probe 扫描 ----------------134    print("② 扫 n_probe：Recall / 候选数 / 耗时的三角关系（残差 PQ）：")135    rows = []136    for n_probe in (1, 2, 4, 8, 16, 32, 64, 128, 256):137        preds, cand_tot = [], 0138        t0 = time.perf_counter()139        for qi in range(len(Q)):140            p, nc = idx.search(Q[qi], k, n_probe)141            preds.append(p)142            cand_tot += nc143        dt = (time.perf_counter() - t0) / len(Q)144        rec = float(np.mean([len(set(p.tolist()) & set(t.tolist())) / k145                             for p, t in zip(preds, truth)]))146        rows.append([n_probe, f"{cand_tot / len(Q):.0f}",147                     f"{cand_tot / len(Q) / N:.2%}",148                     f"{rec:.3f}", f"{dt * 1e3:.2f} ms"])149    table(["n_probe", "平均候选数", "占全库比例", "Recall@10", "单查询耗时"], rows)150    rec_last = float(rows[-1][3])151    print("\n  正文的「n_probe 换 Recall 与 Latency」在这里能看到前半段：")152    print(f"  n_probe=1 只扫约 {rows[0][2]} 的库，召回 {rows[0][3]}；")153    print("  加到 4 个表召回到 7 成，之后再加大几乎不再涨——Top-10 的近邻")154    print("  基本都在最近的几个表里，剩下的漏召回来自 PQ 量化误差")155    print(f"  （这里 n_probe 取到 {rows[-1][0]}，召回上限约 {rec_last}）。")156    print("  「召回随 n_probe 单调大涨」的前提是近邻分散在多个表里：")157    print("  n_list 越小、表越长，近邻越容易落在同一张表，曲线反而更平。")158    print("  下面 ④ 的 n_list 对比可以看到这一点。\n")159160    # ---------------- 3) 残差 vs 直接 PQ ----------------161    print("③ 残差编码 vs 直接编码（同 n_list、同编码长度、同 n_probe）：")162    print("   （n_probe 取 4/8/16，n_list/m 与 ①② 一致）")163    idx_direct = IVFADC(n_list=n_list, m=m, K=K, residual=False, seed=SEED)164    idx_direct.train(X)165    rows = []166    for n_probe in (4, 8, 16):167        out = []168        for model in (idx_direct, idx):169            pr = []170            for qi in range(len(Q)):171                p, _ = model.search(Q[qi], k, n_probe)172                pr.append(p)173            out.append(float(np.mean([len(set(p.tolist()) & set(t.tolist())) / k174                                      for p, t in zip(pr, truth)])))175        # 两边各自的重建误差176        rows.append([n_probe, f"{out[0]:.3f}", f"{out[1]:.3f}",177                     f"{(out[1] - out[0]) / max(out[0], 1e-9):+.1%}"])178    table(["n_probe", "直接 PQ Recall", "残差 PQ Recall", "相对提升"], rows)179180    # 量化误差对比（只看被分到的那个粗中心上的残差）181    rec_direct = idx_direct.pq.reconstruct(idx_direct.codes)182    rec_res = idx.pq.reconstruct(idx.codes) + idx.centroids[idx.assign]183    err_d = float(((rec_direct - X) ** 2).sum(1).mean())184    err_r = float(((rec_res - X) ** 2).sum(1).mean())185    print(f"\n  重建 MSE：直接 PQ {err_d:.2f} → 残差 PQ {err_r:.2f}"186          f"（降低 {(1 - err_r / err_d):.0%}）")187    print("  粗聚类先把「这个向量大概在哪」记下来，残差只描述局部细节，")188    print("  范数更小、分布更集中，同样的 K 个中心就能表示得更细。")189    print("  正文「残差通常范数更小」的说法在这里得到定量验证。\n")190191    # ---------------- 4) n_list 的影响 ----------------192    print("④ n_list 的影响（固定 n_probe=4、m=16，其余不变）：")193    rows = []194    for nl in (16, 64, 256, 512):195        model = IVFADC(n_list=nl, m=m, K=K, residual=True, seed=SEED)196        model.train(X)197        pr, cand_tot = [], 0198        for qi in range(len(Q)):199            p, nc = model.search(Q[qi], k, 4)200            pr.append(p)201            cand_tot += nc202        rec = float(np.mean([len(set(p.tolist()) & set(t.tolist())) / k203                             for p, t in zip(pr, truth)]))204        rows.append([nl, f"{N / nl:.0f}", f"{cand_tot / len(Q):.0f}", f"{rec:.3f}"])205    table(["n_list", "平均表长", "平均候选数", "Recall@10"], rows)206    print("\n  n_list 是一个「表短 vs 跨表漏」的权衡：表越短，同样 probe 4 次")207    print("  扫的候选越少；但近邻落在未访问表里的概率也越高。本数据上召回先升后降，")208    print("  说明存在一个性价比最优点（这里是 256 附近），而不是越大越好。\n")209210    print("⑤ 一句话：IVF 负责「少扫」（粗聚类定向），PQ 负责「扫得快」")211    print("   （查表代替浮点距离），残差负责「压得准」。三者合起来才是")212    print("   工业界能扛十亿级的 IndexIVFPQ；OPQ 只是再给残差加一个旋转。")213214215if __name__ == "__main__":216    main()