Intervening to learn and compose disentangled representations

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

Keywords: disentanglement, generative models, causality, deep learning, out-of-distribution, composition

TL;DR: We introduce a new module that can be attached to the head of any decoder block that learns to process concept information by implicitly inverting linear representations from an encoder.

Abstract: In causal representation learning, 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 also learn disentangled latent structure that enables multi-concept interventions and out-of-distribution (OOD) composition. This is accomplished by adding a simple decoder-only module to the head of an existing decoder block that can be arbitrarily complex. The module learns to process concept information by implicitly inverting linear representations from an encoder. Inspired by the notion of intervention in causal graphical models, 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 disentangled representations that can be composed for OOD generation. To further validate our proposed approach, we show how our module approximates an identifiable concept model by establishing an identifiability result that extends existing work on identifying structured representations in nonlinear models.

Submission Number: 3