结构体 faiss::LocalSearchQuantizer

struct LocalSearchQuantizer : public faiss::AdditiveQuantizer

以下两篇论文中描述的 LSQ/LSQ++ 的实现

Revisiting additive quantization Julieta Martinez, et al. ECCV 2016

LSQ++: Lower running time and higher recall in multi-codebook quantization Julieta Martinez, et al. ECCV 2018

此实现主要从 Julieta Martinez 的 Julia 实现翻译而来:(https://github.com/una-dinosauria/local-search-quantization, https://github.com/una-dinosauria/Rayuela.jl)

训练后的代码存储在 codebooks 中,在 PQ 和 RQ 中称为 centroids

公共类型

enum Search_type_t

编码搜索的执行方式和向量的编码方式。

enumerator ST_decompress

解压缩数据库向量

enumerator ST_LUT_nonorm

使用 LUT,不包括范数(适用于 IP 或归一化向量)

enumerator ST_norm_from_LUT

从查找表计算范数 (成本为 O(M^2))

enumerator ST_norm_float

使用 LUT,并将 float32 范数与向量一起存储

enumerator ST_norm_qint8

使用 LUT,并存储 8 位量化范数

enumerator ST_norm_qint4
enumerator ST_norm_cqint8

使用 LUT,并存储非均匀量化范数

enumerator ST_norm_cqint4
enumerator ST_norm_lsq2x4

使用 2x4 位 lsq 作为范数量化器(用于快速扫描)

enumerator ST_norm_rq2x4

使用 2x4 位 rq 作为范数量化器(用于快速扫描)

公共函数

LocalSearchQuantizer(size_t d, size_t M, size_t nbits, Search_type_t search_type = ST_decompress)
LocalSearchQuantizer()
~LocalSearchQuantizer() override
virtual void train(size_t n, const float *x) override

训练量化器

参数:

x – 训练向量,大小为 n * d

virtual void compute_codes_add_centroids(const float *x, uint8_t *codes, size_t n, const float *centroids = nullptr) const override

编码一组向量

参数:
  • x – 要编码的向量,大小为 n * d

  • codes – 输出代码,大小为 n * code_size

  • n – 向量数

  • centroids – 要添加到 x 的质心,大小为 n * d

void update_codebooks(const float *x, const int32_t *codes, size_t n)

根据编码更新代码本

参数:
  • x – 训练向量,大小为 n * d

  • codes – 编码后的训练向量,大小为 n * M

  • n – 向量数

void icm_encode(int32_t *codes, const float *x, size_t n, size_t ils_iters, std::mt19937 &gen) const

使用迭代条件模式 (icm) 和给定的码本对向量进行编码。

参数:
  • codes – 输出码,大小为 n * M

  • x – 要编码的向量,大小为 n * d

  • n – 向量数

  • ils_iters – 迭代局部搜索的迭代次数

void icm_encode_impl(int32_t *codes, const float *x, const float *unaries, std::mt19937 &gen, size_t n, size_t ils_iters, bool verbose) const
void icm_encode_step(int32_t *codes, const float *unaries, const float *binaries, size_t n, size_t n_iters) const
void perturb_codes(int32_t *codes, size_t n, std::mt19937 &gen) const

向码添加一些扰动

参数:
  • codes – 要扰动的码,大小为 n * M

  • n – 向量数

void perturb_codebooks(float T, const std::vector<float> &stddev, std::mt19937 &gen)

向码本添加一些扰动

参数:
  • T – 模拟退火的温度

  • stddev – 训练数据中每个维度的标准差

void compute_binary_terms(float *binaries) const

计算二元项

参数:

binaries – 二元项,大小为 M * M * K * K

void compute_unary_terms(const float *x, float *unaries, size_t n) const

计算一元项

参数:
  • n – 向量数

  • x – 要编码的向量,大小为 n * d

  • unaries – 一元项,大小为 n * M * K

float evaluate(const int32_t *codes, const float *x, size_t n, float *objs = nullptr) const

计算重建误差的辅助函数

参数:
  • codes – 编码后的码,大小为 n * M

  • x – 要编码的向量,大小为 n * d

  • n – 向量数

  • objs – 如果不为空,则将每个向量的重建误差存储到其中,大小为 n

void compute_codebook_tables()
uint64_t encode_norm(float norm) const

将范数编码为 norm_bits 位

uint32_t encode_qcint(float x) const

通过非均匀标量量化对范数进行编码

float decode_qcint(uint32_t c) const

通过非均匀标量量化对范数进行解码

void set_derived_values()

训练范数量化器。

void train_norm(size_t n, const float *norms)
inline virtual void compute_codes(const float *x, uint8_t *codes, size_t n) const override

量化一组向量

参数:
  • x – 输入向量,大小为 n * d

  • codes – 输出代码,大小为 n * code_size

void pack_codes(size_t n, const int32_t *codes, uint8_t *packed_codes, int64_t ld_codes = -1, const float *norms = nullptr, const float *centroids = nullptr) const

将一系列代码打包成位紧凑格式

参数:
  • codes – 要打包的代码,大小为 n * code_size

  • packed_codes – 输出的位紧凑代码

  • ld_codes – 代码的前导维度

  • norms – 向量的范数(大小为 n)。如果需要但未提供,将会被计算

  • centroids – 要添加到 x 的质心,大小为 n * d

virtual void decode(const uint8_t *codes, float *x, size_t n) const override

解码一组向量

参数:
  • codes – 要解码的代码,大小为 n * code_size

  • x – 输出向量,大小为 n * d

virtual void decode_unpacked(const int32_t *codes, float *x, size_t n, int64_t ld_codes = -1) const

解码一组非打包格式的向量

参数:
  • codes – 要解码的代码,大小为 n * ld_codes

  • x – 输出向量,大小为 n * d

template<bool is_IP, Search_type_t effective_search_type>
float compute_1_distance_LUT(const uint8_t *codes, const float *LUT) const
void decode_64bit(idx_t n, float *x) const

用于解码 64 位字中的代码的函数

virtual void compute_LUT(size_t n, const float *xq, float *LUT, float alpha = 1.0f, long ld_lut = -1) const

计算内积查找表。用于质心搜索函数。

参数:
  • xq – 查询向量,大小 (n, d)

  • LUT – 查找表,大小 (n, total_codebook_size)

  • alpha – 计算 alpha * 内积

  • ld_lut – LUT 的前导维度

void knn_centroids_inner_product(idx_t n, const float *xq, idx_t k, float *distances, idx_t *labels) const

精确 IP 搜索

void compute_centroid_norms(float *norms) const

对于L2搜索,我们需要质心的L2范数

参数:

norms – 输出范数表,大小为 total_codebook_size

void knn_centroids_L2(idx_t n, const float *xq, idx_t k, float *distances, idx_t *labels, const float *centroid_norms) const

精确的 L2 搜索,使用预先计算的范数

公共成员

size_t K

每个码本的代码数量

size_t train_iters = 25

训练中的迭代次数

size_t encode_ils_iters = 16

编码中局部搜索的迭代次数

size_t train_ils_iters = 8

训练中局部搜索的迭代次数

size_t icm_iters = 4

icm中的迭代次数

float p = 0.5f

温度因子

float lambd = 1e-2f

正则化因子

size_t chunk_size = 10000

一次编码的向量数量

int random_seed = 0x12345

随机数生成器的种子

size_t nperts = 4

每个代码中的扰动数量

如果非 NULL,则使用此编码器进行编码 (由对象拥有)

lsq::IcmEncoderFactory *icm_encoder_factory = nullptr
bool update_codebooks_with_double = true
size_t M

码本的数量

std::vector<size_t> nbits

每个步骤的位数

std::vector<float> codebooks

码本

std::vector<uint64_t> codebook_offsets

码本#i 存储在码本表 codebook_offsets[i]:codebook_offsets[i+1] 行中,码本表的大小为 total_codebook_size * d

size_t tot_bits = 0

总位数(索引 + 范数)

size_t norm_bits = 0

为范数分配的位数

size_t total_codebook_size = 0

码本中向量的数量

bool only_8bit = false

是否所有 nbits 都等于 8 (使用更快的解码器)

bool verbose = false

训练期间是否显示详细信息?

bool is_trained = false

是否已经训练

std::vector<float> norm_tabs

ST_norm_lsq2x4 和 ST_norm_rq2x4 的辅助数据,存储用于 4 位快速扫描的码本条目的范数

IndexFlat1D qnorm

存储和搜索范数

std::vector<float> centroid_norms

所有码本条目的范数(大小为 total_codebook_size)

std::vector<float> codebook_cross_products

所有码本条目与之前的码本的点积,大小为 sum(codebook_offsets[m] * 2^nbits[m], m=0..M-1)

size_t max_mem_distances = 5 * (size_t(1) << 30)

使用 beam search 的范数和距离矩阵可能会变得很大,因此使用此选项来控制可以分配的内存量

Search_type_t search_type

还决定了代码中包含的内容。

float norm_min = NAN

范数量化的最小值/最大值

float norm_max = NAN
size_t d

输入向量的大小

size_t code_size

每个索引向量的字节数