Sparsity regularization via tree-structured environments for disentangled representations

TL;DR

This paper introduces Tree-Based Regularization (TBR), a method for causal representation learning that leverages sparsity across related environments to accurately infer latent variables, with applications in biological data analysis.

Contribution

The paper proposes TBR, a novel regularization technique that models sparse changes across environments to improve latent variable inference in causal systems.

Findings

TBR outperforms existing methods in recovering latent variables.

TBR accurately identifies causal variables in both simulated and real gene expression data.

Theoretical guarantees are provided under certain assumptions about sparse changes.

Abstract

Many causal systems such as biological processes in cells can only be observed indirectly via measurements, such as gene expression. Causal representation learning -- the task of correctly mapping low-level observations to latent causal variables -- could advance scientific understanding by enabling inference of latent variables such as pathway activation. In this paper, we develop methods for inferring latent variables from multiple related datasets (environments) and tasks. As a running example, we consider the task of predicting a phenotype from gene expression, where we often collect data from multiple cell types or organisms that are related in known ways. The key insight is that the mapping from latent variables driven by gene expression to the phenotype of interest changes sparsely across closely related environments. To model sparse changes, we introduce Tree-Based Regularization (TBR), an objective that minimizes both prediction error and regularizes closely related environments to learn similar predictors. We prove that under assumptions about the degree of sparse changes, TBR identifies the true latent variables up to some simple transformations. We evaluate the theory empirically with both simulations and ground-truth gene expression data. We find that TBR recovers the latent causal variables better than related methods across these settings, even under settings that violate some assumptions of the theory.