关于 RAG 的一些可能优化方案的研究报告

AI 算法 研究
关于 RAG 的一些可能优化方案的研究报告

本文叙述了一种在 RAG 上可能的优化方案。

研究环境

本实验与研究的运行环境为 iOS,使用 Minis Agent 进行辅助,内接模型为 Deepseek V4 Flash。数据分析等在 Minis 内置的 ISH 终端环境完成,但是由于 ISH 模拟X86过慢,所以模型训练等评估在浏览器端完成,通过编写 HTML 内接 JS 代码,使用 JS 运行以调用苹果原生的芯片能力。

数据方面,我们采用了知乎上抓取的 7~8 篇张量分析相关文章作为基础数据库,然后从中进行统计分析,根据统计分析结果进行同分布生成,生成 1 万级别数据。当然,实际上这个级别的数据还远远不需要优化,但是由于 ISH 模拟层过慢,因此再大的数据可能较难支撑,所以我们主要以收获量与总量的比例,来估算加速比。

前期准备

一、抓取知乎上关于张量分析的相关文章,共 7~8 篇。
二,使用 256 字符为一个 chunk,前后重叠 32 字符进行分段。
三,使用嵌入模型 MiniBLOM 进行嵌入,并且保存。
四,对其进行了统计分析,然后使用 PC 降维以后的提供主要方差的指标进行同分布生成数据。
五,由于我们并不部署在真实生产环境下,因此使用 Agent 生成了 50 条与我们的数据库相关的查询,又生成了 50 条与数据库不太相关的查询,并进行了嵌入与向量保存。

一、IVF 预尝试

我们先使用了 IVF 算法进行优化,其具体操作是,先将全部的数据进行 kmeans 区分为多个质心,然后在查找时,我们会先查质心,或者说先在质心的向量空间当中进行一次查询,寻找到排名靠前的质心,然后再以这些质心的簇作为搜索空间进行搜索。
然而我们在 RAG 的基础上做了稍许改进,根据统计分析发现,在我们的真实数据上,有效维度只有 32 维,剩下的 300 多维向量只提供少量方差。因此我们最后选取了前 16 维的 PC 压缩,对 16 D 向量进行了 IVF 的 Kmeans 提出质心。
然而实际数据并不优良,我们需要一个较大的收割比才能保证高的 recall 率。具体数据如下:

在我们的合成数据(10 万条 384D)上测试,使用 PCA 16D 降维 + K-means 聚类:

S1 probe (16D) S2 probe (32D) recall 收割量
50 8 0.58 ~9.5%
80 10 0.81 ~19%
100 16 0.89 ~39%

要拿到 90%+ recall 需要收割近 40% 的数据,这已经失去了索引的意义。

对 PCA 16D 降维的分析

有了以上结果以后,我们对比了 benchmark,发现正常的 IVF 在不使用 16D 降维情况下,可以达到 比我们数据小得多的收割比,以及高约 5%~10% 的 最高Recall 率。由此,我们对我们的 16D 降维做了一定的分析。
我们怀疑 16 D 降维可能导致保序性变差。因此我们进行了如下计算:使用我们制造的 100 条查询进行测试,计算 384 维的真实排序与 16D 降维以后的排序两个排序之间的曲线下面积,记 AUC,AUC 越接近于一,说明保序性越好。当然,该统计计算在 295 条向量在真实数据库上进行,并没有使用扩充数据。结果如下:

指标 PCA 16D
全排序 AUC@295 0.81
top-10 AUC@10 0.49

这说明 PCA 16D 的全局排序保序性尚可(AUC=0.81),但在前 10 位的精细排序上几乎和随机差不多(AUC=0.49)。这正是 IVF 需要大 probe 才能拿到高 recall 的根本原因——降维破坏了局部保序性。

这也说明,直接使用 PCA 降维确实保留了大部分的方差信息,但是对保序性仍有破坏。因此我们尝试进行模型降维法。

模型降维法

一般的查询当中,我们不可能需要极高的 Top K 量。一般而言,K 在 10 是常用的,在 50是可能的,50 以上则缺少大规模应用的行业典型。因此我们尝试采用模型降维法来降低收割比。具体而言,我们需要设置一个 K 值,我们的优化目标应当是使得对于我们的向量库而言,对于不同的查询向量,需要包含真实 Top K 的数据所应当选取的向量数的期望。此处说期望的原因是因为对于不同的查询,这个截断数是不同的。我们的具体训练方案是:

采用线性投影矩阵 W ∈ ℝ^{16×384},对向量进行 16D 降维后 L2 归一化。优化目标为 triplet margin loss:

L(W) = Σ Σ max(0, sim_W(q, neg) - sim_W(q, pos) + margin)
        q  (pos,neg)

其中 pos ∈ 全空间 top-K(相关项),neg ∉ top-K(非相关项),margin=0.05。SGD 优化 (lr=0.003),每 25 轮评估一次 cutoff,500 轮训练。训练在浏览器中完成(JS JIT,相比 iSH 模拟加速 ~10×)。

当然,由于实际上我们无法真实获得期望这个事情,所以我们只能拿一定的查询分布进行蒙特卡洛近似。但由于我们又没有大量的真实数据,因此这里我们取我们的合成查询进行近似;此外,我们选用在真实的 295 条数据上进行 K 等于 30 的训练。经过 500 个 epoch 的训练,数据如下:

指标 PCA 16D (基准) W 模型 (16D) 改善
test cutoff (top-30 最差排名) 94.7 71.8 ↓24%
test recall@2K 90.6% 94.6% ↑4%
收割百分比 32.1% 24.3% ↓7.8pp

W 模型在 cutoff 和 recall 上均优于 PCA 16D,收割比从 32% 降至 24%。

模型降维法的相关检验

分布不变性检验

在上述过程中,我们取用了 295 个向量的真实数据库进行训练。但实际上,我们需要把训练出来的模型应用到更大规模的同分布数据库上。因此,我们必须需要验证,在一组同分布但是不同的数据下,该模型是否具有泛化性。我们在扩充数据集中均匀采样了 295 个数据,并且经过统计检验,验证其同分布性。之后再进行了交叉验证,即在同分布的合成数据上重新训练一个模型,然后将两个模型进行交叉检验,数据如下:

四路交叉验证结果:

组合 test cutoff 收割%
① W_real → 真实 DB (匹配) 71.8 24.3%
② W_real → 合成 DB (交叉) 78.9 26.7%
③ W_synth → 合成 DB (匹配) 72.0 24.4%
④ W_synth → 真实 DB (交叉) 65.3 22.1%

真实↔合成交叉差异仅 2~4pp,分布不变性确认。有趣的是,W_synth(合成数据训练)在真实数据上表现反而最好(cutoff=65.3, 22.1%),说明合成数据的多样性提升了泛化能力。

分布偏差鲁棒性

实际的RAG 使用当中,我们会不断地加入新的数据内容。这会导致实际分布不断偏移。因此,我们有必要测量,当数据分布存在偏移时,模型在其上的表现,来评估什么时候我们需要重新采样进行模型的重训。

具体的测试方案是,我们先取一组与原数据同分布的数据,然后在其上以一个方向偏差性不断地加入新的数据。随着加入数据的多少,我们可以测量出模型在上面的表现。当然,对于不同的方向偏差性,我们也都需要给出对应的测量结果。经过测量,数据如下:

以 2000 条基库 + 2000 条漂移数据(50% 污染),在 PCA 空间沿不同方向施加 δ 个标准差的偏置:

方向 δ=0.5 δ=1 δ=2 δ=5 δ=10
PC1 +0.1% -0.2% -0.8% +0.0% -1.1%
PC2 -1.2% -2.0% -4.5% -7.2% -5.4%
Random +0.3% +0.1% -0.3% -0.2% -0.7%
PC32 (尾部) -0.3% -0.4% +0.5% +2.3% +4.4%

所有方向 δ≤10 时 cutoff 恶化均不超过 5%,最差的 PC32 方向 δ=10 仅 +4.4%。Random 方向(最现实场景)几乎完全不受影响。

这个结果说明,我们模型其实在数据的偏移分布上具有较大的鲁棒性,当数据有一定的偏差时,即仍可以较好工作。但是随着偏差产生,我们也需要相应地略微扩大截取的百分比。当然,如果希望在分布偏移的数据上仍然有较好的表现性,那么我们仍然需要重采样与重训练。

收割比随规模变化关系

在之前的测验当中,数据大多从 295 这个规模的数据量给出,分布不变性验证了模型可以在更大规模的同分布数据上试验,因此我们也应当测验,在不同规模的数据下,如果我们仍应截取 TOP 30 所需要的收割量与总数据量的比例变化如何。我们从合成数据当中均匀采样了多组数据,并进行同分布检验。经过检验以后,再进行实验测量。经过测量,结果如下:

Ndb Cutoff 收割%
295 78.9 26.7%
500 111.6 22.3%
1000 170.2 17.0%
2000 254.1 12.7%
4000 399.1 10.0%
6000 484.6 8.1%
10000 727.1 7.3%
20000 1143.0 5.7%

从结果可以看出,随着数据量的上升,总收割量确实不断上升,但是收割比却在不断下降。收割比与数据规模的关系可以进行拟合,拟合结果如下:

cutoff = 2.16 × N^0.63

即 cutoff 增长速度为 N^0.63,显著慢于线性(N^1),收割比例按 N^(-0.37) 衰减。这意味着规模越大,W 的相对效率越高——20000 条时收割仅 5.7%。

密度难度曲线

我们的合成查询样本当中,有 50 个强相关样本和 50 个弱相关样本。经过统计分析,我们可以发现,强相关样本落在密度较大区域,弱相关样本落在密度较小区域。而且对于截取 TOP 30 所需要的截断量,对于弱相关样本所需要的截断数明显大于强相关样本所需要的截断数。因此我们可以以截取 TOP 30 所需要的截断数作为一个难度指标,测量查询所落在的区域密度与难度的关系曲线。在最开始的设想中,该曲线可能能用于自适应化查询。也就是说,当有查询时,可以先粗略估计其周围密度,然后根据密度自适应查表选择截断数。在真实数据库上测量数据如下:

α (推离程度) maxSim (密度) cutoff 收割%
0 (原位) 1.000 63.4 21.5%
0.02 0.927 68.9 23.4%
0.05 0.688 68.5 23.2%
0.10 0.430 112.1 38.0%
0.15 0.278 125.0 42.4%
0.30 0.158 180.0 61.0%
0.50 0.122 231.2 78.4%

然而,为了进行自适应查询,我们还需要保证该曲线是分布不变的。使用交叉测验后,数据如下:

对真实 DB 加不同程度高斯噪声后观察曲线变化:

噪声 on-manifold cutoff 收割% vs 原始 KS
0 (基准) 63.4 21.5% —
0.02 66.3 22.5% 0.139
0.05 91.2 30.9% 0.468
0.10 189.8 64.3% 0.776

可以观测到,密度难度曲线在同分布的不同数据上有明显差异。进一步的统计实验表明,密度难度曲线对数据极度敏感,甚至两个数据的略微同分布区别都可能会导致其急剧变化。因此,当前的密度难度曲线暂时不能用于自适应化的查询。

总体开销估算

在我们的模型降维方法当中,我们是在 16D 空间当中进行的全量搜索,进行截断以后,再在完整的全维度空间当中进行收割。16D 空间与 384D 空间的计算开销相比,大体可以忽略。但是根据截断比与规模的变化曲线。当大规模时,例如上亿级的数据规模,16D 空间的全扫计算开销可能反而上涨,导致总体开销反而变大。综合而言,可做以下估算:

总时间 = 16D 全扫 + (收割比 × N) 条 × 384D 重排。收割比按幂律 cutoff = a × N^0.63 计算。

规模 收割比 16D 全扫 重排开销 总计 (SIMD)
1 万 7.5% 0.002ms 0.003ms 0.005ms
10 万 4.2% 0.02ms 0.02ms 0.04ms
100 万 2.4% 0.2ms 0.1ms 0.3ms
1000 万 1.3% 2ms 0.5ms 2.5ms
1 亿 0.7% 17ms 3ms 20ms
10 亿 0.4% 170ms 16ms 186ms

关键转折点:

  • 千万级以内:SIMD 总耗时 <3ms,RAG 场景完全可以接受
  • 亿级:SIMD ~20ms,可接受但需优化
  • 十亿级:SIMD ~186ms,16D 全扫是瓶颈(占 90%+),需在 16D 空间构建索引(比 384D 快 24×)
  • 由于收割比随 N 增大而下降,重排开销增长慢于 16D 全扫,因此核心矛盾是 16D 扫描速度而非重排。