CASTLE: Regularization via Auxiliary Causal Graph Discovery

TL;DR

CASTLE introduces a causality-aware regularization method that jointly learns causal graphs within neural networks, leading to improved out-of-sample prediction accuracy on synthetic and real datasets.

Contribution

The paper proposes CASTLE, a novel regularization technique that incorporates causal structure learning directly into neural network training, enhancing generalization by focusing on causal relationships.

Findings

CASTLE outperforms benchmark regularizers in predictive accuracy.

It effectively learns causal graphs embedded in neural networks.

Theoretical bounds support improved generalization.

Abstract

Regularization improves generalization of supervised models to out-of-sample data. Prior works have shown that prediction in the causal direction (effect from cause) results in lower testing error than the anti-causal direction. However, existing regularization methods are agnostic of causality. We introduce Causal Structure Learning (CASTLE) regularization and propose to regularize a neural network by jointly learning the causal relationships between variables. CASTLE learns the causal directed acyclical graph (DAG) as an adjacency matrix embedded in the neural network's input layers, thereby facilitating the discovery of optimal predictors. Furthermore, CASTLE efficiently reconstructs only the features in the causal DAG that have a causal neighbor, whereas reconstruction-based regularizers suboptimally reconstruct all input features. We provide a theoretical generalization bound for our approach and conduct experiments on a plethora of synthetic and real publicly available datasets demonstrating that CASTLE consistently leads to better out-of-sample predictions as compared to other popular benchmark regularizers.