# CASTLE: Regularization via Auxiliary Causal Graph Discovery

arXiv: https://arxiv.org/abs/2009.13180 · HTML: https://mlsanity.com/paper/2009.13180/

## Key facts

- Cited by 41 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 28 Sep 2020.

## Built on by (most cited first, 30 of 41)

- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [D'ya like DAGs? A Survey on Structure Learning and Causal Discovery](https://mlsanity.com/paper/2103.02582/) (2021)
- [DAGMA: Learning DAGs via M-matrices and a Log-Determinant Acyclicity Characterization](https://mlsanity.com/paper/2209.08037/) (2022)
- [MIRACLE: Causally-Aware Imputation via Learning Missing Data Mechanisms](https://mlsanity.com/paper/2111.03187/) (2021)
- [Unsuitability of NOTEARS for Causal Graph Discovery](https://mlsanity.com/paper/2104.05441/) (2021)
- [Causal Deep Learning](https://mlsanity.com/paper/2303.02186/) (2023)
- [CausPref: Causal Preference Learning for Out-of-Distribution Recommendation](https://mlsanity.com/paper/2202.03984/) (2022)
- [Simultaneous Missing Value Imputation and Structure Learning with Groups](https://mlsanity.com/paper/2110.08223/) (2021)
- [Ordering-Based Causal Discovery for Linear and Nonlinear Relations](https://mlsanity.com/paper/2410.05890/) (2024)
- [GFlowCausal: Generative Flow Networks for Causal Discovery](https://mlsanity.com/paper/2210.08185/) (2022)
- [Deep Causal Learning: Representation, Discovery and Inference](https://mlsanity.com/paper/2211.03374/) (2022)
- [Incorporating Causal Graphical Prior Knowledge into Predictive Modeling via Simple Data Augmentation](https://mlsanity.com/paper/2103.00136/) (2021)
- [Matching Learned Causal Effects of Neural Networks with Domain Priors](https://mlsanity.com/paper/2111.12490/) (2021)
- [On the Generalization and Adaption Performance of Causal Models](https://mlsanity.com/paper/2206.04620/) (2022)
- [Embracing the black box: Heading towards foundation models for causal discovery from time series data](https://mlsanity.com/paper/2402.09305/) (2024)
- [Markov Equivalence and Consistency in Differentiable Structure Learning](https://mlsanity.com/paper/2410.06163/) (2024)
- [Causal Discovery and Knowledge Injection for Contestable Neural Networks (with Appendices)](https://mlsanity.com/paper/2205.09787/) (2022)
- [Boosting Differentiable Causal Discovery via Adaptive Sample Reweighting](https://mlsanity.com/paper/2303.03187/) (2023)
- [Selecting Treatment Effects Models for Domain Adaptation Using Causal Knowledge](https://mlsanity.com/paper/2102.06271/) (2021)
- [Discovering Dynamic Causal Space for DAG Structure Learning](https://mlsanity.com/paper/2306.02822/) (2023)
- [When Graph Neural Network Meets Causality: Opportunities, Methodologies and An Outlook](https://mlsanity.com/paper/2312.12477/) (2023)
- [The Robustness of Differentiable Causal Discovery in Misspecified Scenarios](https://mlsanity.com/paper/2510.12503/) (2025)
- [Invariant Structure Learning for Better Generalization and Causal Explainability](https://mlsanity.com/paper/2206.06469/) (2022)
- [CMVAE: Causal Meta VAE for Unsupervised Meta-Learning](https://mlsanity.com/paper/2302.09731/) (2023)
- [Towards Learning and Explaining Indirect Causal Effects in Neural Networks](https://mlsanity.com/paper/2303.13850/) (2023)
- [Coarse-to-Fine Learning of Dynamic Causal Structures](https://mlsanity.com/paper/2602.22532/) (2026)
- [Navigating causal deep learning](https://mlsanity.com/paper/2212.00911/) (2022)
- [Data-Driven Estimation of Heterogeneous Treatment Effects](https://mlsanity.com/paper/2301.06615/) (2023)
- [A Review and Roadmap of Deep Causal Model from Different Causal Structures and Representations](https://mlsanity.com/paper/2311.00923/) (2023)
- [Enhancing the Performance of Neural Networks Through Causal Discovery and Integration of Domain Knowledge](https://mlsanity.com/paper/2311.17303/) (2023)

## Abstract (arXiv)

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.