LoRIF: Low-Rank Influence Functions for Scalable Training Data Attribution¶
会议: NeurIPS 2026
arXiv: 2601.21929
代码: https://github.com/epfl-mlo/LoRIF
领域: 可解释性
关键词: 训练数据归因, 影响函数, 低秩分解, Gauss-Newton曲率, 随机投影
一句话总结¶
针对大规模模型训练数据归因中投影梯度存储I/O高昂与Hessian逆矩阵内存爆炸的双重瓶颈，本文提出LoRIF，利用梯度张量的样本级与全局低秩特性，结合低秩投影因子存储与截断SVD加速曲率求解，在保留甚至超越全秩归因精度的同时实现最高20倍的存储削减与查询加速。
研究背景与动机¶
训练数据归因（Training Data Attribution, TDA）旨在追踪模型对特定查询样本的预测行为究竟源于哪些训练样本。随着现代语言模型扩展至数十亿甚至上万亿token的训练规模，TDA在模型除错、高质量数据清洗、数据投毒防御、越狱审计以及版权溯源中展现出不可替代的价值。然而，基于影响函数（Influence Functions）等具备严密统计与优化理论支撑的TDA方法，在工业级前沿大模型上始终面临巨大的落地阻碍：精确计算逆Hessian向量积需要反复遍历海量训练集反传梯度，而现有的二阶近似扩展方法在数百万量级的数据集上要么由于磁盘I/O耗尽查询时延，要么因显存溢出而被迫牺牲归因分辨率。
现有最前沿的扩展范式主要依赖随机投影与阻尼Gauss-Newton曲率近似（如TRAK、LoGRA与TrackStar）。这类方法先离线收集所有训练样本在各线性层上的双边随机投影梯度，建立单次索引后供下游查询多次复用。然而，该框架存在不可调和的质量-可扩展性权衡（Quality-Scalability Tradeoff）：一方面，投影维度 \(D\) 决定了子空间的表达能力，增大 \(D\) 是提升归因保真度（如反事实子集重训相关性LDS指标）的核心前提；另一方面，随着有效投影维度 \(D\) 增加，每个样本需要存储并读取 \(O(D)\) 维度的浮点稠密向量，使得在海量数据检索时单次查询有超过96%的时间完全耗费在磁盘I/O读取上，同时构建 \(D \times D\) 的层级阻尼Gauss-Newton逆矩阵需要 \(O(D^2)\) 的显存与 \(O(D^3)\) 的求逆时间开销，使得工程实操中各层投影维度被迫受限在 \(10^5\) 以内。
本文的核心洞见在于：深度神经网络的梯度在微观（单样本跨token外积）和宏观（跨样本梯度协方差谱分布）两个层面上均天然具备高度集中的低秩特性。核心 idea：在每层随机投影梯度上引入低秩分解因子化存储以击穿I/O瓶颈，并利用截断SVD配合Woodbury矩阵恒等式在低维主子空间内隐式解析求逆曲率项，从而在大幅缩减存储与内存的同时，支持将有效投影维度推向更高量级。
方法详解¶
整体框架¶
LoRIF在标准投影影响函数（LoGRA式流程）的基础上，针对“投影梯度存储与查询I/O”和“逆Hessian曲率矩阵显存开销”两处瓶颈进行了针对性的低秩重构。整个归因框架包含三个阶段：离线梯度抽取与低秩分解存储、曲率子空间截断SVD与逆阵近似、以及在线查询阶段的快速影响分数评估。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["训练样本前向与反向传播<br/>获取激活值 X 与输出梯度 δY"] --> B["双边随机投影<br/>降维至 d1 × d2 矩阵"]
B --> C["1. 样本级低秩分解因子化存储<br/>Block Power Iteration 提取秩 c 因子 u, v"]
C --> D["2. 截断SVD与Woodbury曲率近似<br/>随机化SVD获取秩 r 主奇异系统 Vr, Σr"]
D --> E["3. 快速查询影响分数评估<br/>内积分解 + 对角核投影两步打分"]
E --> F["输出高影响力训练样本排序与归因分数"]在离线构建阶段，对于网络中的每个线性层，模型首先将输入的激活特征和反向传播的梯度张量通过双边随机高斯投影矩阵降维为 \(d_1 \times d_2\) 的投影梯度矩阵 \(\tilde{G}_i\)；随后立即通过分块幂迭代将其分解为低秩因子 \(u_i \in \mathbb{R}^{d_1 \times c}\) 与 \(v_i \in \mathbb{R}^{d_2 \times c}\) 写入高速磁盘；在所有样本梯度因子落盘后，利用随机化SVD自低秩因子批次重构梯度矩阵，提取跨样本梯度矩阵 \(G\) 的前 \(r\) 个右奇异向量与奇异值；在在线查询阶段，对于给定的测试样本，通过Woodbury公式将显存昂贵的密集逆阵乘法转化为低维对角阵缩放与低秩因子内积运算，实现近实时的归因排序。
关键设计¶
1. 样本级低秩分解因子化存储：解决存储容量与查询I/O瓶颈 针对每个样本在每层存储全尺寸投影梯度矩阵 \(\tilde{G}_i \in \mathbb{R}^{d_1 \times d_2}\) 带来的巨大磁盘占用与I/O读取延迟，本文利用Transformer线性层梯度本质上由序列各token外积叠加而成、能量高度集中于少数方向的物理特性，不对原始矩阵进行稠密拉平，而是通过少量步骤的分块幂迭代（Block Power Iteration）提取其最佳秩 \(c\) 近似： $\(\tilde{G}_i \approx u_i v_i^\top, \quad u_i \in \mathbb{R}^{d_1 \times c}, \; v_i \in \mathbb{R}^{d_2 \times c}\)$ 该因子化使每个样本每层的存储参数量从 \(d_1 d_2\) 个浮点数急剧压缩至 \(c(d_1 + d_2)\)。在 \(c=1\) 时，存储压缩比逼近 \(\min(d_1, d_2)/2\)（通常达到数十倍）。更重要的是，在实际归因查询中，加载训练集梯度特征所耗费的I/O传输量获得等比例骤降，从根本上消除了制约检索吞吐的核心瓶颈。
2. 截断SVD与Woodbury曲率近似：打破高维Hessian逆阵内存墙 针对形成和存储层级阻尼Gauss-Newton逆矩阵 \((G^\top G + \lambda I)^{-1} \in \mathbb{R}^{D \times D}\) 所需的 \(O(D^2)\) 显存和 \(O(D^3)\) 求逆计算量，本文发掘了深度模型梯度协方差矩阵普遍具备的尖峰谱结构（Spiked Spectrum）。通过对跨样本聚合梯度矩阵 \(G \in \mathbb{R}^{N \times D}\) 执行截断秩 \(r\) 的随机化SVD（\(r \ll \min(N, D)\)），得到前 \(r\) 个主奇异分量 \(G \approx U_r \Sigma_r V_r^\top\)，此时曲率矩阵近似为 \(H \approx V_r \Sigma_r^2 V_r^\top + \lambda I\)。借助Woodbury矩阵逆恒等式，其解析逆可显式展开为： $\((V_r \Sigma_r^2 V_r^\top + \lambda I)^{-1} = \frac{1}{\lambda} I - \frac{1}{\lambda^2} V_r \left(\Sigma_r^{-2} + \frac{1}{\lambda} I_r\right)^{-1} V_r^\top\)$ 其中核心的逆项 \((\Sigma_r^{-2} + \frac{1}{\lambda} I_r)^{-1}\) 退化为极其轻量的 \(r \times r\) 对角矩阵。这使得曲率存储开销从 \(O(D^2)\) 降低至 \(O(Dr)\)，运行期内存开销大幅缓解，使得在固定硬件显存内将有效投影维度 \(D\) 提升一到两个数量级成为可能。
3. 快速查询影响分数评估：结合因子化内积与子空间投影的高效推理 在计算测试样本 \(x_{te}\) 对训练样本 \(x_{tr}\) 的影响值时，将Woodbury逆展开式代入影响函数定义式 \(\mathcal{I}(x_{tr}, x_{te}) \approx g_{te}^\top (G^\top G + \lambda I)^{-1} g_{tr}\)，得到双项分离形式： $\(\mathcal{I}(x_{tr}, x_{te}) \approx \frac{1}{\lambda} g_{te}^\top g_{tr} - \frac{1}{\lambda^2} (g'_{te})^\top \left(\Sigma_r^{-2} + \frac{1}{\lambda} I_r\right)^{-1} g'_{tr}\)$ 其中第一项为梯度的Frobenius无向内积 \(\langle \tilde{G}_{te}, \tilde{G}_{tr} \rangle_F\)，可直接通过低秩因子的高效矩阵乘法 \((u_{te}^\top u_{tr})(v_{te}^\top v_{tr})\) 在 \(O(c^2(d_1 + d_2))\) 复杂度内完成，无需重建 \(d_1 \times d_2\) 稠密阵；第二项中的 \(g'_{tr} = V_r^\top g_{tr} \in \mathbb{R}^r\) 和 \(g'_{te} = V_r^\top g_{te} \in \mathbb{R}^r\) 为向曲率主子空间的紧凑投影，在对角矩阵加权下仅耗费 \(O(r)\) 运算量。计算流水线使得单批次查询能以毫秒级流式处理数百万级训练集。
损失函数 / 训练策略¶
归因过程依赖于预训练或微调完成的模型检查点，不涉及对原始模型权重的额外优化。关键超参数选择与执行策略如下： 1. 阻尼系数 \(\lambda\)：在全特征谱不可直接获得的情况下，设定 \(\lambda = 0.1 \times \text{mean}(\Lambda_{r+p})\)，取随机化SVD中计算出的前 \(r+p\) 个奇异值平方的均值（过采样因子 \(p=10\)）。 2. 分解迭代步数：样本级秩分解的块幂迭代在 \(c=1\) 时固定设为 8 步（\(c>1\) 时设为 16 步）；截断SVD的幂迭代固定设为 3 步，经验证增加步数对归因质量无明显增益。 3. 超参数调优准则：在给定硬件存储与计算预算下，优先选择最小的缩放因子 \(f\)（即最大化有效投影维度 \(D = IO/f^2\)），其次维持 \(c=1\) 的极简低秩结构，最后将截断秩 \(r\) 增加至显存容量或指标饱和区间（典型取 \(2^8 \sim 2^{14}\)）。
实验关键数据¶
主实验¶
论文在三种不同规模的模型与数据场景下展开评估：GPT2-small（124M参数，WikiText-103共23.3万切片）、OLMo-3-7B（7B参数，Dolci Think SFT共220万条指令数据）、以及Apertus-70B（70B参数，公开SFT共380万条样本）。评测指标包括反映反事实子集重训输出相关性的线性数据建模分数（Linear Datamodeling Score, LDS，GPT2-small适用）以及衡量检索出的top-k支持样本对目标单步梯度更新概率提升效果的Tail-patch分数（大模型适用）。
表 1：GPT2-small（233k训练切片）在不同存储约束下的主结果对比（各区间配置与代表性基线）
| 方法 | 缩放因子 \(f\) | 秩 \(c\) | 截断秩 \(r\) | LDS ↑ | 梯度存储 ↓ | 查询时延 ↓ |
|---|---|---|---|---|---|---|
| EK-FAC (参数空间上下文基线) | – | – | – | \(0.3936 \pm 0.013\) | 2.1 GiB | 20 hr |
| RepSim (表示相似度基线) | – | – | – | \(0.0103 \pm 0.009\) | 340.7 MiB | 0.15 s |
| GradDot (高存储) | 8 | – | – | \(0.0576 \pm 0.010\) | 574.9 GiB | 209.8 s |
| TrackStar (高存储) | 8 | – | – | \(0.2623 \pm 0.012\) | 574.9 GiB | 241.1 s |
| LoGRA (高存储) | 8 | – | – | \(0.3123 \pm 0.012\) | 574.9 GiB | 211.1 s |
| LoRIF (高存储，高维大秩) | 4 | 32 | \(2^{14}\) | \(0.3428 \pm 0.012\) | 511.0 GiB | 287.1 s |
| LoGRA (中等存储) | 16 | – | – | \(0.2139 \pm 0.011\) | 143.7 GiB | 56.3 s |
| LoRIF (极低存储配置) | 4 | 1 | \(2^{12}\) | \(0.2073 \pm 0.013\) | 16.6 GiB (降8.7×) | 25.5 s (降2.2×) |
| LoGRA (低存储) | 32 | – | – | \(0.1370 \pm 0.010\) | 35.9 GiB | 16.4 s |
| LoRIF (超紧凑存储) | 16 | 1 | \(2^{11}\) | \(0.1392 \pm 0.011\) | 4.1 GiB (降8.8×) | 2.93 s (降5.6×) |
表 2：OLMo-3-7B 与 Apertus-70B 前沿大模型上的大规模归因结果（Tail-patch 分数）
| 评测模型 | 方法 | 缩放因子 \(f\) | 秩 \(c\) | 截断秩 \(r\) | Tail-patch (%) ↑ | 梯度存储 ↓ | 查询时延 ↓ |
|---|---|---|---|---|---|---|---|
| OLMo-3-7B | RepSim | – | – | – | \(0.144 \pm 0.04\) | 16.99 GiB | 7.26 s |
| OLMo-3-7B | GradDot | 128 | – | – | \(0.411 \pm 0.04\) | 1638.6 GiB | 781.4 s |
| OLMo-3-7B | LoGRA | 128 | – | – | \(0.546 \pm 0.04\) | 1638.6 GiB | 782.1 s |
| OLMo-3-7B | LoRIF (轻量匹配) | 128 | 1 | \(2^8\) | \(0.495 \pm 0.05\) | 80.9 GiB (降20.3×) | 35.6 s (降22.0×) |
| OLMo-3-7B | LoRIF (高维扩张) | 16 | 1 | \(2^8\) | \(0.986 \pm 0.07\) | 647.4 GiB (省60%) | 469.7 s (快1.7×) |
| Apertus-70B | LoGRA | 512 | – | – | \(0.177 \pm 0.04\) | 1815.1 GiB | 1087.9 s |
| Apertus-70B | LoRIF (轻量高质) | 256 | 1 | \(2^9\) | \(0.262 \pm 0.03\) | 333.7 GiB (降5.4×) | 222.5 s (降4.9×) |
| Apertus-70B | LoRIF (高维极致) | 64 | 1 | \(2^9\) | \(0.391 \pm 0.04\) | 1334.7 GiB (省26%) | 843.8 s (快1.3×) |
消融实验¶
表 3：解耦 LoRIF 两大低秩组件的消融分析（GPT2-small，展示两模块对显存与I/O瓶颈的互补贡献）
| 实验配置 | 缩放因子 \(f\) | 秩 \(c\) | 截断秩 \(r\) | LDS ↑ | 存储容量 | 单样本存储 | 查询时延 | 状态/限制 |
|---|---|---|---|---|---|---|---|---|
| 完整 LoRIF (高维紧凑) | 4 | 1 | \(2^{12}\) | 0.2073 | 16.6 GiB | 0.07 MiB | 25.5 s | 正常运行 |
| 去掉截断 SVD (保留低秩分解) | 4 | 1 | – | – | – | – | – | 显存溢出 (OOM) |
| 去掉低秩分解 (保留截断 SVD) | 4 | – | \(2^{12}\) | 0.3525 | 2297.7 GiB | 10.1 MiB | 931.6 s | I/O极重，磁盘膨胀138× |
| 完整 LoRIF (高精度大秩) | 4 | 32 | \(2^{14}\) | 0.3428 | 511.0 GiB | 2.25 MiB | 287.1 s | 正常运行 |
| 去掉截断 SVD (大秩配置) | 4 | 16/32 | – | – | – | – | – | 显存溢出 (OOM) |
| 完整 LoRIF (中等维度) | 16 | 1 | \(2^{11}\) | 0.1392 | 4.1 GiB | 0.02 MiB | 2.93 s | 正常运行 |
| 去掉截断 SVD (中等维度) | 16 | 1 | – | 0.1414 | 4.1 GiB | 0.02 MiB | 6.57 s | 曲率显存可容纳但延时增高 |
| 去掉低秩分解 (中等维度) | 16 | – | \(2^{11}\) | 0.2080 | 143.7 GiB | 0.63 MiB | 66.9 s | 存储与时延增大35× / 23× |
关键发现¶
- 两模块高度互补、缺一不可：消融实验直接证明，去掉截断SVD会导致在高维度（\(f=4\)）下无法构建显存高达数十GB的稠密Hessian逆矩阵而直接OOM；去掉低秩分解虽然利用SVD保持了计算可行性，但磁盘存储飙升至2TB以上，单次查询时延增加数十倍（瓶颈完全被磁盘I/O锁死）。
- 空间预算优先分配给投影维度 \(D\) 而非局部秩 \(c\)：在存储预算固定的前提下，增大有效投影维度 \(D\)（缩小 \(f\)）对归因质量的提振效果显著优于增加每个样本的分解秩 \(c\)。将 \(c\) 压缩至最激进的 \(c=1\) 时，已能捕获大部分梯度能量，释放出的空间让 \(D\) 大幅扩张，进而取得全局最优的Pareto前沿。
- 奇异值权重自衰减机制解释了截断SVD的有效性：理论推导与谱分析表明，逆Hessian修正系数 \(w_i = \sigma_i^2 / (\lambda(\lambda + \sigma_i^2))\) 随着奇异值平方单调递减；在小奇异值方向上，不仅训练集梯度自身均方能量极小（\(\sigma_i^2/N\)），修正权重更以 \(O(\sigma_i^3/\lambda^2)\) 速度衰减，使得截断微小奇异方向对归因排序的扰动几乎为零。
亮点与洞察¶
- 将低秩压缩技术精准作用于归因管线的关键短板：不同于前人简单把低秩用于模型微调（LoRA）或查询梯度缓冲，LoRIF首次辨明投影TDA的根本瓶颈是“样本级离线I/O”与“跨样本曲率显存”，将因子分解与Woodbury逆代数对应嵌入，逻辑严密而闭环。
- 打破投影TDA的维度天花板：以往方法因显存与存储限制，在70B模型上只能退守到极小的投影维度（如 \(f=512\)），丢失了深层语义细节；LoRIF将可行维度扩大数倍（支持 \(f=64\)），直接让70B模型上的Tail-patch分数从 0.177% 跃升至 0.391%。
- 对通用投影与草图管线的高度可迁移性：LoRIF的样本级因子分解不绑定特定投影矩阵，可无缝叠加在任何线性层双边投影或稀疏草图（如TrackStar、GRASS）之上；截断SVD与Woodbury恒等式同样适用于所有基于Gauss-Newton的曲率修正系统。
局限与展望¶
- 离线索引构建仍需完整梯度遍历：虽然LoRIF显著缩减了持久化存储和复用查询延迟，但构建初始归因索引时仍需在训练集上反传一遍梯度。对于长序列微调语料，前序抽取开销依然可观。
- 评测语料局限于SFT微调集而非完整预训练全量集：受算力与数据体量限制，超大模型实验主要在数百万条的SFT指令数据上验证，尚未覆盖数十万亿token的完整海量预训练语料库。
- 超参数选择依赖经验试探：当前 \((f, c, r)\) 的最优搭配依赖于经验预算网格调优，尚未建立端到端预测特定模型架构与数据规模下最优参数组合的解析泛化理论。
相关工作与启发¶
- vs LoGRA / TrackStar: 均属于双边随机投影与Gauss-Newton曲率近似范式。LoGRA直接存储 \(d_1 \times d_2\) 稠密梯度并显式构建 \(D \times D\) 逆矩阵，存储与显存开销随维度呈二次方爆炸；LoRIF通过引入秩 \(c\) 因子存储和秩 \(r\) 截断SVD，突破了维度扩张瓶颈，在中大模型上取得20倍的存储与I/O压缩。
- vs EK-FAC / ASTRA: 属于原始参数空间内的Kronecker积曲率近似方法，虽然理论精度上限高，但每次处理查询批次时必须在线反复重算全部训练样本梯度，无法支持低延迟、高频次的检索式归因；LoRIF走“一次索引、秒级复用”路线，更契合实际大模型排查需求。
- vs RepSim / 文本检索: 基于浅层隐空间表征或文本相似度的方法极其迅速，但完全脱离了模型参数损失面与反向梯度优化动力学，无法有效预测训练样本剔除后对模型输出的因果改变；LoRIF在保持极低查询延迟的同时，提供了严格的二阶影响理论保真度。
评分¶
- 新颖性: ⭐⭐⭐⭐ [首次将样本级与曲率级双重低秩结构深度集成于投影影响函数，数学转化简洁巧妙]
- 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 124M 到 70B 全参数梯队与数百万样本集，兼顾反事实重训LDS、Tail-patch及盲测LLM-as-a-judge]
- 写作质量: ⭐⭐⭐⭐⭐ [架构清晰，瓶颈与化解方案论证透彻，理论谱衰减证明与消融分析极其扎实]
- 价值: ⭐⭐⭐⭐⭐ [首次让基于梯度的影响函数TDA在70B超大模型上具备工程可用性与交互式查询潜力]