Stochastic Amortization: A Unified Approach to Accelerate Feature and Data Attribution

Authors: ian  covert
chanwoo  kim
su-in  lee
james  zou
tatsunori  hashimoto
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 21:44:22

Abstract: Many tasks in explainable machine learning, such as data valuation and feature attribution, perform expensive computation for each data point and are intractable for large datasets. These methods require efficient approximations, and although amortizing the process by learning a network to directly predict the desired output is a promising solution, training such models with exact labels is often infeasible. We therefore explore training amortized models with noisy labels, and we find that this is inexpensive and surprisingly effective. Through theoretical analysis of the label noise and experiments with various models and datasets, we show that this approach tolerates high noise levels and significantly accelerates several feature attribution and data valuation methods, often yielding an order of magnitude speedup over existing approaches.

Main paper: 
 introduction: 
 Many tasks in explainable machine learning (XML) perform some form of costly computation for every data point in a dataset. For example, common tasks include assessing individual data points' impact on a model's accuracy [33] , or quantifying each input feature's influence on individual model predictions [68] . Many of these techniques are prohibitively expensive: in particular, those with game-theoretic formulations have exponential complexity in the number of features or data points, making their exact calculation intractable [89, 4] .Accelerating these methods is therefore a topic of great practical importance. This has been addressed primarily with Monte Carlo approximations [96, 16, 74] , which are faster than brute-force calculations but can be slow to converge and impractical for large datasets. Alternatively, a promising idea is to amortize the computation, or to approximate each data point's output with a learned model, typically a deep neural network [2] . For example, in the feature attribution context, we can train an explainer model to predict Shapley values that describe how each feature affects a classifier's prediction [50] .There are several reasons why amortization is appealing, particularly with neural networks: similar data points often have similar outputs, pretrained networks extract relevant features and can be efficiently fine-tuned, and if the combined training and inference time is low then amortization can be faster than computing the object of interest (e.g., data valuation scores) for the entire dataset. However, it is not obvious how to train such amortized models, because standard supervised learning requires a dataset of ground truth labels that can be intractable to generate. Our goal here is therefore to explore efficiently training amortized models when exact labels are costly. Our main insight is that Figure 1 : Diagram of stochastic amortization. Left: using a dataset with noisy labels ã(b) (e.g., images and data valuation estimates), we can train an amortized model that accurately estimates the true outputs a(b) (e.g., data valuation scores). Right: the default approach of running an expensive approximation algorithm for each example (e.g., a Monte Carlo estimator with many samples [33] ).amortization is surprisingly effective with noisy labels: we train with inexpensive estimates of the true labels, and we find that this is theoretically justified when the estimates are unbiased.We refer to this approach as stochastic amortization (Figure 1 ), and we find that it is applicable to a variety of XML tasks. In particular, we show that it is effective for feature attribution with Shapley values [68] , Banzhaf values [11] and LIME [83] ; for several formulations of data valuation [33, 35, 60, 103] ; and to data attribution with datamodels [47] . Our experiments demonstrate significant speedups for several of these tasks: we find that amortizing across an entire dataset with noisy labels is often more efficient than current per-example approximations, especially for large datasets, and that amortized feature and data attribution models generalize well to unseen examples.Our contributions in this work are the following:• We present the idea of stochastic amortization, or training amortized models with noisy labels.We analyze the role of noise from the label generation process and show theoretically that it is sufficient to use unbiased estimates of the ground truth labels (Section 3). We find that non-zero bias in the labels leads to learning an incorrect function, but that variance in the labels plays a more benign role of slowing optimization.• We identify a range of applications for stochastic amortization in XML. Our theory only requires unbiased estimation of the task's true labels, and we find that such estimators exist for several feature attribution and data valuation methods (Section 4).• Experimentally, we test multiple estimators for Shapley value feature attributions and find that amortization works when the labels are unbiased (Section 5). We also verify that amortization is effective for Banzhaf values and LIME. For data valuation, we apply amortization to Data Shapley and show that it allows us to scale this approach to larger datasets than in previous works.• Throughout our experiments, we also analyze the scaling behavior with respect to the amount of training data and the quality of the noisy labels. In general, we find that amortization is more efficient than per-example computation when the training set used for amortization contains at least a moderate number of data points (e.g., >1K for data valuation).Overall, our work shows the potential for accelerating many computationally intensive XML tasks with the same simple approach: amortizing with noisy, unbiased labels. 
background: 
 We first introduce the basic idea of amortization in ML, which we discuss in a general setting before considering any specific XML tasks. Consider a scenario where we repeat similar computation for a large number of data points. We represent this general setup with a context variable b ∈ B and a per-context output a(b) ∈ A. For example, these can be an image and its data valuation score, or an image and its feature attributions. Amortization can be used with arbitrary domains, but we assume Euclidean spaces, or A ⊆ R m and B ⊆ R d , because many XML tasks involve real-valued outputs.The computation performed for each context b ∈ B can be arbitrary as well. In some cases a(b) is an intractable expectation, in which case we would typically approximate it with a Monte Carlo estimator [96, 33] . In other situations it is the solution to an optimization problem [91, 2] , in which case we can define a(b) via a parametric objective h : A × B → R:EQUATIONWe do not require a specific formulation for a(b) in this work, but we will see that both the expectation and optimization setups provide useful perspectives on our proposal of training with noisy labels.In situations with repeated computation, our two options are generally to (i) perform the computation separately for each b ∈ B, or (ii) amortize the computation by predicting the output with a learned model. We typically implement the latter with a neural network a(b; θ), and our goal is to train it such that a(b; θ) ≈ a(b). In training these models, the main challenge occurs when the per-instance computation is costly: specifically, it is not obvious how to train a(b; θ) without a dataset of ground truth solutions a(b), which can be too slow to generate for many XML methods [68, 33] . We address this challenge in Section 3, where we prove that amortization tolerates training with noisy labels.The general idea of amortized computation captures many tasks arising in physics, engineering, control and ML [2] . For example, amortization is prominent in variational inference [56, 82] , meta learning [38, 29, 80] and reinforcement learning [45, 64, 39] . In the XML context, many recent works have explored amortization to accelerate costly per-datapoint calculations [12, 108, 49, 50, 18, 19] . Some of these works are reviewed by [13] , and we offer a detailed overview in Appendix A.For feature attributions, two works propose predicting Shapley values by training with a custom weighted least squares loss [50, 18] ; our simpler approach can use any unbiased estimator and resembles a standard regression task. Two other works suggest modeling feature attributions with supervised learning [87, 14] ; these recommend training with exact or high-quality labels, whereas we recognize the potential to use noisy labels that can be generated orders of magnitude faster. Concurrently, Zhang et al. [110] proposed training with a custom estimator for Shapley value feature attributions; our work is similar but derives stochastic amortization algorithms for a range of settings, including the use of any unbiased estimator (Section 3) and usage for various XML tasks including data valuation (Section 4).For data valuation, two works consider predicting data valuation scores with supervised learning [35, 36] , but these also use near-exact labels that limit the applicability to large datasets. Concurrently, Li and Yu [65] propose a family of data valuation estimators and a learning-based approach analogous to [50] but for data valuation; our approach uses a simpler training loss that works with any unbiased estimator, and unlike [65] its memory usage does not scale with the dataset size. Separately, another line of work focuses on accelerating the model retraining step underlying most data valuation methods [57, 104, 107] , and these are complementary to our approach. Finally, while there are works that accelerate data attribution with datamodels [78, 27] , we are not aware of any that use amortization. 2More broadly, our proposal to train amortized models with noisy labels can be viewed as a version of stochastic optimization, or training with noisy gradients. This fundamental idea is widely used in machine learning [5] , but to our knowledge we are the first to study the broad applicability of noisy labels for accelerating diverse XML tasks. 
stochastic amortization: 
 We now discuss how to efficiently train amortized models with noisy labels. Following Section 2, we present this as a general approach before focusing on a specific XML task. One natural idea is to treat amortization like a standard supervised learning problem: we can parameterize a model a(b; θ), adopt a distribution p(b) over the context variable, and then train our model with the following objective,EQUATION)This approach is called regression-based amortization [2] because it reduces the problem to a simple regression task. The challenge is that this approach cannot be used when we lack a large dataset of exact labels a(b), which is common for computationally intensive XML methods (Section 4).A relaxation of this idea is to train the model with inexact labels (see Figure 1 ). We assume that these are generated by a noisy oracle ã(b), which is characterized by a distribution of outputs for each context b ∈ B. For example, the noisy oracle could be a statistical estimator of a data valuation score [33] . With this, we can train the model using a modified version of Eq. ( 2), where we consider the loss in expectation over both p(b) and the noisy labels ã(b):EQUATIONIt is not immediately obvious when this approach is worthwhile: if the noisy oracle is too inaccurate we will learn the wrong function, so it is important to choose the oracle carefully. We find that there are two properties of the oracle that matter, and these relate to its systematic error and noise level, or more intuitively its bias and variance. We denote these quantities as follows for a specific value b,B(ã | b) = ∥a(b) -E[ã(b) | b]∥ 2 , N(ã | b) = E ∥ã(b) -E[ã(b) | b]∥ 2 | b ,and based on these we can also define the global measures B(ã) ≡ E p [B(ã | b)] and N(ã) ≡ E p [N(ã | b)]for the distribution over context variables p(b). 3 These terms are useful because they reveal a relationship between the two amortization objectives. In general, the objectives are related by the following two-sided bound (see proof in Appendix D):EQUATIONThis relationship shows that reducing Lreg (θ) towards its minimum value N(ã) is similar to training with L reg (θ), only with a disconnect introduced by the bias B(ã). The bias represents a source of irreducible error, because in the limit Lreg (θ) -N(ã) → 0 we have L reg (θ) = B(ã). On the other hand, when B(ã) = 0 we can see that L reg (θ) = Lreg (θ) -N(ã), which means that training with the noisy loss is equivalent and will recover the correct function asymptotically. This last equality is easy to see given an unbiased noisy oracle ã(b), but the more general relationship in Eq. ( 4) emphasizes how non-zero bias can be problematic and lead to learning an incorrect function.Aside from the bias, the variance plays a role as well, not in determining the function we learn but in making the model's optimization unstable or require more noisy labels. To illustrate the role of variance, we present a theoretical result considering the simplest case of a linear model a(b; θ) trained with SGD, which shows that high label noise slows convergence (see proof in Appendix D). T (T +1) θ t at step T satisfies E[ Lreg ( θT )] -N(ã) ≤ 4 Tr(Σ p ) N q (ã) + 4λ max (Σ q )D 2 λ min (Σ p )(T + 1) , where N q (ã) ≡ E q [N(ã | b)]is the noisy oracle's norm-weighted variance, and λ max (•), λ min (•) are the maximum and minimum eigenvalues.The bound in Theorem 1 shows that noise slows convergence by its presence in the numerator, and interestingly, it appears in the form of a weighted version N q (ã) that puts more weight on values with large norm ∥b∥; this is a consequence of assuming a linear model, but we expect the general conclusion of label variance slowing convergence to hold even for neural networks. As a corollary, we can see that the rate in Theorem 1 applies directly to L reg (θ) when the noisy oracle is unbiased (see Appendix D). We note that very high noise levels can in principle prevent effective optimization, and this can be mitigated by either reducing the noisy oracle's variance or taking more steps T .Overall, our analysis shows that amortization with noisy labels is possible, although perhaps more difficult to optimize than training with exact labels. We next show that unbiased estimates are available for many XML tasks (Section 4), and we later find that this form of amortization is consistently effective with the noise levels observed in practice (Section 5), even providing better accuracy than per-example estimation in compute-matched comparisons. As a shorthand, we refer to training with the noisy objective Lreg (θ) in Eq. (3) as stochastic amortization.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
applications to explainable ml: 
 We now consider XML tasks that can be accelerated with stochastic amortization. Rather than using generic variables b ∈ B and a(b) ∈ A, this section uses an input variable x ∈ X , a response variable y ∈ Y, and a model f or a measure of its performance. As we describe below, each application of stochastic amortization is instantiated by a noisy oracle that generates labels for the given task.One of the most common tasks in XML is feature attribution, which aims to quantify each feature's influence on an individual prediction. The Shapley value has gained popularity because of its origins in game theory [89, 68] , and like many feature attribution methods is based on querying the model while removing different feature sets [17] . Given a model f and input x that consists of d separate features x = (x 1 , . . . , x d ), we assume that we can calculate the prediction f (x S ) ∈ R for any feature set S ⊆ [d]. 4 With this setup, the Shapley values ϕ i (x) ∈ R for each feature i ∈ [d] are defined as:EQUATIONThese scores satisfy several desirable properties [89] , but they are impractical to calculate due to the exponential summation over feature subsets. Our goal is therefore to learn an amortized model ϕ(x; θ) ∈ R d to directly predict feature attribution scores, and for this we require a noisy oracle.Many recent works have studied efficient Shapley value estimation [10] , and we first consider noisy oracles derived from Eq. ( 5), which defines the attribution as the feature's expected marginal contribution. There are several unbiased statistical estimators that rely on sampling feature subsets or permutations [6, 96, 77, 74, 58] , and following Section 3 we can use any of these for stochastic amortization. Our experiments use the classic permutation sampling estimator [96, 74] , which approximates the values ϕ i (x) as an expectation across feature orderings. We defer the precise definition of this noisy oracle to Appendix E, along with the other estimators used in our experiments.Next, we also consider noisy oracles derived from an optimization perspective on the Shapley value. A famous result from Charnes et al. [7] shows that the Shapley values are the solution to the following problem (with abuse of notation we discard the solution's intercept term),EQUATIONwhere we use a least squares weighting kernel defined as µ -1 (S) = d |S| |S|(d -|S|). Several works have proposed approximating Shapley values by solving this problem with sampled subsets, either using projected gradient descent [92] or analytic solutions [68, 16] . Among these, we use KernelSHAP [68] and SGD-Shapley [92] as noisy oracles in our experiments. The first is an Mestimator whose bias shrinks as the number of sampled subsets grows [99] , so we expect it to lead to effective amortization; the latter has been shown to have non-negligible bias [10] , so our theory in Section 3 suggests that it should lead to learning an incorrect function when used for amortization.Next, we consider two alternative feature attribution methods: Banzhaf values [4, 11] and LIME [83] . These are closely related to Shapley values and are similarly intractable [26, 40] , but we find that they offer statistical estimators that can be used for stochastic amortization.First, Banzhaf values assign the following scores to each feature for a prediction f (x) [4] :EQUATIONThese differ from Shapley values only in their choice of weighting function, and they admit a range of similar statistical estimators. One option is the MSR estimator from [103] , which is unbiased and re-uses all model evaluations for each feature attribution estimate. We adopt this as a noisy oracle in our experiments (see a precise definition in Appendix E), but several other options are available.Second, LIME defines its attribution scores ϕ i (x) as the solution to the following optimization problem, given a weighting kernel π(S) and penalty term Ω [83]: 5 arg minEQUATIONAs our noisy oracle for LIME, we use the popular approach of solving the above problem for subsets sampled according to π(S). Similar to KernelSHAP [68] , this is an M-estimator whose bias shrinks to zero as the sample size grows [99] , so we expect it to lead to successful amortization.Aside from these methods, other costly feature interpretation methods rely on unbiased statistical estimators and can be amortized in a similar fashion [98, 62, 32] . We leave further investigation of these methods to future work.Next, data valuation aims to quantify how much each training example affects a model's accuracy. We consider labeled examples z = (x, y) and a training dataset D = {z i } n i=1 , and we analyze each data point's value by fitting models to subsampled datasets D T ⊆ D with T ⊆ [n] and calculating a measure of the model's performance v(D T ) ∈ R (e.g., its 0-1 accuracy). This general approach was introduced by Ghorbani and Zou [33] , who defined the Data Shapley scores ψ(z i ) ∈ R as follows:EQUATIONSubsequent work generalized the approach in different ways, which we briefly summarize before considering amortization. For example, Wang and Jia [103] used the Banzhaf value rather than Shapley value, and Kwon and Zou [60] considered the case of arbitrary semivalues [75] . These correspond to adopting a different weighting over subsets in Eq. ( 9), and the general case can be written as follows for a normalized weighting function w(k):EQUATIONNext, another extension is the case of distributional data valuation. Ghorbani et al. [35] incorporate an expectation over the original dataset D, which they show leads to well defined scores even for data points z = (x, y) outside the training set. Given a distribution over datasets of size |D| = n -1 and a weighting function w(k), this version defines the score ψ(z) ∈ R for arbitrary z as follows:EQUATIONWhen using any of these methods in practice, the scores are difficult to calculate due to the intractable expectation across datasets. However, a crucial property they share is that they can all be estimated in an unbiased fashion, and these estimates can therefore be used as noisy labels for stochastic amortization. We focus on Data Shapley and Distributional Data Shapley in our experiments [33, 35] , and we use the Monte Carlo estimator from Ghorbani and Zou [33] as a noisy oracle (see the precise 5 Following [17] , we only consider the version of LIME with a binary interpretable representation. 6 For this to be a valid expectation, semivalues require that n-1 definition in Appendix E). Doing so allows us to train an amortized valuation model ψ(z; θ) ∈ R that accelerates valuation within our training dataset, and that can also be applied to external data, e.g., when selecting informative new data points for active learning [36] .Finally, Appendix B discusses amortization for the datamodels data attribution technique [47] . This method measures how much each training data point z i ∈ D affects the prediction for an inference example x ∈ X , and we show that the scores are equivalent to a simple expectation that can be estimated in an unbiased fashion. The datamodels scores can therefore be amortized by adopting these estimates as noisy labels, but we leave further investigation of this approach to future work. 
experiments: 
 Our experiments apply stochastic amortization to several of the tasks discussed in Section 4. We consider both feature attribution and data valuation, for image and tabular datasets, and using multiple architectures for our amortized models, including fully-connected networks (FCNs), ResNets [43] and ViTs [24] . Full details are provided in Appendix F, including our exact models and hyperparameters.Our goal in each experiment is to perform feature attribution or data valuation for an entire dataset, and to compare the accuracy of stochastic amortization to running existing estimators on each point. We adopt a noisy oracle for each task (e.g., a Monte Carlo estimator of data valuation scores), we then fit an amortized network with one noisy label per training example, and our evaluation focuses on the accuracy of the amortized predictions relative to the ground truth. Our ground truth is obtained by running the noisy oracle to near-convergence for a large number of samples: for example, we run KernelSHAP [68] for feature attribution with 1M samples, and the TMC estimator [33] for data valuation with 10K samples. We test amortization when using different numbers of training examples, and for both training and unseen external data to evaluate the model's generalization. We find that amortization often denoises and strongly improves upon the noisy labels, leading to a significant accuracy improvement for the same computational budget. We first consider Shapley value feature attributions. This task offers a diverse set of noisy oracles, and we consider three options: Ker-nelSHAP [68] , permutation sampling [74] and SGD-Shapley [92] . Among these, our theory from Section 3 suggests that the first two will be effective for amortization, while the third may not because it is not an unbiased estimator. We follow the setup from [18] and implement our amortized network with a pretrained ViT-B architecture [24] , and we use the ImageNette dataset [46] with 224 × 224 images partitioned into 196 patches of size 14 × 14.As a first result, Figure 2 compares our training targets to the amortized model's predictions. The predicted attributions are significantly more accurate than the labels, even for the noisiest setting with just 512 KernelSHAP samples. Figure 3 (left) quantifies the improvement from amortization, and we observe similar results for both KernelSHAP and permutation sampling (see Appendix G): in both cases the error is significantly lower than that of the noisy labels, and it remains lower and improves as the labels become more accurate. To contextualize our amortized model's accuracy, we find that the error is similar to that of running KernelSHAP for 10-40K samples, even though our labels use an order of magnitude fewer samples (see Appendix G).We also find that the model generalizes to external data points (see Appendix G). In addition, we show results for SGD-Shapley, where we confirm that it leads to poor amortization results due to its non-negligible bias (see Appendix G).Next, we investigate the compute trade-off between calculating attributions separately (e.g., with KernelSHAP) and using amortization. Figure 3 (center-right) shows two results regarding this tradeoff. First, we measure the error as a function of FLOPs, where we account for the cost of generating labels and training the amortized model (see Appendix F). The FLOPs incurred by training are negligible compared to the KernelSHAP estimates, and we find that stopping at any time to fit an amortized model yields significantly better estimates. This suggests that amortization is an inexpensive final denoising step, regardless of how much compute was used for the noisy estimates.Second, we test the effectiveness of amortization for different dataset sizes. We match the compute between the two scenarios, using 2440 KernelSHAP samples for per-example computation 7 and 2257 for amortization to account for the cost of training. This compute-matched comparison shows that amortization achieves lower estimation error for datasets ranging from 250-10K data points (Figure 3 right); it becomes more effective as the dataset grows, but it is useful even for small datasets.Finally, Appendix G shows a comparison between stochastic amortization and FastSHAP [50, 18] , an existing approach to amortized Shapley value estimation. We observe similar estimation accuracy in compute-matched comparisons, and find that both methods are significantly more accurate than per-example estimation (similar to Figure 3 center). Appendix G also show results for amortizing Banzhaf values and LIME: we find that amortization is more difficult for these methods due to the inconsistent scale of attributions between inputs, but we nonetheless observe an improvement in our amortized estimates versus the noisy labels.Next, we consider data valuation with Data Shapley. For our noisy oracle, we obtain training labels by running the TMC estimator with different numbers of samples [33] . We first test our approach with the adult census and MiniBooNE particle physics datasets [25, 84] , and following prior work we conduct experiments using versions of each dataset with different numbers of data points [53] .Our valuation model must predict scores for each example z = (x, y), so we train FCNs that output scores for all classes and use only the relevant output for each data point.As a first result, Figure 4 (left-center) shows the estimation accuracy for the MiniBooNE dataset when using 1K and 10K data points. The noisy estimates converge as we use more Monte Carlo samples, but we see that the amortized estimates are always more accurate in terms of both squared error and correlation with the ground truth. The improvement is largest for the noisiest estimates, where the amortized predictions have correlation >0.9 when using only 50 samples. Amortization is more beneficial for the 10K dataset, which suggests that training with more noisy labels can be a substitute for high-quality labels. Appendix G shows similar results with the adult census dataset.Next, Figure 4 (right) considers the role of training dataset size for the estimates with 50 Monte Carlo samples. For both the adult census and MiniBooNE datasets, we see that the benefits of amortization are small with 250 data points but grow as we approach 10K data points. This number of samples is enough to maintain 0.9 correlation with the ground truth when using amortization, whereas the raw estimates become increasingly inaccurate. Stochastic amortization is therefore promising to scale data valuation beyond previous works, which typically focus on <1K data points [34, 60, 103] .Finally, we consider Distributional Data Shapley [35] , which is similar to the previous experiments but defines valuation scores even for points outside the training dataset; this allows us to test the generalization to unseen data. We use the CIFAR-10 dataset [59] , which contains 50K training examples, and for our valuation model we train a ResNet-18 that outputs scores for each class.Similar to the previous experiments, Figure 5 (left-center) evaluates the estimation accuracy for different numbers of Monte Carlo samples. We observe that the distributional scores converge faster because we use a smaller maximum dataset cardinality (see Appendix F), but that amortization still provides a significant benefit. The improvement is largest for the noisiest estimates, which in this case use just 5 samples: amortization achieves correlation 0.81 with the ground truth, versus 0.58 for the Monte Carlo estimates. The improvement is consistent across several measures of accuracy, including squared error, Pearson correlation and Spearman correlation (see Appendix G).Next, we study generalization and the role of dataset size. We focus on the noisiest estimates with 5 Monte Carlo samples, because this low-sample regime is most relevant for larger datasets with millions of examples. We train the valuation model with different portions of the 50K training set, and we measure the estimation accuracy for both internal and unseen external data. Figure 5 (right) shows large improvements in squared error with as few as 1K data points, and we also observe improvement in correlation when we use at least 5K data points (10% of the dataset, see Appendix G). We additionally observe a small generalization gap, suggesting that the valuation model can be trained with a subset of the data and reliably applied to unseen examples.Lastly, we test the usage of amortized valuation scores in downstream tasks. Following [53] , the tasks we consider are identifying mislabeled examples and improving the model by removing low-value examples. These results are shown in Appendix G, and we find that the amortized estimates identify mislabeled examples more reliably than Monte Carlo estimates, and that filtering the dataset based on our estimates leads to improved performance. 
conclusion: 
 This work explored the idea of stochastic amortization, or training amortized models with noisy labels. Our main finding is that fast, noisy supervision provides substantial compute and accuracy gains over existing XML approximations. This approach makes several feature attribution and data valuation methods more practical for large datasets and real-time applications, and it may have broader applications to amortization beyond XML [2] . Our proposal has certain limitations, including that stochastic amortization may become ineffective with sufficiently high noise levels, and that it is difficult to know a priori how much compute is necessary for label generation to achieve a desired error level in the amortized predictions.Our work suggests multiple directions for future research. One direction is to study the trade-off between using a larger number of noisy labels or a smaller number of more accurate labels, which is a key difference from prior work that uses near-exact labels for amortization [14] . Other directions include scaling to datasets with millions of examples to test the limits of noisy supervision, leveraging more sophisticated data valuation estimators [106] , using alternative model retraining primitives [57, 63, 37] , and exploring amortization for other methods like datamodels (discussed in Appendix B). 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments? Answer: [Yes] Justification: Information on compute resources is provided in Appendix F. Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? Answer: [Yes] Justification: Our proposed method aims to make computationally challenging XML tasks feasible, so our research promotes transparency and interpretability in ML. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: The potential positive and negative societal impacts of our work are discussed in Appendix H. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our paper does not release a new dataset or model.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: The open-source packages and datasets we used are cited in Appendix F.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [Yes] Justification: The code we provide contains usage instructions and scripts for running experiments.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: Our paper does not involve crowdsourcing experiments or research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: Our paper does not involve human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: This section provides a more detailed review of amortization in XML. Many XML tasks involve analyzing individual data points, e.g., to determine the most important features or concepts in a model, or a data point's influence on the model's accuracy or an individual prediction. Among these methods, many require expensive per-example algorithms, and a trend in recent years has been to amortize this computation using deep learning. We discuss these works below, grouping them into several main XML problems that they address.Within this discussion, we differentiate works not only based on their goal (e.g., feature attribution or data valuation) but also based on how they calculate the object of interest. For those that perform per-example computation, we distinguish between methods that solve a parametric optimization problem of the form a(b) = arg min a ′ h(a ′ ; b) (e.g., with gradient descent) from those that exploit an analytic solution, either by directly calculating the result or using a Monte Carlo approximation.For those that use amortization, we distinguish those that use regression-based amortization with L reg (θ) from those that perform objective-based amortization with L obj (θ) (see Appendix C for more details on objective-based amortization [2] ).Table 1 summarizes the methods we discuss, including the various tasks they solve, the context and output domains (A, B) for each problem, and the computational approach. This perspective highlights the significant role of amortization in XML, and it also reveals opportunities for new applications, some of which we discuss in Section 4. Some of the same works are discussed by Chuang et al. [13] , but we cover a wider range of methods, and we adopt the framing of Amos [2] by outlining the various parametric optimization problems and different amortization approaches.Feature attribution. These methods aim to quantify the importance of a model's input features, typically for individual predictions (such methods are called local rather than global feature attributions, see Lundberg et al. [69] ). The context is therefore a data example x ∈ X , and the output is a vector of attributions in R d when we have d features. Feature attribution algorithms can be efficient to calculate, particularly when they are based on simple propagation rules [93, 97, 90, 95, 88, 111, 3] or a transformer's attention values [1] . However, another family of approaches are based on feature removal [17] , and querying the model with many feature subsets is typically less efficient.Among the feature removal-based methods, one of the most famous is Shapley value feature attributions [68, 89] . Two related methods are LIME [83] and Banzhaf values [11, 4] , see our discussion of their similarity in Section 4. Many works have focused on efficiently calculating these attributions, because their computational complexity is exponential in the number of features. As discussed by Chen et al. [10] , there are many stochastic estimators derived from an analytic expression for the attributions [6, 96, 77, 16, 74, 103, 58] , and these are typically unbiased Monte Carlo estimators. Others are based on solving an optimization problem, for example with either M-estimation [83, 68, 105, 16] or gradient descent [92] .Lastly, there are also methods based on amortization, which offer the possibility of real-time feature attribution without a significant loss in accuracy. Two works consider objective-based amortization, which avoids the need for labels during training [50, 18] . Others consider regression-based amortization [86, 87, 14] , and these recommend using high-quality labels; for example, Chuang et al. [14] use exact labels when possible. Our perspective is most similar to the concurrent work by Zhang et al. [110] , in that we highlight the potential for regression-based amortization with inexpensive, noisy labels; but our work is more general in that we highlight the potential to use noisy labels from any unbiased estimator, and the applicability of this approach to other XML tasks like data valuation.Instance-wise feature selection. A related set of methods select the most important features for individual predictions. The context in this case is an input x ∈ X , and the output is a set of feature indices S ∈ P([d]) where we use P(•) to denote the power set. The optimal subset can be defined either based on the prediction f (x S ) or f (x [d]\S ) for a specific class [31] , or based on the deviation from the original prediction f (x) [12] . Among these methods, some solve the underlying optimization problem on a per-input basis [31, 30] , and others solve it using objectivebased amortization [21, 12, 108, 49] . These methods all require differentiating through discrete subsets, so they employ various continuous relaxations (e.g., the Concrete distribution [71] ), as well as other tricks to constrain or penalize the subset size. These methods are sometimes thought of as attribution methods due to the continuous relaxations [31, 30] , but we describe them as feature x ∈ X R d Štrumbelj and Kononenko [96] Okhrati and Lipani [77] Mitchell et al. [74] Lundberg and Lee [68] Simon and Vincent [92] Covert and Lee [16] Schwarzenberg et al. [87] Chuang et al. [14] Zhang et al. [110] Jethani et al. [50] Covert et al. [18] Leave-one-out attribution Zeiler and Fergus [109] Schwab and Karlen [86] LIME attribution Ribeiro et al. [83] Banzhaf value attribution Covert et al. [17] Chen and Jordan [11] Instance-wise selectionx ∈ XChen et al. [12] Yoon et al. [108] Jethani et al. [49] Image maskingFong and Vedaldi [31] Fong et al. [30] Dabkowski and Gal [21] Dynamic feature selectionx S ∈ X × P([d]) [d]Ma et al. [70] He et al. [42] Chattopadhyay et al. [9] Covert et al. [19] Counterfactual explanationx ∈ X X Wachter et al. [102] Mahajan et al. [72] Verma et al. [101] Data Shapley zi ∈ D R Ghorbani and Zou [33] Ghorbani et al. [35, 36] Li and Yu [65] Beta Shapley Kwon and Zou [60] Data Banzhaf Wang and Jia [103] Datamodels (zi, x) ∈ D × X R selection methods because the final scores for each feature are restricted to the range [0, 1] and are made continuous mainly for optimization purposes.Dynamic feature selection. Next, other works select features separately for each prediction to achieve high accuracy with a small feature acquisition budget. Each selection is made based on a partially observed input, so the context is a feature subset x S , which we write as belonging to the domain X × P([d]), and the output is an index i ∈ [d]. There are several ways to define the optimal selection at each step, but one common approach is to define it as the feature with maximum conditional mutual information, or i * = arg max i I(y; x i | x S ), where the response variable y and unobserved features x i are random variables and x S has a fixed observed value. Given access to this objective, it is trivial to solve by enumerating the possible indices; however, the mutual information is typically unavailable in practice and must be approximated. One approach is therefore to fit a proxy for the mutual information (e.g., via a generative model) and then optimize this proxy for each selection [70, 81, 8, 44 ]. An alternative is to learn a network that predicts the optimal selection at each step: He et al. [42] do so using imitation learning, which resembles regression-based amortization, and Chattopadhyay et al. [9] and Covert et al. [19] do so with an optimization-based view of the mutual information I(y; x i | x S ), which is objective-based amortization.Counterfactual explanation. The goal of counterfactual explanations (also known as recourse explanations) is to identify small changes to an input that cause a desired change in a model's prediction. The context is an input x ∈ X , and the output is a modified version x ′ ∈ X that is typically not too different from the original input. This family of methods was introduced by Wachter et al. [102] and is generally framed as an optimization problem involving the prediction f (x ′ ) and a measure of the perturbation strength between x ′ and x. Verma et al. [100] provide a review of these methods and the various choices for the optimization formulation. When computing these explanations, the most common approach is to solve the problem separately for each input, e.g., using gradient descent [102] . Other works have explored learning models that directly output the modified example x ′ , and these are typically implemented using objective-based amortization [72, 94, 101, 13] .Data valuation. The goal of data valuation is to assign scores to each training example that represent their contribution to a model's performance. The context is therefore a labeled training example z ∈ D, and the output is a real-valued score. These methods are typically not defined via an optimization problem, but instead as a measure of the example's expected impact on performance across a distribution of preceding datasets [33, 35, 60, 103] . Existing methods therefore rely on analytic expressions for the valuation scores, which require Monte Carlo approximation because they are intractable. Two works have considered predicting data valuation scores given a dataset of near-exact estimates [35, 36] , which is regression-based amortization. Our work is similar, but we use inexpensive estimates to reduce the cost of amortization and scale more efficiently to larger datasets. Concurrently, Li and Yu [65] proposed a learning-based approach analogous to FastSHAP for feature attribution [50] , but whose memory requirements scale with the dataset size and limit applicability to large datasets. Besides these methods, there are also data valuation approaches that can be calculated without any approximations [51, 54, 61] .Data attribution. Finally, data attribution methods aim to quantify the effect of training examples on individual predictions. The context is therefore a training example z ∈ D paired with an inference example x ∈ X , and the output is a real-valued attribution score. One classic approach to this problem is influence functions [15] , which use a gradient-based approximation to avoid the cost of retraining [57, 37, 63] . Another class of methods involve measuring the effect of training with subsampled datasets [28, 47] . For datamodels, the main existing approximation algorithm is based on solving a global least squares problem [47] , which does not correspond to any of the computational approaches listed in Table 1 ; notably, it is not a form of amortization because there is no model that can predict the attribution given a new tuple (z, x). As an alternative to the global least squares problem, TRAK calculates similar scores to datamodels using a gradient-based approximation [78, 27] . We are not aware of any work investigating amortization for datamodels, but our results in Appendix B show how to approximate the scores based on their analytic solution, and how to amortize the computation by adopting Monte Carlo estimates as noisy training labels.The datamodels technique [47] aims to quantify how much each training data point z i ∈ D affects the prediction for an inference example x ∈ X . Similar to data valuation, the inference example's output given a training dataset D T is represented by a function v x (D T ), and the attribution scores ζ(z i , x) ∈ R are then defined as the solution to a joint least squares problem that can be solved after training many models with different datasets, as we show below. 8 For the inference example's output v x (D T ), Ilyas et al. [47] focus on the loss rather than the raw prediction, but our perspective can accommodate any definition of this output.Our insight on the datamodels technique is twofold. First, we show that the datamodels scores are equal to a simple expectation and can be estimated in an unbiased fashion. Second, we show that these calculations can be amortized by using the noisy Monte Carlo estimates as training targets. Our findings rely on a slight reformulation of the datamodels scores, where we deviate from Ilyas et al. [47] by using an intercept term and a subtly different dataset distribution: our distribution is biased towards size nq for a value q ∈ (0, 1) rather than using a fixed size, which was also considered in recent work by Saunshi et al. [85] . In the following result, we use the shorthandζ(x) ≡ [ζ(z 1 , x), . . . , ζ(z n , x)]for the vector of data attributions. Given a probability q ∈ (0, 1), we also define the weighting function u(k) = q k (1 -q) n-1-k . With this setup, our reformulation is the following (see the proof in Appendix D). Proposition 1. Given a subset distribution that includes each data point z i ∈ D with probability q ∈ (0, 1), the data attribution scores defined byζ(x) ≡ arg min a∈R n+1 E T   a 0 + i∈T a i -v x (D T ) 2  can be expressed as the following expectation:ζ(z i , x) = T ⊆[n]\i u(|T |) (v x (D T ∪ {z i }) -v x (D T )) .Perhaps surprisingly, this corresponds to a game-theoretic semivalue and reduces to the Banzhaf value when q = 1/2 [73] . Based on this perspective, we can estimate the score ζ(z i , x) without solving the regression problem from Ilyas et al. [47] . For example, we can simply sample k datasets D T from the distribution with q ∈ (0, 1), and then calculate the empirical average as follows:ζ(z i , x) = 1 k k j=1 v x (D Tj ∪ {z i }) -v x (D Tj \ {z i }).Following Recall that many tasks with repeated computation have an optimization view where a(b) ≡ arg min a ′ h(a ′ ; b) [91, 2] (see Section 2). For completeness, this section describes an interpretation of stochastic amortization from this optimization perspective. An alternative to regression-based amortization when we have an objective h(a ′ ; b) that defines a(b) is to train a(b; θ) directly with the h objective:EQUATIONThis approach provides an alternative when exact labels are costly, but it can be unappealing because (i) the task may not offer a natural objective h(a ′ ; b), and (ii) minimizing h(a ′ ; b) may seem disconnected from the error ∥a(b; θ) -a(b)∥. We discuss these issues below and how they are related to stochastic amortization.For issue (i), certain problems like Shapley values have a natural optimization characterization (see Section 4.1), but this is not always the case: for example, data valuation scores are framed as an expectation rather than via optimization (see Section 4.3 or [33] ), so they lack an optimization view.We always have the trivial optimization perspective h(a ′ ; b) = ∥a ′ -a(b)∥ 2 , but this is not useful because it reduces L obj (θ) to L reg (θ) and requires the exact outputs a(b). Regardless of whether the task offers a natural objective h(a ′ ; b), stochastic amortization can be viewed as defining a new objective that does not require the exact outputs: given a noisy oracle ã(b), stochastic amortization is equivalent to L obj (θ) with the objective h(a ′ ; b) = E ∥a ′ -ã(b)∥ 2 , and if the oracle is unbiased then the optimal predictions are a(b; θ) = a(b).Next, issue (ii) is a concern when h(a ′ ; b) is available but lacks a clear connection to the estimation error, which is directly reflected by L reg (θ). The squared error is in many cases a more meaningful accuracy measure, and it is commonly used to evaluate estimation accuracy for feature attribution and data valuation [50, 103, 10] . Certain works on objective-based amortization have considered nonconvex objectives h(a ′ ; b) where the exact output a(b) is not well-defined [2] , but we can understand the potential disconnect by focusing on a class of well behaved objectives. In particular, if we assume that h(a ′ ; b) is α-strongly convex and β-smooth in a ′ for all b, then the L reg (θ) and L obj (θ) objectives bound one other as follows,EQUATIONwhere we define L * obj ≡ E[h(a(b); b)] (see proof in Appendix D.1). Eq. (13) shows that by minimizing L obj (θ), we effectively minimize both upper and lower bounds on L reg (θ). However, these bounds can be loose: for example, because we have L reg (θ) ≤ 2 α L obj (θ) and the strong convexity constant α shrinks to zero for Shapley values as the number of features grows [18] , optimizing L obj (θ) may become less effective in high dimensions. Stochastic amortization resolves this potential disconnect between L obj (θ) and L reg (θ) as follows: if we use an unbiased noisy oracle ã(b), then training with Lreg (θ) is equivalent to using L obj (θ) with h(a ′ ; b) = E ∥a ′ -ã(b)∥ 2 , so we have α = β = 1 in Eq. ( 13) and the regression objectives L reg (θ) and Lreg (θ) are equal up to a constant L * obj = N(ã).This section provides proofs for our claims in the main text. Appendix D.1 shows proofs for our results related to stochastic amortization in Section 3 and Appendix C, and then Appendix D.2 shows the proof for Proposition 1 related to datamodels.We first derive the inequality in Eq. ( 4) from the main text, which is the following:Lreg (θ) -N(ã) -B(ã) 2 ≤ L reg (θ) ≤ Lreg (θ) -N(ã) + B(ã) 2 .Proof. We begin by decomposing a per-input version of the noisy oracle loss Lreg (b; θ), which takes an expectation over the noisy oracle ã(b) for a fixed b:Lreg (b; θ) = E ∥a(b; θ) -ã(b)∥ 2 | b = ∥a(b; θ) -E[ã(b) | b]∥ 2 + E[∥ã(b) -E[ã(b) | b]∥ 2 | b] = ∥a(b; θ) -E[ã(b) | b]∥ 2 + N(ã | b).Next, we can also decompose a per-input version of the original regression loss L reg (b; θ) as follows using triangle inequality:L reg (b; θ) = ∥a(b; θ) -a(b)∥ 2 ≤ (∥a(b; θ) -E[ã(b) | b]∥ + ∥a(b) -E[ã(b) | b]∥) 2 = Lreg (b; θ) -N(ã | b) + B(ã | b) 2 .Taking both sides of the inequality in expectation over p(b), we arrive at the following upper bound:L reg (θ) ≤ E Lreg (b; θ) -N(ã | b) + B(ã | b) 2 .One side of the bound in Eq. ( 4) comes from developing the square and applying Cauchy-Schwarz to the cross term:L reg (θ) ≤ E Lreg (b; θ) -N(ã | b) + B(ã | b) 2 = E Lreg (b; θ) -N(ã | b) + E [B(ã | b)] + 2E B(ã | b) Lreg (b; θ) -N(ã | b) ≤ Lreg (θ) -N(ã) + B(ã) + 2 B(ã) Lreg (θ) -N(ã) = Lreg (θ) -N(ã) + B(ã) 2 .This proves the upper bound in Eq. ( 4). For the lower bound, we return to the decomposition of the per-input noisy oracle loss Lreg (b; θ) and apply triangle inequality as follows:Lreg (b; θ) = ∥a(b; θ) -E[ã(b) | b]∥ 2 + N(ã | b) ≤ (∥a(b; θ) -a(b)∥ + ∥a(b) -E[ã(b) | b]∥) 2 + N(ã | b) = L reg (b; θ) + B(ã | b) 2 + N(ã | b).Rearranging terms, we haveL reg (b; θ) ≥ Lreg (b; θ) -N(ã | b) -B(ã | b) 2 ,and applying the same logic as above with Cauchy-Schwarz, we arrive at:L reg (θ) ≥ Lreg (θ) -N(ã) -B(ã) 2 .Putting the two results together, this yields the two-sided bound in Eq. ( 4):Lreg (θ) -N(ã) -B(ã) 2 ≤ L reg (θ) ≤ Lreg (θ) -N(ã) + B(ã) 2 .Next, we prove our SGD result for stochastic amortization in Theorem 1. E[ Lreg ( θT )] -N(ã) ≤ 4 Tr(Σ p ) N q (ã) + 4λ max (Σ q )D 2 λ min (Σ p )(T + 1) ,where N q (ã) ≡ E q [N(ã | b)] is the noisy oracle's norm-weighted variance, and λ max (•), λ min (•) are the maximum and minimum eigenvalues.Proof. Given our assumption about the noisy oracle, we can re-write the noisy objective Lreg (θ) as follows:Lreg (θ) = E ∥(θ -θ)b∥ 2 + N(ã) = Tr (θ -θ) ⊤ Σ p (θ -θ) + N(ã).It is clear that the minimizer is θ = θ and that the minimum achievable error is Lreg ( θ) = N(ã).Because the objective composes a linear function with a convex function, it is convex in θ and we can analyze its convergence using a standard SGD result. Specifically, because we assume that ∥ θ∥ F ≤ D, we consider a projected SGD algorithm where each step is followed by a projection into the D-ball, or θ ← θ • min(1, D/∥θ∥ F ). Note that it is common to assume a bounded solution space when proving SGD's convergence [76, 22, 5] .Before proceeding to the main convergence result, we must derive several properties of this objective and its stochastic gradients. The first is the strong convexity constant, and we begin by noting that the gradient is the following (see Section 2.5 of Petersen et al. [79] for a review of matrix derivatives):∇ Lreg (θ) = 2(θ -θ)E[bb ⊤ ] = 2(θ -θ)Σ p .For the strong convexity constant, we require a value α > 0 such that the following is satisfied for all parameters θ, θ ′ :EQUATIONFor the first two terms on the right side of the inequality, we can write:Lreg (θ ′ ) + Tr (θ -θ ′ ) ⊤ ∇ Lreg (θ ′ ) = N(ã) + Tr (θ ′ -θ)Σ p (θ ′ -θ) ⊤ + 2 Tr (θ ′ -θ)Σ p (θ -θ ′ ) ⊤ = N(ã) + Tr (θ -θ)Σ p (θ -θ) ⊤ -Tr (θ ′ -θ)Σ p (θ ′ -θ) ⊤ .By using the minimum eigenvalue of Σ p , we can write the following,Tr (θ ′ -θ)Σ p (θ ′ -θ) ⊤ ≥ λ min (Σ p )∥θ -θ ′ ∥ 2 F ,and therefore satisfy Eq. ( 14) as follows:Lreg (θ) = N(ã) + Tr (θ -θ)Σ p (θ -θ) ⊤ = Lreg (θ ′ ) + Tr (θ -θ ′ ) ⊤ ∇ Lreg (θ ′ ) + Tr (θ ′ -θ)Σ p (θ ′ -θ) ⊤ ≥ Lreg (θ ′ ) + Tr (θ -θ ′ ) ⊤ ∇ Lreg (θ ′ ) + λ min (Σ p )∥θ -θ ′ ∥ 2 F .We can therefore conclude that Lreg (θ) is 2λ min (Σ p )-strongly convex.The next property to derive is the expected gradient norm. When running SGD, we make updates using the following stochastic gradient estimate:g(θ) = 2 (θb -ã(b)) b ⊤ .We require an upper bound on the stochastic gradient norm, which can be written as follows:EQUATION)For the first term in Eq. ( 15), which represents the gradient variance, we have:E ∥g(θ) -∇ Lreg (θ)∥ 2 F = 4E ∥(θb -ã(b))b ⊤ -(θ -θ)Σ p ∥ 2 F .To find a simple expression for this term, we first consider the expectation over the distribution ã(b) with fixed b ∈ B, which isolates label variation due to the noisy oracle:E ã|b ∥(θb -ã(b))b ⊤ -(θ -θ)Σ p ∥ 2 F = E ã|b ∥(ã(b) -θb)b ⊤ ∥ 2 F + ∥(θ -θ)(bb ⊤ -Σ p )∥ 2 F = N(ã | b) • ∥b∥ 2 + ∥(θ -θ)(bb ⊤ -Σ p )∥ 2 F .When we take this in expectation over p(b), the first term can be understood as a norm-weighted average of the noisy oracle's conditional variance:E N(ã | b) • ∥b∥ 2 = p(b)N(ã | b) • ∥b∥ 2 db = Tr(Σ p ) p(b) • ∥b∥ 2 Tr(Σ p ) N(ã | b)db = Tr(Σ p )N q (ã).The second term can be rewritten as follows using Σ p and Σ q :E ∥(θ -θ)(bb ⊤ -Σ p )∥ 2 F = E Tr (bb ⊤ -Σ p )(θ -θ) ⊤ (θ -θ)(bb ⊤ -Σ p ) = E Tr (θ -θ)(bb ⊤ -Σ p )(bb ⊤ -Σ p )(θ -θ) ⊤ = Tr (θ -θ)E (bb ⊤ -Σ p )(bb ⊤ -Σ p ) (θ -θ) ⊤ = Tr (θ -θ)(Tr(Σ p )Σ q -Σ 2 p )(θ -θ) ⊤ .For the deterministic gradient upper bound in Eq. ( 15), we have:∥∇ Lreg (θ)∥ 2 F = 4 Tr (θ -θ)Σ 2 p (θ -θ) ⊤ .Putting these results together, we can upper bound the expected gradient norm for any parameter values ∥θ∥ F ≤ D as E ∥g(θ)∥ 2 F = 4 Tr(Σ p )N q (ã) + 4 Tr(Σ p ) Tr (θ -θ)Σ q (θ -θ) ⊤ ≤ 4 Tr(Σ p )N q (ã) + 4 Tr(Σ p )λ max (Σ q )∥θ -θ∥ 2 F ≤ 4 Tr(Σ p )N q (ã) + 16 Tr(Σ p )λ max (Σ q )D 2 .Using our results for the objective's strong convexity and expected gradient norm, we can now invoke a standard SGD convergence result. Following Theorem 6.3 from Bubeck et al. [5] , if we make T updates with the specified step size, we have the following upper bound on the expected objective value:EQUATIONNext, we prove a simple consequence of this result, which is that the rate from Theorem 1 applies to the original objective L reg (θ) when we assume an unbiased noisy oracle. Corollary 1. Following the setup from Theorem 1, if the noisy oracle is unbiased, then the averaged iterate at step T satisfiesE[L reg ( θT )] ≤ 4 Tr(Σ p )N q (ã) + 16 Tr(Σ p )λ max (Σ q )D 2 λ min (Σ p )(T + 1) .If the oracle is also noiseless, or if we assume access to the exact oracle ã(b) = a(b), then the averaged iterate at step T satisfiesEQUATIONProof. The first result follows from combining Theorem 1 with the relationship between L reg (θ) and Lreg (θ): if we assume the noisy oracle is unbiased, or B(ã) = 0, then we have L reg (θ) = Lreg (θ) -N(ã), which implies the first inequality. The second inequality follows from setting N q (ã) = 0 in the upper bound.Next, we provide a proof for Eq. ( 13) regarding the connection between regression-and objectivebased amortization, which is the following:α 2 L reg (θ) ≤ L obj (θ) -L * obj ≤ β 2 L reg (θ).Proof. This result relies on well known properties from convex optimization [41] . Consider a fixed context variable b ∈ B. Strong convexity with α > 0 means that for all a, a ′ ∈ A we have:h(a; b) ≥ h(a ′ ; b) + (a -a ′ ) ⊤ ∇ a h(a ′ ; b) + α 2 ∥a -a ′ ∥ 2 .Considering the optimal value a ′ = a(b), this simplifies to:EQUATIONNext, smoothness with β > 0 means that for all a, a ′ ∈ R d we have:∥∇ a h(a; b) -∇ a h(a ′ ; b)∥ 2 ≤ β∥a -a ′ ∥.Lemma 3.4 in Bubeck et al. [5] shows that β-smoothness also implies the following:h(a; b) -h(a ′ ; b) -(a -a ′ ) ⊤ ∇ a h(a ′ ; b) ≤ β 2 ∥a -a ′ ∥ 2 .Considering the optimal value a ′ = a(b), this simplifies to:EQUATION)The bounds in Eq. ( 13) follow from substituting a for predictions from a model a(b; θ), and considering Eq. ( 17) and Eq. ( 18) in expectation across the distribution p(b).Before proving our main claim for datamodels in Proposition 1, we first prove a more general result. This version considers an arbitrary symmetric distribution p(T ) over datasets, rather than the specific distribution parameterized by q ∈ (0, 1) used in Proposition 1. By a symmetric distribution, we mean one that satifies p(T) = p(T ′ ) whenever |T | = |T ′ |.As setup for our derivation, notice that when we assume a symmetric distribution p(T ), we can define the following probabilities that are identical for all indices i, j ∈ [n]:p 1 ≡ Pr(i ∈ T ) for i ∈ [n] p 2 ≡ Pr(i, j ∈ T ) for i ̸ = j p 3 ≡ Pr(i ∈ T, j / ∈ T ) for i ̸ = j.Note that we have p 1 = p 2 + p 3 . For example, given a uniform distribution p(T ) over subsets with size |T | = k, which is used by Ilyas et al. [47] , we have:p 1 = n-1 k-1 n k = k n p 2 = n-2 k-1 n k = k(k -1) n(n -1) p 3 = n-2 k-1 n k = k(n -k) n(n -1).Alternatively, if we have each data point included independently with probability q ∈ (0, 1), which is used by Saunshi et al. [85] and Proposition 1, then we have p 1 = q, p 2 = q 2 , and p 3 = q(1 -q).Our more general claim about datamodels with a symmetric distribution p(T ) is the following. Lemma 1. Given a symmetric distribution p(T ), the data attribution scores defined byζ(x) ≡ arg min a∈R n+1 E p(T )   a 0 + i∈T a i -v x (D T ) 2   can be expressed as the expectation ζ(z i , x) = E[c i (T )v x (D T )],where we define the weighting function c i (T ) as follows:c i (T ) ≡ 1 p 3 1(i ∈ T ) - p 2 -p 2 1 p 3 + n(p 2 -p 2 1 ) |T | - p 1 p 3 p 3 + n(p 2 -p 2 1 ).Proof. We can write the problem's partial derivatives as follows, both for the intercept term a 0 and the coefficients a i for i ∈ [n]:∂ ∂a 0 E a 0 + a ⊤ 1 T -v x (D T ) 2 = 2 a 0 + E 1 ⊤ T a -E [v x (D T )] = 2 a 0 + p 1 1 ⊤ n a -E [v x (D T )] ∂ ∂a i E a 0 + a ⊤ 1 T -v x (D T ) 2 = 2 E 1(i ∈ T )(1 ⊤ T a + a 0 ) -E [1(i ∈ T )v x (D T )] = 2 p 1 a 0 + p 2 1 ⊤ n a + (p 1 -p 2 )a i -p 1 E [v x (D T ) | i ∈ T ] .We use the shorthand notation v ∈ R n+1 for a vector with entries v0= E[v x (D T )] and vi = E[v x (D T ) | i ∈ T ] for i ∈ [n].Combining this with the partial derivatives, we can derive an analytic solution a * ∈ R n+1 by setting the derivative to zero,1 p 1 1 ⊤ n p 1 1 n p 2 1 n 1 ⊤ n + p 3 I n a * - 1 0 0 p 1 I n v = 0,which yields the following equation for the solution:a * = 1 p 1 1 ⊤ n p 1 1 n p 2 1 n 1 ⊤ n + p 3 I n -1 1 0 0 p 1 I n v.To find the required matrix inverse, we can combine the Sherman-Morrison formula with the formula for block matrix inversion. The formula for block matrix inversion is the following [79] :A B C D -1 = A -1 + A -1 B(D -CA -1 B) -1 CA -1 -A -1 B(D -CA -1 B) -1 -(D -CA -1 B) -1 CA -1 (D -CA -1 B) -1 .We are mainly interested in the lower rows of this matrix because we do not use the learned intercept term. For the lower right matrix, we have the following:(D -CA -1 B) -1 = p 2 1 n 1 ⊤ n + p 3 I n -p 2 1 1 n 1 ⊤ n -1 = (p 2 -p 2 1 )1 n 1 ⊤ n + p 3 I n -1 = 1 p 3 I n - p 2 -p 2 1 p 3 (p 3 + n(p 2 -p 2 1))1 n 1 ⊤ n .Next, for the lower left vector, we have the following:-(D -CA -1 B) -1 CA -1 = - 1 p 3 I n - p 2 -p 2 1 p 3 (p 3 + n(p 2 -p 2 1))1 n 1 ⊤ n p 1 1 n = - p 1 p 3 1 n + np 1 p 3 p 2 -p 2 1 p 3 + n(p 2 -p 2 1)1 n = - p 1 p 3 + n(p 2 -p 2 1 )1 n .This yields the following solution for the optimal coefficients a * i with i ∈ [n]:a * i = p 1 p 3 vi - p 1 (p 2 -p 2 1 ) p 3 (p 3 + n(p 2 -p 2 1 )) j∈[n] vj - p 1 p 3 + n(p 2 -p 2 1 )v0 .Based on this solution, we can design a function c i (T ) so that we have the expectationE[c i (T )v x (D T )] = a * i .We define the function as follows,c i (T ) ≡ 1 p 3 1(i ∈ T ) - p 2 -p 2 1 p 3 + n(p 2 -p 2 1 ) |T | - p 1 p 3 p 3 + n(p 2 -p 2 1 ), and it can be verified that this satisfies the required expectation.A similar derivation to Lemma 1 is possible when we omit an intercept term in the datamodels optimization problem, but we do not show the result here.Finally, we prove the special case of Lemma 1 considered in Proposition 1. Proposition 1. Given a subset distribution that includes each data point z i ∈ D with probability q ∈ (0, 1), the data attribution scores defined byζ(x) ≡ arg min a∈R n+1 E T   a 0 + i∈T a i -v x (D T ) 2  can be expressed as the following expectation:ζ(z i , x) = T ⊆[n]\i u(|T |) (v x (D T ∪ {z i }) -v x (D T )) .Proof. Consider the weighting function c i (T ) introduced in the proof for Lemma 1. In this case, we can use the fact that p 2 = p 2 1 and write the weighting function as follows:c i (T ) = 1 p 3 (1(i ∈ T ) -p 1 ) = 1-p1 p3 i ∈ T -p1 p3 i / ∈ T.Using the fact that p(T ) = q |T | (1 -q) n-|T | and the coefficient values (1 -p 1 )/p 3 = q -1 and -p 1 /p 3 = -(1 -q) -1 , we arrive at the result in Proposition 1.This section describes the statistical estimators used for each XML task, which are briefly introduced in Section 4. Each estimator serves as a noisy oracle for the given task, which allows us to train amortized models with inexpensive supervision. Shapley values. We use three statistical estimators for Shapley value feature attributions. The simplest is permutation sampling, where we average each feature's contribution across a set of sampled orderings [6, 96] . For this approach, we use ρ to denote a permutation of the indices [d], where we let ρ(i) ⊆ [d] \ {i} denote the elements appearing before i in the permutation. The Shapley value can be understood as the marginal contribution f (x ρ(i)∪{i} ) -f (x ρ(i) ) averaged across all possible permutations [89] , so our first estimator involves sampling k permutations ρ 1 , . . . , ρ k , and then calculating the following average for each feature:EQUATIONNext, we consider two estimators based on the Shapley value's least squares view. Recall that the Shapley value is the solution to the following problem (where we discard the intercept term),EQUATION)where we use a weighting kernel defined as µ -1 (S) = d |S| |S|(d -|S|) [7] . The first estimator we use based on this perspective is KernelSHAP [68] , which solves this problem using k subsets S j ⊆ [d] sampled according to µ(S). The problem is convex but constrained due to the weighting terms µ([d]) = µ(∅) = ∞, so it must be solved via the KKT conditions; we refer readers to Covert and Lee [16] for the closed-form solution. We also consider the SGD-Shapley approach from Simon and Vincent [92] : rather than solving Eq. ( 20) exactly given sampled subsets, this approach solves the problem iteratively with projected stochastic gradient descent. Unlike the other estimators, SGD-Shapley has been shown to have non-negligible bias [10] . We used a custom implementation for SGD-Shapley, and we used an open-source implementation of permutation sampling and KernelSHAP. 9Banzhaf values. When calculating Banzhaf value feature attributions, we adapt the maximum sample re-use (MSR) estimator from Wang and Jia [103] , which was originally used for data valuation but is equally applicable to feature attribution. For a single example x ∈ X , we generate predictions using k subsets S j ⊆ [d] sampled uniformly at random. We then split the subsets into those that include or exclude each feature i ∈ [d], and we estimate the Banzhaf value as follows:EQUATION)The re-use of samples across all features makes this estimator more efficient than independent Monte Carlo estimates [103] , and re-using samples in this way is simpler for Banzhaf values than for other methods like Shapley values [20, 58] . We used a custom implementation for this approach.LIME. Similar to KernelSHAP, the most popular estimator for LIME feature attributions is based on approximately solving its weighted least squares problem. Following the problem formulation in Section 4.2, which we simplify by omitting the regularization term Ω, we sample k subsets S j ⊆ [d] uniformly at random and solve the following importance sampling version of the objective: φ1 (x), . . . , φd (x) = arg minEQUATION)When doing so, we discard the intercept term, we ensure that k is large enough to avoid singular matrix inversion, and we use the default weighting kernel π(S) for images in the official LIME implementation. 10 We opt to sample subsets uniformly rather than according to π(S) because this approach is used in the official implementation, and because the weighting kernel is similar to sampling subsets uniformly at random [66] . We used a custom implementation of this approach.Data valuation. For the various data valuation methods discussed in Section 4.3, we use the simplest unbiased approximation, which is a Monte Carlo estimator that averages the performance difference across a set of sampled datasets. Given a labeled data point z, we sample k datasets D j ⊆ D from the appropriate distribution and calculate the following empirical average:EQUATION)This is similar to the TMC algorithm from Ghorbani and Zou [33] , but we do not employ truncation; our approach can be used with truncation, but the noisy labels and valuation model predictions would both become biased. Similar to previous works, we also adopt a minimum subset cardinality to avoid training models with an insufficient number of data points. We implemented this approach using the OpenDataVal package [53] . More sophisticated estimators are available, like those that exploit stratification [106] , re-use samples across all data points [103] , or assume sparse attribution scores [52] , but we leave exploration of these estimators to future work.Model architectures. For the feature attribution experiments, we used the ViT-Base architecture [24] for the classifier, and we used a modified version of the architecture for the amortized attribution model: following the approach from Covert et al. [18] , we added an extra self-attention layer and three fully-connected layers that operate on each token, so that the output contains an attribution score for each class. When estimating feature attributions, we make predictions with subsets of patches by setting the held-out patches to zero, and the classifier was fine-tuned with random masking to accommodate missing patches [17, 48, 50, 18] .For the data valuation experiments, we used a FCN with two hidden layers of size 128 for the tabular datasets, and we used a ResNet-18 architecture [43] for CIFAR-10. Valuation scores are defined for labeled examples z = (x, y), so the valuation model must account for both the features x and the class y when making predictions; rather than passing y as a model input, our architecture makes predictions simultaneously for all classes, and we only use the estimate for the relevant class. When generating noisy labels using the TMC estimator [33] , we trained a logistic regression model on raw input features for the tabular datasets, and for CIFAR-10 we trained a logistic regression on pre-trained ResNet-50 features whose dimensionality was reduced with PCA.Hyperparameters. For the feature attribution experiments, we optimized the models using the AdamW optimizer [67] with a linearly decaying learning rate schedule. The maximum learning rate was tuned using the validation loss, we trained for up to 100 epochs, and we selected the best model based on the validation loss. Due to the small scale of the noisy labels for Banzhaf and LIME feature attributions, we re-scaled the labels during training to improve the model's stability (see Appendix G.2), but this re-scaling was not necessary for Shapley values.For the data valuation experiments, we optimized the models using Adam [55] with a cosine learning rate schedule. The maximum learning rate, the number of training epochs and the best model from the training run were determined using the validation loss. Due to the small scale of the noisy labels, we found it helpful to re-scale them during training to have standard deviation approximately equal to 1. The cost of performing multiple training runs is negligible compared to generating the noisy training labels, so it is not reflected in our plots comparing amortization to the Monte Carlo estimator (e.g., Figure 4 ).Pretraining. We found that training the amortized models was faster and more stable when we initialized using pretrained architectures. For the feature attribution experiments, we initialized the model using the existing ViT-Base classifier weights. For data valuation, we initialized from a classifier trained using the entire dataset, where we used a FCN for the tabular datasets and a ResNet-18 for CIFAR-10. When adapting these models to their respective amortization tasks, the feature attribution models had freshly initialized output layers (one self-attention and three fully-connected), and the data valuation models had identical output layers that we re-initialized to zero.Validation loss. For the feature attribution experiments, we calculated the validation loss using independent estimates generated with same estimator and the same number of samples used for the noisy training labels. For the data valuation experiments, we implemented the validation loss using independent Monte Carlo estimates of the data valuation scores. Our independent estimates use only 10 samples for the tabular datasets, and just one sample for CIFAR-10. Amortization provides a significant benefit over per-example calculations even after accounting for the cost of the validation loss.Noisy oracle details. The estimators used for each task are described in detail in Appendix E. For the feature attribution experiments, the permutation sampling and KernelSHAP estimators only require the number of samples as hyperparameters, and we tested multiple values in our experiments (e.g., Figure 3 ). For SGD-Shapley, we tuned the algorithm in several ways to improve its performance: we used a constant rather than decaying learning rate, we took a uniform average over iterates, we calculated gradients using multiple subsets, and we used the paired sampling trick from Covert and Lee [16] to reduce the gradient variance. We also tuned the learning rate to a value that did not cause divergence for any examples (5e-4). As described in Appendix G.1, we observed that SGD-Shapley often had the lowest squared error among the three Shapley value estimators (Figure 11 ), but that it did not lead to successful amortization due to its non-negligible bias (Figure 10 ).For the data valuation experiments, models were trained on each subsampled dataset by fitting a logistic regression model, either on raw input features for the tabular datasets or on pre-trained ResNet-50 features for CIFAR-10. For the Data Shapley experiments in Section 5.2 with tabular datasets, we sampled datasets without replacement, we used a minimum cardinality of 5 points and a maximum cardinality equal to the number of points n. For the Distributional Data Shapley experiments in Section 5.3 with CIFAR-10, we sampled datasets with replacement, and we used a minimum cardinality of 100 and maximum cardinality of 1000.Ground truth. In all our experiments, the ground truth is obtained by running an existing perinstance estimator for a large number of samples, and we use enough samples to ensure that it has approximately converged to the exact result. For the Shapley value feature attribution experiments, we run KernelSHAP for 1M samples. For Banzhaf values, we run the MSR estimator for 1M samples.For LIME, we run the least squares estimator for 1M samples. For the data valuation experiments, we run the Monte Carlo estimator for 10K samples. These estimators are costly to run for a large number of samples, so we do so for only a small portion of the dataset: we calculate the ground truth for 100 examples for the feature attribution experiments. For the data valuation experiments, we use 250 examples for the tabular datasets and 500 for CIFAR-10.Metrics. The performance metrics used throughout our experiments are related to the estimation accuracy, which we evaluate using our ground truth values. Other works have evaluated Shapley value feature attributions against other methods [50, 18] or used data valuation scores in a range of downstream tasks [53] , but our focus is on efficient and accurate estimation. For the feature attribution experiments, we used the squared error distance, which we calculate across all attributions; we used Pearson and Spearman correlation, which are calculated individually for each flattened vector of attribution scores and then averaged across data points; and we used sign agreement, which is averaged across all attributions. For the data valuation experiments, we used squared error distance, which we normalized so the mean ground truth valuation score has error equal to 1 (i.e., we report the error divided by the variance in ground truth scores); and we used Pearson correlation, Spearman correlation and sign agreement, which were calculated using the vectors of estimated and ground truth valuation scores.Datasets. We used publicly available, open-source datasets for our experiments. For the feature attribution experiments, we used the ImageNette dataset, a natural image dataset consisting of ten ImageNet classes [46, 23] . For the data valuation experiments, we used two tabular datasets from the UCI repository: the MiniBooNE particle classification dataset [84] and the adult census income classification dataset [25] . We obtained these using the OpenDataVal package [53] FLOPs profiling. When profiling compute for our feature attribution experiments (see Figure 3 ), we first measured the following constants using the DeepSpeed FLOPs profiler, all using the ViT-Base architecture:• The classifier's forward pass requires 17,563,067,904 FLOPs ≈ 17.6 GFLOPs per prediction.• The amortized model's forward pass requires 21,335,153,664 FLOPs ≈ 21.3 GFLOPs per prediction, due to the model having an extra transformer block for fine-tuning.• The amortized model has 104,730,000 ≈ 105M trainable parameters.Next, for the FLOPs comparison shown in Figure 3 we calculated total FLOPs for the two approaches as follows. For KernelSHAP, the total FLOPs = 17.6 GFLOPs × (number of subset samples) × (number of training datapoints). For the amortized models, the total FLOPs are the sum of multiple terms:1 In other words, the amount of compute for training an amortized model for up to 50 epochs translates to obtaining predictions for just 182.5 subset samples per datapoint, which is not enough to meaningfully improve the estimates by running KernelSHAP for more iterations. This explains why the lines for amortized models appear almost vertical in Figure 3 when we compared amortization to KernelSHAP. Our procedure for calculating FLOPs is similar to other works that profile computation when training neural networks. Figure 10 shows an expanded version of Figure 3 , where we can see the benefits of amortization for KernelSHAP and permutation sampling across four metrics. We also see that SGD-Shapley does not lead to successful amortization, because the predictions are worse than the noisy labels across all four metrics. This is perhaps surprising, because Figure 11 shows that SGD-Shapley has the lowest squared error among the three noisy oracles. The crucial issue with SGD-Shapley is that its estimates are not unbiased (see Section 3), an issue that has been shown in prior work [10] .Finally, Figure 13 provides a comparison between stochastic amortization and FastSHAP [50, 18] , one of the main existing approaches to amortized Shapley value estimation. For stochastic amortization, we use our previous results with permutation sampling as the noisy oracle and five different numbers of samples. For FastSHAP, we train the same ViT-Base architecture following the approach from Covert et al. [18] , using 32 subset samples per gradient step. We monitor FastSHAP's estimation accuracy at the end of each epoch while training for a total of 100 epochs. We observe that FastSHAP and stochastic amortization achieve similar error at each total FLOPs budget, where both methods incur FLOPs from training the amortized model, stochastic amortization incurs FLOPs upfront when generating noisy labels, and FastSHAP incurs FLOPs during training while sampling subsets for each gradient step. FastSHAP is slightly more accurate at the largest computational budget, and both are significantly more accurate than per-sample estimation with KernelSHAP. Due to their similar performance, the main advantage of our approach is its simplicity (a standard regression objective rather than the custom FastSHAP objective), and the flexibility to use any unbiased estimator as the noisy oracle; future work may find that stochastic amortization is more effective with other noisy oracles that we did not try here. As we discussed in Section 4.2, Banzhaf value and LIME feature attributions are two XML tasks closely related to Shapley values that can be amortized in a similar fashion. Both are intractable to calculate exactly, but they have (approximately) unbiased estimators that can be used as noisy labels for stochastic amortization. Appendix E describes these estimators, namely the MSR estimator for Banzhaf values [103] and the least squares estimator for LIME [83] . Prior work has also shown that these approaches generate similar outputs [66] , although the standard computational approaches are quite different.In amortizing these methods, one trend we observed is that Banzhaf value and LIME feature attributions have a very different scale from Shapley values. Figure 14 shows that they not only have smaller norm, but that the distribution is concentrated at small magnitudes with a long tail of larger magnitudes. This is troublesome, because our theory is focused on the squared error ∥a(b; θ) -a(b)∥ 2 (see Section 3), which is dominated by the small number of examples with large magnitudes. As a result, an amortized attribution model can appear to train well by accurately estimating attributions with a large norm, and predicting the remaining attributions to be roughly zero; in doing so it may fail to provide the correct relative feature ordering for most examples, which is more important for practical usage of feature attributions. We plotted the L2 norm of the 100 ground truths for each feature attribution.The issue described above is precisely what occurs when we amortize Banzhaf value feature attributions, as shown in the top row of Figure 15 . We observe that the squared error from our predictions is lower than that of the noisy labels, which is consistent with our theory from Section 3 and our results for Shapley values (Appendix G.1). However, the amortized estimates perform worse than the noisy labels on the remaining metrics, particularly for the correlation scores that evaluate the relative feature ordering in each example's attributions (see details for the metrics in Appendix F).To alleviate this issue, we experimented with a per-label normalization that eliminates the long tail of large attribution norms that dominate training: we simply normalized each example's attributions for each class to have a norm of 1. The results are shown in the bottom row of Figure 15 , where we see that the predictions have higher correlation scores than the noisy labels. The squared error is higher for the amortized predictions, and the sign agreement is roughly the same. One issue with this heuristic is that the normalized noisy labels are biased: the normalized noisy label is not an unbiased estimator of the normalized exact label, because the normalization constant is not known a priori and must be calculated using the noisy label. However, the results show that amortization can to some extent denoise inexact labels even with this non-zero bias.Figure 16 shows similar results as Figure 15 but with LIME feature attributions. We observe the same improvement in squared error from amortization, and a similar degradation in the correlation metrics (Figure 16 top). We then apply the same per-label normalization trick, and we observe a similar modest improvement in the correlation metrics, at least for noisier settings of the least squares estimator (Figure 16 bottom).Overall, these results show that for the purpose of amortization, the consistency in scale of Shapley values is an unexpected advantage over Banzhaf values and LIME. This property is due in part to the Shapley value's efficiency axiom [89] , which guarantees that the Shapley values sum to the difference between the prediction with all features and no features [68] . In comparison, LIME and Banzhaf values focus on marginal contributions involving roughly half of the features, which in many cases are near zero when the prediction is saturated; for example, previous works have observed that forFor the data valuation experiments, our first additional result compares amortization to the Monte Carlo estimator when using the MiniBooNE and adult datasets with different numbers of data points.The results are shown in Figure 17 and Figure 18 , where we use datasets ranging from 1K to 10K points, and we see that the benefits of amortization increase with the size of the dataset. We also provide more detailed performance metrics for CIFAR-10: Figure 19 shows an expanded version of Figure 5 , and we observe that amortization improves upon the Monte Carlo estimator across all four performance metrics.Next, we consider the use of CIFAR-10 data valuation scores in two downstream tasks. First, we attempt to identify mislabeled examples, which we expect should have large negative valuation scores.Figure 20 shows how quickly each method identifies negative examples when we sort the scores from lowest to highest: our amortized estimates that train with just 5 samples provide the highest accuracy, outperforming Monte Carlo estimates that use as many as 100 samples. This result is consistent across different levels of label noise, where we randomly flip either 10%, 25% or 50% of the labels. Our proposed method aims to make many computationally challenging XML tasks feasible and scalable to large datasets. We expect our research to contribute to a better understanding of ML models, enhancing their transparency and explainability. Additionally, our research could help understand model fairness properties if our method is used to identify undesirable or unfair dependencies of ML models on input features related to protected attributes like race or gender. The potential downside is that that if users overly trust our method and rely solely on it for these tasks without careful usage, it could lead to misunderstanding of models and result in harmful outcomes.NeurIPS Paper Checklist Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.Datamodels [47] performs data attribution by fitting a linear regression model, but this is not amortization because the model cannot predict attributions for new data points (see Appendix B).N(ã | b) is equal to the trace of the conditional covariance Cov(ã | b), and N(ã) = Tr(E[Cov(ã | b)]).There are many ways to do so[17], e.g., we can set features to their mean or use masked self-attention.This is the default number of samples for 196 features in the official repository: https://github.com/ shap/shap.As we noted in the main text, although datamodels performs data attribution by fitting a linear regression model, it is not a form of amortization because the resulting model cannot predict attribution scores ζ(z, x) for new datapoints. The attribution scores are given by the model's coefficients, not its predictions.https://github.com/iancovert/shapley-regressionhttps://github.com/marcotcr/limeWe consulted the following resources: (i) https://www.lesswrong.com/posts/jJApGWG95495pYM7C/ how-to-measure-flop-s-for-neural-networks-empirically, (ii) https://www.lesswrong.com/ posts/fnjKpBoWJXcSDwhZk/what-s-the-backward-forward-flop-ratio-for-neural-networks, and (iii) https://www.adamcasson.com/posts/transformer-flops 


Decision: The reviewers were generally positive, and the authors did a good job addressing their concerns during the response period. Please make sure that the changes you promised to the reviewers are incorporated into the current version of the paper.

Reviews: review: paper_summary: paper_summary: This paper proposes a stochastic amortization framework for efficiently estimating feature attribution and data attribution values. The idea is to learn a parameterized model from noisy while unbiased samples of the value to be estimated. In comparison to naive Monte Carlo sampling, this amortized estimation is empirically more efficient when achieving similar estimation accuracy. The authors conducted experiments on a diverse set of feature attribution and data attribution problems.
strength_weakness: strength_weakness: - The theoretical analysis, especially Theorem 1, is not directly relevant to the main argument of this paper, i.e., stochastic amortization is more efficient than naive Monte Carlo. IMO this Theorem 1 is a bit distracting than being helpful for this paper. Furthermore, the theoretical result about the unbiasedness in lines 128 - 132 only concerns the expected loss, $\tilde{\mathcal{L}}_{reg}$, which doesn't directly imply the property of the empirically learned parameterized model. Overall, the advantage of stochastic amortization over Monte Carlo estimation is established empirically instead of theoretically, while the current presentation reads a bit misleading on this point.
- In most of the experiments, only relatively naive Monte Carlo estimators compared as baselines. However, for certain specific problems, more efficient estimators may be available. For example, in the context of Data Shapley, how does the stochastic amortization compare with Gradient Shapley proposed in Ghorbani and Zou [33], or the compressive-sensing-based method proposed in Jia et al. (2019).

References
- Jia et al. (2019) Towards Efficient Data Valuation Based on the Shapley Value. AISTATS 2019.
questions: questions: See Weaknesses.
limitations: limitations: None.
paper_summary: paper_summary: In this paper, the authors introduce a framework termed stochastic amortization that can accelerate computationally expensive explainable machine learning (XML) tasks by training models with noisy but unbiased labels. They provide theoretical analysis that shows that unbiased noisy labels allow learning the correct function. Empirically, the authors demonstrate the effectiveness of this approach on several XML tasks, including Shapley value feature attributions, Banzhaf values, LIME, and data valuation, achieving speedups over traditional per-example computation and improvements over simply using noisy labels.
strength_weakness: strength_weakness: While the concept of amortization for Shapley value prediction is not entirely new, this paper makes a good contribution through its use of noisy labels and the accompanying theoretical analysis. I don't find any major weaknesses that would significantly detract from the paper's value, and I recommend acceptance. Reflecting my limited familiarity with some aspects of the relevant literature, I set my confidence to 3.
questions: questions: - For applying stochastic amortization to a new XML task, how should practitioners determine an appropriate error level from a practical perspective?
- Regarding line 99, "while there are works that accelerate data attribution with datamodels, we are not aware of any that use amortization". Could datamodels themselves be considered a form of amortized optimization (linear regression model)?
- What guidelines can you provide for defining the architecture for amortization (e.g., ResNet) in practice? Did you observe significant differences with various architectures? Should the architecture be similar to the base model?
- In Figure 4 (right), why does the correlation decrease for MC as the number of data points increases (intuitively, I expected correlation to remain relatively flat for both MC and amortized approaches, assuming a constant number of samples per point)?
limitations: limitations: The authors adequately addressed the limitations of the proposed approach.
paper_summary: paper_summary: This paper presents a fast prediction explanation approach for machine unlearning models that approximates traditional explanation approaches by using an neural network. The network is trained with noisy labels (so called stochastic amortization) such that it learns to approximate the prediction explanations of various approaches, including Shaply, LIME, and others. Experiments covers multiple explanation approaches from feature explanation to data instance explanation, showing effectiveness of the approximation.
strength_weakness: strength_weakness: It is concerning to use neural networks to explaining the predictive behaviour of another machine learning model. In fact, how do we know whether the explanation itself is trustworthy? If the prediction explanation is given to an stakeholders, is it possible to distinguish the potential concern between target machine learning model's trustworthiness and the approximative explainer? This concern needs to be addressed further to let the proposed approach practically useful.

Partial experimental results show the proposed model may not be a good approximation approach; the results can be interpreted as the approximation fails to provide correct explanation to feature importance (Section 5.1). It causes me wonder if there is a practical value for the proposed approach if it is fast but not accurate.
questions: questions: How do you prove the proposed approximation is, in general, faithful to the ground-truth explanation results from the actual explanation models? 
Why would the proposed approach to be useful if the explanation is off from the ground-truth?
limitations: limitations: I don't see limitations. But, as the proposed approach is about XAI, it should be sound in terms of improving the trustworthiness of prediction instead of introducing another layer of uncertainty.
paper_summary: paper_summary: The paper proposes using stochastic amortization to speed up feature attribution and data attribution. This is applicable when the attribution technique is expensive to compute exactly (e.g. LIME), and when unbiased estimators of the attributions exist. Specifically, the paper proposes to, for a given machine learning model (e.g. Resnet) and a dataset of inputs (e.g. cifar) on which we want to compute attribution (e.g. Shapley feature values), obtain noisy estimates of the attributions (e.g. by sampling permutations), and then train a least square regression model to predict these noisy estimates from the input. The paper evaluates this method against the ground-truth attribution values for Shapley values on Imagenette,  data valuation shapley on adult census and MiniBooNE datasets, and Distributional data valuation on the CIFAR-10 dataset.
strength_weakness: strength_weakness: - The related work section could benefit with a more direct and concrete comparison with prior work. For example, the paper says that "while there are works that accelerate data attribution with datamodels, we are not aware of any that use amortization". It is not clear to me in what "amortization" mean here exactly --- in what way does your work perform amortization that datamodels do not?
- In addition, the experimental results lack baselines that compare with prior amortization work. The related work section describes a number of prior works (citations 50, 86, 18, 14) that all seem to be about computing Shapley values with amortization. 
- Figures are missing error bars
questions: questions: - S5.1: would be interested to look at fig 3 right with a log scale on the y axis. Amortization seems to benefit with increasing dataset size, and I wonder what the scaling looks like
- Related to that, a key hyperparameter in your amortization set up is to choose between spending compute to get fewer but more accurate labels vs more noisy labels. Is there a way of determining an optimal trade-off?
- How does the paper's method compare with prior works? This is a crux for me.
limitations: limitations: The authors adequately discussed technical limitations of their amortization approach in the appendix.

Comments: title: Author response
Author response
Author response
Thanks for the response
Response
Rebuttal by Authors (cont.)
comment: Thanks for your answers. While my concerns of faithfulness of the NN approximation are not fully address, I think the paper may add value to the literature. Hence, I bumped my rating.
Thank you for the clarification. I encourage the authors to make that connection with prior work clear in the paper (i.e. that the key observation is that stochastic amortization can tolerate more noise than previous methods assumed). I believe, already, a lot of the empirical tests in the current is aimed at studying how much noise can be tolerated.

Because you have committed to a more direct comparison with prior work, and have shown empirical tests of prior baselines, I will raise my score accordingly.
Thanks very much for reading our rebuttal and for your response. We'll make sure to add that clarification about datamodels to the paper, especially since hGpm had the same question.
Thanks very much for reading our rebuttal and for your response. We understand the confusion about Theorem 1, thanks for pointing out that sentence in the abstract. We'll make sure to clarify that and add a remark after Theorem 1 to avoid confusion for other readers.
Thanks for your response, we hope we resolved the concerns mentioned in your review. As for whether our technique is a simplification of prior works – you could describe it that way, that's basically correct regarding [86, 14]: the fact that amortization tolerates high noise levels enables our more practical approach, although the mathematical justification makes it perhaps less conceptually simple than training with exact targets. As for FastSHAP, our approach is simpler, but it’s not exactly a simplification of [50, 18] in the sense that the objective is derived differently. (FastSHAP is based on a weighted least squares view of SVs, whereas our approach is based on SVs being the expectation of a Monte Carlo estimator.)
Thank you for the detailed response. Just an additional clarification about the comparison with prior work - would you agree that the technique you use is a simplification of previous techniques, which is enabled by the observation that stochastic amortization can tolerate more noise than previous methods assumed?
I appreciate the authors' response. I would encourage the authors to at least add a remark after Theorem 1 clarifying the gap between Theorem 1 and the main claims of the advantage of the proposed method (which are evaluated empirically). This will help avoid potential confusions by readers like myself. 

PS: Part of the confusion comes from the paper abstract: "Through theoretical analysis of the label noise and experiments with various models and datasets, we show that this approach significantly accelerates several feature attribution and data valuation methods", which sets up an expectation that Theorem 1 is directly about the advantage of the proposed method.

Overall I think this is a solid paper worth publication at NeurIPS. And this assessment has already been reflected by my previous rating.
Thank you for the comments. I read the author's response and other reviewers' comments. The clarification on the datamodels was helpful, and it would be great if you could add this explicit answer somewhere (I have read Appendix B but was still unsure). Overall, I am satisfied that the authors addressed all my questions.
> Related to that, a key hyperparameter in your amortization set up is to choose between spending compute to get fewer but more accurate labels vs more noisy labels. Is there a way of determining an optimal trade-off?

That’s a good question, and one that we didn’t try to answer conclusively here. Our goal was to show that training with noisy labels works, that it tolerates surprisingly high noise levels, and that it’s applicable to many XML tasks (including both feature attribution and data valuation). Investigating the quality-quantity tradeoff in noisy labels is an interesting question and a natural subject for future work, but addressing it thoroughly would be a bit too much work for our current paper. We are happy to mention this in the conclusion and provide some preliminary thoughts on how one should approach it.

Intuitively, the optimal point can’t be at either extreme: a very small number of exact labels would be unlearnable, as would a large number of extremely noisy labels. Finding the optimal point requires reasoning about how well neural networks learn as a function of the dataset size and label noise, which is dataset- and model-dependent and likely best to approach empirically. One idea is to 1) measure the estimation error as a function of dataset size and number of Monte Carlo samples (i.e., run stochastic amortization under a range of settings), 2) fit the observed errors using a simple analytic function (in the style of recent scaling laws), and 3) use this to estimate the optimal trade-off (sort of like Chinchilla’s compute-optimal training). This approach seems reasonable and straightforward, but it would require many training runs and is beyond the scope of our current paper.