HydraViT: Adaptive Multi-Branch Transformer for Multi-Label Disease Classification from Chest X-ray Images

TL;DR

HydraViT introduces an adaptive multi-branch transformer model that improves multi-label disease classification from chest X-ray images by effectively capturing long-range context and focusing on critical regions, outperforming existing methods.

Contribution

The paper presents HydraViT, a novel transformer-based architecture with multi-branch outputs for enhanced multi-label chest X-ray disease classification, addressing heterogeneity and co-occurrence challenges.

Findings

HydraViT outperforms attention-guided methods by 1.2%

It surpasses region-guided methods by 1.4%

It exceeds semantic-guided methods by 1.0% in classification accuracy

Abstract

Chest X-ray is an essential diagnostic tool in the identification of chest diseases given its high sensitivity to pathological abnormalities in the lungs. However, image-driven diagnosis is still challenging due to heterogeneity in size and location of pathology, as well as visual similarities and co-occurrence of separate pathology. Since disease-related regions often occupy a relatively small portion of diagnostic images, classification models based on traditional convolutional neural networks (CNNs) are adversely affected given their locality bias. While CNNs were previously augmented with attention maps or spatial masks to guide focus on potentially critical regions, learning localization guidance under heterogeneity in the spatial distribution of pathology is challenging. To improve multi-label classification performance, here we propose a novel method, HydraViT, that synergistically combines a transformer backbone with a multi-branch output module with learned weighting. The transformer backbone enhances sensitivity to long-range context in X-ray images, while using the self-attention mechanism to adaptively focus on task-critical regions. The multi-branch output module dedicates an independent branch to each disease label to attain robust learning across separate disease classes, along with an aggregated branch across labels to maintain sensitivity to co-occurrence relationships among pathology. Experiments demonstrate that, on average, HydraViT outperforms competing attention-guided methods by 1.2%, region-guided methods by 1.4%, and semantic-guided methods by 1.0% in multi-label classification performance.