Interpretable Variational Autoencoder with Stabilized Tree Regularization

ICLR 2026 Conference Submission17905 Authors

Keywords: Interpretability, High-dimentional Data, VAE, SHAP, Decision Tree

Abstract: With the growth of digital health, bioinformatics and healthcare now produce massive high-dimensional (HD) datasets that challenge both prediction and interpretability. This work introduces the Tree-Regularized Interpretable Variational Autoencoder (TRI-VAE), which couples a VAE with a surrogate decision tree to impose rule-consistent structure on the latent space. TRI-VAE aligns embeddings to soft leaf distributions for a cluster-aware representation learning, and employs a SHAP-based attribution scheme tailored to HD settings to select salient features and harmonize feature-level explanations with path-based rules. A tree-regularizer, optimized via a learned average-path-length surrogate, promotes compact and balanced trees; stability-controlled tree updates further preserve assignment consistency over training. Across public (TCGA-LIHC, TUEP) and private (PPH) datasets, TRI-VAE delivers competitive predictive performance while yielding faithful, human-readable explanations. An LLM-assisted evaluation protocol with clinician review supports the accessibility and reliability of the extracted rules and attributions, advancing trustworthy AI for medical data analysis.

Primary Area: interpretability and explainable AI

Submission Number: 17905