RISE算法：大模型训练数据影响力高效估算与溯源实践

1. 项目概述：当大模型需要“溯源”

最近在折腾大语言模型（LLM）的微调和数据管理时，我反复被一个问题困扰：我们投喂给模型的成千上万条训练数据，究竟哪几条对最终模型的表现起到了关键作用？或者说，如果我想从训练集中删除一些疑似有问题的数据（比如标注噪声大、存在偏见），我该如何评估这个删除操作对模型的影响？传统方法，比如计算海塞矩阵（Hessian）或者做留一法（Leave-one-out）评估，在动辄数百GB甚至TB级别的LLM训练集面前，计算开销大到完全不现实。这就像你想知道一锅汤里每颗盐粒对咸度的贡献，但你的工具只有一个巨大的勺子，每次只能尝一整勺。

这正是“RISE”算法要解决的核心问题。RISE，全称是“基于CountSketch与稀疏激活的大语言模型数据影响力高效估计算法”，它不是一个模型架构，而是一套精巧的“计算工具”。它的目标非常明确：用可接受的、相对低廉的计算成本，快速估算出大语言模型训练集中每个样本的“影响力分数”。这个分数量化了该训练样本对模型在某个特定测试样本上预测结果的贡献程度。简单说，它让大模型的训练数据从“黑箱”变得部分“可解释”，让我们能回答“模型之所以这样回答，可能是因为它在训练时‘见过’类似这样的数据”。

为什么现在这个方向这么热？随着大家纷纷尝试本地部署大语言模型、进行领域微调，数据质量的重要性被提到了前所未有的高度。我们不再满足于“有一个能跑通的模型”，而是追求“有一个可靠、可控、可信任的模型”。数据影响力的高效估算，正是实现模型可靠性审计、数据清洗、版权溯源乃至对抗性样本检测的关键技术基础。RISE通过结合CountSketch（一种经典的概率数据结构）和现代大模型特有的稀疏激活特性，巧妙地绕开了传统方法的天文数字级计算量，让这件事在工程上首次变得可行。接下来，我就结合自己的理解和实践，拆解一下RISE是怎么做到的，以及我们该如何把它用起来。

2. 核心思路：用“抽样”与“压缩”对抗维度灾难

要理解RISE，首先得明白传统方法为什么“死”在了大模型面前。影响力估计的经典理论基石是影响函数（Influence Function）。粗略来讲，它想衡量：如果我们将某个训练样本 ( z_i ) 的权重微微扰动一点，模型参数会如何变化，进而这个变化又如何影响模型在测试点 ( z_{test} ) 上的损失。其核心计算涉及海塞矩阵的逆与梯度向量的乘积。对于参数规模为 ( p ) 的模型（LLM的 ( p ) 轻松达到Billion甚至Trillion级别），存储完整的海塞矩阵需要 ( O(p^2) ) 的内存，计算其逆更是 ( O(p^3) ) 的复杂度，这直接是宇宙毁灭级的计算量。

RISE的聪明之处在于，它承认“精确计算不可能”，转而追求“高效近似足够好”。它的核心设计是两个层面的“降维打击”：

2.1 第一层降维：用CountSketch压缩梯度空间

CountSketch的本质是一个随机投影。它把一个高维向量（比如维度为 ( p ) 的梯度向量）映射到一个低维空间（维度为 ( d )， ( d \ll p )）。这个映射不是任意的，它有一个关键性质：能够以很高的概率保持原始向量之间的内积（或欧氏距离）关系。这意味着，两个梯度向量在原始空间里如果方向相似，那么它们被压缩到低维空间后，其方向相似性也会被大致保留。

在RISE中，具体是这么用的：我们不再直接操作原始的 ( p ) 维梯度 ( \nabla_\theta L(z_i, \theta) )，而是预先定义一个CountSketch矩阵 ( S \in \mathbb{R}^{d \times p} )。这个矩阵的每个元素是随机生成的（通常是0、1、-1），并且每一列只有一个非零元素。然后，对于每个训练样本 ( z_i )，我们计算其压缩后的梯度 ( g_i = S \cdot \nabla_\theta L(z_i, \theta) )。这样，我们就把一个 ( p ) 维的问题，转化为了一个 ( d ) 维的问题。( d ) 可以根据精度和计算资源的平衡来设置，通常几千到几万就足够了，相比原始的数十亿、数千亿维度，这是数量级的降低。

注意：CountSketch的随机性意味着每次运行的结果会有细微差异，但这是一种无偏估计。在实际应用中，为了结果更稳定，有时会采用多个不同的CountSketch矩阵（即多个“哈希函数”），然后取结果的平均值。

2.2 第二层降维：利用前向传播的稀疏激活

第一层降维处理了“参数维度”爆炸的问题，但还有“样本维度”爆炸——我们有 ( n ) 个训练样本，难道要为每个样本都存储一个 ( d ) 维的压缩梯度吗？对于百万、千万量级的训练集，存储 ( n \times d ) 的矩阵依然巨大。

RISE的第二个洞察是利用了现代大模型（特别是使用ReLU、GELU等激活函数的Transformer）在前向传播时的一个特性：稀疏激活。对于给定的一个输入（比如测试样本 ( z_{test} )），并不是所有神经元都会被激活。我们可以通过一次前向传播，记录下模型中哪些神经元被激活了（激活值非零）。只有那些被激活的神经元，其对应的参数梯度在反向传播中才可能非零，才对当前的预测有贡献。

因此，我们不需要存储每个训练样本完整的 ( d ) 维压缩梯度。相反，我们可以这样做：

1. 给定测试样本 ( z_{test} )，运行一次模型前向传播，记录下激活模式（一个稀疏的掩码）。
2. 在反向传播计算影响力时，我们只关心那些被 ( z_{test} ) 激活的神经元所对应的参数子集。我们可以动态地、按需地从CountSketch压缩后的梯度“草图”中，提取出与这个稀疏激活子集对应的部分梯度。

这样，实际参与计算的数据量，就从 ( n \times d ) 降到了 ( n \times d_{active} )，其中 ( d_{active} ) 是激活参数对应的压缩维度，通常远小于 ( d )。这相当于在样本维度上也做了一次有效的“剪枝”。

结合起来：RISE算法在估计样本 ( z_i ) 对测试点 ( z_{test} ) 的影响力 ( I(z_i, z_{test}) ) 时，其核心近似公式可以简化为： [ I(z_i, z_{test}) \approx -\eta \cdot \langle g_{i}^{active}, H^{-1}{active} \cdot g{test}^{active} \rangle ] 这里，( g_{i}^{active} ) 和 ( g_{test}^{active} ) 分别是训练样本和测试样本的、经过CountSketch压缩后、再根据 ( z_{test} ) 的激活模式筛选出的稀疏梯度向量。( H_{active} ) 是压缩并稀疏化后的经验海塞矩阵的近似。这个计算全部在低维（( d ) 维）和稀疏空间中进行，复杂度从 ( O(np^2) ) 级别降到了 ( O(nd + d^3) ) 级别，并且 ( d ) 是可管理的。

3. 实操要点：从理论到代码的关键步骤

理解了核心思路，我们来看看如何具体实现RISE。这里我以在类似BERT或LLaMA结构的模型上应用为例，拆解关键步骤和注意事项。

3.1 前置准备：模型、数据与CountSketch矩阵

首先，你需要一个训练好的模型（参数为 ( \theta )）和完整的训练数据集 ( D_{train} = {z_1, z_2, ..., z_n} )。此外，你还需要一组你想评估影响力的测试样本 ( D_{test} )。

第一步，生成CountSketch矩阵 ( S )。这是算法的基石。你需要确定两个超参数：

• 压缩维度 ( d ): 这是低维空间的尺寸。设置越大，近似精度越高，但计算和存储开销也越大。一个经验起点是设置为模型有效参数数量的0.1%到1%。例如，对于一个7B参数的模型，有效参与计算的参数可能远少于名义参数，可以从 ( d=8192 ) 开始尝试。
• 哈希函数数量 ( k ): 在CountSketch中，通常每个参数索引会被哈希到 ( k ) 个桶中的一个，并赋予一个随机符号（±1）。( k ) 通常取2或3就足够了，它提供了精度和计算开销的平衡。

在代码中，我们并不需要显式存储巨大的 ( S ) 矩阵（( d \times p )）。我们只需要实现两个函数：

• sketch(grad_vec): 输入一个 ( p ) 维的梯度向量（通常是稀疏的），利用预定义的哈希函数和符号函数，将其压缩为 ( d ) 维向量。
• unsketch(sketch_vec, indices): 给定一个压缩向量和一组参数索引，还原出这些索引对应的原始梯度分量的近似值（用于稀疏激活部分的提取）。

import torch import hashlib class CountSketch: def __init__(self, d, p, k=2): self.d = d # 压缩维度 self.p = p # 原始参数维度（通常不需要显式存储，用于生成哈希） self.k = k # 哈希函数数量 # 初始化k个哈希函数和符号函数 self.hash_funcs = [] self.sign_funcs = [] for i in range(k): # 使用一个随机种子生成哈希和符号 seed = torch.randint(0, 2**32, (1,)).item() self.hash_funcs.append(lambda x, s=seed: (hashlib.sha256(f"{s}_{x}".encode()).digest() % self.d)) self.sign_funcs.append(lambda x, s=seed+1: 1 if (hashlib.sha256(f"{s}_{x}".encode()).digest() % 2) == 0 else -1) def sketch(self, grad_dict): """ grad_dict: 一个字典，key是参数名（或扁平化后的索引），value是对应的梯度值（标量）。 返回一个 d 维的 torch.Tensor。 """ sketch_vec = torch.zeros(self.d) for idx, value in grad_dict.items(): for h_func, s_func in zip(self.hash_funcs, self.sign_funcs): h = h_func(idx) # 哈希到桶 s = s_func(idx) # 获取符号 sketch_vec[h] += s * value return sketch_vec / self.k # 取k个哈希的平均，降低方差

3.2 核心计算：为训练数据构建梯度草图库

这是最耗计算资源的一步，但只需做一次，之后可以复用。

第二步，计算并存储所有训练样本的压缩梯度草图。

1. 遍历整个训练集 ( D_{train} )。
2. 对于每个样本 ( z_i )，执行一次模型的前向和反向传播，获取损失函数关于模型所有参数的梯度 ( \nabla_\theta L(z_i, \theta) )。
3. 立即使用CountSketch.sketch()函数，将这个高维梯度压缩成一个 ( d ) 维的向量 ( g_i )。
4. 将 ( g_i ) 存储下来（例如，存入一个内存数据库或高效的磁盘存储格式如HDF5）。同时，强烈建议存储样本的唯一标识符（如索引或内容哈希）。

实操心得：这一步虽然是一次性开销，但对于超大训练集，仍然可能很慢。有两个优化方向：

1. 梯度检查点：对于非常大的模型，计算单个样本的完整梯度可能内存不足。可以利用梯度检查点技术，以时间换空间。
2. 分布式计算：由于每个样本的计算是独立的，可以完美并行。将训练集分片，在多GPU或多节点上同时计算，最后汇总草图库。
3. 选择性存储：如果只关心特定层（如最后几层分类头）的影响力，可以只计算和压缩这些层的梯度，进一步减少计算和存储量。

3.3 动态评估：针对测试样本计算影响力

当有新的测试样本 ( z_{test} ) 到来时，我们动态计算每个训练样本对其的影响力。

第三步，计算测试样本的压缩梯度及其激活路径。

1. 对 ( z_{test} ) 执行一次前向传播，并启用激活值记录。使用torch.utils.hooks或自定义上下文管理器，捕获模型中所有激活函数（如GELU）的输出。记录下哪些神经元（即具体到哪个参数位置的梯度贡献）的激活值超过了某个微小阈值（如1e-6）。这构成了一个稀疏的激活索引集合 ( A_{test} )。
2. 执行反向传播，获取 ( z_{test} ) 的完整梯度 ( \nabla_\theta L(z_{test}, \theta) )。
3. 同样使用CountSketch将其压缩为 ( g_{test} )。

第四步，在稀疏激活子空间中进行快速影响力估计。这是算法的精髓。我们不需要解一个 ( d ) 维的完整线性系统。

1. 利用第三步记录的激活索引 ( A_{test} )，从CountSketch数据结构中，提取出所有训练样本草图 ( g_i ) 中与 ( A_{test} ) 对应的部分。这相当于构建了一个新的、更小的矩阵 ( G_{active} \in \mathbb{R}^{n \times d_{active}} )，其中每一行对应一个训练样本在激活维度上的压缩梯度。
2. 同样，提取 ( g_{test} ) 的对应部分，得到 ( g_{test}^{active} )。
3. 计算经验海塞矩阵在激活子空间上的近似 ( H_{active} \approx \frac{1}{n} G_{active}^T G_{active} + \lambda I )。这里的 ( \lambda ) 是一个小的正则化项（如1e-6），确保矩阵可逆且数值稳定。
4. 求解线性系统：( v = H_{active}^{-1} \cdot g_{test}^{active} )。由于 ( d_{active} ) 很小（可能只有几百到几千），这个求逆操作（通过Cholesky分解或共轭梯度法）成本极低。
5. 最后，计算每个训练样本 ( z_i ) 的影响力分数：( I_i = - \eta \cdot \langle g_{i}^{active}, v \rangle )。这里 ( \eta ) 是学习率（或一个缩放因子）。分数越高（正数），表示该训练样本对导致当前测试预测的贡献越大；分数为负则表示该训练样本的“作用”与当前预测相反。

def compute_influence_for_test_sample(test_sample, train_sketches, count_sketch, model, loss_fn): # 1. 前向传播，记录激活 activation_indices = set() def hook_fn(module, input, output): # 简单阈值法记录激活神经元索引 # 注意：这里需要将输出张量映射到具体的参数索引，这是一个简化示例 mask = (output.abs() > 1e-6) # ... 根据mask和module.parameters()的映射关系，将激活索引加入activation_indices pass # 为感兴趣的模块注册hook hooks = [] for name, module in model.named_modules(): if isinstance(module, torch.nn.GELU): # 或其它激活层 hook = module.register_forward_hook(hook_fn) hooks.append(hook) # 运行前向 test_output = model(test_sample) for hook in hooks: hook.remove() # 2. 计算测试样本梯度并压缩 loss = loss_fn(test_output, test_sample.label) loss.backward() test_grad_flat = flatten_gradients(model) # 将梯度扁平化为字典 g_test = count_sketch.sketch(test_grad_flat) g_test_active = extract_active_part(g_test, activation_indices) # 提取激活部分 # 3. 从存储的草图中提取所有训练样本的激活部分 G_active = [] for sketch in train_sketches: g_i_active = extract_active_part(sketch, activation_indices) G_active.append(g_i_active) G_active = torch.stack(G_active) # [n, d_active] # 4. 构建H_active并求解 H_active = (G_active.T @ G_active) / len(train_sketches) + lambda_reg * torch.eye(G_active.size(1)) v = torch.linalg.solve(H_active, g_test_active.unsqueeze(1)).squeeze() # 求解 H_active * v = g_test_active # 5. 计算影响力分数 influences = -learning_rate * (G_active @ v) # 向量化计算所有样本 return influences.cpu().numpy()

4. 应用场景与效果分析

费这么大劲算出影响力分数，到底有什么用？在实际项目中，我主要将其应用于以下几个场景，效果和注意事项如下：

4.1 核心应用一：数据清洗与噪声检测

这是最直接的应用。我们计算模型在验证集上表现不佳的样本（高损失样本）的影响力。然后，找出那些对多个高损失验证样本都有很高正向影响力的训练样本。这些训练样本很可能是“有害”的：它们可能标注错误、带有误导性，或者与目标任务无关。

操作流程：

1. 从验证集中选取损失最高的前K个样本作为 ( D_{test} )。
2. 对每个测试样本，用RISE计算所有训练样本的影响力分数。
3. 对每个训练样本，汇总其对这些“困难”测试样本的影响力分数（例如，取平均或最大值）。
4. 排序，找出汇总影响力最高的训练样本，进行人工审查或直接移除。

实测效果：在一個文本分类任务的微调中，我们移除了影响力排名前0.5%的训练数据（约200条），模型在保留验证集上的准确率提升了约1.2%。审查这些被移除的数据，发现其中超过70%确实存在明显的标注错误或歧义。这比随机移除同样数量的数据效果要好得多。

注意事项：影响力分数高的样本不一定是“坏”样本，也可能是“关键但困难”的正样本。因此，在自动移除前，建议对高分样本进行小批量的人工抽样检查，确认其性质。可以设定一个影响力阈值，只移除高于阈值且经过抽查确认有问题的样本。

4.2 核心应用二：模型预测溯源与解释

当模型对一个输入做出令人惊讶或关键的预测时（例如，在医疗诊断或金融风控中），我们可以使用RISE来追溯这个预测最可能来源于训练数据中的哪些样本。这极大地增强了模型的可解释性和可信度。

操作流程：

1. 给定一个需要解释的模型预测（测试样本 ( z_{test} )）。
2. 运行RISE，得到所有训练样本相对于 ( z_{test} ) 的影响力分数 ( I_i )。
3. 展示影响力分数最高的前N个训练样本（例如，前5条）。这些样本的内容，就是模型做出当前预测的“主要依据”。

效果分析：在一個问答系统中，对于模型给出的一个特定答案，我们通过RISE找到了几条相关的训练问答对。分析发现，模型之所以能给出精准答案，正是因为它在训练时“见过”高度相似的问题和答案组合。这为算法工程师和领域专家提供了宝贵的调试和信任依据。同时，如果发现模型依据的是一条质量不高的数据，则提示我们需要清洗数据。

4.3 核心应用三：高效数据价值评估与核心集选择

有时我们想从一个巨大的候选数据池中，挑选出一个小的、有代表性的子集（核心集）来进行高效训练或主动学习。RISE可以帮助我们评估每个候选数据的“价值”。

操作流程（基于影响力的核心集选择）：

1. 用一个在少量干净数据上预训练的模型作为起点。
2. 将庞大的未标注候选数据集视为“训练集”，将一个小而精的验证集作为 ( D_{test} )。
3. 使用RISE（可能需要近似，因为候选数据没有标签）或类似方法，估算每个候选数据对验证集损失的“预期影响力”。
4. 选择那些预期能最大程度降低验证集损失（即影响力分数最负，因为损失降低）的候选样本进行标注和加入训练。

这种方法比随机选择或基于不确定性的选择（如熵）更能直接瞄准提升模型整体性能的目标数据点。

4.4 性能与精度权衡

RISE的优势是效率，但代价是近似误差。我们需要明确其精度边界：

• 近似误差来源：主要来自CountSketch的随机投影（可通过增加压缩维度 ( d ) 和哈希函数数 ( k ) 来降低）以及对海塞矩阵的低秩近似。
• 与精确方法的对比：在计算资源允许的小规模问题上（如逻辑回归、小CNN），可以将RISE的结果与精确计算的影响函数结果对比。实验表明，在合适的 ( d ) 设置下（如 ( d ) 为参数数量的0.5%），RISE计算出的影响力排名（Top-K）与精确结果的吻合度（如Spearman相关系数）可以超过0.9，完全能满足排序和筛选的需求。
• 计算开销：主要开销在“构建梯度草图库”阶段，复杂度为 ( O(n \cdot (T_{forward} + T_{backward}) + n \cdot d) )，其中 ( T ) 是单次前向/反向传播时间。一旦草图库建成，针对单个测试样本的影响力查询开销仅为 ( O(n \cdot d_{active} + d_{active}^3) )，速度极快。

5. 常见问题与实战避坑指南

在实际部署RISE的过程中，我踩过不少坑，这里总结一下最常见的问题和解决方案。

5.1 内存与存储爆炸

问题：即使经过压缩，存储数百万个 ( d ) 维（如8192维）的梯度草图，内存或磁盘占用依然巨大。解决方案：

• 量化存储：将草图向量从FP32转换为FP16甚至INT8进行存储，在计算时再转换回来。这对最终的影响力排序结果影响微乎其微，但能减少50%-75%的存储空间。
• 分片加载：不要一次性将所有草图加载进内存。设计一个索引，在计算时按需加载批次数据。计算G_active @ v这个矩阵-向量乘法时，可以分批进行。
• 考虑更激进的压缩：对于超大规模数据集，可以结合产品量化（Product Quantization）等更高级的向量压缩技术来存储草图，进一步牺牲少量精度换取存储空间。

5.2 梯度爆炸与数值不稳定

问题：大模型的梯度可能非常大或非常小，导致CountSketch压缩时出现数值下溢或上溢，或者使 ( H_{active} ) 矩阵病态。解决方案：

• 梯度裁剪：在计算单个样本梯度后、进行压缩前，对梯度进行全局范数裁剪（如torch.nn.utils.clip_grad_norm_）。这不会改变梯度的方向，但能稳定数值范围。
• 增加正则化：在计算 ( H_{active} ) 时，务必加上一个足够大的正则化项 ( \lambda I )。( \lambda ) 可以从1e-6开始尝试，如果求解线性系统时仍然报错（如非正定），逐步增大到1e-4或1e-3。
• 使用双精度：在计算 ( H_{active} ) 和求解线性系统时，使用torch.double精度，虽然慢一点，但能极大增强数值稳定性。

5.3 激活路径记录不准确

问题：如何准确、高效地将前向传播中的神经元激活映射到参数梯度索引？这是一个工程难点。解决方案：

• 使用Hook记录张量：在PyTorch中，通过前向Hook记录激活张量。关键是要建立一个从激活张量中的位置到模型参数扁平化向量中的索引的映射表。这需要在模型初始化后，遍历一次所有参数，建立参数名、形状与全局索引的映射关系。
• 仅关注关键层：实践表明，对于Transformer模型，中间层（特别是后几层）的激活对最终预测的影响最大。可以只在这些层的输出上注册Hook，忽略嵌入层和靠前的层，这能大幅减少需要处理的激活数量，且对结果影响不大。
• 阈值选择：激活阈值不宜过小，否则会记录大量接近零的噪声激活；也不宜过大，否则会漏掉重要信号。可以通过观察激活值的分布直方图，选择一个分位数（如95%分位数）作为阈值。

5.4 影响力分数的解释与校准

问题：计算出的影响力分数绝对值大小没有直接物理意义，只有相对排名有意义。不同测试样本之间的分数也无法直接比较。解决方案：

• 始终使用标准化或排序后的结果：不要直接比较原始分数。对于单个测试样本，关注的是训练样本影响力的相对排名。对于跨测试样本的汇总（如找对多个坏样本都有害的数据），可以先将每个测试样本下的影响力分数进行标准化（如减去均值除以标准差），然后再进行求和或平均。
• 结合领域知识进行验证：定期对高影响力样本进行人工审查。如果发现算法找出的“重要”样本在领域专家看来无关紧要，或者找出的“有害”样本其实是优质数据，就需要反思模型、损失函数或数据处理流程是否存在更深层次的问题。RISE是一个强大的工具，但它反映的是当前模型与当前数据之间的关系，如果模型本身有缺陷，RISE的结果也会出现偏差。

5.5 分布式计算的同步难题

问题：在多个GPU/节点上并行计算训练样本的梯度草图时，如何高效汇总？解决方案：

• 草图的可加性：CountSketch的一个美妙性质是，多个数据点的草图可以线性相加。也就是说，如果每个工作节点计算了自己分片内数据草图的和，那么主节点只需要将这些“和草图”相加，就能得到全局数据草图的和。这对于计算海塞矩阵的近似 ( \frac{1}{n} G^T G ) 非常有用。
• 设计两阶段流程：
  1. Map阶段：每个工作节点独立计算自己分片内每个样本的草图，并计算两个聚合量：a) 本分片所有草图的和（一个 ( d ) 维向量），b) 本分片草图的外积和（一个 ( d \times d ) 矩阵，即 ( \sum_{i \in shard} g_i g_i^T )）。
  2. Reduce阶段：主节点收集所有工作节点的和向量与外积和矩阵，分别相加，得到全局的和向量与全局的外积和矩阵。全局海塞近似矩阵即为(全局外积和矩阵) / n。 这样，通信开销只有 ( O(d + d^2) )，与数据量 ( n ) 无关，非常适合分布式计算。