Efficient Ensembles Improve Training Data Attribution

TL;DR

This work targets the efficiency gap in training data attribution (TDA) for large-scale models by moving beyond costly fully independent ensembles of gradient-based methods. It introduces two efficient ensemble strategies—Dropout Ensemble and LoRA Ensemble—that replace expensive independent retraining with cost-effective variants, achieving substantial reductions in training time, serving time, and space while preserving attribution efficacy as measured by LDS. Empirical results across datasets (MNIST-10, CIFAR-2, MAESTRO) and models (MLPs, ResNet-9, Music Transformer) demonstrate that the proposed methods outperform naive ensembles and maintain strong TDA performance, including in generative settings. These approaches advance the Pareto frontier of TDA methods, enabling scalable, accurate data attribution in large-scale and edge contexts, with potential applicability to a wide range of gradient-based TDA techniques.

Abstract

Training data attribution (TDA) methods aim to quantify the influence of individual training data points on the model predictions, with broad applications in data-centric AI, such as mislabel detection, data selection, and copyright compensation. However, existing methods in this field, which can be categorized as retraining-based and gradient-based, have struggled with the trade-off between computational efficiency and attribution efficacy. Retraining-based methods can accurately attribute complex non-convex models but are computationally prohibitive, while gradient-based methods are efficient but often fail for non-convex models. Recent research has shown that augmenting gradient-based methods with ensembles of multiple independently trained models can achieve significantly better attribution efficacy. However, this approach remains impractical for very large-scale applications. In this work, we discover that expensive, fully independent training is unnecessary for ensembling the gradient-based methods, and we propose two efficient ensemble strategies, DROPOUT ENSEMBLE and LORA ENSEMBLE, alternative to naive independent ensemble. These strategies significantly reduce training time (up to 80%), serving time (up to 60%), and space cost (up to 80%) while maintaining similar attribution efficacy to the naive independent ensemble. Our extensive experimental results demonstrate that the proposed strategies are effective across multiple TDA methods on diverse datasets and models, including generative settings, significantly advancing the Pareto frontier of TDA methods with better computational efficiency and attribution efficacy.

Figures (9)

• Figure 1: Dropout Ensemble consists of two steps: (1) train $I$ models ($\{\Theta^{(i)}\}_{i=1}^I$) independently; (2) get $D$ dropout-masked models ($\{f^{(d)}_{\Theta^{(i)}}\}_{d=1}^D$) for each $i = 1,\ldots,I$.
• Figure 2: LoRA Ensemble consists of two steps: (1) train $I$ models ($\{\Theta^{(i)}\}_{i=1}^I$) independently; (2) get $L$ LoRA fine-tuned models ($\{\Theta_{\text{LoRA}}^{(i, l)}\}_{l=1}^L$) for each $i = 1,\ldots,I$.
• Figure 3: The LDS of naive independent ensemble and Dropout Ensemble with different numbers of dropout-masked passes ($D$) and independently trained models ($I$). We apply the ensemble methods to the TDA method, TRAK. There are four experiment settings: MLP classifiers trained on MNIST and CIFAR-2 (top row); ResNet9 trained on CIFAR-2 (bottom-left); and Music Transformer trained on MAESTRO (bottom-right). The $x$-axis indicates the training time cost measured by the number of independently trained models ($I$). The $y$-axis indicates the attribution efficacy measured by LDS.
• Figure 4: The LDS of naive independent ensemble and Dropout Ensemble on more TDA methods, IF, Grad-Dot, and Grad-Cos. The experiments are performed on MLP classifiers trained on MNIST. The plot setup is similar as Figure \ref{['fig:droptrak-lds']}.
• Figure 5: The LDS of Dropout Ensemble and its variant Dropout Ensemble (forward-only) when applied to TRAK on different dataset and model settings. The $x$-axis indicates the serving time cost measured by the running time on a single A40 GPU. The $y$-axis indicates the attribution efficacy measured by LDS. The points in the plot correspond to different numbers of dropout masked models ($D$). The number of independently trained models ($I$) is fixed to 5.

Theorems & Definitions (1)

• Definition D.1: Linear datamodeling score