LucidAtlas: Learning Uncertainty-Aware, Covariate-Disentangled, Individualized Atlas Representations

TL;DR

LucidAtlas is a novel method for extracting individualized, uncertainty-aware, and covariate-disentangled representations from high-dimensional data, particularly useful in medical imaging for understanding population and individual variations.

Contribution

It introduces LucidAtlas, a versatile atlas representation that captures covariate effects and uncertainty, with methods for interpretability and validation on medical datasets.

Findings

Demonstrates generalizability on two medical datasets.

Provides robust covariate interpretation and uncertainty estimation.

Highlights importance of interpretable models in scientific discovery.

Abstract

The goal of this work is to develop principled techniques to extract information from high dimensional data sets with complex dependencies in areas such as medicine that can provide insight into individual as well as population level variation. We develop $\texttt{LucidAtlas}$, an approach that can represent spatially varying information, and can capture the influence of covariates as well as population uncertainty. As a versatile atlas representation, $\texttt{LucidAtlas}$ offers robust capabilities for covariate interpretation, individualized prediction, population trend analysis, and uncertainty estimation, with the flexibility to incorporate prior knowledge. Additionally, we discuss the trustworthiness and potential risks of neural additive models for analyzing dependent covariates and then introduce a marginalization approach to explain the dependence of an individual predictor on the models' response (the atlas). To validate our method, we demonstrate its generalizability on two medical datasets. Our findings underscore the critical role of by-construction interpretable models in advancing scientific discovery. Our code will be publicly available upon acceptance.