Learning disentangled representations for explainable chest X-ray classification using Dirichlet VAEs

Rachael Harkness, Alejandro F. Frangi, Kieran Zucker, Nishant Ravikumar

Research output: Chapter in Book/Report/Conference proceedingConference contributionpeer-review

Abstract

This study explores the use of the Dirichlet Variational Autoencoder (DirVAE) for learning disentangled latent representations of chest X-ray (CXR) images. Our working hypothesis is that distributional sparsity, as facilitated by the Dirichlet prior, will encourage disentangled feature learning for the complex task of multi-label classification of CXR images. The DirVAE is trained using CXR images from the CheXpert database, and the predictive capacity of multi-modal latent representations learned by DirVAE models is investigated through implementation of an auxiliary multi-label classification task, with a view to enforce separation of latent factors according to class-specific features. The predictive performance and explainability of the latent space learned using the DirVAE were quantitatively and qualitatively assessed, respectively, and compared with a standard Gaussian prior-VAE (GVAE). We introduce a new approach for explainable multi-label classification in which we conduct gradient-guided latent traversals for each class of interest. Study findings indicate that the DirVAE is able to disentangle latent factors into class-specific visual features, a property not afforded by the GVAE, and achieve a marginal increase in predictive performance relative to GVAE. We generate visual examples to show that our explainability method, when applied to the trained DirVAE, is able to highlight regions in CXR images that are clinically relevant to the class(es) of interest and additionally, can identify cases where classification relies on spurious feature correlations.

Original languageEnglish
Title of host publicationMedical Imaging 2023
Subtitle of host publicationImage Processing
EditorsOlivier Colliot, Ivana Isgum
PublisherSPIE
ISBN (Electronic)9781510660335
DOIs
Publication statusPublished - 2023
EventMedical Imaging 2023: Image Processing - San Diego, United States
Duration: 19 Feb 202323 Feb 2023

Publication series

NameProgress in Biomedical Optics and Imaging - Proceedings of SPIE
Volume12464
ISSN (Print)1605-7422

Conference

ConferenceMedical Imaging 2023: Image Processing
Country/TerritoryUnited States
CitySan Diego
Period19/02/2323/02/23

Keywords

  • chest X-ray images
  • Dirichlet distribution
  • disentanglement
  • explainability
  • Multi-label classification
  • variational autoencoders

Fingerprint

Dive into the research topics of 'Learning disentangled representations for explainable chest X-ray classification using Dirichlet VAEs'. Together they form a unique fingerprint.

Cite this