使用LSH 进行特征提取

news2024/10/7 10:24:02

局部敏感哈希(LSH)通常用于近似最近邻算法(ANN) 操作(向量搜索)。LSH的特性也可以在以矢量为输入的神经网络模型中得到利用(例如，各种的音频、视频和文本嵌入等内容信号)。

通常情况下，特定领域模型中输入的流形是复杂的（非i. i. d）。这种复杂性使得使用计算密集型操作的多层感知机来分离这些流形非常困难。学习复杂映射的经典方案是记忆结果，而不是学习函数。如何记忆向量图？最直接的方法就是嵌入向量。但是我们需要离散的对象来计算嵌入，而向量不是离散的。那么怎么把向量嵌入算法应用到向量输入中呢？对向量进行哈希运算，在哈希运算后，附近的点必须保持“附近”状态。这就是LSH的做法，所以我LSH运算顶部的嵌入可以作为浅层特征提取器。

“局部敏感哈希”（Locality Sensitive Hashing，简称LSH）是一种用于解决这类问题的近似搜索技术。它的主要思想是将相似的数据点映射到同一个"哈希"桶中，从而可以在特定的桶中进行搜索，而不必对整个数据集进行线性搜索。虽然这种方法不保证找到确切的最近邻，但它在高维数据中提供了一种高效的近似搜索方法。

LSH的核心概念如下：

局部敏感性函数（Locality Sensitive Function）：这是一个函数，它能够将相似的数据点映射到相同的哈希桶中，但也不是那么严格，因此即使有些数据点被映射到相同的桶中，它们也不一定是真正相似的。局部敏感性函数的设计取决于所处理的数据类型和相似性度量。
哈希桶（Hash Bucket）：数据点通过局部敏感性函数映射到不同的哈希桶中。相似的数据点可能被映射到相同的桶，从而提供了搜索的起点。
哈希表（Hash Table）：哈希桶构成了一个哈希表，通过在哈希表中进行搜索，可以快速定位具有相似性的数据点。

LSH的性能取决于局部敏感性函数的设计和哈希桶的构建。这涉及到在保持相似性的同时，将数据点映射到不同的桶，以及在哈希表中组织和检索数据。LSH通常用于解决近似最近邻搜索（Approximate Nearest Neighbor Search，ANN）问题，其中目标是在给定查询点的情况下，找到与其相似度较高的数据点。

选择LSH算法和将LSH桶转换为嵌入的方式非常重要。所以这里只介绍一种只有方向感知的算法（忽略向量的大小），它基于这个简单的LSH算法：

 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 
 class CosineVectorEmbedding(nn.Module):
     """
     LSH based vector indexer for highly non-linear ops
     """
 
     def __init__(self, inp_dim: int, emb_dim: int, n_proj: int = 16, num_bins: int = 20):
         super().__init__()
         self.register_buffer(
             'projection_mat',
             F.normalize(torch.randn((inp_dim, n_proj)), p=2.0, dim=0),
             persistent=True,
         )
         resolution = 2.0 / num_bins
         self.register_buffer(
             'grid',
             torch.linspace(-1, 1, num_bins + 1)[:-1] + 0.5 * resolution,
             persistent=True,
         )
         self.register_buffer(
             'pos_offset',
             ((num_bins + 1) * torch.arange(0, n_proj, dtype=torch.long)).long().reshape(-1, 1, 1),
             persistent=True
         )
         self.emb = nn.EmbeddingBag((num_bins + 1) * n_proj, emb_dim)
         self.emb_dim = emb_dim
         self.n_proj = n_proj
 
     def forward(self, x):
         bs, seq_len, emb_dim = x.size()
         z = F.normalize(x, p=2.0, dim=-1) @ self.projection_mat
         z = torch.bucketize(z, self.grid).transpose(0, -1)
         z = (z + self.pos_offset).transpose(0, -1).contiguous()
         return self.emb(z.view(-1, self.n_proj)).reshape(bs, seq_len, self.emb_dim)

为了说明它的有效性，我们将它应用到输入32维的输入内容嵌入的RecSys LLM的训练中。使用从低分辨率到高分辨率的独立级联LSH嵌入（inp_dim = 32，emb_dim = 512，n_proj = 32，num_bins =（1，2，4，8，12，16，20））并将其输出相加。把它与使用一个简单投影进行了对比(使用nn. Linear (32, 512))。