[Paper Review] Neural Decomposition: Functional ANOVA with Variational Autoencoders

Kaspar Märtens, Christopher Yau|arXiv (Cornell University)|Jun 25, 2020

Explainable Artificial Intelligence (XAI)23 references1 citations

TL;DR

This paper introduces Neural Decomposition, a method that integrates functional ANOVA variance decomposition into Conditional Variational Autoencoders (CVAEs) to enable interpretable, feature-level decomposition of data variation into main effects of latent variables (z) and covariates (c), and their non-linear interactions. By imposing functional constraints on the decoder's weight space, the approach achieves identifiability and enables scalable, non-linear variance attribution in high-dimensional data, demonstrated effectively on single-cell genomics data with improved robustness over linear models.

ABSTRACT

Variational Autoencoders (VAEs) have become a popular approach for dimensionality reduction. However, despite their ability to identify latent low-dimensional structures embedded within high-dimensional data, these latent representations are typically hard to interpret on their own. Due to the black-box nature of VAEs, their utility for healthcare and genomics applications has been limited. In this paper, we focus on characterising the sources of variation in Conditional VAEs. Our goal is to provide a feature-level variance decomposition, i.e. to decompose variation in the data by separating out the marginal additive effects of latent variables z and fixed inputs c from their non-linear interactions. We propose to achieve this through what we call Neural Decomposition - an adaptation of the well-known concept of functional ANOVA variance decomposition from classical statistics to deep learning models. We show how identifiability can be achieved by training models subject to constraints on the marginal properties of the decoder networks. We demonstrate the utility of our Neural Decomposition on a series of synthetic examples as well as high-dimensional genomics data.

Motivation & Objective

• To address the interpretability gap in Variational Autoencoders (VAEs), particularly in healthcare and genomics, where black-box latent representations limit insight into feature-level variation.
• To extend classical functional ANOVA (fANOVA) to deep learning by embedding additive and interaction effects of latent variables (z) and covariates (c) into the decoder architecture of CVAEs.
• To enforce identifiability of variance components—main effects and interactions—by constraining the marginal properties of the decoder network through functional constraints in weight space.
• To provide a scalable, non-linear alternative to linear models for decomposing gene expression variation in single-cell RNA-seq data, reducing risk of model misspecification.
• To demonstrate that non-linear interaction effects in gene expression can be accurately captured and attributed to pseudotime and stimulant type in dendritic cell data.

Proposed method

• Proposes Neural Decomposition as a CVAE-based framework where the decoder is structured to explicitly model additive effects of latent variables (z) and fixed covariates (c), along with their non-linear interactions.
• Implements functional constraints on the decoder’s weight space to enforce identifiability, ensuring that variation is not absorbed by one component at the expense of others.
• Uses constrained optimization during training to maintain marginal properties of the decoder, such as zero-mean activation for main effects, to separate variance contributions.
• Applies amortized variational inference to approximate posterior distributions, enabling scalable training on large datasets like single-cell genomics data.
• Employs a deep neural network architecture in the decoder that decomposes the output into distinct components: f(z), f(c), and f(z,c), with shared representations for interaction terms.
• Validates the method using synthetic data and real-world single-cell RNA-seq data from mouse dendritic cells, with pseudotime inferred as the latent variable z.

Experimental results

Research questions

• RQ1 Can functional ANOVA decomposition be effectively adapted to deep generative models like CVAEs to enable interpretable variance attribution in high-dimensional data?
• RQ2 How can identifiability of main effects and interactions be enforced in a neural network-based decoder without sacrificing modeling flexibility?
• RQ3 Does the proposed non-linear decomposition outperform linear models in capturing true biological variation in noisy single-cell genomics data?
• RQ4 To what extent can Neural Decomposition accurately separate and quantify the contributions of pseudotime (z), stimulant type (c), and their interaction to gene expression variation?
• RQ5 Can the method robustly identify biologically relevant interaction effects in single-cell data where linear models may fail due to model misspecification?

Key findings

• Neural Decomposition successfully decomposed gene expression variation in single-cell RNA-seq data into distinct components: main effects of pseudotime (z), main effects of stimulant type (c), and non-linear interaction effects (z,c), with variance fractions quantified for each.
• The method confirmed the presence of strong interaction effects in three key genes—Tnf, Rasgefb1, and Tnfaip3—whose expression patterns were previously detected using linear models, but now validated under a more flexible non-linear framework.
• The inferred latent dimension (z) was significantly correlated with capture time, confirming its role as a valid proxy for pseudotime in the dendritic cell differentiation process.
• Genes with strong additive z effects showed mixed clustering in UMAP visualizations, while those with strong c or interaction effects exhibited clear segregation between LPS and PAM-stimulated cells, indicating correct attribution of biological signals.
• The non-linear model provided a more accurate and robust variance decomposition than prior linear approaches, reducing the risk of false positives due to model misspecification in high-noise single-cell data.
• The constrained optimization approach successfully maintained identifiability, ensuring that each component (z, c, interaction) captured distinct sources of variation without redundancy.