Test Time Training for Supervised Causal Learning

2026-05-28

AI Summary

Supervised causal learning often suffers from poor generalization under distribution shifts in real-world scenarios, struggling to bridge the performance gap between synthetic and real data. This work proposes TTT-SCL, a novel framework that introduces test-time training into supervised causal learning for the first time. By dynamically constructing an aligned training set for each test instance, TTT-SCL enables instance-adaptive causal discovery. The method integrates a scoring function to design an efficient training set generation module and establishes a theoretical connection with score-based causal discovery approaches. Experiments demonstrate that TTT-SCL significantly outperforms existing supervised causal learning and traditional causal discovery methods across synthetic, semi-real, and real-world datasets, effectively enhancing out-of-distribution generalization and robustness.

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 reveal three limitations of previous SCL practices: a significant performance gap between synthetic benchmarks and real-world data, fragility to distribution shifts, and failure in compositional generalization, 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 sets explicitly aligned with any specific test instance. We demonstrate the correlation between TTT-SCL and score-based methods, and design an efficient module for generating training sets based on the classic scoring function. Experiments on synthetic benchmarks, pseudo-real and real-world datasets demonstrate that TTT-SCL significantly outperforms existing SCL and traditional causal discovery methods.

Problem

• Supervised Causal Learning
• out-of-distribution generalization
• distribution shift
• compositional generalization
• causal discovery

Innovation

• Test-Time Training
• Supervised Causal Learning
• Out-of-Distribution Generalization
• Causal Discovery
• Dynamic Training Set Generation