Intervening to learn and compose causally disentangled representations

Alex Markham, Isaac Hirsch, Jeri A. Chang, Liam Solus, Bryon Aragam

Keywords: causal disentanglement, generative models, compositional generalization, concept learning

TL;DR: We propose a new context module that fine-tunes arbitrary black-box generative models for causal disentanglement, concept learning, and compositional generalization.

Abstract: In designing generative models, it is commonly believed that in order to learn useful latent structure, we face a fundamental tension between expressivity and structure. In this paper we challenge this view by proposing a new approach to training arbitrarily expressive generative models that simultaneously learn causally disentangled concepts. This is accomplished by adding a simple \emph{context module} to an arbitrarily complex black-box model, which learns to process concept information by implicitly inverting linear representations from the model's encoder. Inspired by the notion of intervention in a causal model, our module selectively modifies its architecture during training, allowing it to learn a compact joint model over different contexts. We show how adding this module leads to causally disentangled representations that can be composed for out-of-distribution generation on both real and simulated data. The resulting models can be trained end-to-end or fine-tuned from pre-trained models. To further validate our proposed approach, we prove a new identifiability result that extends existing work on identifying structured representations.

Submission Number: 88