From Attribution to Action: Jointly ALIGNing Predictions and Explanations

TL;DR

ALIGN addresses the challenge that high-quality, task-relevant explanation signals are essential for reliable EGL but are hard to obtain. By jointly training a learnable masker with a classifier in an iterative framework, ALIGN produces soft masks that align with the model’s saliency, boosting accuracy and generalization under domain shifts. The approach is theoretically justified via bounds on generalization error and empirically validated on VLCS and Terra Incognita, where ALIGN outperforms strong baselines in both predictive and interpretability metrics. The results demonstrate that annotation-free, task-driven guidance can yield more faithful explanations and more robust generalization than static or non-task-aligned supervision.

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.

Figures (5)

• Figure 1: Segmentations as guidance may emphasize background, limiting reliability for task-specific explanations.
• Figure 2: Impact of mask quality for predictions.
• Figure 3: Overview of the proposed ALIGN framework.
• Figure 4: Two case studies from the VLCS dataset.
• Figure 5: Mask quality evaluation results on VLCS dataset. Upper: human assessment; Lower: visualization examples.

Theorems & Definitions (7)

• Lemma 1
• Lemma 2: MSE Discrepancy Bound under Domain Shift
• Lemma 3: Cross-Entropy Stability to Small Shifts
• Lemma 4: In-domain errors concerning Feature Inclusion
• Proof 1: Proof of Lemma \ref{['lemma:mse_discrepancy']}
• Proof 2: Proof Sketch of Lemma \ref{['lemma:ce_discrepancy']}
• Proof 3: Proof of Lemma \ref{['lemma:in_domain_error_feature_inclusion']}