Attention Consistency Regularization for Interpretable Early-Exit Neural Networks
Yanhua Zhao
TL;DR
The paper tackles interpretability in early-exit neural networks, where intermediate predictions may rely on different features than the final decision, risking inconsistent explanations. It proposes Explanation-Guided Training (EGT), a multi-objective framework that adds an attention-consistency loss to align early-exit attention maps with the final exit, formalized as $L_{\text{total}} = L_{\text{cls}} + \alpha L_{\text{consistency}}$ and $L_{\text{consistency}} = \frac{1}{4} \sum_{i=1}^{4} d_{\cos}(\tilde{A}_i, A_5)$ with $d_{\cos}(x,y) = 1 - \frac{x^\top y}{\|x\|\|y\|}$. Experiments on a real-world 9-class dataset show that EGT achieves up to $98.97\%$ overall accuracy with a $1.97\times$ inference speedup and up to $18.5\%$ improvement in attention consistency over a baseline. This work offers a practical path toward trustworthy, resource-efficient AI by producing consistent, interpretable explanations across all exit points without sacrificing performance.
Abstract
Early-exit neural networks enable adaptive inference by allowing predictions at intermediate layers, reducing computational cost. However, early exits often lack interpretability and may focus on different features than deeper layers, limiting trust and explainability. This paper presents Explanation-Guided Training (EGT), a multi-objective framework that improves interpretability and consistency in early-exit networks through attention-based regularization. EGT introduces an attention consistency loss that aligns early-exit attention maps with the final exit. The framework jointly optimizes classification accuracy and attention consistency through a weighted combination of losses. Experiments on a real-world image classification dataset demonstrate that EGT achieves up to 98.97% overall accuracy (matching baseline performance) with a 1.97x inference speedup through early exits, while improving attention consistency by up to 18.5% compared to baseline models. The proposed method provides more interpretable and consistent explanations across all exit points, making early-exit networks more suitable for explainable AI applications in resource-constrained environments.