# Efficient Ensembles Improve Training Data Attribution

Junwei Deng
University of Illinois Urbana-Champaign
<EMAIL_ADDRESS>
&Ting-Wei Li11footnotemark: 1
University of Illinois Urbana-Champaign
<EMAIL_ADDRESS>
&Shichang Zhang
University of California, Los Angeles
<EMAIL_ADDRESS>
&Jiaqi Ma
University of Illinois Urbana-Champaign
<EMAIL_ADDRESS>
These authors contributed equally to this work.

###### 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.

_Keywords_ training data attribution, data-centric AI, influence function

##  1 Introduction

Training data plays an increasingly crucial role in modern artificial
intelligence (AI) models (Kaplan et al., 2020). Consequently, data-centric AI
emerges as a vital paradigm, emphasizing the collection, curation, and
understanding of training data. _Training Data Attribution_ (TDA) is a family
of methods that assess the influence of each training sample on a model’s
output. Numerous TDA methods have been developed and applied to a wide range
of data-centric AI applications, such as mislabel detection (Koh and Liang,
2017), data selection (Engstrom et al., 2024), and copyright compensation
(Deng and Ma, 2023), thereby gaining increasing popularity.

However, accurately attributing the training data influence on very large-
scale AI applications remains an open challenge. Existing TDA methods can be
generally categorized into two groups: _retraining-based methods_ and
_gradient-based methods_ (Hammoudeh and Lowd, 2024). Retraining-based methods
involve the systematic retraining of the model with and without specific
training samples to observe changes in the model’s output (Ghorbani and Zou,
2019; Jia et al., 2019; Feldman and Zhang, 2020; Ilyas et al., 2022; Wang and
Jia, 2023). Such methods often require thousands of model retraining,
sometimes even growing with the size of the training data, to achieve
satisfactory performance, which makes them infeasible for moderately large
models. Gradient-based methods, on the other hand, estimate the influence by
tracking the gradients of data samples (Koh and Liang, 2017; Yeh et al., 2018;
Pruthi et al., 2020). These methods are typically computationally more
efficient as they do not require the training of multiple models. But
empirically they can be brittle in handling complex non-convex models (Basu et
al., 2020; Bae et al., 2022) and be sensitive to the randomness associated
with model initialization and training dynamics (Søgaard et al., 2021).

Recent studies have shown that gradient-based TDA methods can be significantly
improved by ensembling tens of models independently trained with different
random seeds, leading to state-of-the-art attribution efficacy on modern
neural network models (Søgaard et al., 2021; Park et al., 2023). Aggregating
these independently trained models helps to mitigate the randomness introduced
by training dynamics, such as random initializations and stochastic
optimization. However, despite its impressive performance, scaling such a
naive independent ensemble approach further for very large models remains
challenging due to the significant computational costs associated with
training multiple models.

In this work, we hypothesize that training models independently is unnecessary
for the purpose of TDA ensemble, and we propose two efficient ensemble
strategies alternative to the naive independent ensemble approach. Our first
strategy, Dropout Ensemble, is motivated by dropout (Srivastava et al., 2014),
a common deep learning module initially designed to efficiently approximate
ensemble. Dropout Ensemble reduces the computational costs by replacing the
independently trained models in the naive ensemble approach with multiple
dropout-masked models using the same original model. Our second strategy, LoRA
Ensemble, is motivated by an efficient fine-tuning technique, LoRA (Hu et al.,
2021). Similar to Dropout Ensemble, LoRA Ensemble reduces the computational
costs by replacing the independently trained models with LoRA fine-tuned
models from the same original model, which is particularly suitable for
generative Transformer models (Vaswani et al., 2017).

We evaluate the proposed Dropout Ensemble and LoRA Ensemble with extensive
experiments. Our experiments span various datasets (MNIST (LeCun et al.,
1998), CIFAR (Krizhevsky et al., 2009), and MAESTRO (Hawthorne et al., 2018)),
model architectures (Multi-Layer Perceptrons (MLP), Residual Neural Networks
(ResNet) (He et al., 2016), and Transformers (Vaswani et al., 2017)), and TDA
methods (TRAK (Park et al., 2023), Influence Function (Koh and Liang, 2017),
and Grad-Dot/Grad-Cos (Charpiat et al., 2019)). Compared to the naive
independent ensemble approach, Dropout Ensemble and LoRA Ensemble can
significantly reduce the training time, serving time, and space costs by
respectively up to 80%, 60%, and 80% while maintaining similar TDA efficacy.

We summarize the contributions of this work as follows:

  * •

We demonstrate that fully independent training is not a strict requirement for
effective TDA with ensembles, which opens up a promising direction for
developing more efficient and effective TDA methods.

  * •

We draw intriguing connections among dropout, LoRA fine-tuning, and ensemble,
which leads to two novel efficient ensemble strategies for TDA: Dropout
Ensemble and LoRA Ensemble.

  * •

Our proposed Dropout Ensemble and LoRA Ensemble achieve significant efficiency
improvement across a diverse range of machine learning models, datasets, and
TDA methods, advancing the state-of-the-art of TDA.

##  2 Related Work

TDA methods quantify the influence of each training sample on the model
predictions by assigning a TDA score to the sample. These methods can be
categorized into retraining-based ones and gradient-based ones (Hammoudeh and
Lowd, 2024), and our work focuses on the latter.

Retraining-based TDA methods.  Retraining-based methods compute TDA scores by
systematically retraining the model with and without specific training samples
to quantify their influence on the model’s output. These methods are
computationally expensive due to the requirement of the large amount of model
retrainings. For example, Leave-One-Out (LOO) influence tuk (1958) measures
the prediction difference between the model trained on the full training
dataset and models trained on subsets with only one specific sample dropped.
Data Shapley (Ghorbani and Zou, 2019; Jia et al., 2019), Beta-Shapley (Kwon
and Zou, 2022), and Data Banzhaf (Wang and Jia, 2023) extends the LOO idea to
consider data interactions for more equitable TDA scores, but they require
retraining models on all possible subsets of the training dataset. Similarly,
DataModels (Ilyas et al., 2022) tries to learn the model predictions when the
model is trained on each subset of the training dataset, which requires a
nontrivial amount of model retraining. While, in practice, sampling or
approximation will be used to reduce the number of model retrainings, these
methods typically still require thousands of or more retrainings to achieve
satisfactory attribution efficacy. In summary, the high demand for retraining
makes these TDA methods computationally inefficient and limits their practical
applicability to even moderately large models.

Gradient-based TDA methods.  Gradient-based methods are another group of TDA
methods that usually provide closed-form TDA scores using gradients. Since the
seminal work of influence function by Koh and Liang (2017), gradient-based
methods have become increasingly popular due to their scalability. The
influence function (Koh and Liang, 2017) and subsequent studies (Guo et al.,
2021; Barshan et al., 2020; Schioppa et al., 2022; Kwon et al., 2023) obtain
TDA scores by approximating the effect of upweighting a training sample on the
loss function. Moreover, Representer Point Selection (Yeh et al., 2018)
decomposes the pre-activation of a neural network as a linear combination of
training samples. TracIn (Pruthi et al., 2020) traces the loss changes on the
test points during the training process. TRAK (Park et al., 2023) uses the
neural tangent kernel with random projection to assess influence. These
gradient-based methods significantly reduced the computational cost compared
to retraining-based methods. The downside is that they typically rely on the
convexity assumption and Taylor approximation to calculate TDA scores. These
requirements lead to performance degradation on non-convex neural networks and
sensitivity to the randomness inherent in model initialization and training.

Ensembling for gradient-based TDA methods.  Recent studies have shown the
effectiveness of ensembling for improving TDA scores computed with gradient-
based methods Søgaard et al. (2021); Park et al. (2023), which mitigates their
typical issues associated with non-convexity and sensitivity to randomness.
Ensembling normally applies the TDA method to many independently trained
models. Either averaging the final TDA scores Søgaard et al. (2021) or
aggregating some intermediate terms for score calculation Park et al. (2023).
Besides independently trained models, Park et al. (2023) suggest that model
checkpoints at different stages of a single training can be used for
ensembling as well. All these ensembling methods, though effective, also
require a non-trivial amount of ensembles to perform well. Empirical studies
show that their TDA performance suffers significantly when ensemble size is
limited (Park et al., 2023). Consequently, the ensemble size, and thus the
cost associated with each ensemble, poses a significant barrier to the
effective use of ensembling in current TDA methods.

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$.

##  3 Method

Following our hypothesis that independently trained models are unnecessary for
ensembling TDA methods, we propose efficient ensemble strategies alternative
to the naive independent ensemble.

###  3.1 Preliminaries

We start by formalizing the TDA problem and the naive independent ensemble
method for TDA.

The TDA problem.  We have a training set
$\mathcal{S}=\\{x_{1},\ldots,x_{n}\\}$ where each
$x_{i}\in\mathcal{X}_{\text{train}}$, a test set
$\mathcal{T}=\\{x_{1},\ldots,x_{m}\\}$ where each
$x_{i}\in\mathcal{X}_{\text{test}}$, and a trained model output function
$f_{\Theta}$ that is parameterized by $\Theta$. Here
$\mathcal{X}_{\text{train}}$ and $\mathcal{X}_{\text{test}}$ are the space of
the training and test data respectively. We consider both supervised learning
and generative modeling settings. For the supervised learning setting,
$\mathcal{X}_{\text{train}}$ and $\mathcal{X}_{\text{test}}$ are typically
identical, and each element of them corresponds to a data point with both the
feature and label. For the generative modeling setting,
$\mathcal{X}_{\text{train}}$ refers to the space of training data (e.g., text
sequence segments for autoregressive language models) while
$\mathcal{X}_{\text{test}}$ refers to the space of model generation. For the
model output function $f_{\Theta}$, typically we will use the model learned
from the training set $\mathcal{S}$ (i.e., the model parameters will be chosen
as
$\Theta^{*}=\operatornamewithlimits{argmin}_{\Theta}\sum_{x_{i}\in\mathcal{S}}\mathcal{L}(x_{i};f_{\Theta})$
for some loss function $\mathcal{L}$). However, this is not always the case in
practice (Pruthi et al., 2020). For a training set $\mathcal{S}$ and any test
sample $x\in\mathcal{T}$, a TDA method $\tau$ derives the TDA scores
$\tau(x,\mathcal{S};f_{\Theta})\in\mathbb{R}^{n}$ to quantify the influence of
each training data point in $\mathcal{S}$ on the model learned from
$\mathcal{S}$. More specifically, the $i$-th element,
$\tau(x,\mathcal{S};f_{\Theta})_{i}\in\mathbb{R}$, is a real-valued score
indicating the importance of $x_{i}$ on the model output on the test sample
$x$. In the supervised classification setting, the “model output on $x$”
usually refers to the loss, (log-)likelihood, or logit for the model
predicting the correct class of $x$ (Koh and Liang, 2017; Park et al., 2023);
in the generative setting, it typically refers to the (log-)likelihood for the
model generating $x$ (Deng and Ma, 2023).

The naive independent ensemble method for TDA.  To mitigate the randomness led
by the training dynamics of non-convex deep learning models, the naive
independent ensemble method first trains a set of $I$ independent models,
$\\{\Theta^{(i)}\\}_{i=1}^{I}$, and then derives ensembled TDA scores
$\tau_{\text{ens}}(x,\mathcal{S};\\{f_{\Theta^{(i)}}\\}_{i=1}^{I})\in\mathbb{R}^{n}$
based on the set of models.

The ensembled TDA scores could be simply obtained by averaging over the TDA
scores for individual models, i.e.,

$\tau_{\text{ens}}(x,\mathcal{S};\\{f_{\Theta^{(i)}}\\}_{i=1}^{I})=\frac{1}{I}\sum_{i=1}^{I}\tau(x,\mathcal{S};f_{\Theta^{(i)}}).$
(1)

They could also come from more sophisticated aggregation over the individual
models. For example, the TRAK method (Park et al., 2023) works as the
following:

$\tau_{\text{TRAK}}(x,\mathcal{S};\\{f_{\Theta^{(i)}}\\}_{i=1}^{I})=\left(\frac{1}{I}\sum_{i=1}^{I}\mathbf{Q}_{f_{\Theta^{(i)}}}\right)\left(\frac{1}{I}\sum_{i=1}^{I}\phi_{f_{\Theta^{(i)}}}\left(\Phi_{f_{\Theta^{(i)}}}^{\top}\Phi_{f_{\Theta^{(i)}}}\right)^{-1}\Phi_{f_{\Theta^{(i)}}}^{\top}\right),$
(2)

where $\Theta^{(i)}$ are parameters of models independently trained on the
training data; $\mathbf{Q}_{f_{\Theta^{(i)}}}$ is a diagonal matrix with each
diagonal element corresponding to the “one minus correct-class probability” of
a training data point under model $\Theta^{(i)}$; $\phi_{f_{\Theta^{(i)}}}$ is
the vector gradient of $f_{\Theta^{(i)}}(x)$ with respect to $\Theta^{(i)}$;
and $\Phi_{f_{\Theta^{(i)}}}$ is the matrix of the vector gradients of
$f_{\Theta^{(i)}}(x_{j})$ stacked over the training samples $x_{j}\in
S$.111Some details of the exact TRAK implementation are omitted here.

###  3.2 Computational costs of TDA methods

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$.

While recent research has shown that ensembling gradient-based TDA methods can
achieve decent attribution efficacy with tens of indepedently trained models
(as opposed to hundreds or even thousands of model retraining in retraining-
based TDA methods) (Park et al., 2023), it is still impractical for very
large-scale or edge-device applications due to the increased time and space
costs. To better quantify the time and space costs associated with TDA
methods, we categorize the costs into three parts: training time cost, serving
time cost, and storage cost. We provide a detailed explanation for these costs
as follows:

  * •

Training time cost: This refers to the computational time incurred by
processing and training models to be used by TDA ensembles. This is the major
computational challenge for ensembling TDA methods.

  * •

Serving time cost: This refers to the remaining computational time to deploy
TDA ensembles after model training. Typically, ensembling gradient-based TDA
methods derives different TDA scores using multiple trained models and then
aggregates them. The serving time cost will include all computational costs
incurred during this process.

  * •

Space cost: This refers to the parameters of the trained models that are
stored by TDA ensembles. These model parameters are required to run forward
and backward passes on top of the models for TDA score calculation. The space
costs are generally proportional to the model size.

###  3.3 Dropout Ensemble

General methodology.  Dropout, a common module used in many modern deep
learning models, is initially designed as an efficient approximation of
ensemble (Srivastava et al., 2014). Motivated by this idea, we propose a
simple, efficient ensemble strategy, Dropout Ensemble, for TDA methods.
Instead of performing the TDA ensemble over a large number of independently
trained models, the proposed method utilizes multiple dropout masks on the
same model to perform the TDA ensemble.

In practice, Dropout Ensemble consists of two steps as illustrated in Figure
1. In the first step, we train $I$ independent models,
$\\{\Theta^{(i)}\\}_{i=1}^{I}$, as shown by the  green arrows in Figure 1.
Next, as shown by the  blue arrows in Figure 1, for each model
$\Theta^{(i)},i\in\left\\{1,\ldots I\right\\}$, we obtain $D$ variants of the
model with different dropout masks, which are denoted as
$\\{f_{\Theta^{(i)}}^{(d)}\\}_{d=1}^{D}$. For any TDA method, Dropout Ensemble
calculates the ensembled TDA scores through
$\tau_{\text{ens}}(x,\mathcal{S};\\{f_{\Theta^{(i)}}^{(d)}\\}_{1\leq i\leq
I,1\leq d\leq D})$.

In comparison to the naive independent ensemble method, Dropout Ensemble can
significantly reduce the training time and space costs. As we will show in
Section 4.2, this method allows us to replace the expensive independently
trained models with dropout-masked models that incur no additional training
cost or model parameters.

TDA-method-specific optimization.  It is possible to further optimize the
proposed ensemble strategy when applying it to a particular TDA method. For
example, we develop a variant of Dropout Ensemble tailored for the TRAK method
to reduce the serving time cost of Dropout Ensemble. Recall that in Eq. (2),
the quantities $Q$’s only involve the forward pass of the models on the
training data, while $\phi$’s and $\Phi$’s require the backward pass on the
training data. We find that, perhaps a bit surprisingly, if we only use the
dropout-masked models to calculate $Q$’s while using the original models to
calculate $\phi$’s and $\Phi$’s, we can get similar attribution efficacy.
Concretely, directly applying Dropout Ensemble on TRAK leads to

$\tau_{\text{ens}}\left(x,\mathcal{S};\\{f_{\Theta^{(i)}}^{(d)}\\}_{\begin{subarray}{c}1\leq
i\leq I\\\ 1\leq d\leq D\end{subarray}}\right)=\left(\frac{1}{I\cdot
D}\sum_{i=1}^{I}\sum_{d=1}^{D}\mathbf{Q}_{f^{(d)}_{\Theta^{(i)}}}\right)\left(\frac{1}{I\cdot
D}\sum_{i=1}^{I}\sum_{d=1}^{D}\phi_{f^{(d)}_{\Theta^{(i)}}}\left(\Phi_{f^{(d)}_{\Theta^{(i)}}}^{\top}\Phi_{f^{(d)}_{\Theta^{(i)}}}\right)^{-1}\Phi_{f^{(d)}_{\Theta^{(i)}}}^{\top}\right),$

while in the optimized version, we have

$\tau_{\text{ens}}\left(x,\mathcal{S};\\{f_{\Theta^{(i)}}^{(d)}\\}_{\begin{subarray}{c}1\leq
i\leq I\\\ 1\leq d\leq D\end{subarray}},\\{f_{\Theta^{(i)}}\\}_{1\leq i\leq
I}\right)=\left(\frac{1}{I\cdot
D}\sum_{i=1}^{I}\sum_{d=1}^{D}\mathbf{Q}_{f^{(d)}_{\Theta^{(i)}}}\right)\left(\frac{1}{I}\sum_{i=1}^{I}\phi_{f_{\Theta^{(i)}}}\left(\Phi_{f_{\Theta^{(i)}}}^{\top}\Phi_{f_{\Theta^{(i)}}}\right)^{-1}\Phi_{f_{\Theta^{(i)}}}^{\top}\right).$

In this case, we only need to evaluate the forward pass on the dropout-masked
models and reduce the serving time cost by avoiding the backward pass. We call
this variant of our method as Dropout Ensemble (fowrad-only).

###  3.4 LoRA Ensemble

For large-scale generative models, especially Transformer models (Vaswani et
al., 2017), LoRA adapters have shown great success for efficiently fine-tuning
the models (Hu et al., 2021). Our second efficient ensemble strategy, LoRA
Ensemble, is motivated by the LoRA techniques. In particular, we propose to
replace the independently trained models in the naive ensemble method with
LoRA fine-tuned models.

Similar to Dropout Ensemble, LoRA Ensemble also consists of two steps as
illustrated in Figure 2. The first step trains $I$ independent models,
$\\{\Theta^{(i)}\\}_{i=1}^{I}$ ( green arrows in Figure 2), while the second
step further trains $L$ LoRA fine-tuned models,
$\\{\Theta_{\text{LoRA}}^{(i,l)}\\}_{l=1}^{L}$, for each $\Theta^{(i)}$,
$i\in\left\\{1,\ldots,I\right\\}$ ( brown arrows in Figure 2).

In comparison to Dropout Ensemble, LoRA Ensemble will introduce a small amount
of additional training time cost and model parameters for LoRA fine-tuning.
However, using LoRA adapters can reduce the serving time cost of TDA compared
to either dropout-masked models in Dropout Ensemble or independent models in
the naive method. This leads to a unique advantage for LoRA Ensemble if the
serving time cost is critical among the computational costs.

##  4 Experiments

In this section, we empirically evaluate the efficiency and efficacy of the
proposed Dropout Ensemble and LoRA Ensemble.

###  4.1 Experimental setup

We conduct extensive experiments on a wide range of settings.

TDA methods.  We apply the proposed methods to various gradient-based TDA
methods, including influence function (IF) (Koh and Liang, 2017), Grad-
Dot/Grad-Cos (Charpiat et al., 2019), and TRAK (Park et al., 2023). The
detailed descriptions and implementations of these algorithms are provided in
Appendix A.

Datasets and models.  We consider models with different architectures trained
on diverse datasets: (1) a three-layer MLP classifier trained on the MNIST-10
dataset (LeCun et al., 1998), (2) a three-layer MLP classifer and a ResNet-9
classifier (He et al., 2016) trained on the CIFAR-2 dataset (a two-class
subset of the CIFAR-10 dataset (Krizhevsky et al., 2009)), and (3) a Music
Transformer (Huang et al., 2018) trained on the MAESTRO dataset (Hawthorne et
al., 2018). Notably, the former two are supervised classification settings,
while the last one is a generative modeling setting. We sample 5000 training
samples and 500 test samples from MNIST-10 and CIFAR-2 datasets. For MAESTRO
dataset, we sample 5000 training samples and 178 generated samples. More
detailed setups for each dataset and model are listed in Appendix B. MNIST-10
dataset holds CC BY-SA 3.0 license. CIFAR-10 dataset holds CC-BY 4.0 license.
MAESTRO dataset holds CC BY-NC-SA 4.0 license.

Evaluation metric for TDA efficacy.  We utilize the linear datamodeling score
(LDS) (Park et al., 2023) to evaluate the efficacy of the TDA methods
augmented by different ensembling methods. Intuitively, LDS measures the rank
correlation between the TDA scores among training samples and the change of
model outputs by removing a subset of training samples and retraining the
model from scratch. A higher LDS value corresponds to a better alignment
between the TDA scores and the influence of the training samples on the model
outputs, thus better TDA quality. We refer the reader to Appendix D for the
exact definition of LDS.

Evaluation metrics for TDA efficiency.  As introduced in Section 3.2, we
measure the TDA efficiency in terms of the training time cost, serving time
cost, and space cost. For both training and serving time costs, we measure the
wall-clock time on a single A40 GPU. For the space cost, although ideally one
would measure it by memory usage, this can be challenging because memory usage
is sensitive to the specific implementation and parallelization. Therefore, we
measure the space cost by the total parameter count instead. For Dropout
Ensemble, we will also use the number of independently trained models ($I$) as
a measure of the training time cost and space cost, as both of the two costs
are proportional to $I$ in this case. More details about the measurements are
provided in Appendix E.

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 3. 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.

###  4.2 Dropout Ensemble

Improvement over the training time cost.  To illustrate the improvement over
training time cost comparing Dropout Ensemble to the naive ensemble, we first
report the results on the TRAK method across four experiment settings. As can
be seen in Figure 3, across all four experiment settings, increasing the
number of dropout-masked passes ($D$) results in significant LDS improvement
for a fixed number of independently trained models ($I$). Recall that Dropout
Ensemble does not incur any additional training time cost for a fixed $I$.
This implies that Dropout Ensemble can significantly reduce the training time
cost for achieving the same level of attribution efficacy as measured by LDS.
For example, Dropout Ensemble with $I=5$ can reach higher LDS than naive
independent ensemble with $I=25$ for MLP on MNIST and MLP/ResNet9 on CIFAR-2,
which achieves a 80% reduction on training time cost. For Music Transformer on
MAESTRO, Dropout Ensemble with $I=1$ can reach higher LDS than naive
independent ensemble with $I=10$, which achieves a 90% reduction.

We find that Dropout Ensemble works similarly well for other TDA methods. In
Figure 4, we report the results on IF, Grad-Dot, and Grad-Cos. We only
experiment on the setting of MLP classifiers trained on MNIST, as these TDA
methods do not have meaningfully good efficacy on more complex settings.
Similar to the TRAK experiments, we observe that Dropout Ensemble with $I=1$
could match the LDS of naive independent ensemble with $I=10$ for IF and Grad-
Dot, and match that with $I=5$ for Grad-Cos. Therefore, we expect that the
proposed Dropout Ensemble can be generalized to various gradient-based TDA
methods.

Improvement over the space cost.  The improvement of Dropout Ensemble over the
space cost is self-evident, which can also be measured by $I$, the number of
independently trained models. According to Figure 3, for example, applying
Dropout Ensemble on TRAK (in comparison to naive independent ensemble) can
lead to a 80% or 90% reduction in space cost for different model and dataset
settings.

Dropout Ensemble (forward-only) improves the serving time cost.  Figure 5
shows the results of Dropout Ensemble (forward-only), a variant of Dropout
Ensemble specifically tailored for TRAK. In comparison to the vanilla Dropout
Ensemble, this optimized variant can achieve comparable TDA accuracy with
around 50% less serving time.

We also report additional ablation studies and results of Dropout Ensemble in
Appendix F.

((a)) Training time cost improvement.

((b)) Serving time cost improvement.

((c)) Total time cost improvement.

((d)) Space cost improvement.

Figure 6: The LDS of naive ensemble and LoRA Ensemble with respect to
different cost measurements. Here, we apply the ensemble methods to TRAK on
Music Transformer trained on the MAESTRO dataset. The x-axis of Figure (6, 6,
6) indicates training/serving time costs (running time on a single A40), while
that of Figure 6 specifies the space cost (total parameter count). The y-axis
of all figures is the attribution efficacy measured by LDS. LoRA Ensemble has
significantly fewer costs in all aspects than naive ensemble for achieving
similar LDS.

###  4.3 LoRA Ensemble

To demonstrate the efficiency of LoRA Ensemble, we only consider TRAK Park et
al. (2023) as the main TDA method in this section since it is the most
effective (in terms of LDS) among the four gradient-based TDA methods
discussed in Section 4.2. We experiment LoRA Ensemble under the setting of
Music Transformer on MAESTRO because this setting aligns better with large-
scale real-world generative settings where LoRA adapters are more prevalent.

As shown in Figure 6, LoRA Ensemble outperforms naive independent ensemble
since it can reach similar LDS as the latter with significantly less
computational costs in terms of all types of measurements. For example, LoRA
Ensemble with ($I=1,L=3$) achieves similar LDS as naive ensemble with $I=5$,
while having reduction in training time, serving time, and total parameter
count respectively for $78.4\%$, $60.5\%$, and $79.6\%$.

##  5 Conclusion

We present Dropout Ensemble and LoRA Ensemble as efficient alternatives to the
naive independent ensemble approach for improving gradient-based TDA methods.
The proposed strategies significantly reduce training time (up to 80%),
serving time (up to 60%), and space cost (up to 80%), while maintaining
similar attribution efficacy in comparison to the naive ensemble. Empirical
results from our extensive experiments show that the proposed efficient
ensembles can remarkably advance the Pareto frontier of TDA methods with
better computational efficiency and TDA efficacy. Notably, we demonstrate the
proposed methods work well on a generative modeling setting. In the future, we
will utilize our methods in more real-world generative settings that were
blocked by high computational costs or low effectiveness of TDA methods.

## References

  * tuk (1958) Bias and confidence in not quite large samples.  _Ann. Math. Statist._ , 29:614, 1958.
  * Bae et al. (2022) Juhan Bae, Nathan Ng, Alston Lo, Marzyeh Ghassemi, and Roger B Grosse.  If influence functions are the answer, then what is the question?  _Advances in Neural Information Processing Systems_ , 35:17953–17967, 2022.
  * Barshan et al. (2020) Elnaz Barshan, Marc-Etienne Brunet, and Gintare Karolina Dziugaite.  Relatif: Identifying explanatory training samples via relative influence.  In _International Conference on Artificial Intelligence and Statistics_ , pages 1899–1909. PMLR, 2020.
  * Basu et al. (2020) Samyadeep Basu, Philip Pope, and Soheil Feizi.  Influence functions in deep learning are fragile.  _arXiv preprint arXiv:2006.14651_ , 2020.
  * Charpiat et al. (2019) Guillaume Charpiat, Nicolas Girard, Loris Felardos, and Yuliya Tarabalka.  Input similarity from the neural network perspective.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Deng and Ma (2023) Junwei Deng and Jiaqi Ma.  Computational copyright: Towards a royalty model for ai music generation platforms.  _arXiv preprint arXiv:2312.06646_ , 2023.
  * Engstrom et al. (2024) Logan Engstrom, Axel Feldmann, and Aleksander Madry.  Dsdm: Model-aware dataset selection with datamodels.  _arXiv preprint arXiv:2401.12926_ , 2024.
  * Feldman and Zhang (2020) Vitaly Feldman and Chiyuan Zhang.  What neural networks memorize and why: Discovering the long tail via influence estimation.  _Advances in Neural Information Processing Systems_ , 33:2881–2891, 2020.
  * Ghorbani and Zou (2019) Amirata Ghorbani and James Zou.  Data shapley: Equitable valuation of data for machine learning.  In _International conference on machine learning_ , pages 2242–2251. PMLR, 2019.
  * Guo et al. (2021) Han Guo, Nazneen Rajani, Peter Hase, Mohit Bansal, and Caiming Xiong.  FastIF: Scalable influence functions for efficient model interpretation and debugging.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_. Association for Computational Linguistics, 2021.  doi: 10.18653/v1/2021.emnlp-main.808.
  * Hammoudeh and Lowd (2024) Zayd Hammoudeh and Daniel Lowd.  Training data influence analysis and estimation: A survey.  _Machine Learning_ , pages 1–53, 2024.
  * Hawthorne et al. (2018) Curtis Hawthorne, Andriy Stasyuk, Adam Roberts, Ian Simon, Cheng-Zhi Anna Huang, Sander Dieleman, Erich Elsen, Jesse Engel, and Douglas Eck.  Enabling factorized piano music modeling and generation with the maestro dataset.  _arXiv preprint arXiv:1810.12247_ , 2018.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 770–778, 2016.
  * Hu et al. (2021) Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.  Lora: Low-rank adaptation of large language models.  _arXiv preprint arXiv:2106.09685_ , 2021.
  * Huang et al. (2018) Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Noam Shazeer, Curtis Hawthorne, Andrew M Dai, Matthew D Hoffman, and Douglas Eck.  Music transformer: Generating music with long-term structure.  _arXiv preprint arXiv:1809.04281_ , 2018.
  * Ilyas et al. (2022) Andrew Ilyas, Sung Min Park, Logan Engstrom, Guillaume Leclerc, and Aleksander Madry.  Datamodels: Predicting predictions from training data.  _arXiv preprint arXiv:2202.00622_ , 2022.
  * Jia et al. (2019) Ruoxi Jia, David Dao, Boxin Wang, Frances Ann Hubis, Nick Hynes, Nezihe Merve Gürel, Bo Li, Ce Zhang, Dawn Song, and Costas J Spanos.  Towards efficient data valuation based on the shapley value.  In _The 22nd International Conference on Artificial Intelligence and Statistics_ , pages 1167–1176. PMLR, 2019.
  * Kaplan et al. (2020) Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei.  Scaling laws for neural language models.  _arXiv preprint arXiv:2001.08361_ , 2020.
  * Koh and Liang (2017) Pang Wei Koh and Percy Liang.  Understanding black-box predictions via influence functions.  In _International conference on machine learning_ , pages 1885–1894. PMLR, 2017.
  * Krizhevsky et al. (2009) Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * Kwon and Zou (2022) Yongchan Kwon and James Zou.  Beta shapley: a unified and noise-reduced data valuation framework for machine learning.  In _International Conference on Artificial Intelligence and Statistics_ , pages 8780–8802. PMLR, 2022.
  * Kwon et al. (2023) Yongchan Kwon, Eric Wu, Kevin Wu, and James Zou.  Datainf: Efficiently estimating data influence in lora-tuned llms and diffusion models.  _arXiv preprint arXiv:2310.00902_ , 2023.
  * LeCun et al. (1998) Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner.  Gradient-based learning applied to document recognition.  _Proceedings of the IEEE_ , 86(11):2278–2324, 1998.
  * Park et al. (2023) Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclerc, and Aleksander Madry.  Trak: Attributing model behavior at scale.  _arXiv preprint arXiv:2303.14186_ , 2023.
  * Pruthi et al. (2020) Garima Pruthi, Frederick Liu, Satyen Kale, and Mukund Sundararajan.  Estimating training data influence by tracing gradient descent.  _Advances in Neural Information Processing Systems_ , 33:19920–19930, 2020.
  * Schioppa et al. (2022) Andrea Schioppa, Polina Zablotskaia, David Vilar, and Artem Sokolov.  Scaling up influence functions.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 36, pages 8179–8186, 2022.
  * Søgaard et al. (2021) Anders Søgaard et al.  Revisiting methods for finding influential examples.  _arXiv preprint arXiv:2111.04683_ , 2021.
  * Spearman (1904) C. Spearman.  The proof and measurement of association between two things.  _The American Journal of Psychology_ , 15(1):72–101, 1904.
  * Srivastava et al. (2014) Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov.  Dropout: a simple way to prevent neural networks from overfitting.  _The journal of machine learning research_ , 15(1):1929–1958, 2014.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Wang and Jia (2023) Jiachen T Wang and Ruoxi Jia.  Data banzhaf: A robust data valuation framework for machine learning.  In _International Conference on Artificial Intelligence and Statistics_ , pages 6388–6421. PMLR, 2023.
  * Yeh et al. (2018) Chih-Kuan Yeh, Joon Kim, Ian En-Hsu Yen, and Pradeep K Ravikumar.  Representer point selection for explaining deep neural networks.  _Advances in neural information processing systems_ , 31, 2018.

##  Appendix A TDA methods

In this section, we provide details on the TDA methods we perform efficient
ensembling on. Additionally, we also introduce the ensembling aggregation
methods, i.e., the way to aggregate
$\tau_{\text{ens}}(x,\mathcal{S};\\{f_{\Theta^{(i)}}\\}_{i=1}^{I})$, for each
method. The notation will follow Section 3.1.

#### Tracing with the Randomly-projected After Kernel (TRAK).

This is a state-of-the-art gradient-based TDA method provided by Park et al.
(2023). It natively introduces ensembling in their definition.

$\displaystyle\tau_{\text{TRAK}}(x,\mathcal{S};\\{f_{\Theta^{(i)}}\\}_{i=1}^{I})=\left(\frac{1}{I}\sum_{i=1}^{I}\mathbf{Q}_{f_{\Theta^{(i)}}}\right)\left(\frac{1}{I}\sum_{i=1}^{I}\phi_{f_{\Theta^{(i)}}}\left(\Phi_{f_{\Theta^{(i)}}}^{\top}\Phi_{f_{\Theta^{(i)}}}\right)^{-1}\Phi_{f_{\Theta^{(i)}}}^{\top}\right),$

Detailed notation description is described in Section 3.1. We use 2048 as the
random projection dimension for TRAK. For each independently trained model, we
train the model on half sampled training set following Park et al. (2023).

#### Influence functions based on the conjugate gradients (IF).

First proposed by Koh and Liang (2017), the definition of IF is

$\displaystyle\tau_{\text{IF}}(x,\mathcal{S};\\{f_{\Theta^{(i)}}\\}_{i=1}^{I})=\left[\frac{1}{I}\sum_{i=1}^{I}g_{f_{\Theta^{(i)}}}(x_{j})^{\top}H_{f_{\Theta^{(i)}}}^{-1}g_{f_{\Theta^{(i)}}}(x):x_{j}\in\mathcal{S}\right],$

where $g_{f_{\Theta^{(i)}}}(x)$ is the vector gradient of model output
$f(x;\Theta^{(i)})$ with respect to the parameters $\Theta^{(i)}$,
$H^{-1}_{f_{\Theta^{(i)}}}$ is the hessian matrix with respect to the training
set, and $g_{f_{\Theta^{(i)}}}(x_{j})^{\top}$ is the vector gradient to the
training sample. The product of the first two terms are an inverse-hessian-
vector-product problem, we implement conjugate gradients approach to solve it.

#### Grad-Dot.

Grad-Dot is proposed by Charpiat et al. (2019). We simply calculate Grad-Dot
multiple times on trained parameters
$\Theta^{(i)},i\in\left\\{1,\ldots,K\right\\}$ and take the average value.

$\displaystyle\tau_{Grad-
Dot}(x,\mathcal{S};\\{f_{\Theta^{(i)}}\\}_{i=1}^{I})=\left[\frac{1}{I}\sum_{i=1}^{I}g_{f_{\Theta^{(i)}}}(x)^{\top}g_{f_{\Theta^{(i)}}}(x_{j}):x_{j}\in\mathcal{S}\right],$

where $g_{f_{\Theta^{(i)}}}(x)^{\top}$ is the vector gradient of model output
$f(x;\Theta^{(i)})$ with respect to the parameters $\Theta^{(i)}$, and
$g_{f_{\Theta^{(i)}}}(x_{j})^{\top}$ is the gradient to the training sample.

#### Grad-Cos.

Similar to Grad-Dot,

$\displaystyle\tau_{Grad-
Cos}(z,\mathcal{S};\\{f_{\Theta^{(i)}}\\}_{i=1}^{I})=\left[\frac{1}{I}\sum_{i=1}^{I}\frac{g_{f_{\Theta^{(i)}}}(x)^{\top}}{\|g_{f_{\Theta^{(i)}}}(x)\|}\frac{g_{f_{\Theta^{(i)}}}(x_{j})}{\|g_{f_{\Theta^{(i)}}}(x_{j})\|}:x_{j}\in\mathcal{S}\right],$

where $g_{f_{\Theta^{(i)}}}(x)^{\top}$ is the vector gradient of model output
$f(x;\Theta^{(i)})$ with respect to the parameters $\Theta^{(i)}$, and
$g_{f_{\Theta^{(i)}}}(x_{j})^{\top}$ is the gradient to the training sample.

##  Appendix B Detailed experiment setup

#### MLP on MNIST-10.

For the MNIST-10 (LeCun et al., 1998) experiment, we sampled 5000 training
samples and 500 testing samples. We used a 3-layer MLP with hidden layer sizes
equal to 128 and 64 and placed dropout layers after the first two linear
layers with a rate of 0.1, which resulted in a total of about 0.11M
parameters. We employed an SGD optimizer with learning rate 0.01, momentum
0.9, and batch size 64 to train this MLP classifier for 100 epochs.

#### MLP/ResNet-9 on CIFAR-2.

For CIFAR-2 experiment, we construct the CIFAR-2 dataset by sampling from the
CIFAR-10 dataset that only include the “cat” and “dog” classes (same as the
setting in Park et al. (2023),). To incorporate a diverse set of model
architectures, we train an MLP and a CNN-based model on this dataset. We also
only consider a subset of the CIFAR-2 dataset with a training size equal to
5000 and a testing size equal to 500. Firstly, we use a 3-layer MLP with
hidden layer sizes equal to 120 and 84 and place dropout layers after the
first two linear layers with a rate of 0.1, which results in a total of about
0.38M parameters. We employ an SGD optimizer with a learning rate of 0.01,
momentum of 0.9, and batch size 64 to train this MLP classifier for 50 epochs.
Additionally, we consider a standard ResNet-9 model (He et al., 2016) with
dropout layers being placed after all convolution layers, which has roughly
4.83M trainable parameters. We train this model for 50 epochs.

#### MusicTransformer on MAESTRO.

For the MAESTRO experiment, we use the MIDI and Audio Edited for Synchronous
TRacks and Organization (MAESTRO) dataset (v2.0.0) Hawthorne et al. (2018) and
construct a Music Transformer following the original setting in (Huang et al.,
2018). Specifically, the number of layers equals to 6, the number of
independent heads equal to 8, the input feature size is 512 and the dimension
of the feedforward network is 1024. For data processing, we follow the basic
experiment setup used by (Deng and Ma, 2023). To be more specific, we define a
vocabulary set of size equal to 388, which includes “NOTE ON” and “NOTE OFF”
events for 128 different pitches, 100 “TIME SHIFT” events, and 32 “VELOCITY”
events. The raw data is pre-processed as sequences of about 90K events. Due to
computational constraints, we train a Music Transformer model on a subset of
the official training set in the MAESTRO dataset with a size of 5000. For
training, the batch size has been set to 64, and the model is trained by a
classic seq2seq loss function. We employ an Adam optimizer with a learning
rate equal to 1e-4, $\beta_{1}=0.9$ and $\beta_{2}=0.98$. We apply zero warm-
up steps since the dataset size is comparably small, and we train the model
for 20 epochs. For music event generation, we use 178 samples from the
official testing dataset as prompts to generate music with a single event,
which is used to evaluate the TDA methods.

##  Appendix C Efficient ensemble setup

###  C.1 Dropout Ensemble

Dropout is applied to the same layers as the model training stated in Appendix
B. Dropout rate is set to 0.1 for all experiment settings.

###  C.2 LoRA Ensemble

This method is only applied to MusicTransformer trained on the MAESTRO
dataset. According to Section 3.4, we first train $I$ independent models on
the full training dataset with 10 epochs using different model initialization,
and all the other remaining settings follow Appendix B. The LoRA adapters used
in this experiment all have rank $r=8$, alpha $\alpha=8$, and trainable
biases. No dropouts are applied for LoRA parameters. We augment LoRA adapters
to only the $W_{q}$ and $W_{v}$ matrices in the self-attention modules within
all the layers of the MusicTransformer. Then, we define fine-tuning datasets
for each LoRA adapter as random training data subsets with a size equal to
2500 (i.e., half of the original training dataset size). We further fine-tune
each of the aforementioned $I$ independent models using $L$ different LoRA
adapters for an additional 10 epochs on these random training data subsets,
which results in a total of $I\cdot L$ LoRA fine-tuned models.

##  Appendix D Linear datamodeling score (LDS)

Park et al. (2023) proposed the linear datamodeling score (LDS), aiming at
probing the TDA method’s ability to make counterfactual predictions based on
the attribution score derived from the learned model output function
$f_{\Theta}$ and the corresponding dataset to train $f_{\Theta}$. Because most
TDA methods are assumed to be additive222If a TDA method is additive, then it
defines an attribution score that the overall influence of a group is the sum
of the individual influence in the group., the TDA scores can be used to
predict the model output function learned from a subset of training data in a
summation form. Formally, the attribution-based output predictions of the
model output function $f_{\Theta_{\mathcal{S}^{\prime}}}$ is defined as
follows:

$g_{\tau}(x,\mathcal{S}^{\prime};\mathcal{S})\triangleq\sum_{i:x_{i}\in\mathcal{S}^{\prime}}\tau(x,\mathcal{S};f_{\Theta})_{i},$
(3)

where $\mathcal{S}$ is the training set,
$\mathcal{S}^{\prime}\subseteq\mathcal{S}$ is a subset of $\mathcal{S}$ and
$f_{\Theta_{\mathcal{S}^{\prime}}}$ is the model output function with
$\Theta_{\mathcal{S}^{\prime}}$ learned from $\mathcal{S}^{\prime}$.
Intuitively, $g_{\tau}(x,\mathcal{S}^{\prime};\mathcal{S})$ computes the
overall attribution of the subset $\mathcal{S}^{\prime}$ on example $x$, which
should be a powerful indicator of the model prediction on $x$ (i.e.,
$f_{\Theta_{\mathcal{S}^{\prime}}}(x)$) if the TDA method works well. The
linear datamodeling score (LDS) is defined to measure the predictive power of
$g_{\tau}(x,\mathcal{S}^{\prime};\mathcal{S})$ and can be formalized as
follows:

######  Definition D.1 (Linear datamodeling score).

Given a training set $\mathcal{S}$, a model output function $f_{\Theta}$, and
a corresponding TDA method $\tau$. Let
$\\{\mathcal{S}_{1},\ldots,\mathcal{S}_{m}:\mathcal{S}_{j}\subseteq\mathcal{S}\\}$
be $m$ randomly sampled subsets of $\mathcal{S}$, each of size $\alpha\times
n$ for some fixed $\alpha\in(0,1)$. The linear datamodeling score (LDS) of
$\tau$ for a specific example $x$ is defined as

$LDS(\tau,x)\triangleq\boldsymbol{\rho}(\\{f_{\Theta_{\mathcal{S}_{j}}}(x):j\in[m]\\},\\{g_{\tau}(x,\mathcal{S}_{j};\mathcal{S}):j\in[m]\\}),$

where $\boldsymbol{\rho}$ is the Spearman rank correlation (Spearman, 1904),
$f_{\Theta_{\mathcal{S}_{j}}}$ is the model output function with
$\Theta_{\mathcal{S}_{j}}$ learned from $\mathcal{S}_{j}$ and
$g_{\tau}(x,\mathcal{S}_{j};\mathcal{S})$ is defined in Eq (3).

To compute LDS for our experiment settings, we use 50 models that are
independently trained on random subsets with size half of the full dataset
(i.e., we set $m=50$ and $\alpha=0.5$ in Definition D.1).

##  Appendix E Wall-clock time measurements

In this paper, we use the wall-clock time on a single A40 GPU to measure the
computational costs if the number of independently trained models (i.e., the
value of $I$) can not precisely demonstrate the costs, e.g., the training time
cost and serving time cost of LoRA Ensemble.

Here, we define several components that dominate the wall-clock time of
different ensemble methods.

  * •

TTrain: The time to train a model from scratch.

  * •

TTrain, Base: The time to train a base model from scratch for LoRA tuning.
Normally speaking, TTrain, Base $<$ TTrain.

  * •

TTrain, LoRA: The time to fine-tune for one LoRA adapter.

  * •

TServing: The time to calculate the TDA scores for one trained model after
model training.

  * •

TServing, Forward-only: The time to calculate the TDA scores for one trained
models after model training with shared gradients. (will only be used by
Dropout Ensemble(forward only). Normally speaking, TServing, Forward-only $<$
TServing, Forward-only.

  * •

TServing, LoRA: The time to calculate the TDA scores after model training for
one LoRA adapter. Normally speaking, TServing, LoRA $<$ TServing.

The total computational cost is approximated and summarized in t:

  * •

Training time cost (naive independent ensemble/Dropout Ensemble/Dropout
Ensemble (forward only)): $I\times$ TTrain.

  * •

Training time cost (LoRA Ensemble): $I\times$ TTrain, Base \+ $I\times
L\times$ TTrain, LoRA.

  * •

Serving time cost (naive independent ensemble): $I\times$ TServing.

  * •

Serving time cost (Dropout Ensemble): $I\times D\times$ TServing.

  * •

Serving time cost (Dropout Ensemble (forward only)): $I\times$ TServing \+
$I\times(D-1)\times$ TServing, Forward-only.

  * •

Serving time cost (LoRA Ensemble): $I\times L\times$ TServing, LoRA.

##  Appendix F Additional experiment for Dropout Ensemble

###  F.1 Ablation experiment to random projection

Here, we examine the root of the Dropout Ensemble’s performance through an
ablation experiment. There are no random factors other than dropout for IF,
Grad-Dot, or Grad-Cos, while there is another random factor, i.e., random
projection, in TRAK. We perform $D$ dropout-masked passes with dropout
enabled, i.e., “Dropout Ensemble’́ and $D$ dropout-masked passes with dropout
disabled, i.e., “Only Random Projection”, in Table 1 and calculate the LDS.
Random projection does contribute to the accuracy improvement, but the
improvement will saturate when $D$ is large.

|  $D$ $I$ | 1 | 3 | 5
---|---|---|---|---
Naive Independent Ensemble | N/A | 0.122 | 0.217 | 0.275
Dropout Ensemble | 10 | 0.249 | 0.399 | 0.457
Only Random Projection | 0.210 | 0.335 | 0.398
Dropout Ensemble | 25 | 0.316 | 0.458 | 0.502
Only Random Projection | 0.217 | 0.351 | 0.413
Table 1: Ablating the contribution of test-time dropout. The experiment is
carried out on MNIST+MLP. Note that $D$ is the number of dropout-masked
passes, and $I$ is the number of independently trained models.

###  F.2 Intermediate checkpoints

Here, we show that Dropout Ensemble ensemble can also be applied to
intermediate checkpoints and improve the accuracy. Park et al. (2023) states
that intermediate checkpoints with fewer epochs can be used for ensembling to
reduce the training time cost. We save multiple checkpoints from different
epochs of the same training process ($I=1$).

Ensembling methods |  $D$ #ckpts | 1 | 3 | 5
---|---|---|---|---
Naive Independent Ensemble | N/A | 0.122 | 0.170 | 0.188
Dropout Ensemble | 10 | 0.249 | 0.318 | 0.342
25 | 0.316 | 0.359 | 0.373
Table 2: The TDA efficacy of Dropout Ensemble on intermediate checkpoints from
different epochs of the same training process ($I=1$). Note that $D$ is the
number of dropout-masked passes, and #ckpts is the number of intermediate
checkpoints used.

##  Appendix G Memory costs of Dropout Ensemble and LoRA Ensemble

Here we record the peak memory usage at serving time for Dropout Ensemble
(Table 3) and LoRA Ensemble (Table 4) applied on TRAK algorithm. The memory of
vanilla Dropout Ensemble is the same as naive independent ensemble. The
meomory of Dropout Ensemble (Forward Only) is larger than vanilla Dropout
Ensemble because some of the cached terms. The memory usage of LoRA Ensemble
is slightly lower than naive independent ensemble because of the reduction in
parameter size.

Datasets and Models |  method variants $D$ | 3 | 10 | 25
---|---|---|---|---
MNIST+MLP | Dropout Ensemble | 342M
Dropout Ensemble (Forward Only) | 636M | 1956M | 4787M
CIFAR2+MLP | Dropout Ensemble | 344M
Dropout Ensemble (Forward Only) | 638M | 1966M | 4797M
CIFAR2+ResNet9 | Dropout Ensemble | 477M
Dropout Ensemble (Forward Only) | 785M | 2087M | 4877M
MAESTRO+MusicTransformer | Dropout Ensemble | 538M
Dropout Ensemble (Forward Only) | 634M | 1529M | 3421M

Table 3: The peak memory usage of Dropout Ensemble and its alternative on TRAK with different numbers of dropout-masked passes ($D$) and the number of independently trained models fixed to 5 ($I=5$).  Ensembling Methods |  $L$ $I$ | 1 | 3 | 5 | 10
---|---|---|---|---|---
Naive Independent Ensemble | N/A | 431M | 538M | 538M | 538M
LoRA Ensemble | 3 | 404M | 404M | 404M | 404M
LoRA Ensemble | 10 | 404M | 404M | 404M | 404M
LoRA Ensemble | 25 | 404M | 404M | 404M | 404M
Table 4: The peak memory usage of LoRA Ensemble and naive independent ensemble
applied on TRAK for MusicTransformer trained on MAESTRO dataset. Note that $I$
represents the number of independently trained models, and $L$ is the number
of LoRA adapters augmented on each model.

##  Appendix H Space costs of Dropout Ensemble and LoRA Ensemble

Here we record the parameter count to present the space cost for Dropout
Ensemble (Table 5) and LoRA Ensemble (Table 6).

$I$(w/ any $D$) settings |  MNIST +MLP |  CIFAR-2 +MLP |  CIFAR-2 +ResNet-9 |  MAESTRO +MusicTransformer
---|---|---|---|---
1 | 0.11M | 0.38M | 4.83M | 13.11M
3 | 0.33M | 1.14M | 14.48M | 39.35M
5 | 0.55M | 1.89M | 24.13M | 65.58M
10 | 1.09M | 3.79M | 48.25M | 131.16M
25 | 2.73M | 9.48M | 120.63M | 327.89M
Table 5: The space cost (total parameter count) of different experiment settings under different numbers of independently trained models ($I$). Note that Dropout Ensemble will not incur any additional storage cost with more dropout-masked passes (i.e., the space cost is fixed with respect to $D$). $I$ $L$ |  Naive Independent ensemble ($L=0$) | 3 | 10 | 25
---|---|---|---|---
1 | 13.11M | 13.41M | 14.09M | 15.57M
3 | 39.35M | 40.23M | 42.29M | 46.72M
5 | 65.58M | 67.05M | 70.49M | 77.87M
10 | 131.76M | 134.11M | 140.99M | 155.73M
Table 6: The space cost (total parameter count) of MusicTransformer under
different numbers of independently trained models ($I$) and different numbers
of LoRA adapters ($L$). LoRA Ensemble only add marginal space cost (at most a
$18.7\%$ increment for $D=25$ across all $I$) compared to naive independent
ensembling.