Hybrid Explanation-Guided Learning for Transformer-Based Chest X-Ray Diagnosis

TL;DR

This work addresses spurious correlations in transformer-based chest X-ray diagnosis by introducing Hybrid Explanation-Guided Learning (H-EGL), which jointly leverages self-supervised and human-guided attention constraints. The framework combines a Discriminative Attention Learning (DAL) module with a human-AI alignment component, optimized via a composite loss $\mathcal{L}_{\text{H-EGL}} = \mathcal{L}_{\text{CE}} + \alpha \mathcal{L}_{\text{HA}} + \beta \mathcal{L}_{\text{DAL}}$, within a DWARF ViT-based encoder-decoder. Evaluated on ChestXDet with chest X-ray images, H-EGL outperforms state-of-the-art EGL baselines (KAD, GAIN) in AUC, F1, and MCC, and exhibits reduced generalization gaps and improved robustness to noise. The approach yields attention maps better aligned with expert pathology regions, offering a scalable, annotation-efficient path toward robust human-AI aligned vision models for clinical imaging.

Abstract

Transformer-based deep learning models have demonstrated exceptional performance in medical imaging by leveraging attention mechanisms for feature representation and interpretability. However, these models are prone to learning spurious correlations, leading to biases and limited generalization. While human-AI attention alignment can mitigate these issues, it often depends on costly manual supervision. In this work, we propose a Hybrid Explanation-Guided Learning (H-EGL) framework that combines self-supervised and human-guided constraints to enhance attention alignment and improve generalization. The self-supervised component of H-EGL leverages class-distinctive attention without relying on restrictive priors, promoting robustness and flexibility. We validate our approach on chest X-ray classification using the Vision Transformer (ViT), where H-EGL outperforms two state-of-the-art Explanation-Guided Learning (EGL) methods, demonstrating superior classification accuracy and generalization capability. Additionally, it produces attention maps that are better aligned with human expertise.

Figures (3)

• Figure 1: Overview of the Hybrid Explanation-Guided Learning (H-EGL) framework. H-EGL is a hybrid Explanation-Guided Learning (EGL) approach designed for Vision Transformers (ViTs). The model integrates visual inputs (e.g., chest X-rays) and textual information (e.g., disease labels such as Atelectasis, Cardiomegaly) via a ViT-based encoder-decoder architecture, employing a frozen text encoder and trainable ViT components. It generates both class-wise predictions and attention maps $A$. The H-EGL framework combines two key components: (i) a self-supervised module, Discriminative Attention Learning (DAL), which encourages class-distinctive attention patterns, and (ii) a supervised Human-AI Alignment module. The total training objective incorporates: (i) Cross-Entropy Loss$L_{CE}$ for classification, (ii) Human-AI Alignment Loss$L_{HA}$ to align attention maps with expert annotations (e.g., bounding boxes), and (iii) Discriminative Attention Loss$L_{DAL}$, which enforces separation among class-specific attention maps. The figure illustrates attention distinctiveness (e.g., similarity $s_{ij}$ between attention maps $A_i$ and $A_j$) in a four-class setting, though the approach generalizes to any number of classes.
• Figure 2: Robustness of models under different levels of perturbed inputs. The noises applied are normally distributed with mean zero and different standard deviation levels ($\sigma=0, 0.03, 0.05, 0.1$). Note: H-EGL and DAL curves closely overlap.
• Figure 3: Comparison of attention maps from the baseline model ($\alpha=0$, $\beta$=0), DWARF ($\beta=0$), DAL ($\alpha=0$), and H-EGL model. Abnormal pathologies are highlighted with white bounding boxes, while arrows indicate areas with visible changes across models.