GRASP: 面向可扩展预训练数据归因的几何感知残差对齐

# GRASP：面向可扩展预训练数据归因的几何感知残差对齐
来源：https://arxiv.org/html/2606.06892

###### 摘要

可扩展的数据归因方法通常为单个训练样本分配孤立的效用分数。这种普遍的加性假设从根本上无法捕捉关键的子集动态，包括数据冗余和互补覆盖。在本文中，我们将归因重新定义为子集级别的反事实效用预测，并引入GRASP，一种交互感知的替代模型。基于理论平滑性下界，GRASP通过二次几何惩罚显式建模子集交互。为了实现预训练规模的高效性，同时避免依赖隐藏的Oracle调优，我们将低维特征草图与严格有下限置信区间选择协议相结合。大量的子集重训练评估表明，GRASP显著优于现有的可扩展基线方法。在反事实子集保真度方面，它将任务级秩相关性提高了一倍以上，同时将前期工件构建成本降低了近一个数量级。下游诊断进一步表明，这种评分机制可迁移至语言模型策展和跨领域视觉选择，为优化大规模预训练语料库奠定了坚实基础。

---

**GRASP: Geometry-aware Residual Alignment for Scalable Pretraining Data Attribution**

Yue Min¹, Ruining Chen², Yujun Li¹  
¹Wizard Quant, ²中国科学技术大学  
通讯作者：[email protected] (https://arxiv.org/html/2606.06892v1/mailto:[email protected]), [email protected] (https://arxiv.org/html/2606.06892v1/mailto:[email protected])

## 1 引言

大语言模型的能力和失败模式根本上由其预训练语料库决定（Raffel等，2020 (https://arxiv.org/html/2606.06892#bib.bib11)；Penedo等，2024 (https://arxiv.org/html/2606.06892#bib.bib20)）。因此，数据策展已成为模型性能的核心驱动力。这一过程包括提取、过滤、去重和领域混合。然而，这些结构性决策通常依赖于粗略的启发式方法或成本高昂的重训练试验。数据归因通过将模型行为与数据管道连接起来，填补了这一方法论空白。这种连接使得在固定计算预算下，能够进行有针对性的审计、污染分析和有原则的数据选择。

![图1：GRASP架构概览。GRASP通过结合加性相关性核心（底部）和交互感知几何惩罚（顶部）来预测数据干预的反事实效用。通过将LLM隐藏状态和残差处理为紧凑的直和草图（φi）和稳定的相关性分数，该方法高效地捕捉子集冗余和互补覆盖。然后将标准化组件线性聚合，以评估整个子集，无需组合重训练。](https://arxiv.org/html/2606.06892/something.png)

尽管有这些前景，但将归因扩展到指导实际策展时，揭示了当前范式的一个根本缺陷。现有的可扩展方法，包括影响函数（Koh和Liang，2017 (https://arxiv.org/html/2606.06892#bib.bib1)）、TracIn（Pruthi等，2020 (https://arxiv.org/html/2606.06892#bib.bib13)）、TRAK（Park等，2023 (https://arxiv.org/html/2606.06892#bib.bib16)）、DataInf（Kwon等，2024 (https://arxiv.org/html/2606.06892#bib.bib18)）和LESS（Xia等，2024 (https://arxiv.org/html/2606.06892#bib.bib19)），都通过为每个训练样本分配一个标量相关性分数来运作。然而，实际的数据策展涉及保留或移除连贯的数据切片，而不是孤立的句子。当从业者使用逐点分数来评估一个数据子集时，他们隐含地将该子集的效用视为个体贡献的简单总和。这种加性假设在数学上具有限制性。在包含近似重复和重叠领域的现实语料库中，训练样本会相互作用。选择高度对齐但相同的样本会产生递减的回报，而多样化的样本则提供互补的覆盖。尽管基于联盟的估值方法，如Data Shapley（Ghorbani和Zou，2019 (https://arxiv.org/html/2606.06892#bib.bib6)）和数据建模（Ilyas等，2022 (https://arxiv.org/html/2606.06892#bib.bib23)；Wang等，2025 (https://arxiv.org/html/2606.06892#bib.bib12)），形式上捕捉了这些子集依赖关系，但精确的联盟评估对于常规的预训练管道来说计算成本高昂。这造成了一个关键的方法论空白：可扩展的归因需要一个可复用的替代模型，该模型显式建模训练-训练交互，而无需组合重训练。因此，归因方法的评估必须从单个样本排序结构性转向*子集级别的反事实效用预测*。一个严格的归因方法必须根据重训练后的经验效用准确地对候选数据干预进行排序。

为了应对这一挑战，我们提出了GRASP，一种用于可扩展预训练数据归因的几何感知残差对齐替代模型。我们的方法基于加权子集干预下效用改进的一步平滑性下界。这一理论基础自然产生了一个线性目标相关性项和一个二次几何惩罚项。展开这个二次表达式正式捕捉了加性评分所忽略的相互依赖关系。为了预训练规模部署，GRASP用低维特征草图替换了精确的高维更新方向。为了保证实验纪律，所有组件权重和保留/删除模式都在测试评估之前在开发环境上选择，从而将可扩展子集预测与隐藏的每任务调优分开。

总之，我们的主要贡献是：

- • 我们将预训练数据归因形式化为*子集级别的反事实效用预测*，直接将目标与实际语料库干预对齐，以克服加性逐点评分的局限性。
- • 我们推导出GRASP，一种基于一步平滑性下界的交互感知替代模型。其二次几何惩罚显式建模数据冗余和互补覆盖，无需组合重训练。
- • 我们实现了具有紧凑特征草图的GRASP，以实现预训练规模的高效性。它将工件构建时间减少了近一个数量级，在5秒内对100,000个子集进行评分，并将任务级LDS秩相关性相比现有可扩展基线提高了一倍以上。

## 2 问题设定

我们研究*子集级别的反事实效用预测*：给定一个预训练语料库、一个评估目标和候选数据干预，目标是根据重训练后的实现效用对干预措施进行排序。

### 2.1 预训练数据归因

令 D={1,...,n} 索引一个预训练语料库 {zi}i∈D。一个数据干预保留了一个子集 S⊆D，在重训练后产生参数 θS。对于一个目标 e，我们将效用定义为 Ue(θ)=−Le(θ)，其中 Le 是目标损失。因此真实的子集效用为：

Ue(S)=Ue(θS)=−Le(θS). (1)

预训练归因旨在为这个集合函数 S↦Ue(S) 提供一个可扩展的替代函数 Ûe:2^D→R。虽然传统归因为每个样本 i 分配一个标量分数 ai,e∈R（Koh和Liang，2017 (https://arxiv.org/html/2606.06892#bib.bib1)；Park等，2023 (https://arxiv.org/html/2606.06892#bib.bib16)；Xia等，2024 (https://arxiv.org/html/2606.06892#bib.bib19)），但一个子集级别的替代模型预测整个干预 S 的效用。

### 2.2 通过LDS的反事实保真度

遵循数据建模（Ilyas等，2022 (https://arxiv.org/html/2606.06892#bib.bib23)；Park等，2023 (https://arxiv.org/html/2606.06892#bib.bib16)），我们使用线性数据建模分数（LDS）来评估反事实保真度。给定一个由 B 个采样干预组成的族 B={Sb}_{b=1}^B，LDS 测量目标 e 的预测效用与实现效用之间的秩相关性：

ρe=Spearman_b(Ûe(Sb), Ue(Sb)). (2)

当依赖逐点分数时，子集预测器退化为一个加性基线：

Ûe^add(S)=∑_{i∈S} ai,e. (3)

这种模块化形式假设样本之间独立，因此对于包含重复和领域簇的语料库显得脆弱。因此，我们寻求一个包含结构性子集交互项的可扩展替代模型，以改善 ρe。

## 3 方法

GRASP 预测候选数据干预的效用，而无需对每次查询进行重训练。它首先计算训练-目标对的残差对齐相关性分数，然后通过添加几何感知的训练-训练交互对整个子集进行评分，最后在开发环境上选择一个固定协议。图1（https://arxiv.org/html/2606.06892#S1.F1）总结了这一流程。

### 3.1 GRASP相关性分数

为了评估数据干预的效用，GRASP 首先需要一个可扩展的度量来测量训练样本与目标样本之间的对齐程度。我们并非计算精确但难以处理的影响函数 (H^{-1}∇ℓ)，而是从因果语言模型的最后一层几何中推导出一个高效的代理。令 hi 和 ξi=pi−yi 表示给定 token 的隐藏状态和下一 token 残差，其中局部梯度具有精确的外积形式 ∇_W ℓ_tok = h ξ^⊤。对于包含多个 token 的序列，精确梯度需要对所有这些外积求和。为了保持线性存储复杂度，我们将序列池化为一个单一的摘要表示，并采用秩-1代理：∇_W ℓ_i ≈ h_i ξ_i^⊤。尽管这一简化在形式上丢弃了序列内的协方差，但它有效地捕捉了序列的宏观对齐。在 Kronecker 分解的广义 Gauss-Newton (GGN) 近似下（Martens和Grosse，2015 (https://arxiv.org/html/2606.06892#bib.bib34)），逆经验 Fisher 或 GGN 矩阵近似为一个 Kronecker 乘积 F^{-1}≈Σ_h^{-1}⊗D_ξ。因此，预条件梯度的内积优雅地分解为独立的隐藏和残差匹配：

⟨∇_W ℓ_i, F^{-1}∇_W ℓ_e⟩≈(h_i^⊤ Σ_h^{-1} h_e)(ξ_i^⊤ D_ξ ξ_e). (4)

受此分解启发，对于一个具有表示 h_e, ξ_e 的目标评估单元 e，我们使用全局训练统计对隐藏状态进行白化：

h̃_i = Σ_h^{-1/2}(h_i - h̄), h̃_e = Σ_h^{-1/2}(h_e - h̄), (5)

其中 Σ_h^{-1/2} 包含了 Tikhonov 正则化。然后定义分解的相似度为：

H_{i,e} = ⟨h̃_i, h̃_e⟩, (6)
R_{i,e} = ξ_i^⊤ D_ξ ξ_e. (7)

这里，D_ξ = diag(γ+ε)^{-α}（α∈{1,1/2}）作为残差的稳定化逆协方差预条件子，其中 γ 通过稀疏的 top-k 残差草图估计，以最小化内存占用。严格遵循局部影响需要直接的双线性乘积 H_{i,e} R_{i,e}。然而，在实践中，隐藏状态内积表现出极端重尾分布，这通常会破坏子集评估的稳定性。为了使公式更加稳健，我们明确偏离精确的泰勒展开，引入一个单调阻尼函数：

a_i,e^{GRASP}=R_{i,e} · sgn(H_{i,e}) log(1+|H_{i,e}|). (8)

这个奇函数、严格递增的对数变换保留了表示匹配的序数排序和方向支持（符号），同时安全地截断了几何异常值的不成比例杠杆。最后，一个加权干预子集 T 的一阶相关性通过加性方式计算：

A_e^{GRASP}(T)=∑_{i∈T} w_i a_i,e^{GRASP}. (9)

GRASP 采用 A_e^{GRASP}(T) 作为其基础线性预测器，随后通过几何感知的交互惩罚进行增强，以建模子集冗余。

### 3.2 交互感知的子集效用

令 g_e = ∇_θ L_e(θ) 为参考检查点处的目标梯度。如附录A（https://arxiv.org/html/2606.06892#A1）中推导，加权训练目标的隐式微分产生了示例 i 的影响引导更新 u_i = H^{-1}∇_θ ℓ_i(θ)，直至一个共享尺度。一个标准的标量归因方法通过这种一阶对齐加性地给子集 S 打分：

Û_e^add(S)=∑_{i∈S} a_{i,e}, a_{i,e}≈⟨g_e, u_i⟩. (10)

然而，这种模块化形式假设严格独立，从根本上无法捕捉子集内的冗余或互补性。为了正式建模子集动态，我们必须检查组合更新的几何结构，这自然会产生一个曲率惩罚。

**定理 1（一步交互下界）**。固定参数 θ 和目标 e。令 g_e = ∇_θ L_e(θ)，步长 η≥0，更新方向 u_i∈R^{dim(θ)}。对于一个加权集 T，权重为 w_i，定义组合更新 D_T = ∑_{i∈T} w_i u_i。如果 L_e 在 θ 和 θ-η D_T 之间（关于欧几里得范数）是 β_e-光滑的，则目标效用改进 Δ_e(T)=U_e(θ-η D_T)−U_e(θ) 满足：

Δ_e(T) ≥ η ∑_{i∈T} w_i ⟨g_e, u_i⟩ - (β_e η^2)/2 ‖∑_{i∈T} w_i u_i‖². (11)

证明（附录A (https://arxiv.org/html/2606.06892#A1)）限定了完整重训练的反事实效用。关键的是，展开二次惩罚项将成对交互分离出来：

‖∑_{i∈T} w_i u_i‖² = ∑_{i∈T} w_i² ‖u_i‖² + 2 ∑_{i<j} w_i w_j ⟨u_i, u_j⟩.

因此，具有正对齐更新方向（⟨u_i, u_j⟩ > 0）的示例会承担严格为正的非对角惩罚。在离散 c

（注：由于原文在`离散 c`处截断，后续内容待补充。目前按给定文本翻译处理。）