LLM 的大规模文本去重技术
LLM document 的去重在大规模场景要做好还是比较复杂。
两种思路，粗筛走 MinHashLSH，精筛走 Embedding 去重。
MinHashLSH 路线
朴素思想
- 直接对每个 doc 做 n-gram (然后把一个 doc 看作是多个 n-gram 的集合)
- 定义 K 个 Hash function，对每个 n-gram 做 hash，拿到 doc 到 hash 的集合
- 每个 hash function 对所有 n-gram 做 hash，取最小值，得到一个 length=k 的 minhash 集合
- LSH 分桶，k 维切成 b 个 band，每个 band 整体 hash 一次，对应一个 band_hash
- 按 band_idx 和 band_hash 进行 groupBy，相同的 band_hash 表示两个 doc 至少在一个 band 里相似，加入比较的 candidates，没有相同 band_hash 的不需要再处理了
- 对相同 band_hash 的 doc 做一个连通图处理，假设边是 (u,v)，接下来通过多轮的连通处理，首先定一批节点，每个节点负责一批 doc_id，每一轮：
- (u,v) 的边发送给 u 和 v 的两个节点
- 每个节点内会产生新的连通，即新的 (u,v)，同时每个节点的数据不断累积，doc_id 对应的簇(doc_id 集合)会越来越大
- 连通多轮之后（一般 2-3 轮就够了），做决策，哪些 doc_id 要留下
- 遍历原始数据，留下不重复的 doc_id
用 codex 做了一个流程图：
这里要注意一个点：最后每个 doc_id 等价于多个 band_hash 值的集合，所以判断两者 doc_id 重复是通过 Jaccard 相似度来判断。而通常为了计算性能考虑，不一定会把一个 doc_id 的所有连通图都打通来计算完整的 Jaccard，比如 BigCode(https://huggingface.co/blog/zh/dedup) 里只要有一个 band_hash 相同就认为这是重复的。所以这里会引入一个概率关系：
| 真实 Jaccard | 被判为候选重复的概率 |
|---|---|
| 0.50 | 2.4% |
| 0.60 | 14% |
| 0.70 | 51% |
| 0.80 | 94% |
| 0.85 | 99.6% |
Spark
Spark 是最简单直接，也是早期最多人用的方法。基于 (band_idx, band_hash) 进行 groupBy，一样的值放在一起比较，阿里云的 Serverless Spark 实现参考 https://help.aliyun.com/zh/emr/emr-serverless-spark/use-cases/minhash-lsh-based-large-scale-text-duplication-scheme，但是这个连通图涉及到数据 shuffle，不建议。
NeMo Curator(NVIDIA 开源，基于 Ray)
ref: https://docs.nvidia.com/nemo/curator/curate-text/process-data/deduplication/fuzzy
- 本质也是 follow 朴素思想，但是换成了 GPU 库，LSH 的时候每次只处理了 5 个 band（中间结果落 parquet)
- minhash 换成了 cuDF，连通图换成了 cuGraph 的多 GPU 弱连通算法（同时对重复的校验没有那么严格）
Data-Juicer 实现（Ray 对朴素思想的实现 + 热点优化）
- 遍历数据，算 band hash 然后先每个 actor 内自行合并一次，得到一个单机的 2 级连通图，让所有 doc_id 指向最小的 doc_id
- 按 doc_id 做一次 reshuffle，每个 actor 有自己的一批 doc_id，每个 edge 发送到两个 actor
- 每个 actor 经过多轮 reshuffle 后，会有一个 doc_id->簇 的映射，此时决定哪些是重复值
这种 case 对于 doc_id 簇过大的情况怎么办（比如 Licence Text）：有一个 proxy 的思路(rebalance)，大致思想是不要把完整簇放到一个节点里，每个节点放一个属于自己的小簇，再把每个小簇的根节点作为 proxy 连到真正的根节点，同时每次真正根节点的变化可能带来 actor 的通信。
SimHash
以文本为例：
- 对文本分词或提取特征。
- 给每个特征计算普通哈希。
- 根据词频、TF-IDF 等给特征赋权重。
- 将各特征的哈希按位加权投票。
- 每一位的总值大于 0 记为 1，否则记为 0，得到固定长度的 SimHash。
但是 simhash 这里有个问题，因为 bit 数比较少，很容易受到一些关键词的影响（比如代码里的 Python、 README 等），所以当前大模型文本里用 MinHashLSH 还是要多一些。
Embedding 路线
Kmeans
比较容易的是复用 Faiss 的方法（同时也支持 GPU），Faiss 里的 Kmeans 原理如下：
- 抽样数据，训一个 Kmeans Cluster，得到 k 个 cluster centroids
- 遍历旧数据，算 cluster，分发，在每一个 centroids 里做 embedding 去重（相当于是一个矩阵乘法）
- 如果 k 个 centroids 还是比较大，继续做 k 个 cluster（两次肯定是够了）
Embedding Similarity
另一个就是构建分布式的 Embedding Service，通过建索引的方式，常见索引有两类：
1. HNSW：类似跳表数组
- 思路：Embedding 写入时，动态构建 Graph，与此同时根据写入的 emb 数量，会有多层的子 subgraph 结构，产生多个 Layer，每个 level=n+1 的 subgraph 可以认为是 level=n 这层 graph 的若干个中心点的连接
- 特点：写入慢（建图），查询快（跳表结构直接快速定位）；
2. IVF：倒排索引的思路
- 思路：Embedding 写入时，会抽样跑 kmeans cluster，训出若干个 centroids，然后对于每个 cluster 有一批 embedding 数据，类似一个倒排索引的样子，然后查询时直接定位到 cluster 再从这个倒排里算 similarity
- 特点：写入快（一旦训完 kmeans 就是一个追加写），查询慢（遍历倒排）且有热点问题
IVF 有很多变种，上面讲的是最常见的 IVF-FLAT，可以想象，如果单 cluster 数量过多，带来的查询时间会线性增长。另一种是 IVF-PQ，基本思想是在 IVF-FLAT 基础上，每个 cluster 的倒排数组里，对于每个 embedding，将 m-dim 的 embedding 切成 n 段，每一段单独做 kmeans cluster，然后得到一个固定大小 K 的 codebook，每个 embedding 可以当做是对 K 个 codebook 的距离(length=n)。
- IVF-PQ = IVF-Flat + 每个 cluster 里存储一个 PQ 码本（这样大大减小了 similarity 计算量）
IVF-PQ 的查询很快（多级索引），但是同时也丢失了 embedding 原值，此时要在磁盘上维护一个 vector_id -> embedding 的映射（用 memmap 来读，性能也很快）。
实测
今天用 FAISS 单机 DB 测了一下，14M 个 document embedding(384d)，
- IVF-FLAT: 训 kmeans 1min，写入 5min 左右（50k/s)，查询用 nprobe=4 的话 p99 4ms；
- HNSW：写入 52min，查询是 0.7ms；
场景：
- 如果是一次性离线去重，和 IVF 一样的思路，聚类后在聚类内部两两比较就可以了；
- 如果是不断有新数据进来，可以依赖 Embedding Cluster 来完成去重；
补充
datasketch
一个很有意思的开源项目，支持了各种去重能力以及外挂 embedding DB 的能力，适合单机玩玩，数据量不大的话这个就够了。