English

Attention Consistency Regularization for Interpretable Early-Exit Neural Networks

Machine Learning 2026-02-05 v2 Artificial Intelligence

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.

Keywords

Cite

@article{arxiv.2601.08891,
  title  = {Attention Consistency Regularization for Interpretable Early-Exit Neural Networks},
  author = {Yanhua Zhao},
  journal= {arXiv preprint arXiv:2601.08891},
  year   = {2026}
}

Comments

2 pages, 1 figure

R2 v1 2026-07-01T09:03:23.635Z