[Paper Review] Amortized learning of neural causal representations

Nan Rosemary Ke, Jane X. Wang+3 more|arXiv (Cornell University)|Aug 21, 2020

Bayesian Modeling and Causal Inference23 references4 citations

TL;DR

This paper introduces Causal Relational Networks (CRN), a neural network framework that learns continuous, amortized representations of causal graphs to enable fast adaptation to new causal structures. By leveraging meta-learning with intervention data and a differentiable graph decoder, CRN achieves 90% edge accuracy on synthetic data, significantly outperforming baselines and demonstrating the effectiveness of continuous causal representations for scalable, transferable causal discovery.

ABSTRACT

Causal models can compactly and efficiently encode the data-generating process under all interventions and hence may generalize better under changes in distribution. These models are often represented as Bayesian networks and learning them scales poorly with the number of variables. Moreover, these approaches cannot leverage previously learned knowledge to help with learning new causal models. In order to tackle these challenges, we represent a novel algorithm called extit{causal relational networks} (CRN) for learning causal models using neural networks. The CRN represent causal models using continuous representations and hence could scale much better with the number of variables. These models also take in previously learned information to facilitate learning of new causal models. Finally, we propose a decoding-based metric to evaluate causal models with continuous representations. We test our method on synthetic data achieving high accuracy and quick adaptation to previously unseen causal models.

Motivation & Objective

• To address the scalability and transferability limitations of traditional Bayesian network-based causal structure learning.
• To enable fast adaptation to new causal graphs by leveraging previously learned knowledge through meta-learning.
• To develop a differentiable, continuous representation of causal graphs that supports efficient inference and generalization.
• To propose a decoding-based metric for evaluating causal models with continuous representations, overcoming the challenge of comparing implicit, continuous structures to ground-truth graphs.

Proposed method

• Uses a meta-learning setup where each episode samples a new, distinct causal graph and generates intervention data from it.
• Employs a relational neural network encoder to process intervention data and produce a continuous belief state representing the causal structure.
• Introduces a graph decoder that reconstructs the adjacency matrix from the continuous belief state to enable evaluation via edge accuracy.
• Applies unsupervised contrastive learning via a relational autoencoder to learn meaningful representations without explicit supervision on graph structure.
• Uses a stop-gradient operation to decouple the training of the belief state from the graph decoder, enabling end-to-end training with a supervised signal.
• Leverages attention mechanisms in the relational encoder to model complex dependencies between variables in the causal graph.

Experimental results

Research questions

• RQ1 Can continuous, neural representations of causal graphs enable faster and more scalable causal discovery compared to discrete Bayesian network search?
• RQ2 Does amortized learning via meta-learning allow the model to generalize to new, unseen causal structures after training on diverse graphs?
• RQ3 Is the continuous belief state sufficient to reconstruct accurate causal graphs, or is explicit supervision from a decoder necessary?
• RQ4 How effective is the proposed decoding-based metric in evaluating continuous causal representations compared to traditional structure matching?
• RQ5 To what extent does the relational architecture improve performance over simpler baselines like MLPs or LSTMs?

Key findings

• The Causal Relational Network (CRN) achieved 90% edge prediction accuracy on synthetic causal graphs after training, significantly outperforming the LSTM baseline, which plateaued at around 70% accuracy.
• The model demonstrated rapid adaptation within episodes, improving edge accuracy as more intervention samples were observed, indicating effective amortized inference.
• Removing the belief state from the graph decoder reduced accuracy to below 60%, proving that the continuous representation contains essential information for causal structure recovery.
• Introducing a supervised signal to the graph decoder increased edge accuracy from 90% to 95%, showing that the decoder benefits from gradient feedback.
• The relational autoencoder was essential: without it, the model failed to learn meaningful representations, as evidenced by poor performance on the baseline without the autoencoder.
• The decoding-based evaluation metric successfully captured the model’s ability to recover ground-truth causal structures, providing a scalable alternative to full intervention prediction comparison.