开源工具库dattri让数据溯源更高效,支持多种AI模型的训练样本影响评估。
$\texttt{dattri}$: A Library for Efficient Data Attribution
- 统一API接口,几行代码即可接入PyTorch流程
- 封装常用计算模块,如Hessian向量积等低层操作
- 内置预训练模型与评测框架,覆盖生成式AI场景
数据溯源方法旨在量化单个训练样本对人工智能模型预测的影响。随着大规模AI模型的发展,数据溯源在提升模型性能与安全性方面应用广泛。然而,尽管近期涌现出大量新方法,仍缺乏一个全面的工具库来支持方法开发、基准测试与部署。本文提出dattri,一个开源数据溯源库,具备三大创新设计:第一,提供统一且易用的API,用户仅需修改少量代码即可将不同溯源方法集成到PyTorch机器学习流程中;第二,模块化封装数据溯源中常用的底层函数,如Hessian-向量积、逆Hessian-向量积或随机投影,降低新方法研发门槛;第三,提供包含预训练模型与真实标注的综合评测框架,覆盖多种基准设置,包括生成式AI场景。我们已实现多种先进高效的数据溯源方法,可应用于大规模神经网络模型,并将持续更新。基于dattri,我们开展了跨多类方法的全面公平基准分析。源代码见https://github.com/TRAIS-Lab/dattri。
原文摘要 · Abstract (English)
Data attribution methods aim to quantify the influence of individual training samples on the prediction of artificial intelligence (AI) models. As training data plays an increasingly crucial role in the modern development of large-scale AI models, data attribution has found broad applications in improving AI performance and safety. However, despite a surge of new data attribution methods being developed recently, there lacks a comprehensive library that facilitates the development, benchmarking, and deployment of different data attribution methods. In this work, we introduce $\texttt{dattri}$, an open-source data attribution library that addresses the above needs. Specifically, $\texttt{dattri}$ highlights three novel design features. Firstly, $\texttt{dattri}$ proposes a unified and easy-to-use API, allowing users to integrate different data attribution methods into their PyTorch-based machine learning pipeline with a few lines of code changed. Secondly, $\texttt{dattri}$ modularizes low-level utility functions that are commonly used in data attribution methods, such as Hessian-vector product, inverse-Hessian-vector product or random projection, making it easier for researchers to develop new data attribution methods. Thirdly, $\texttt{dattri}$ provides a comprehensive benchmark framework with pre-trained models and ground truth annotations for a variety of benchmark settings, including generative AI settings. We have implemented a variety of state-of-the-art efficient data attribution methods that can be applied to large-scale neural network models, and will continuously update the library in the future. Using the developed $\texttt{dattri}$ library, we are able to perform a comprehensive and fair benchmark analysis across a wide range of data attribution methods. The source code of $\texttt{dattri}$ is available at https://github.com/TRAIS-Lab/dattri.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。