提出低秩影响函数方法,让大规模模型训练数据溯源更高效
LoRIF: Low-Rank Influence Functions for Scalable Training Data Attribution
- 用低秩分解存储梯度投影,减少内存和读取开销
- 通过截断SVD和Woodbury恒等式降低海森逆矩阵内存消耗
- 在700亿参数模型上实现20倍存储压缩,适合大模型可解释性研究
训练数据溯源(TDA)旨在识别哪些训练样本最影响模型预测。影响函数方法基于梯度理论,但计算梯度在大规模数据集上存在可扩展性挑战。现有主流策略如随机投影(如TRAK、LoGRA)仍面临两大瓶颈:(i) 存储和加载所有N个训练样本的投影梯度,查询延迟受I/O主导;(ii) 构造D×D的海森逆近似,需O(D²)内存。二者均随投影维度D增长,但提升D才能提高溯源质量——形成质量与可扩展性的权衡。本文提出低秩影响函数(LoRIF),利用梯度的低秩结构解决双重瓶颈。首先,存储投影梯度的秩-c因子而非完整矩阵,使每层每样本的存储与查询开销从O(D)降至O(c√D);其次,采用截断SVD结合Woodbury恒等式,在r维子空间中近似海森逆项,内存从O(D²)降至O(Dr)。在0.1B至70B参数模型、百万级样本数据集上,LoRIF相比LoGRA实现最高20倍存储压缩与查询加速,同时保持或超越其溯源质量。该方法使基于梯度的TDA在前沿规模下成为现实。
原文摘要 · Abstract (English)
Training data attribution (TDA) identifies which training examples most influenced a model's prediction. Influence function methods are a theoretically grounded family of TDA methods and exploit gradients. To overcome the scalability challenge arising from gradient computation, the most popular strategy is random projection (e.g., TRAK, LoGRA). However, this still faces two bottlenecks when scaling to large training sets and high-quality attribution: \emph{(i)} storing and loading projected per-example gradients for all $N$ training examples, where query latency is dominated by I/O; and \emph{(ii)} forming the $D \times D$ inverse Hessian approximation, which costs $O(D^2)$ memory. Both bottlenecks scale with the projection dimension $D$, yet increasing $D$ is necessary for attribution quality -- creating a quality--scalability tradeoff. We introduce \textbf{LoRIF} (\textbf{Lo}w-\textbf{R}ank \textbf{I}nfluence \textbf{F}unctions), which exploits low-rank structures of gradient to address both bottlenecks. First, we store rank-$c$ factors of projected per-example gradients rather than full matrices, reducing storage and query-time I/O from $O(D)$ to $O(c\sqrt{D})$ per layer per sample. Second, we use truncated SVD with the Woodbury identity to approximate the inverse Hessian term in an $r$-dimensional subspace, reducing memory from $O(D^2)$ to $O(Dr)$. On models from 0.1B to 70B parameters trained on datasets with millions of examples, LoRIF achieves up to 20$\times$ storage reduction and query-time speedup compared to LoGRA, while matching or exceeding its attribution quality. LoRIF makes gradient-based TDA practical at frontier scale.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。