[论文] 循环 Transformer 的步级数据归因

Source: arXiv - 2602.10097v1

概览

论文 “Step‑resolved data attribution for looped transformers” 解决了现代可解释性工具中的一个盲点：它们可以告诉你 哪些 训练样本重要，但无法说明这些样本在模型的循环推理过程中的 何时 起作用。通过展开循环（递归）Transformer 的计算——例如，GPT 风格的模型对同一块进行 τ 次迭代——作者们提出了一种细粒度的影响估计器，能够揭示训练样本发挥作用的确切迭代步骤。

关键贡献

• Step‑Decomposed Influence (SDI)： 对 TracIn influence 估计器的全新扩展，生成长度为 τ 的轨迹，为每一次循环迭代分配单独的影响分数。
• 基于 TensorSketch 的实现： 使 SDI 能在 transformer 规模下运行，且无需显式生成每个样本的梯度，大幅降低内存和计算开销。
• 在循环 GPT‑style 模型上的实证验证： 表明 SDI 在匹配完整梯度基线的同时（误差 ≤ 5 %），能够扩展至数十亿参数。
• 广泛的适用性： 展示了 SDI 可用于数据调试、课程设计以及探究算法任务（如排序、奇偶性）中的潜在推理步骤。

方法论

1. 展开递归图： 共享的 transformer 块被应用 τ 次，产生 τ 个不同的“步骤”。
2. 分解 TracIn： 传统 TracIn 计算测试样本梯度与所有训练步骤中训练样本梯度之和的点积。SDI 用 τ 个部分和的向量代替求和，得到一个影响 轨迹 ([I_1, I_2, …, I_τ])。
3. TensorSketch 压缩： 为了避免存储每个样本的梯度（成本过高），作者使用 TensorSketch 算法将梯度哈希投影到低维草图中。草图是可加的，因此可以通过在草图空间的简单内积恢复逐步影响。
4. 评估流程：
   • 在合成算法数据集（例如复制、加法、排序）上训练循环 transformer 模型。
   • 对一组测试查询和训练样本池计算 SDI。
   • 与完整梯度基线（精确的每样本梯度）和经典 TracIn 进行比较。

结果与发现

• 逐步洞察： 对于排序任务，影响最高的峰值恰好出现在模型执行“比较‑交换”操作的迭代阶段，证实了 SDI 能准确定位推理阶段。
• 可扩展性： 在单个 8 GPU 节点上对 1.3 B 参数的循环 GPT‑style 模型（τ = 12）进行实验，而完整梯度基线则需要 > 200 GB 的 GPU 内存。

Practical Implications

• Debugging training data: 开发者现在可以问“哪些训练示例导致模型在这个特定查询上失败，以及在什么推理步骤？” 这对于发现仅影响后期推理阶段的错误标记或对抗性示例非常宝贵。
• Curriculum learning: 通过观察逐步影响概况，可以安排先教授早期推理步骤的训练示例，然后逐步引入对后期重要的示例，可能加速收敛。
• Model auditing & compliance: 监管框架日益要求模型决策的可追溯性。SDI 提供了将决策链接回具体数据点以及确切内部计算步骤的具体审计轨迹。
• Improved probing tools: 构建用于潜在推理的探测分类器的研究者现在可以在影响峰值出现的步骤上条件化探测，从而产生更干净、更易解释的信号。

限制与未来工作

• 假设固定循环计数 (τ)： SDI 的轨迹长度等于展开的步数；需要自适应决定何时停止的模型（例如 early‑exit transformers）将需要动态处理。
• Sketch 近似误差： 虽然在报告的实验中可以忽略不计，但 TensorSketch 引入的偏差在极深循环（τ ≫ 20）或梯度高度稀疏时可能变得显著。
• 聚焦于合成算法任务： 未对真实世界的 NLP 基准（如代码生成、对话）进行评估；将 SDI 扩展到这些领域是一个未解的方向。
• 与现有工具的集成： 当前实现是研究原型；将 SDI 打包为流行库（PyTorch‑Lightning、Hugging Face）的插件将降低采纳门槛。

底线： 步分解影响（Step‑Decomposed Influence）为循环 Transformer 的内部工作打开了新视角，使开发者能够追踪训练样本何时起作用。凭借其可扩展的基于 Sketch 的引擎，它弥合了学术可解释性研究与实际生产级模型调试之间的鸿沟。

作者

• Georgios Kaissis
• David Mildenberger
• Juan Felipe Gomez
• Martin J. Menten
• Eleni Triantafillou

论文信息

• arXiv ID: 2602.10097v1
• 分类: cs.LG, cs.AI
• 出版日期: 2026年2月10日
• PDF: Download PDF