$\texttt{LucidAtlas}$: Learning Uncertainty-Aware, Covariate-Disentangled, Individualized Atlas Representations

Abstract

Interpreting how covariates influence spatially structured biological variation — for example, how pediatric airway geometry changes along the airway and across a growing population — remains a key challenge in developing models suitable for clinical application. We present $\texttt{LucidAtlas}$, a versatile framework for modeling and interpreting spatially varying information with associated covariates. To address the limitations of neural additive models when analyzing dependent covariates, we introduce a marginalization approach that enables accurate explanations of how combinations of covariates shape the learned atlas. $\texttt{LucidAtlas}$ integrates covariate interpretation, spatial representation, individualized prediction, population distribution analysis, and out-of-distribution detection into a single interpretable model. We validate its effectiveness on a synthetic spatiotemporal dataset, the OASIS brain volume dataset, and a pediatric airway shape dataset. Our findings underscore the critical role of by-construction interpretable models in advancing scientific discovery. The implementation is publicly available at https://github.com/****.