From Attribution to Action: Jointly ALIGNing Predictions and Explanations

TL;DR

This paper introduces ALIGN, a joint training framework for classifiers and maskers that enhances interpretability and generalization in vision models by producing task-relevant explanations without relying on noisy external supervision.

Contribution

ALIGN is a novel iterative training method that jointly optimizes classifiers and maskers to produce high-quality, task-relevant explanations, improving interpretability and robustness.

Findings

ALIGN outperforms six baselines on VLCS and Terra Incognita benchmarks.

ALIGN improves explanation quality in terms of sufficiency and comprehensiveness.

ALIGN enhances both interpretability and generalization in vision models.

Abstract

Explanation-guided learning (EGL) has shown promise in aligning model predictions with interpretable reasoning, particularly in computer vision tasks. However, most approaches rely on external annotations or heuristic-based segmentation to supervise model explanations, which can be noisy, imprecise and difficult to scale. In this work, we provide both empirical and theoretical evidence that low-quality supervision signals can degrade model performance rather than improve it. In response, we propose ALIGN, a novel framework that jointly trains a classifier and a masker in an iterative manner. The masker learns to produce soft, task-relevant masks that highlight informative regions, while the classifier is optimized for both prediction accuracy and alignment between its saliency maps and the learned masks. By leveraging high-quality masks as guidance, ALIGN improves both interpretability and generalizability, showing its superiority across various settings. Experiments on the two domain generalization benchmarks, VLCS and Terra Incognita, show that ALIGN consistently outperforms six strong baselines in both in-distribution and out-of-distribution settings. Besides, ALIGN also yields superior explanation quality concerning sufficiency and comprehensiveness, highlighting its effectiveness in producing accurate and interpretable models.