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

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.

Strengths

- The paper clearly explains its motivations and approach - Using the amortized model improves performance over the initial noisy labels used for training, clearly showing the benefits of using an amortized method - The paper provides some theoretical arguments for using unbiased estimators.

Weaknesses

- 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

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

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately discussed technical limitations of their amortization approach in the appendix.

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.

Strengths

Using network to approximate prediction explanation of various explanation methods seem an interesting idea. Using noisy labels to train the approximation model is also well justified with corresponding proofs. Experiments covers the various scenarios from feature to instance based explanation approaches.

Weaknesses

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

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?

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

3

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.

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.

Strengths

- The paper is well-written, with clear motivations and a thorough discussion of related work. The technical details appear sound and well-presented. - The proposed approach offers a simple yet effective framework for accelerating computationally expensive XML methods, potentially enabling their application to large-scale datasets. - The authors provide code (with good documentation) and detailed experimental information in the Appendix for reproducability. - Limitations of the proposed approach are well discussed, demonstrating a balanced perspective. - Comprehensive experiments across multiple XML tasks (Shapley values, Banzhaf values, LIME, and data valuation) consistently demonstrate the benefits of stochastic amortization over per-example computation.

Weaknesses

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

- 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)?

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors adequately addressed the limitations of the proposed approach.

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.

Strengths

- This paper tackles the computational efficiency issue in a set of feature attribution and data attribution problems, which is an important problem that could have significant impact to the area of explainable machine learning (XML). - This paper abstracts the paradigm of learning a parameterized function to estimate the attribution values in some XML methods with a more general amortized estimation framework, making it possible to extend this approach to more XML methods that conventionally rely on naive Monte Carlo estimators. - The authors conducted comprehensive experiments in diverse settings to demonstrate the effectiveness of the stochastic amortization.

Weaknesses

- 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

See Weaknesses.

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

None.

Rebuttal by Authors (cont.)

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

Response

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.