Causal-DFM: Diffusion-Based Causal-Invariant Dynamic Factor Model for Stock Return Prediction

Abstract

Stock return prediction is highly sensitive to market regime shifts, such as bull-bear transitions, financial crises, and monetary policy cycles, where statistical relationships between observed features and returns can change abruptly. While recent deep learning-based factor models achieve strong performance in relatively stable periods, they often rely on correlations that are specific to historical market environments, causing learned patterns to break down when the market enters a different regime and leading to severe out-of-distribution (OOD) degradation. To address this challenge, we propose the Causal-Invariant Dynamic Factor Model (Causal-DFM), a causality-inspired factor-based framework that encourages compact and stable return-relevant representations while suppressing environment-dependent variations, thereby improving robustness under distribution shift. Concretely, Causal-DFM couples a conditional diffusion factor prior that captures regime-varying market dynamics with an Information-Bottleneck-regularized Transformer state encoder that distills a compact, return-predictive representation, suppressing spurious regime-specific signals while preserving the most informative content for prediction. The diffusion component provides a flexible generative mechanism for systematic market factors, while the Information Bottleneck regularization is implemented via an auxiliary return-prediction proxy that encourages the learned representation to be highly predictive yet capacity-controlled. We further present an information-theoretic analysis that motivates this objective and clarifies the conditions under which the learned representation can align with stable return mechanisms across environments. Empirical results on real-world stock datasets demonstrate that Causal-DFM consistently outperforms strong baselines, with particularly pronounced gains during regime transitions and crisis periods, highlighting its effectiveness in mitigating regime-induced OOD failures. The source code is available at https://github.com/bjtu-yzh/Causal-DFM.