[论文解读] TRAK: Attributing Model Behavior at Scale

Sung‐Min Park, Kristian Georgiev+3 more|arXiv (Cornell University)|Mar 24, 2023

COVID-19 diagnosis using AI被引用 10

一句话总结

TRAK 是一种数据归因方法，具有可扩展性并且对大规模、可微分模型有效，仅需要少量训练模型即可实现强归因。

ABSTRACT

The goal of data attribution is to trace model predictions back to training data. Despite a long line of work towards this goal, existing approaches to data attribution tend to force users to choose between computational tractability and efficacy. That is, computationally tractable methods can struggle with accurately attributing model predictions in non-convex settings (e.g., in the context of deep neural networks), while methods that are effective in such regimes require training thousands of models, which makes them impractical for large models or datasets. In this work, we introduce TRAK (Tracing with the Randomly-projected After Kernel), a data attribution method that is both effective and computationally tractable for large-scale, differentiable models. In particular, by leveraging only a handful of trained models, TRAK can match the performance of attribution methods that require training thousands of models. We demonstrate the utility of TRAK across various modalities and scales: image classifiers trained on ImageNet, vision-language models (CLIP), and language models (BERT and mT5). We provide code for using TRAK (and reproducing our work) at https://github.com/MadryLab/trak .

研究动机与目标

• 阐明将模型预测追溯到训练数据的重要性。
• 识别现有数据归因方法中的计算权衡。
• 介绍 TRAK 作为适用于大规模非凸模型的可扩展、有效归因方法。
• 在多模态和多种体系结构下展示 TRAK 的有效性。

提出的方法

• 引入 TRAK（Tracing with the Randomly-projected After Kernel）作为数据归因方法。
• 仅利用少量训练模型来执行归因。
• 目标与需要训练成千上万模型的归因方法的性能相匹配。
• 将 TRAK 应用于跨多模态的大规模、可微分模型。
• 提供用于使用 TRAK 和复现结果的代码。

实验结果

研究问题

• RQ1 在大规模非凸环境下，数据归因是否既可扩展又有效？
• RQ2 TRAK 与需要大量训练模型的现有归因方法相比如何？
• RQ3 TRAK 是否在不同模态和模型家族（视觉、视觉-语言、语言）中都有效？

主要发现

• TRAK 在比基于采样的基线需要更少训练模型的情况下实现了有效的数据归因。
• 在实验中，TRAK 与训练成千上万模型的方法的性能相当或接近。
• TRAK 展示了在图像分类器（ImageNet）、视觉-语言模型（CLIP）和语言模型（BERT、mT5）中的实用性。
• 该工作提供了可复现结果并将 TRAK 应用于其他模型的代码。