Multi-Head Explainer: A General Framework to Improve Explainability in CNNs and Transformers

TL;DR

The paper presents MHEX, a modular framework that improves the explainability and accuracy of CNNs and Transformers by integrating attention, supervision, and unified representations, with demonstrated success in medical imaging and text tasks.

Contribution

Introduces MHEX, a versatile framework that enhances model interpretability and performance with minimal modifications across CNNs and Transformers.

Findings

Improves classification accuracy on benchmark datasets.

Produces detailed and interpretable saliency maps.

Easily integrates into existing architectures.

Abstract

In this study, we introduce the Multi-Head Explainer (MHEX), a versatile and modular framework that enhances both the explainability and accuracy of Convolutional Neural Networks (CNNs) and Transformer-based models. MHEX consists of three core components: an Attention Gate that dynamically highlights task-relevant features, Deep Supervision that guides early layers to capture fine-grained details pertinent to the target class, and an Equivalent Matrix that unifies refined local and global representations to generate comprehensive saliency maps. Our approach demonstrates superior compatibility, enabling effortless integration into existing residual networks like ResNet and Transformer architectures such as BERT with minimal modifications. Extensive experiments on benchmark datasets in medical imaging and text classification show that MHEX not only improves classification accuracy but also produces highly interpretable and detailed saliency scores.