Who Gets Credit or Blame? Attributing Accountability in Modern AI Systems

Technical Summary: Who Gets Credit or Blame? Attributing Accountability in Modern AI Systems

Problem Formulation
Modern AI systems are developed through sequential stages—such as pretraining, fine-tuning, and alignment—where each stage updates the model based on specific data and optimization dynamics. This modular structure complicates accountability: when a deployed model exhibits specific behaviors (beneficial, harmful, or surprising), it is unclear which development stage is responsible. Existing methods for causal responsibility analysis are often limited to small-scale, discrete settings; learning dynamics research is typically descriptive rather than attributive; and data attribution methods often assume an "average model" or simple optimizers, failing to account for complex, model-specific training processes involving learning rate schedules, momentum, and weight decay. The authors formulate the accountability attribution problem as tracing model behavior back to specific stages of the development process.

Methodology
The authors propose a general framework for accountability attribution based on the potential outcomes formalism. The core question is counterfactual: How would the model’s behavior have changed if the updates from a particular stage had not occurred?

1. Causal Effect Definition: For a development process with parameters evolving from $\theta_0$ to $\theta_K$, a stage $S$ is defined as a set of update steps. The treatment $T^S$ indicates whether these steps occurred ($T^S=1$) or were skipped ($T^S=0$). The accountability of stage $S$ is defined as the causal effect $\tau_{K,S}$ on the final model performance $\gamma(x, \theta_K)$:
   $$\tau_{K,S} = \gamma(x, \theta_K) - \gamma(x, \theta_K(0^S))$$
   where $\theta_K(0^S)$ is the counterf final parameter state if stage $S$ were skipped.

2. AA-Score Estimator: To estimate this causal effect without retraining the model, the authors derive efficient estimators using Taylor approximations and state trajectory interpolation.

   • They define a propagator matrix $M_k$ that captures the influence of optimization dynamics (including SGD with momentum and weight decay) on the state trajectory.
   • The effect of a single step $t$ on the final state is propagated from the initial effect at $t+1$ using these matrices.
   • The effect on model behavior (the Accountability Attribution Score or AA-Score) is computed as a dot product between the gradient of the performance function at the final parameters and the estimated parameter change caused by the stage.
   • The total effect of a stage is the sum of the effects of its constituent steps, justified by the linearity of the first-order approximation.

3. Error Analysis: The authors provide an error bound for the estimator, showing that accuracy depends on the magnitude of the initial state effects, the Lipschitz continuity of the loss Hessian, and the stability of the optimization dynamics. Attribution is most reliable for late-stage interventions or stable regimes.

Key Contributions

• Problem Definition: Formalizing accountability attribution as tracing model behavior to training stages via counterfactual queries.
• Framework: A general framework based on potential outcomes that accounts for complex optimization dynamics (learning rates, momentum, weight decay).
• Efficient Estimators: Derivation of AA-Score estimators that quantify stage effects without retraining, scaling to deep models by computing reusable representations of stage influence.
• Empirical Validation: Demonstrating the utility of the framework in identifying stages responsible for spurious correlations, domain generalization, and noisy label degradation.

Results
The authors evaluate AA-Score on vision (MNIST, CELEBA, Chest X-rays) and language (CIVILCOMMENTS) tasks:

1. MNIST Experiments:

   • Positive/Negative Stages: AA-Score accurately identifies a single inserted data point as having a high positive effect on recognizing that digit and identifies stages with mislabeled data as having negative effects.
   • Distributional Shifts: In a 3-stage setup with rotated images, AA-Score correctly attributes high positive effects to the "in-distribution" stage for each test set.
   • Optimization Dynamics: AA-Score reflects the influence of hyperparameters; for example, reducing learning rate or momentum in a later stage reduces the AA-Score for that stage, consistent with their reduced impact on parameters.
   • Correlation with Ground Truth: AA-Score shows high correlation (avg 0.9456) with counterfactual results obtained by retraining, slightly outperforming the adapted DVEmb baseline.

2. Spurious Correlations (CELEBA & CIVILCOMMENTS):

   • On CELEBA (blonde hair prediction), AA-Score identifies stages where the model learns spurious correlations with gender. Skipping these high-effect stages reduces the model’s reliance on the confounding feature.
   • On CIVILCOMMENTS (toxicity detection), AA-Score identifies stages where demographic identity terms are spuriously correlated with toxicity. Skipping top positive-effect stages reduces the likelihood of predicting toxicity based on these confounders.

3. Complex Multi-Stage Settings:

   • On CELEBA, CIVILCOMMENTS, and Chest X-ray datasets, AA-Score achieves higher correlation with retraining counterfactuals than DVEmb, particularly in complex multi-stage settings, due to its explicit modeling of optimization dynamics.

4. Robustness and Efficiency:

   • Backdoor Attacks: AA-Score can identify poisoned batches in dispersed backdoor attacks, though discriminative power decreases when poisoning is fully mixed across all batches.
   • Runtime/Memory: The method involves logging gradients/optimizer states during training and computing attributions post-training. It scales roughly linearly with dataset and model size, making it practical for small-to-medium models, though scaling to foundation models remains a challenge addressed via structured approximations.

Significance and Limitations
The paper positions accountability attribution as a practical tool for model analysis, debugging, and auditing. It provides a principled way to assign credit or blame to specific development stages, aiding in the identification of harmful behaviors (e.g., bias, spurious correlations) and beneficial ones (e.g., robustness).

Limitations noted by the authors include:

• Reliance on first-order approximations, which may be insufficient when higher-order effects are significant.
• High computational cost for large-scale models due to Hessian and propagator matrix computations, though structured approximations (layer-wise, block-wise) are proposed as mitigation strategies.
• The requirement for recorded trajectories, optimizer states, and hyperparameters, which may not always be available in practice.