Causal Local States: Scalable Simultaneous Causal Network Inference and Forecasting for Dynamical Systems

Jonas Braun; Fabian Fischbach; Daniel Köglmayr; Sebastian Baur; Christoph Räth

cs.LGAug 18, 2026

Causal Local States: Scalable Simultaneous Causal Network Inference and Forecasting for Dynamical Systems

Authors: Jonas Braun, Fabian Fischbach, Daniel Köglmayr, Sebastian Baur, Christoph Räth

Organizations: Ludwig-Maximilians-Universität München, Faculty of Physics, Geschwister-Scholl-Platz 1, 80539 Munich, Germany · Institut für Theoretische Physik, Universität zu Köln, Zülpicher Strasse 77, 50937 Cologne, Germany · Institut für KI-Sicherheit, Deutsches Zentrum für Luft- und Raumfahrt (DLR), St. Augustin & Ulm, Germany · Entrox Systems, Munich, Germany ·

Abstract

Machine learning methods predict many real-world systems with remarkable accuracy, but they are typically treated as black boxes that offer no insight into which interactions drive the dynamics. Causal discovery methods reconstruct the interaction network from observational data, but without regard to whether the inferred structure supports prediction. Existing approaches combining both tasks rely on a single global hyperparameter, such as a causal threshold or a fixed neighborhood size, which cannot recover the structure of heterogeneous systems. Here we introduce causal local states (CLS), a framework that simultaneously infers an approximate Granger-causal interaction network and forecasts the system dynamics. For each node independently, we select the smallest set of neighbors that allows a predictive model to forecast the node near-optimally, and the resulting neighborhoods are then combined for a forecast of the full system. On three benchmarks of increasing difficulty, we achieve reconstruction of the underlying networks with high fidelity and forecasts on par with a model that is supplied with the true network, providing a step toward explainable and scalable forecasting of complex systems.

Figures & tables

Figure 1: Schematic illustration of the CLS ( CLS ) framework. a , The input consists of high-dimensional time series data generated by a system with an unknown interaction network. b , A causal measure is used to rank candidate neighbors for a core node and generate neighborhood candidates. c , A wrapper-based optimization procedure selects the optimal neighborhood candidate and subsequently removes redundant features through backward elimination. d , The inference process yields an approximation of the causal neighborhood of the core node. e , After repeating a-d for all core nodes, the inferred local neighborhoods are combined into a global neighborhood matrix. f , The CLS framework performs one-step prediction using the inferred neighborhoods by training a model $g_{i}$ on only the data present in the neighborhood $\mathcal{N}_{i}^{*}$ . g , One-step predictions are recursively fed back and concatenated to generate multi-step predictions of the full system dynamics. Due to the independent neighborhood inference and prediction procedure for each node, the CLS framework is naturally parallelizable and scalable to high-dimensional systems.

Figure 2: Analysis of composite (non-interacting) Lorenz63–Rössler systems using CLS ( CLS ). a , Summary graph of a single Lorenz63–Rössler attractor pair. b , Exhaustive wrapper validation for a single attractor pair. All non-empty neighborhood candidates of a core node are evaluated, and their wrapper MSE ( MSE ) is shown on a logarithmic axis, ordered by decreasing error. Bars are colored by how many of the core node’s two same attractor neighbors the candidate contains: red, neither; blue and orange, exactly one (one color per neighbor, identified in the sub-panel legends); green, both. c-e , CLS analysis for a single attractor pair system. c , TE ( TE ) matrix computed during the filtering step. d , Inferred neighborhood matrices before and after feature elimination. e , Short-term prediction performance of CLS compared with standard NGRC ( NGRC ). f , Inferred neighborhood matrix after feature elimination for a system of 15 coupled attractor pairs (90 variables), shown for seed 0. g,h , CLS performance across different system sizes ( $N\in[1,15]$ ), using TE as the filter and NGRC as the wrapper model. Every system size was run for $n=10$ seeds, and every seed forecast from 10 disjoint starting points; markers give the mean over seeds and error bars $\pm 1$ standard deviation across them. g , Neighborhood inference accuracy: TPR ( TPR ) on the left axis (circles; grey dashed, before elimination; green solid, after) and FPR ( FPR ) on the right axis (squares; grey dash-dotted, before elimination; red solid, after). h , Prediction quality as the mean VPS ( VPS ) at $\epsilon=0.3$ , for CLS (blue circles, solid line) and for a monolithic NGRC trained on the full state vector with linear features only (red squares, dashed line). All experiment parameters can be extracted from Table 1 .

Figure 3: CLS ( CLS ) analysis of a 40-dimensional Lorenz96 system. a , TE ( TE ) matrix computed during the filtering step. b , Ground-truth neighborhood matrix defined by the governing equations. c , Inferred neighborhood matrix prior to feature elimination. d , Final neighborhood matrix after feature elimination. e , Histogram showing the distribution of calculated TE values split between true parents of a node and non-parents. f , Ground-truth system trajectory. g–i , Short-term predictions using the true adjacency matrix ( g ), the inferred matrix before feature elimination ( h ), and the inferred matrix after feature elimination ( i ). j , VPS ( VPS ) at $\epsilon=0.3$ expressed as VPT ( VPT ) in maximum Lyapunov times for CLS using the three neighborhood matrices. Each of the 40 grid nodes contributes one point, its VPS averaged over 100 forecasts (10 seeds $\times$ 10 disjoint launch points). All experiment parameters can be extracted from Table 2 .

Figure 4: CLS ( CLS ) analysis of the Kuramoto model of the UK power grid. a , Ground-truth adjacency matrix. b,c , Neighborhood matrices inferred using Sine-NGRC before ( b ) and after ( c ) feature elimination. d,e , Neighborhood matrices inferred using Kuramoto-NGRC before ( d ) and after ( e ) feature elimination. f–i , Example neighborhood inference process for a single core node using Kuramoto-NGRC as the wrapper model. f, h , Initial neighborhood identified after the filtering step. f shows that multiple spurious neighbors are selected in the first step because their CCM ( CCM ) score is higher than the true neighbors. g, i , Final neighborhood after backward feature elimination. The spurious neighbors are largely removed from the neighborhood. j , Histogram showing the distribution of calculated CCM values split between true parents of a node and non-parents. k , VPS ( VPS ) for both wrapper models and all three neighborhood matrices, on $\sin\theta$ at $\epsilon=0.3$ . Each of the 120 grid nodes contributes one point, its VPS averaged over 100 forecasts (10 seeds $\times$ 10 disjoint launch points). All experiment parameters can be extracted from Table 3 .

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 1 — Figure 1: No global hyperparameter recovers the neighborhoods of the composite Lorenz63–Rössler system. a , The TE ( TE ) matrix of a single attractor pair, from which both reconstructions below are read off. Entry $(i,j)$ is the TE from source node $j$ to target node $i$ , in the convention of Fig. 2 c of the main text, which shows the same matrix without values; cell shading is proportional to the entry, and the black outlines mark the two ground-truth blocks. Diagonal entries are not scored. b , Global neighborhood size: with $q=1$ the true neighbor $R_{z}$ of the core node $R_{x}$ is missed; with $q=5$ , the smallest size that includes $R_{z}$ , the three independent Lorenz variables are admitted as well, since TE ranks them above $R_{z}$ . No $q$ recovers the neighborhood of $R_{x}$ correctly. c , Global threshold: $0.28$ is the highest threshold that still captures the true neighbor $R_{x}$ of the core node $R_{z}$ (top); at this value the neighborhood of $L_{x}$ (bottom) already contains spurious Rössler variables. No single threshold recovers both neighborhoods without introducing spurious edges in one or false negatives in the other.

Figure 2: Edge-importance experiment on the UK power grid with three representative single-edge removals shown. In each case, two reservoir-computing local-state models are trained, one with the true adjacency matrix and one with the selected edge removed; a node is counted as affected if its VPS ( VPS ) change by more than a threshold (

\mathrm{VPS_{th}}=5

). Left: the removed edges and the corresponding affected nodes highlighted on the UK power grid network. Right: node-wise change in VPS for each removal. Some edges affect many nodes, others only a few.

Figure 3: Distribution of the change in VPS ( VPS ) and of the number of affected nodes across all removed edges. The importance of an edge for prediction is highly heterogeneous, so that a criterion aimed at structural recovery rather than forecast skill is suboptimal when the goal is prediction.

Panel of Figure 4 — Figure 4: Hidden confounder inducing a spurious but predictive link. The hidden process $H$ drives $X$ at lag $\tau=1$ and $Y$ at lag $\tau=2$ . There is no direct causal link $X\to Y$ (dotted), yet $X$ is predictive of the future of $Y$ because both are echoes of $H$ . Consequently a spurious $X\to Y$ link improves the forecast of $Y$ .

Figure 5: Sample core-node predictions for a composite Lorenz63–Rössler system. For the core node

L_{x}

, wrapper models are trained on a neighborhood of the true neighbors

L_{y},L_{z}

(green) and on a neighborhood of spurious neighbors

R_{x},R_{y},R_{z}

(red), for each of the three strategies. Under the one-step strategy the core node’s own past dominates, so even the spurious neighborhood reproduces the trajectory; under NGRC ( NGRC ) inference mode the spurious neighborhood fails to track the core trajectory at all.

Panel of Figure 6 — Figure 6: Core-node MSE ( MSE ) for all $32$ candidate neighborhoods of the composite Lorenz63–Rössler system, evaluated for each core node and grouped by how many features come from the same attractor as the core node: none (red), one (orange, the two possible cases distinguished), or both (green). Only NGRC ( NGRC ) inference mode (c) cleanly separates the neighborhoods that contain the true neighbors from those that do not; the one-step strategy (a) barely separates them because the core node dominates the forecast, and the multi-step strategy (b) improves the separation but does not fully resolve it.

System
	Exhaustive wrapper	Single attractor pair	Scaling sweep
	validation (Fig. 2b)	(Fig. 2c–e)	(Fig. 2f–h)
Attractor pairs $N$	$1$	$1$	$1$ – $15$
Lorenz63 $\mathrm{d}t$ ; $\sigma,\rho,\beta$	$0.01$ ; $10.0,\ 28.0,\ 8/3$
Rössler $\mathrm{d}t$ ; $a,b,c$	$0.05$ ; $0.2,\ 0.2,\ 5.7$
Initial condition, every subsystem	drawn from $\mathcal{U}(-3,3)^{3}$ with the seed below

Appendix

Table 1: Simulation and model hyperparameters for the composite Lorenz63–Rössler benchmark. NGRC , NGRC ; CLS , CLS ; VPT , VPT .

Lorenz96 system
Dimension $N$	$40$
Forcing $F$	$5$
Time step $\mathrm{d}t$	$0.05$
Initial condition	$x_{0}=F\cdot\mathbf{1}+10^{-5}\cdot\mathcal{U}(0,1)^{N}$
Random seeds	$0,1,\dots,9$ (ten seeds)
Transient steps (discarded)	$10000$

Appendix

Table 2: Simulation and model hyperparameters for the Lorenz96 benchmark. TE , TE ; NGRC , NGRC ; CLS , CLS ; VPT , VPT .

UK power grid
	Sine-NGRC	Kuramoto-NGRC
Phase lags $\alpha,\ \beta$	$0.02,\ 0.06$
Coupling $\gamma_{1},\ \gamma_{2}$	$0.4,\ 0.4$
$\omega_{\mathrm{abs,low}}$	$0.3$
Time step $\mathrm{d}t$	$8\times 10^{-2}$
Initial phases	$\mathcal{U}(0,2\pi)^{N}$ , drawn with the seed

Appendix

Table 3: Simulation and model hyperparameters for the UK power grid benchmark. The seed fixes both the natural frequencies

\omega_{i}

and the initial phases. NGRC , NGRC ; CLS , CLS ; VPS , VPS ; VPT , VPT .

Time-Series Causal Discovery

Time-series causal discovery seeks to infer which variables influence others over time and when those effects occur.

See all 65 papers

Mask2Cause: Causal Discovery via Adjacency Constrained Causal Attention
May 8, 2026Omar Muhammad, Pasupuleti Dhruv Shivkant, Deepak N. SubramaniTime-Series Causal Discovery Causal Structure Learning
1Indian Institute of Science, Bengaluru, India
Causal-INSIGHT: Probing Temporal Models to Extract Causal Structure
Mar 26, 2026Benjamin Redden, Hui Wang, Shuyan LiCausal Interpretability Causal Discovery
School of Electronics, Electrical Engineering and Computer Science, Queen’s University Belfast, Belfast, United Kingdom
Causal Discovery in Action: Learning Chain-Reaction Mechanisms from Interventions
Mar 23, 2026Panayiotis Panayiotou, Özgür ŞimşekCausal Discovery Causal Identifiability
Department of Computer Science, University of Bath, UK

May 8, 2026J/K move · Enter open · S save

Mask2Cause: Causal Discovery via Adjacency Constrained Causal Attention

Omar Muhammad, Pasupuleti Dhruv Shivkant, Deepak N. Subramani

1Indian Institute of Science, Bengaluru, India

Time-Series Causal Discovery Causal Structure Learning Causal Attention

Open paper AskPDF arXiv

Leveraging deep learning for causal discovery in time series remains challenging because existing neural methods predominantly rely on component-wise architectures that fail to capture shared system dynamics or employ decoupled post-hoc graph extraction that risks overfitting to spurious correlations. We propose

\textbf{Mask2Cause}

, an end-to-end framework that recovers the underlying causal graph directly during the forecasting forward pass. Our approach introduces an Inverted Variable Embedding and an Adjacency-Constrained Masked Attention mechanism, trained with homoscedastic or heteroscedastic objectives to capture causal influences in both mean and variance. Empirical results on diverse benchmarks, from synthetic chaotic dynamics to realistic biological simulations, demonstrate state-of-the-art causal discovery with significantly reduced parameter complexity compared to standard baselines. We further show that inferred causal structures can be used to reduce parameter count of forecasting models by more than 70% on average while maintaining predictive accuracy.

Abstract

Time-Series Causal Discovery

Explore similar work

Mask2Cause: Causal Discovery via Adjacency Constrained Causal Attention

Causal-INSIGHT: Probing Temporal Models to Extract Causal Structure

Causal Discovery in Action: Learning Chain-Reaction Mechanisms from Interventions