Intervening to Learn and Compose Causally Disentangled Representations

TL;DR

This paper introduces a novel training method for generative models that achieves causally disentangled representations by adding a context module, enabling out-of-distribution generation and extending identifiability theory.

Contribution

It proposes a simple yet effective context module that allows arbitrarily expressive models to learn causally disentangled concepts during training.

Findings

The approach produces causally disentangled representations capable of out-of-distribution generation.

The method can be integrated into end-to-end training or fine-tuning of pre-trained models.

A new identifiability theorem extends existing results on structured representations.

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 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.