Decomposing and Editing Predictions by Modeling Model Computation

Abstract: How does the internal computation of a machine learning model transform inputs into predictions? In this paper, we introduce a task called component modeling that aims to address this question. The goal of component modeling is to decompose an ML model's prediction in terms of its components -- simple functions (e.g., convolution filters, attention heads) that are the "building blocks" of model computation. We focus on a special case of this task, component attribution, where the goal is to estimate the counterfactual impact of individual components on a given prediction. We then present COAR, a scalable algorithm for estimating component attributions; we demonstrate its effectiveness across models, datasets, and modalities. Finally, we show that component attributions estimated with COAR directly enable model editing across five tasks, namely: fixing model errors, ``forgetting'' specific classes, boosting subpopulation robustness, localizing backdoor attacks, and improving robustness to typographic attacks. We provide code for COAR at https://github.com/MadryLab/modelcomponents .

• The paper introduces the Coar algorithm, a regression-based method that efficiently assigns attribution scores to individual model components.
• It demonstrates robust performance against methods like Neuron Conductance across diverse architectures and datasets, including ResNet and Vision Transformers.
• The approach enables practical model editing by facilitating precise interventions in model behavior without the need for extensive retraining.

Unveiling Model Predictions via Component Attribution with Coar

Introduction to Component Modeling

In traditional machine learning applications, the opaque nature of model computations, especially within large models, presents significant challenges in understanding and interpreting a model's predictions. Addressing this issue, the concept of component modeling presents itself as a method to decompose model predictions in accordance with individual model components, such as attention heads or convolutional filters.

Task Definition and the Coar Framework

Component modeling aims at building interpretable estimators or models that can predict the output of a model after hypothetical modifications (ablations) to its components. This task helps in both understanding the underlying functioning of model components and leveraging this understanding for practical applications like model editing.

One pivotal approach within component modeling is known as component attribution. This method assigns scores to each model component, predicting the effect of ablating (modifying or removing) a set of components based on the sum of their scores. The Coar algorithm, a key contribution of the discussed work, is a scalable method to compute these attributions efficiently across various architectures and data modalities.

Methodology of Coar

Coar stands for Component Attribution via Regression. It simplifies the estimation of component attributions by treating the problem as a regression analysis where the contributive weights of components toward the model's prediction are directly learned from data. To generate necessary training data, Coar ablates components at random and observes changes in predictions, thus constructing a dataset of component outcomes that it uses to learn attributions.

Empirical Validation

The effectiveness of Coar is demonstrated across models trained on different datasets (such as ImageNet and CIFAR-10) and model architectures like ResNet and Vision Transformers. Results indicate that:

• Coar calculations efficiently predict the impact of component ablations.
• Component attributions offer insights into the contribution of specific parts of the model towards making predictions.
• Coar consistently outperforms existing methods such as Neuron Conductance and Internal Influence in predicting these impacts, demonstrating robust and scalable performance.

Practical Applications and Model Editing

Moving beyond interpretation, Coar's computed attributions facilitate model editing, where direct interventions in a model's behavior are executed without retraining. Through various tasks like fixing specific predictive errors, enhancing robustness against types of noise, and "forgetting" undesired biases, component attributions by Coar enable precise and target-specific modifications to the model's behavior.

Concluding Thoughts

The combination of component modeling and the Coar algorithm represents a significant advancement in interpretable machine learning techniques. By allowing an analytical peek into the often opaque process of model predictions and providing a tool to manipulate these predictions consciously and purposefully, this approach paves the way for more understandable, fair, and controllable AI systems. Future work might explore extending the linear regression-based attributions to non-linear models for enhanced accuracy and investigating other forms of component interventions beyond zeroing out weights, potentially allowing even finer control over model behavior.