Test Time Training for Supervised Causal Learning

ICLR 2026 Conference Submission12059 Authors

Keywords: Supervised causal learning, causal discovery, out of distribution, test time training

Abstract: Supervised Causal Learning (SCL) has shown promise in causal discovery by framing it as a supervised learning problem. However, it suffers from significant out-of-distribution generalization challenges. We reveals three fundamental limitations of previous SCL practices: severe fragility to distribution shifts, failure in compositional generalization, and a stark performance gap between synthetic benchmarks and real-world data, collectively questioning its real-world applicability. To address this, we propose Test-Time Training for Supervised Causal Learning (TTT-SCL), a novel framework that dynamically generates training data explicitly aligned with any specific test instance. We find that the causal similarity between training and test data can be implicitly captured through distributional alignment, which we operationalize via a proposed Alignment of Distribution (AD) metric. To prevent degenerate solutions and enforce causal minimality, we incorporate sparsity constraints into the optimization. Building on this foundation, we introduce Test-time Aligned Causal Training with Informed Construction (TACTIC), the first instantiation of TTT-SCL, which jointly optimizes AD and sparsity via stochastic graph refinement to dynamically generate aligned training data. Extensive experiments on synthetic benchmarks, real-world and pseudo-real dataset demonstrate that TACTIC significantly outperforms existing SCL and traditional causal discovery methods.

Primary Area: causal reasoning

Submission Number: 12059