STRIDE:基于子集扰动的稀疏恢复训练数据归因方法 STRIDE: Training Data Attribution via Sparse Recovery from Subset Perturbations
通过激活空间steering operators和压缩感知,实现高效LLM训练数据归因
前置知识
训练数据归因(Training Data Attribution, TDA)
训练数据归因是量化单个训练样本对模型特定预测的因果影响的技术。核心思想是回答反事实问题:如果某个训练样本从训练集中移除,模型的预测会如何变化?最严谨的方法是Leave-One-Out(LOO)重训练,计算$\delta_x(\{z_i\}) = F_x(S) - F_x(S \setminus \{z_i\})$,但这对LLM来说计算成本太高。现有方法通过参数空间梯度近似来估计这个效应,但存在计算昂贵和局部近似失效的问题。
理解本文需要TDA作为基础概念,因为STRIDE正是要解决TDA的计算可扩展性问题,同时保持因果解释力。
影响函数(Influence Functions)
影响函数源于稳健统计理论,用于估计训练样本对模型参数和预测的影响。在神经网络中,常用参数扰动近似:$\hat{\theta}_{\epsilon, z} \approx \hat{\theta} - \epsilon H^{-1}\nabla_\theta \ell(z; \hat{\theta})$,其中$H$是Hessian矩阵。预测的影响可以通过链式法则计算:$\frac{d\ell(z_{test}; \hat{\theta})}{d\epsilon} = -\nabla_\theta \ell(z_{test}; \hat{\theta})^\top H^{-1}\nabla_\theta \ell(z; \hat{\theta})$。对LLM来说,计算和存储数十亿参数的梯度及二阶导数矩阵是巨大的计算负担。
影响函数是STRIDE要改进的主要baseline之一,理解其局限性(参数空间计算、局部线性假设)有助于理解STRIDE为何转向激活空间方法。
压缩感知(Compressive Sensing)
压缩感知是信号处理理论,表明如果信号是稀疏的(只有少数非零元素),可以从远低于信号维数的线性测量中精确恢复信号。给定测量$y = Mx + e$,其中$x$是k-稀疏向量,$M \in \mathbb{R}^{K \times n}$是测量矩阵($K \ll n$),$e$是噪声,通过求解$\min \|x\|_1 \quad \text{s.t.} \quad \|y - Mx\|_2 \leq \epsilon$可以恢复$x$。关键理论保证是RIP(Restricted Isometry Property)或null-space property。STRIDE将训练数据归因问题映射到这个框架:训练样本影响力向量$w(x)$是稀疏的,子集扰动响应$y_x$是测量,成员关系矩阵$M$是测量矩阵。
压缩感知是STRIDE的核心数学框架,理解稀疏恢复原理和测量矩阵要求(如本文使用的expander图性质)是理解方法理论基础的关键。
线性数据建模评分(Linear Datamodeling Score, LDS)
LDS是评估训练数据归因方法的标准指标,衡量预测的影响力向量与真实重训练子集响应之间的相关性。给定$K$个子集的扰动响应向量$y_x$和成员关系矩阵$M$,LDS计算真实影响力$w_{true} = (M^\top M)^{-1}M^\top y_x$与预测影响力$w_{pred}$之间的Spearman秩相关系数。值越高(接近1)表示归因越准确。LDS的优势是不需要昂贵的LOO重训练,只需要子集级别的重训练即可提供可靠的评估基准。
LDS是论文中所有实验的主要评估指标,理解LDS的含义和计算方式有助于正确解读STRIDE的性能优势。
Expander图(Expander Graph)
Expander图是一类具有强连通性质的稀疏图,在压缩感知中用作测量矩阵可以保证稀疏恢复的唯一性。对于二部图$G = (L, R, E)$,如果每个左顶点的度为$d$,且对于任意小集合$T \subseteq L$($|T| \leq s$),其邻域满足$|\Gamma(T)| \geq (1-\epsilon)d|T|$,则称为$(s, \epsilon)$-expander。STRIDE中,训练样本对应左顶点,子集对应右顶点,边表示样本属于子集。如果二部图是$(2k, \epsilon)$-expander且$\epsilon < 1/6$,则可以通过$\ell_1$最小化唯一恢复任意$k$-稀疏的影响力向量。理论要求$K = O(k \log(n/k))$个子集。
Expander图是STRIDE稀疏恢复理论保证的基础,理解其性质可以解释为什么只需少量子集测量就能恢复单个样本的影响力。
研究动机
现有训练数据归因方法在大规模语言模型(LLM)上面临严重的计算和理论瓶颈。最严谨的重训练方法(如Leave-One-Out或Data Models)需要对每个目标样本重新训练模型,这对数十亿参数的LLM来说计算成本极高,完全不可行。主流的梯度方法(如Influence Functions及其扩展)需要在参数空间跟踪数十亿参数的梯度,并计算Hessian矩阵的逆或近似。例如,DataInf、EKFAC、LOGRA等方法虽然提升了效率,但仍然需要在巨大的参数空间中操作,内存和计算开销随参数数量线性增长。更大的问题是,这些方法严重依赖损失函数局部凸性和模型严格收敛的假设,而深度学习中这些假设常常被违反,导致估计失效。基于表示的方法(如AirRep、GTE)虽然可扩展,但依赖启发式的特征相似性或学习到的评分函数,缺乏重训练级别的因果基础,在预训练目标下特别难以捕捉数据贡献的复杂动力学。
本文的目标是本文的核心目标是开发一种既能保持重训练级别的因果解释力,又能在LLM规模上高效可扩展的训练数据归因方法。具体来说,作者希望绕过参数空间的计算瓶颈,同时避免表示方法的启发式局限性,找到一种既能准确模拟数据子集训练引起的模型行为变化,又能高效恢复单个样本影响力的解决方案。
与已有工作不同的是,本文的独特切入角度是从"参数空间"转向"激活空间"。现有方法要么在参数空间建模(梯度方法),要么在启发式特征空间建模(表示方法)。STRIDE假设训练数据的因果影响可以更稳健、更高效地建模为模型激活空间中的功能性变化。这一视角转变基于近期工作表明低秩激活扰动可以忠实复制微调引起的行为变化。核心洞察是:当LLM在数据子集上训练时,它在内部表示中引起系统性变化。与其计算权重应该如何调整,不如建模激活应该如何被"steer"以产生更新的预测。这消除了对数十亿参数梯度进行计算、存储或求逆的需要,同时保持了强因果联系。
核心方法
STRIDE(Steering-based Training Data Influence Decomposition)将训练数据归因建模为激活空间扰动上的稀疏恢复问题。框架分为两个阶段:首先,离线阶段学习轻量级的低秩\"steering operators\",应用于冻结模型的中间激活,显式模拟在特定训练数据子集上训练时的模型输出概率变化;其次,在线推理阶段,对这些learned operators应用测试查询,生成扰动响应向量,然后通过稀疏线性分解恢复单个训练样本的影响力。直觉上,如果某些训练样本对特定预测很重要,那么这些样本所属的子集的steering operator应该能显著改变模型的预测。通过测量所有子集的扰动响应,并假设影响力是稀疏的(只有少数样本有重要影响),可以用压缩感知技术从子集级别的聚合测量中恢复单个样本的影响力。
STRIDE的核心创新点是将归因问题从参数空间梯度估计转移到激活空间steering,用低秩算子显式拟合真实子集重训练的响应,然后用压缩感知恢复单个样本影响力。与现有方法的本质区别在于:它不在参数空间近似梯度,也不使用启发式特征相似性,而是直接学习激活空间的函数映射来模拟重训练行为。这保持了与实际重训练的强因果联系,同时完全避开了大规模参数空间的计算瓶颈。理论上,STRIDE利用expander图的性质证明了稀疏二进制子集矩阵足以唯一恢复稀疏的影响力向量,测量复杂度为$K = O(k \log(n/k))$,与经典压缩感知的采样率匹配。
方法步骤详情
STRIDE的完整步骤分为离线训练和在线推理两个阶段。离线阶段(训练steering operators):首先构建$K$个训练数据子集,形成二进制成员关系矩阵$M \in \{0,1\}^{K \times n}$,其中$M_{k,i} = 1$表示样本$z_i$属于子集$A_k$。然后学习共享低秩basis network $B_\psi: \mathbb{R}^{d_h} \to \mathbb{R}^{1 \times r}$和子集特定的steering矩阵集合$A = \{a_1, ..., a_K\}$,其中$a_k \in \mathbb{R}^{r \times C}$,$C$是输出类别数。steered logits计算为$\tilde{o}_k(x) = o_x + B_\psi(h_x)a_k$。训练目标包含三个损失:(1)Fidelity Loss:$L_{fid}(A_k) = -\frac{1}{|A_k|}\sum_{x\in A_k}\sum_{t=1}^T \omega_t \log \sigma(o_{x,t} + B_\theta(f_{x,t})a_k)_{y_{x,t}}$,确保steering改善子集内样本的预测,模拟只在该子集上重训练的效果;(2)Stability Loss:$L_{stab}(A_k, R_k) = \frac{1}{|R_k|}\sum_{x\in R_k} D_{KL}(\sigma(o_x) || \sigma(o_x + B_\theta(f_x)a_k))$,惩罚对无关样本预测的偏离,防止全局退化或捷径学习,$R_k$是从$S \setminus A_k$中随机采样的补集;(3)Linearity Loss(LDS正则化):$L_{LDS}(y, \tilde{M}) = \|\tilde{M}(\tilde{M}^\top\tilde{M} + \gamma I)^{-1}\tilde{M}^\top y - y\|_2^2$,鼓励learned steering效果遵循稀疏恢复所需的加性结构,其中$\tilde{M} = MR$,$R \in \mathbb{R}^{n \times q}$是固定随机投影,$q \ll \min\{K, n\}$。总目标$L_{total} = \sum_{k=1}^K (L_{fid}(A_k) + \lambda_{stab}L_{stab}(A_k, R_k)) + \lambda_{LDS}L_{LDS}(y, \tilde{M})$通过联合优化$\psi$和$A$来最小化。在线推理阶段(稀疏恢复):对于测试查询$x$,首先提取冻结基础模型的latent特征$h_x$和logits $o_x$;然后对所有$K$个子集应用steered logits计算$\tilde{o}_k(x) = o_x + B_\psi(h_x)a_k$;测量每个子集的扰动响应$y_{x,k} = \text{Loss}(o_x) - \text{Loss}(\tilde{o}_k(x))$,形成测量向量$y_x \in \mathbb{R}^K$;最后求解稀疏恢复问题$\hat{w}(x) = \arg\min_{w\in\mathbb{R}^n} \frac{1}{2}\|y_x - Mw\|_2^2 + \lambda\|w\|_1$,得到每个训练样本的影响力分数。
技术新颖性
STRIDE的技术新颖性体现在多个层面。首先,视角新颖性:首次将激活空间steering与压缩感知结合用于训练数据归因,完全避开参数空间计算。其次,理论新颖性:证明了稀疏二进制子集矩阵(由expander二部图构造)足以唯一恢复稀疏影响力向量,达到与经典压缩感知相同的$O(k \log(n/k))$测量复杂度,但使用稀疏矩阵而非密集随机矩阵。第三,架构新颖性:steering operators使用共享低秩basis和子集特定steering矩阵的分解,参数效率高,且通过fidelity、stability和linearity三重损失的联合训练确保了因果保真度、预测稳定性和可加性。第四,实用性新颖性:离线训练的operators可以零样本应用于任何新查询,无需额外训练,推理时只需一次批量前向传播和稀疏优化,大幅提升了可扩展性。相比最先进的AirRep方法,STRIDE在1.38B参数模型上快12倍且LDS更高。
实验结果
STRIDE在LLM预训练归因任务上达到了最先进的准确性和效率。在Nanochat模型(286M、537M、897M、1.38B参数)上,STRIDE的LDS相关系数分别为0.1581、0.1792、0.1598、0.1671,显著超越所有baseline。相比之下,LoGRA在286M上为0.1126,在1.38B上仅为0.1139;AirRep在286M上为0.1406,但无法扩展到1.38B规模。计算效率方面,STRIDE在1.38B模型上总运行时间为9.9小时,峰值VRAM为8.41GB;而AirRep在1.38B上预估需要116.1小时,是STRIDE的11.7倍;LoGRA需要52.3小时(仅286M),且随模型大小扩展性差。STRIDE的运行时间主要来自离线训练阶段(9.02小时),推理阶段(0.57小时前向传播 + 0.34小时Lasso)很快。在监督微调(SFT)任务上,STRIDE在Alpaca、Tulu、FLAN、SafeRLHF四个数据集上的LDS分别为0.2426、0.1611、0.1932、0.3995,在三个数据集上超越所有baseline,在SafeRLHF上略低于AirRep(0.4608)但快得多。下游应用验证:在FLAN 100K数据选择任务上,STRIDE的平均unigram F1为49.65(±22.08),与AirRep(49.66)和LoGRA(49.94)相当,远超随机基线(42.97)。在数据污染检测任务上,STRIDE+LoGRA组合达到74.2%的召回率,比单独使用LoGRA(62.1%)提升12.1个百分点,证明激活空间和梯度空间信号互补。定性分析显示STRIDE能识别有意义的训练样本:例如,在证明$\sqrt{2}$无理数时,STRIDE发现模型模仿了训练数据中$\sqrt{3}$证明的结构;在AI合理撒谎问题上,归因追踪到关于感知AI的新闻报道和数据权利的政策文本。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| LLM预训练归因 | LDS (Spearman correlation) | 0.1671 (1.38B) | LoGRA: 0.1139, AirRep: 不可行 | 相对LoGRA提升46.7%,且AirRep无法完成 |
| LLM预训练归因 | Runtime (1.38B, hours) | 9.9 | AirRep: 116.1 | 快11.7倍 |
| SFT归因 | LDS (Alpaca) | 0.2426 | AirRep: 0.2258, LESS: 0.0959 | 相对AirRep提升7.5% |
| SFT归因 | LDS (SafeRLHF) | 0.3995 | AirRep: 0.4608, DSDM: 0.2594 | 相对DSDM提升54.1% |
| 数据选择 (FLAN 100K) | Unigram F1 | 49.65 | Random: 42.97, LoGRA: 49.94 | 相对Random提升15.5% |
| 数据污染检测 (MATH) | Recall | 74.2% (LoGRA+STRIDE) | LoGRA: 62.1% | 提升12.1个百分点 |
局限与改进
作者承认STRIDE作为基于表示的steering方法,其有效性依赖于基础模型内部激活的质量和干预层的选择。方法依赖于局部线性和加性影响力假设,虽然这在标准预训练和指令微调regime下鲁棒成立,但在极端分布偏移或高度非凸记忆现象下可能失效。论文未验证对RL-based训练目标(如RLHF)的扩展,作者承认这是一个开放方向。作者的观察显示,STRIDE在视觉模型(如CIFAR-10上的ResNet-9)上LDS略低于LoGRA,表明方法在卷积架构上的性能可能略有下降。另一个局限是稀疏恢复假设影响力是稀疏的,虽然实验支持这一假设(预训练中sparsity = 1.0,SFT中约60-68%),但在某些场景下可能不成立。最后,steering operators的训练需要选择合适的干预层和秩$r$,论文通过实验选择但没有提供自动选择策略。
独立分析的弱点
STRIDE的独立分析弱点包括:(1)干预层敏感性:方法需要在特定层的激活上进行steering,不同层的选择可能显著影响性能,论文通过实验选择但缺乏理论指导。改进方向:探索多层联合steering或自动化层选择策略,例如通过可学习的层注意力机制自适应地组合不同层的信号。(2)秩的选择:低秩分解中的秩$r$需要手动调优,过小可能表达能力不足,过大可能增加计算和过拟合风险。改进方向:使用自适应秩估计或正则化自动压缩,例如通过奇异值衰减或基于LDS验证的秩搜索。(3)对卷积架构的适应性:在CIFAR-10 ResNet-9上LDS低于LoGRA,表明方法对transformer架构更友好。改进方向:为卷积网络设计专门的基础网络结构,或结合卷积特定的特征表示(如通道注意力)来提升性能。(4)子集数量与质量的权衡:理论要求$K = O(k \log(n/k))$个子集,但实际中需要平衡子集数量(计算成本)和每个子集的大小(训练质量)。改进方向:研究自适应子集采样策略,例如重要性采样或基于梯度的子集构造,以在固定$K$下最大化信息量。
未来方向
作者提出的未来研究方向包括将STRIDE扩展到RL-based训练目标(如RLHF),这对AI安全归因非常重要。基于成果可延伸的方向包括:(1)多任务归因:探索如何在不同任务或模态之间共享steering operators,以实现更高效的大规模归因系统。(2)层级归因:不仅识别影响单个预测的训练样本,还分析样本影响力的层次结构,例如哪些样本影响模型的整体能力,哪些仅影响特定细分领域。(3)主动归因:结合STRIDE与主动学习,自动识别需要归因的关键查询或样本,以优化资源分配。(4)因果关系挖掘:利用STRIDE的影响力分数挖掘训练数据中的因果关系,例如发现数据泄露、偏差传递或虚假信息传播链。(5)实时归因:研究如何将STRIDE部署到实时系统中,实现模型行为的动态监控和调试,例如在生产环境中识别有害输出的根源。(6)理论扩展:进一步探索expander图在稀疏恢复中的最优构造,以及放宽稀疏性假设下的恢复理论。
复现评估
STRIDE的复现评估:作者在项目页面提供了代码和实验细节,但目前未正式开源。实验环境为单个H100-80GB GPU。预训练实验使用Nanochat系列模型(286M-1.38B参数)在Nemotron-ClimbMix数据集上训练,数据集规模约为11.4M训练样本。SFT实验使用Qwen2.5-0.5B模型,在FLAN(100K样本)、Alpaca(51,760样本)、Tulu(100K样本)、SafeRLHF(251,963样本)上微调。数据污染实验在OpenWebText上注入MATH问题。超参数包括:子集数量$K=256$(预训练)或$K=100$(SFT),低秩$r=64$,正则化系数$\lambda$通过验证集调优。复现难度中等:主要挑战是计算资源需求(特别是预训练实验需要多个GPU小时),以及steering operators的训练需要仔细调优三个损失的权重。但离线训练一次后可以用于任意查询,推理成本很低。作者提供了详细的实现细节在附录F,包括低秩算子架构、指令微调token masking、干预层选择、内存优化、超参数设置等,有助于复现。
论文图表