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

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.

• The paper demonstrates that training with noisy, unbiased labels enables efficient approximation of complex attribution methods like Shapley values and Data Shapley.
• The paper shows that amortized models reduce computation time while maintaining robust estimation accuracy across multiple metrics and data domains.
• The paper outlines potential extensions for applying stochastic amortization to broader explainable AI tasks, enabling scalable, real-time analytics.

Amortized Approaches for Efficient Feature and Data Attribution in Machine Learning

The research paper "Stochastic Amortization: A Unified Approach to Accelerate Feature and Data Attribution" proposes a framework for improving the efficiency of various tasks in explainable machine learning (XML), particularly feature attribution and data valuation, by employing a strategy called stochastic amortization. This approach leverages amortized models to predict computationally expensive outputs directly, significantly reducing inference time by replacing costly per-instance calculations with a one-time model training and a subsequent fast inference process.

Theoretical Framework and Methodology

The central premise of the paper is the use of noisy labels to train amortized models. These noisy labels are derived from statistical estimators or approximations that are less resource-intensive than computing the exact outputs for every single data point. The authors establish that, under certain conditions, training with these noisy, unbiased labels can indeed lead to effective models that approximate the desired attributions or valuations accurately.

The theoretical contribution of the paper includes proving that if the noise introduced by the labels is unbiased, the estimators used as labels can still lead to robust, reliable amortized models. The paper demonstrates that even with high variance, the model can learn effectively, albeit with potentially slower convergence rates.

Application to Explainable Machine Learning

1. Shapley Value Feature Attribution: The paper details the use of amortization for Shapley values, a popular feature attribution method. The computational burden associated with exact Shapley value calculation, due to its combinatorial nature, is alleviated using amortized models trained on noisy estimations from permutation sampling or Kernel SHAP methods.
2. Alternative Attribution Methods: The study extends the amortization framework to Banzhaf values and LIME attributions, deriving efficient training targets from unbiased estimators.
3. Data Valuation: The framework is applied to Data Shapley methods to assess the impact of individual training data points on overall model performance. Using amortized models, the authors propose significant reductions in computation compared to existing Monte Carlo-based methods.
4. General Extensions: The paper briefly discusses potential extensions of their framework to datamodels, indicating broader applicability to other data attribution tasks.

Experiments and Results

Empirically, the paper substantiates its claims with experiments conducted on several data domains, including image and tabular data. The use of amortized models led to substantial compute savings, showing improved estimation accuracy across multiple metrics, including squared error and correlation with ground truth, even when using significantly fewer computational resources. A notable experiment highlighted the efficiency gains where training with amortized models using noisy labels was comparable or superior in quality to exhaustive traditional methods but at a fraction of the cost.

Implications and Future Directions

The implications of this work are multifaceted. Practically, it suggests that many existing XML tasks can be conducted more efficiently with negligible loss in explanatory power by using stochastic amortization. This has significant implications for real-world applications where interpretability needs to be balanced with computational feasibility, such as in large-scale AI systems and real-time analytics.

Theoretically, this research invites further exploration into domains where noisy labels could be effectively employed in training without sacrificing model accuracy. It also opens the avenue for the development of better estimators and more robust training frameworks that can handle high label noise.

The paper concludes by suggesting several potential research directions, including scaling the approach to larger datasets, refining estimation techniques for data valuation, and further exploring the integration of amortization with other data influence techniques currently dependent on exact computations or approximations.

In summary, "Stochastic Amortization: A Unified Approach to Accelerate Feature and Data Attribution" provides a comprehensive framework that effectively marries the needs of computational efficiency with fidelity in explainable AI, marking a promising advancement in machine learning interpretability methods.