1 Introduction↩︎

Inferring causal relationships from observational time series is a central challenge in climate science, neuroscience, finance, gene regulatory network reconstruction, industrial root cause analysis, and many other fields [1]. Given a multivariate time series \(\mathbf{x}_t = (x_t^1, \ldots, x_t^d)\) with \(T\) observations, the goal is to recover the lagged causal graph \(G\), where \(G[c, \ell, e] = 1\) indicates that variable \(x^c\) at lag \(\ell\) is a direct cause of variable \(x^e\) at time \(t\).

Two dominant families of methods have emerged for time-series causal discovery, each with complementary strengths and limitations. Score-based methods—Neural Granger Causality [2], NAVAR [3], DYNOTEARS [4], CUTS+ [5]—take a global optimization view, searching for the graph that best fits the data under a scoring criterion. They can capture nonlinear dependencies and scale to moderate dimensionality, but enforce sparsity through continuous penalties (\(L_1\) or Gumbel-Softmax), requiring post-hoc weight thresholding to obtain a binary graph—and model misspecification or poor scoring assumptions can silently degrade the recovered graph. Attention-based methods like TCDF [6] similarly optimize a predictive objective but use attention weights as causal proxies; however, softmax normalization creates artificial competition among candidate parents, and attention thresholds must be dataset-tuned. Constraint-based methods—PCMCI/PCMCI+ [7], [8], CDNOTS [9], SyPI+ [10], LPCMCI [11]—build causal graphs from conditional independence (CI) tests, offering interpretable, statistically principled inference without requiring a global scoring function. However, these tests are sensitive to sample size and conditioning-set dimension: nonlinear CI tests such as CMI [12], KCIT [13] (\(O(n^3)\)), and RCoT [14] (\(O(d_f^2 n)\) for \(d_f\) Fourier features) remain costly at high \(d\), forcing practitioners to rely on fast linear tests that miss nonlinear dependencies. For broader surveys, see [15] and [16].

In this work, we propose GRACE (Gated Refinement for Accurate Causal Edge discovery), a framework for improving constraint-based causal discovery at high dimensionality. GRACE operates in two stages: first, a high-recall constraint-based method such as CDNOTS [9] or PCMCI [7] produces a candidate skeleton; then, a gated neural model refines this skeleton by assigning each candidate edge an independent Hard Concrete gate [17], [18] trained with \(L_0\) regularization, whose gate values concentrate near 0 or 1, yielding a robust binary graph at a natural 0.5 threshold. The gated model learns which skeleton edges genuinely improve prediction, pruning false positives while preserving true causal links—narrowing the gap with expensive nonlinear CI tests at a fraction of the computational cost. Although our experiments focus on lagged effects (\(\tau \geq 1\)), the framework natively supports contemporaneous (lag-0) edges: enabling lag-0 gates allows the same \(L_0\) mechanism to select among instantaneous candidates provided by the skeleton.

Our contributions are:

1. The GRACE framework: A gated refinement stage that converts a high-recall skeleton into a high-precision causal graph via Hard Concrete gates with \(L_0\) regularization. The gating mechanism is independent of the skeleton source: we demonstrate it with both CDNOTS and PCMCI skeletons.

2. An empirically derived regularization rule (Eq. 12 ) that adapts \(\lambda\) to problem dimensions (\(d\), \(T\)) and skeleton density, paired with \(L_0\) normalization that decouples regularization from batch size.

3. Real-world evaluation on the Elbe River network [19], where a temporal bootstrap variant—applying GRACE to random temporal windows and retaining edges that appear consistently—handles non-stationarity, hidden confounders, and variable lags without domain-specific heuristics.

2 Methodology↩︎

2.1 Problem Formulation↩︎

Consider a \(d\)-variate stationary time series \(\{\mathbf{x}_t\}_{t=1}^T\) with \(\mathbf{x}_t \in \mathbb{R}^d\). We assume a structural causal model with maximum lag \(L\): \[x_t^e = f_e\!\left(\{x_{t-\ell}^c : G[c, \ell, e] = 1\}\right) + \epsilon_t^e, \quad \epsilon_t^e \sim \mathcal{N}(0, \sigma_e^2), \label{eq:scm}\tag{1}\] where \(G \in \{0,1\}^{d \times (L+1) \times d}\) is the lagged causal graph and \(f_e\) is the (possibly nonlinear) causal mechanism for effect variable \(e\). The goal is to recover \(G\) from observations \(\{\mathbf{x}_t\}_{t=1}^T\).

The gated refinement stage assumes: (i) Causal sufficiency: there are no unobserved common causes; (ii) Finite lag: all causal effects occur within lags \(0, \ldots, L\); (iii) Additive structure: effects combine additively across parents, though each individual effect \(f_{c,\ell \to e}\) can be nonlinear. Additional assumptions (e.g., stationarity) are inherited from the skeleton discovery algorithm; when using CDNOTS, nonstationarity is handled natively.

2.2 Gated Causal Model↩︎

We construct a predictive model for each effect variable \(e\), where candidate causal inputs are individually gated. The architecture consists of three components (Figure 1): shared encoder, per-effect gated aggregation, and per-effect decoder.

2.2.0.1 Shared encoder.

A linear projection maps each lagged input to a hidden representation shared across all effect variables: \[h_{c,\ell} = \mathbf{W}_{c,\ell}\, x_{t-\ell}^c + b_{c,\ell} \in \mathbb{R}^H, \label{eq:encoder}\tag{2}\] where \(H\) is the encoder hidden dimension, \(\mathbf{W}_{c,\ell} \in \mathbb{R}^{H \times 1}\), and \(b_{c,\ell} \in \mathbb{R}^H\). Sharing the encoder across effects reduces parameters from \(O(d^2 L H)\) to \(O(d L H)\); effect-specificity is provided by the gates.

2.2.0.2 Gated aggregation.

For each effect \(e\), a Hard Concrete gate \(z_{c,\ell,e} \in \{0, 1\}\) independently controls each candidate parent: \[\hat{\mathbf{x}}_e = \sum_{c,\ell} z_{c,\ell,e} \cdot h_{c,\ell}. \label{eq:aggregate}\tag{3}\] Gates are parameterized by learnable log-odds \(\log\alpha_{c,\ell,e}\) (Section 2.3).

2.2.0.3 Per-effect decoders.

Each effect variable has an independent two-layer MLP decoder with SiLU activations that maps the aggregated representation to Gaussian parameters: \[(\mu_e, \log\sigma_e) = \text{MLP}_e(\hat{\mathbf{x}}_e), \quad \text{MLP}_e : \mathbb{R}^H \to \mathbb{R}^2. \label{eq:decoder}\tag{4}\] Using per-effect decoders (rather than a shared decoder) prevents conflation of distinct causal mechanisms—e.g., a quadratic effect \(x^c_{t-1} \to x^e_t\) and a linear effect \(x^c_{t-1} \to x^{e'}_t\) can learn separate response functions.

2.3 Hard Concrete Gates↩︎

We adopt the Hard Concrete distribution [17], which stretches the binary concrete (Gumbel-Softmax) [18], [20] distribution to produce exact zeros and ones.

2.3.0.1 Training (stochastic).

During training, gates are sampled via the reparameterization trick: \[\begin{align} u &\sim \text{Uniform}(\epsilon, 1-\epsilon), \tag{5} \\ s &= \mathop{\mathrm{sigmoid}}\!\left(\frac{\log u - \log(1-u) + \log\alpha}{\tau}\right), \tag{6} \\ \bar{z} &= s \cdot (\zeta - \gamma) + \gamma, \tag{7} \\ z &= \min(\max(\bar{z}, 0), 1), \tag{8} \end{align}\] where \(\tau = 2/3\) is the temperature, and \(\gamma = -0.1\), \(\zeta = 1.1\) define the stretch interval. The stretch beyond \([0,1]\) followed by clamping places positive probability mass on exact 0 and exact 1.

2.3.0.2 Evaluation (deterministic).

At test time, we use the deterministic Hard Concrete estimator from [17]: \[z = \min\!\left(\max\!\left(\mathop{\mathrm{sigmoid}}(\log\alpha) \cdot (\zeta - \gamma) + \gamma,\; 0\right),\; 1\right). \label{eq:hc95det}\tag{9}\] Crucially, the stretch-and-clamp mechanism places positive probability mass on exact zero, so trained gates converge to either \(z=0\) (edge absent) or \(z \approx 1\) (edge present); Appendix 5 confirms this empirically with gate value histograms showing clear bimodal separation between true and false edges. The final causal graph is obtained by thresholding deterministic gate values at \(0.5\)—a natural boundary that falls in the gap between the two modes, making the binary decision robust to the exact threshold choice. Unlike the post-hoc weight thresholding required by \(L_1\) and attention-based methods—where continuous edge scores cluster in a narrow range and results are sensitive to the chosen cutoff—the bimodal gate distribution makes the \(0.5\) boundary a principled default rather than an arbitrary tuning parameter. The gate initialization (\(\log\alpha={-}0.5\)) and \(L_0\) penalty strength (\(\lambda\)) remain as design choices, but these are set once before training rather than tuned after inspecting learned weights.

2.3.0.3 \(L_0\) penalty.

The probability that a gate is nonzero has a closed-form expression: \[\mathbb{P}(z \neq 0) = \mathop{\mathrm{sigmoid}}\!\left(\log\alpha - \tau \log\!\frac{-\gamma}{\zeta}\right). \label{eq:l095penalty}\tag{10}\] This enables an analytic \(L_0\) regularizer without Monte Carlo estimation.

2.3.0.4 Initialization.

We initialize \(\log\alpha = -0.5\) for all gates, yielding a deterministic gate value of approximately \(0.36\) (below the active threshold of \(0.5\)). Gates start mostly closed, requiring edges to demonstrate predictive value to open. This asymmetric initialization reduces false positives compared to a symmetric initialization near \(\log\alpha = 0\).

2.4 Loss Function↩︎

The training objective combines prediction quality with sparsity: \[\begin{align} \mathcal{L} &= \underbrace{\frac{1}{Bd}\sum_{t=1}^B \sum_{e=1}^d \text{NLL}(x_t^e \mid \mu_t^e, \sigma_t^e)}_{\text{prediction}} + \underbrace{\lambda \cdot s \sum_{c,\ell,e} \mathbb{P}(z_{c,\ell,e} \neq 0)}_{\text{L_0 sparsity}} \notag \\ &\quad + \underbrace{\lambda \cdot \lambda_{\text{lag}} \sum_{c,e} \mathop{\mathrm{ReLU}}\!\left(\sum_\ell \mathbb{P}(z_{c,\ell,e} \neq 0) - k\right)}_{\text{lag concentration (optional)}}, \label{eq:loss} \end{align}\tag{11}\] where \(\text{NLL}(x \mid \mu, \sigma) = \frac{(x-\mu)^2}{2\sigma^2} + \log\sigma + \frac{1}{2}\log 2\pi\) is the Gaussian negative log-likelihood, \(B\) is the batch size, and \(d\) is the number of variables.

2.4.0.1 \(L_0\) normalization.

The scaling factor \(s = 1/\lceil N/B \rceil\) normalizes the \(L_0\) penalty by the number of gradient steps per epoch, ensuring that \(\lambda\) values transfer reliably across dataset sizes (Section 2.5).

2.5 Practical Considerations↩︎

2.5.0.1 Batch size and \(L_0\) normalization.

With \(N\) samples and batch size \(B\), there are \(\lceil N/B \rceil\) gradient steps per epoch. Without normalization, the \(L_0\) penalty is applied per step, while the NLL is averaged per sample, creating an implicit coupling: \(\text{effective } \lambda \propto \lambda \cdot N/B\). Our \(L_0\) normalization (factor \(s = 1/\lceil N/B \rceil\) in Eq. 11 ) decouples regularization from batch size. The batch size is set automatically as \(B = \max(32, \min(N/8, 256))\).

2.5.0.2 Lag concentration (optional).

By default, each gate is penalized independently by the \(L_0\) term and the skeleton determines the lag structure—if the CI skeleton identifies edges at multiple lags for a pair, all corresponding gates are learned without additional constraint. When domain knowledge or exploratory analysis (e.g., partial cross-correlation) suggests that causal effects concentrate at a small number of lags, an optional lag concentration penalty (third term in Eq. 11 ) can be enabled by setting \(\lambda_{\text{lag}} > 0\). This penalty discourages more than \(k\) active lags per (cause, effect) pair via a soft hinge: \(k\) lags are “free”; the penalty is linear above \(k\). The lag term omits the per-step normalization factor \(s\): it is a structural penalty on edge multiplicity, not a per-sample regularizer. In our experiments, \(\lambda_{\text{lag}} = 0\) (disabled) since the skeleton already provides the appropriate lag structure.

2.5.0.3 Training details.

We use Adam [21] with learning rate \(10^{-3}\) and weight decay \(10^{-4}\). The encoder hidden dimension is \(H = 64\); each per-effect decoder is a two-layer MLP with hidden size 64 and SiLU activations. The noise scale \(\sigma_e\) is not a fixed parameter but a learned output of the decoder (via softplus), optimized jointly with all other parameters. Training runs for up to 150 epochs with early stopping (patience 20, monitoring training loss).

2.6 CDNOTS-Gated Refinement↩︎

Constraint-based methods like CDNOTS [9] with fast linear CI tests (partial correlation) achieve high recall but suffer from low precision as dimensionality increases: in our benchmarks (Section 3), CDNOTS achieves \(95\%\) TPR but only \(37\%\) precision at \(d=100\)/\(T=1000\), yielding F1 \(=\) 0.53. GRACE uses the gated model from Section 2.2 as a precision filter: starting from this high-recall skeleton, it learns which edges genuinely improve prediction and prunes the rest.

2.6.0.1 Pipeline.

The CDNOTS-Gated refinement operates in two stages:

1. CDNOTS skeleton: Run CDNOTS with partial correlation CI tests at significance level \(\alpha = 0.05\) to obtain a binary skeleton \(S \in \{0,1\}^{d \times d \times (L+1)}\). This produces a high-recall superset of the true causal graph.

2. Gated refinement: Train a single gated model with the skeleton as a hard mask—only edges present in \(S\) have learnable gates; all others are fixed to zero. Gates are initialized at \(\log\alpha = -0.5\) (deterministic gate value \(\approx 0.36\), below the active threshold of \(0.5\)) for skeleton edges, so each edge must demonstrate predictive value to open.

2.6.0.2 \(\lambda\) selection.

We use an empirically derived rule based on problem dimensions and skeleton density: \[\lambda = \max\!\bigl(0.007 + 0.16\,\rho_S,\;\; \tfrac{1}{d} + \tfrac{4}{T} - 0.4\,\rho_S^{0.8}\bigr), \label{eq:auto95lambda}\tag{12}\] where \(\rho_S = |S| / (d^2(L+1) - d)\) is the fraction of candidate edges in the skeleton. The coefficients were fitted by grid search over the oracle-best \(\lambda\) on a held-out calibration set of \({\sim}1{,}000\) random Erdős–Rényi graphs spanning \(d \in \{10, \ldots, 100\}\), \(T \in \{500, \ldots, 2000\}\), and edge densities from sparse to 0.12—deliberately disjoint from the SCP benchmarks used in Section 3, to avoid overfitting to the evaluation data. The main term (\(1/d + 4/T - 0.4\,\rho_S^{0.8}\)) adapts to problem scale: more variables and more data allow weaker regularization, with a sublinear density correction. The floor (\(0.007 + 0.16\,\rho_S\)) ensures sufficient regularization for dense skeletons. This reduces computation to a single model-training run.

2.6.0.3 Why it works.

The combination succeeds because the two components address complementary failure modes. CDNOTS with partial correlation provides a computationally cheap skeleton with very high recall—in practice, most nonlinear causal relationships also induce some linear association that partial correlation can detect—but many false positives at high \(d\). The gated model, constrained to the skeleton, operates in a much smaller search space (\(|S| \ll d^2 L\)), enabling it to learn nonlinear causal mechanisms and prune edges that don’t genuinely improve prediction.

A natural question is why a single \(\lambda\) value works across all edges, rather than requiring per-edge tuning. To build intuition (not a formal guarantee for the nonlinear case), Appendix 6 analyzes a simplified linear Gaussian model and shows that each edge has a critical penalty \(\lambda^*\) proportional to its NLL reduction \(\Delta\)—the improvement in prediction from opening that gate. True causal edges reduce NLL by a constant \(\Delta^{\text{true}} = \Theta(1)\) (determined by signal strength), while false edges can only overfit noise, giving \(\Delta^{\text{false}} = O(1/T)\). The ratio of their critical penalties is therefore \(\lambda^{*\text{true}} / \lambda^{*\text{false}} = T \cdot \text{SNR}\), where \(\text{SNR}\) is the per-edge signal-to-noise ratio (formally defined in Appendix 6 as the conditional SNR). This means a separation interval of valid \(\lambda\) values exists and widens linearly with sample size, so any \(\lambda\) in this interval recovers the correct graph. The analysis does not directly derive the constants in Eq. 12 , but it explains three qualitative properties that the heuristic satisfies: (i) the skeleton pre-filters candidates, reducing the number of false edges that \(\lambda\) must suppress; (ii) the gate initialization below the decision boundary creates inertia against false edges opening; and (iii) the dependence on \(\rho_S\) is justified because denser skeletons admit more false candidates, requiring stronger regularization to keep the expected false positive count bounded.

2.6.0.4 Comparison with nonlinear CI tests.

An alternative to gated refinement is using expensive nonlinear CI tests (RCoT, KCIT) within CDNOTS. Using a nonlinear CI test (RCoT) within the skeleton generally improves F1, since it can detect dependencies that linear partial correlation misses. However, the runtime cost is steep—in our experiments, 40–75\(\times\) slower at \(d \geq 50\). Gated refinement with a fast linear skeleton narrows the F1 gap at a fraction of the cost, requiring only \(O(|S| \cdot H)\) per gradient step. We validate this empirically in Section 3.3.0.3 (full results in Appendix 7).

3 Experiments↩︎

We evaluate GRACE on synthetic benchmarks with known ground-truth causal graphs, focusing on high-dimensional scaling (\(d \in \{20, 30, 50, 70, 100\}\)) with varying sample sizes (\(T \in \{300, 500, 1000, 2000\}\)).

3.1 Setup↩︎

3.1.0.1 Data generation.

We generate random structural causal processes (SCPs) with controlled properties—the standard benchmark class for time-series causal discovery, modeling systems common in climate science, econometrics, and neuroscience [7], [15]. Graphs are sparse (\(\sim\)​1.3–2.2 edges per variable), with lags drawn uniformly from \(1\) to \(L=5\). Dependency coefficients are sampled uniformly from \([0.40, 0.89]\) or \([-0.89, -0.40]\) with equal probability, and functions are drawn from a mixed pool of linear and nonlinear forms: \(\{x, x^2, |x|, \sin(x)\}\). For each dimensionality \(d\), we generate 10 random graphs (seeds 0–9) at \(T_{\max} = 2500\) and slice to each target \(T\) for evaluation, ensuring that results across sample sizes reflect the same underlying causal structure.

3.1.0.2 Metrics.

We report F1 score (harmonic mean of precision and recall), along with precision and true positive rate (TPR/recall) to characterize the precision–recall tradeoff. Structural Hamming Distance (SHD) results are provided in Appendix 8. All results are averaged over 10 random graph seeds; we report mean \(\pm\) standard deviation.

3.2 Baselines↩︎

We compare against representative methods from each major paradigm—constraint-based (PCMCI, CDNOTS), attention-based (TCDF), and score-based (CUTS+). CUTS+ subsumes the score-based family: it outperforms NGC, DYNOTEARS, and NAVAR on its own benchmarks [5], so including all four would be redundant.

• PCMCI [7]: Constraint-based method with partial correlation CI test. Significance level \(\alpha = 0.05\).

• CDNOTS [9]: Nonstationary causal discovery with partial correlation.

• TCDF [6]: Attention-based convolutional approach with permutation testing.

• CDNOTS-Gated (GRACE): Our CDNOTS-Gated refinement pipeline (Section 2.6).

• CUTS+ [5]: Score-based method using Gumbel-Softmax gates on a GRU-based message-passing network, designed explicitly for high-dimensional causal discovery from irregular time series. CUTS+ is the most directly comparable score-based baseline at scale; since it outputs pairwise probabilities without lag resolution, we evaluate both methods at pair-level F1 (lags collapsed) for a fair comparison, setting aside GRACE’s lag-specific output. GRACE substantially outperforms CUTS+ across all configurations (Appendix 9). We use their published default settings with threshold \(0.5\) and set the input window to match the true maximum lag.

3.3 Results: F1 Across Dimensionality and Sample Size↩︎

Table 1 and Figure 2 present F1 scores across all methods, dimensionalities, and sample sizes. We highlight several key findings:

3.3.0.1 CDNOTS-Gated dominates at high dimensionality.

At \(d=50\)/\(T=1000\), CDNOTS-Gated achieves F1 \(=\) 0.915, compared to 0.742 for CDNOTS and 0.691 for TCDF. At \(d=100\), the gap widens further (F1 \(=\) 0.878 vs. at \(T=500\)), with the advantage growing monotonically in \(d\). PCMCI achieves \(>97\%\) recall at all \(d\) but precision collapses from 30% (\(d=20\)) to 6% (\(d=100\))—a common challenge as the hypothesis space grows as \(O(d^2 L)\). TCDF achieves high precision (\(>93\%\)) through permutation testing but limited recall (\({\sim}56\%\)), plateauing around F1 \(\approx\) 0.72. CUTS+ precision collapses at high \(d\): pair-level F1 drops to 0.055 at \(d=100\) (Appendix 9).

3.3.0.2 Precision–recall tradeoff.

Table 2 reveals the mechanism: CDNOTS-Gated preserves most of CDNOTS’s recall (83–88% vs.–95%) while boosting precision from 37–84% to 86–97%. At \(d=70\), precision reaches 97.5%—nearly every predicted edge is correct. Runtime remains practical (\({\sim}4\) min at \(d=100\); Appendix 10).

3.3.0.3 Nonlinear CI skeleton.

Replacing the linear skeleton with RCoT [14] improves F1 slightly (0.944 vs. at \(d=100\)/\(T=2000\)) but at \(75\times\) the cost; at lower \(T\), ParCorr-Gated matches or exceeds RCoT-Gated (Appendix 7).

3.3.0.4 Ablations and generalization.

Without a skeleton, \(L_0\) suppression alone cannot recover the signal at high \(d\) (F1 degrades substantially); the adaptive \(\lambda\) formula consistently outperforms fixed values; lag overspecification is handled gracefully (\(F_1 \geq 0.878\)) via \(L_0\) suppression (Appendix 11). GRACE generalizes across skeleton sources: PCMCI-Gated yields \(3\)–\(5\times\) F1 over raw PCMCI (Appendix 12), and DYNOTEARS-Gated matches CDNOTS-Gated at \(d \leq 50\) (Appendix 13). It also generalizes across non-additive mechanisms (Lorenz-96; Appendix 14) and diverse graph topologies (scale-free, Erdős–Rényi, small-world; Appendix 15).

3.4 Real-World Evaluation: Elbe River Network↩︎

Synthetic benchmarks offer controlled evaluation but cannot capture the compounding assumption violations present in real-world systems. We evaluate GRACE on a river network dataset [19] providing water level measurements along river networks in eastern Germany, where ground truth is determined by physical flow direction (upstream \(\to\) downstream).

3.4.0.1 Challenges.

River data violates standard assumptions: rainfall acts as a time-varying, spatially-varying hidden confounder; causal lags vary with water level (higher flow \(=\) faster propagation); seasonal floods create distributional shifts; and reservoirs disconnect natural flow. These challenges make causal discovery on river data substantially harder than on synthetic benchmarks [19].

3.4.0.2 Setup.

We focus on the Elbe River’s main branch in eastern Germany: 12 gauging stations forming a chain graph with 11 directed edges (upstream \(\to\) downstream), measured at 6-hour resolution over 2019–2023 (\(T=3{,}164\)). We compare PCMCI, CDNOTS, single-run GRACE, and GRACE-Bootstrap. When applied to the full multi-year time series, all skeleton methods achieve perfect recall but extremely poor precision due to ubiquitous confounding—e.g., CDNOTS recovers all 11 true edges but produces 106 false positives (\(F_1=0.17\)).

3.4.0.3 Temporal bootstrap.

We apply GRACE independently to random 1-month windows and retain edges appearing in more than a \(\tau\) fraction of runs—a temporal analogue of stability selection [22]. Because confounders like rainfall are transient and episodic, they inflate correlations only in the windows they affect; true flow edges, by contrast, persist across all hydrological conditions. Windowing thus separates stable causal signal from transient confounding without requiring confounder measurements.

3.4.0.4 Results.

We set \(\tau=0.70\), requiring an edge to appear in more than 70% of windows to be retained—a conservative threshold that does not require ground-truth calibration and is interpretable on its own terms. Table 3 shows that GRACE-Bootstrap (50 runs, \(\tau=0.70\)) recovers 9 of 11 edges with only 1 false positive (\(F_1=0.86\), Precision \(=\) 0.90), reducing CDNOTS’s 106 FPs by 99%. The bootstrap AUROC reaches 0.986, showing that edge frequency provides an excellent ranking even when the binary threshold excludes borderline edges.

Together, the synthetic experiments (Section 3.3) and this real-world evaluation cover complementary axes of difficulty: the former tests scalability under controlled conditions, the latter tests robustness when causal assumptions are simultaneously and severely violated—a combination that no single benchmark can address.

4 Conclusion↩︎

We presented GRACE, a framework that refines constraint-based causal discovery via Hard Concrete gates with \(L_0\) regularization. On synthetic benchmarks, GRACE nearly doubles F1 over raw CDNOTS at \(d=100\) (0.92 vs.) with 97% precision, generalizing across skeleton sources (PCMCI, DYNOTEARS) and graph topologies. On the real-world Elbe River benchmark [19], a temporal bootstrap variant achieves \(F_1 = 0.86\) with AUROC\({} = 0.99\) despite hidden confounders, variable lags, and distributional shifts.

4.0.0.1 Limitations.

On dense graphs (scale-free hubs at \(d \gtrsim 70\)), per-edge SNR decreases and gate values cluster below threshold—AUROC remains high but binary F1 degrades, partially mitigated by more data. Spurious edges from latent confounders may survive if they improve prediction. The linear skeleton may miss purely symmetric nonlinear dependencies (e.g., \(x^2\)) that induce no linear association; a nonlinear CI skeleton mitigates this at higher cost (Appendix 7). Gated refinement trades a small number of true positives for a large precision gain; the recall reduction is minor in our experiments but may be more pronounced at very low sample sizes.

5 Gate Value Distribution↩︎

Figure 3 plots the distribution of deterministic gate values (Eq. 9 ) after training, partitioned by whether each skeleton edge corresponds to a true or false causal link. At \(d=100\), false edges cluster near zero (mass at \(0.0\) and \(0.25\)–\(0.40\)) while true edges concentrate at \(0.55\)–\(0.85\), with almost no overlap at the \(0.5\) decision boundary. At \(d=50\), the separation is somewhat narrower, with minor overlap near \(0.5\), consistent with the slightly lower F1 at this scale. The clear bimodal pattern confirms that Hard Concrete gates learn to separate true from false edges without requiring post-hoc threshold tuning: the \(0.5\) boundary falls squarely in the gap between the two modes.

6 Theoretical Analysis of \(\lambda\)-Sensitivity↩︎

This appendix analyzes the \(\lambda\)-sensitivity of the gated model under a simplified linear Gaussian setting. The results provide qualitative intuition for why a single \(\lambda\) can separate true and false edges, but do not constitute identifiability guarantees for the full nonlinear model with MLP decoders. We view these results as motivation for the design choices in Section 2.6, not as formal recovery proofs.

6.1 Setup and Notation↩︎

Consider a stationary VAR process with \(d\) variables and maximum lag \(L\): \[x^e_t = \sum_{(c,\ell) \in \text{pa}(e)} \beta_{c\ell}^e \, x^c_{t-\ell} + \varepsilon^e_t, \qquad \varepsilon^e_t \stackrel{\text{iid}}{\sim} \mathcal{N}(0, \sigma_e^2), \label{eq:dgp}\tag{13}\] where \(\text{pa}(e)\) denotes the true parent set of variable \(e\) (the set of \((c,\ell)\) pairs with \(\beta_{c\ell}^e \neq 0\)). We analyze the linear encoder case; the nonlinear extension is discussed in the Remark (Nonlinear extension) below.

The per-sample loss for effect \(e\), with weight \(w_{c\ell}^e\) for each gated input, is: \[\ell_e = \frac{1}{2\hat{\sigma}_e^2}\left(x^e_t - \sum_{c,\ell} z_{c\ell e} \, w_{c\ell}^e \, x^c_{t-\ell}\right)^2 + \log\hat{\sigma}_e + \tfrac{1}{2}\log 2\pi. \label{eq:per95sample95loss}\tag{14}\] The full training objective (Eq. 11 ) averages \(\ell_e\) over both \(T\) samples and \(d\) effects—i.e., \(\overline{\text{NLL}} = \frac{1}{d}\sum_e \frac{1}{T}\sum_t \ell_e(t)\)—and adds the \(L_0\) penalty \(\lambda \cdot s \sum_{c,\ell,e} \mathbb{P}(z_{c\ell e} \neq 0)\), where \(s = 1/\lceil T/B \rceil\) is the per-step normalization factor. The \(1/d\) averaging is important: it means each gate’s NLL contribution is scaled by \(1/d\), which enters the critical \(\lambda\) formula below.

6.2 Per-Edge NLL Reduction↩︎

Under the linear Gaussian DGP (Eq. 13 ), assuming the other gates and weights are at their population optima:

6.2.0.1 True edge \(((c,\ell) \in \text{pa}(e))\).

The optimal weight is \(w^{*} = \beta_{c\ell}^e\), and the NLL reduction equals the variance uniquely explained by this edge—the variance of \(x^c_{t-\ell}\) after projecting out all other parents already included: \[\Delta_{c\ell e}^{\text{true}} = \frac{(\beta_{c\ell}^e)^2 \, \widehat{\text{Var}}(x^c_{t-\ell} \mid \widehat{\text{pa}}^e_{-(c,\ell)})}{2\,\sigma_e^2}, \label{eq:delta95true}\tag{15}\] where \(\widehat{\text{Var}}(x^c_{t-\ell} \mid \widehat{\text{pa}}^e_{-(c,\ell)})\) is the empirical variance of the residual of \(x^c_{t-\ell}\) after linear projection onto all other true parents of \(e\). This is the conditional variance formula from multivariate regression: in a regression with correlated regressors, the incremental gain from one predictor depends on its residualized variance, not its marginal variance. By the law of large numbers this converges to \(\text{Var}(x^c_{t-\ell} \mid \text{pa}^e_{-(c,\ell)}) > 0\) (positive unless a true parent is a perfect linear combination of others), so \(\Delta^{\text{true}}\) converges to a positive constant.

6.2.0.2 False edge \(((c,\ell) \notin \text{pa}(e))\).

Opening a spurious gate adds one free parameter that can only fit noise. By standard regression theory, the in-sample NLL reduction from fitting one noise regressor to \(T\) residuals is: \[\Delta_{c\ell e}^{\text{false}} = \frac{\hat{r}_{c\ell e}^2}{2} + O(T^{-1}), \label{eq:delta95false}\tag{16}\] where \(\hat{r}_{c\ell e}\) is the sample correlation between the residual and \(x^c_{t-\ell}\). Under the null hypothesis (no causal effect), \(T \cdot \hat{r}^2 \sim \chi^2(1)\), so \(\hat{r}^2\) is itself \(O(1/T)\) and \[\mathbb{E}[\Delta^{\text{false}}] = \frac{1}{2T}. \label{eq:delta95false95exp}\tag{17}\]

6.3 Critical \(\lambda\) and Separation↩︎

7 Nonlinear CI Test Results↩︎

Section 3.3.0.3 compares the fast linear CI skeleton (ParCorr) with a nonlinear one (RCoT) in terms of F1 accuracy and runtime cost. Here we present the full results. We compare CDNOTS-Gated with the default ParCorr skeleton against CDNOTS-Gated with an RCoT skeleton, reporting both F1 and runtime.

Table 4 compares both F1 and runtime for the two Gated variants. RCoT-Gated outperforms ParCorr-Gated at high \(T\), since RCoT can detect nonlinear conditional dependencies that linear partial correlation misses. At \(d=100\)/\(T=2000\), RCoT-Gated achieves F1 \(=\) 0.944 vs. for ParCorr-Gated. However, at moderate \(T\), the gap narrows or reverses (e.g., 0.916 vs. at \(d=100\)/\(T=1000\)), while the runtime cost is substantial: at \(d=50\)/\(T=1000\), RCoT-Gated requires 4,060s vs.100s for ParCorr-Gated—a 40\(\times\) slowdown. At \(d=100\), the ratio reaches 75\(\times\) (\({\sim}17{,}000\)s vs.\({\sim}228\)s). The practical conclusion is clear: for \(d \geq 50\), ParCorr-Gated is the most practical choice, with RCoT offering diminishing returns at disproportionate cost.

8 SHD Results↩︎

We report the Structural Hamming Distance (SHD)—the total number of edge additions, deletions, and reversals needed to transform the predicted graph into the true graph—to complement the F1 results in the main text. Lower SHD indicates a predicted graph closer to ground truth.

Table 5 confirms the F1 findings: CDNOTS-Gated achieves the lowest SHD at all dimensionalities and sample sizes. The SHD reduction over raw CDNOTS is particularly striking at high \(d\) with sufficient samples: at \(d=100\)/\(T=1000\), CDNOTS-Gated achieves SHD \(=\) 27 vs. for CDNOTS—a 90% reduction in graph edit distance. PCMCI’s SHD grows dramatically with \(d\) (reaching \({\sim}2{,}500\) at \(d=100\)), reflecting its high false positive rate, while TCDF achieves competitive SHD at high \(T\) but is less consistent at low sample sizes.

9 Score-Based Baseline: CUTS+ Comparison↩︎

Since CUTS+ [5] outputs a pairwise probability matrix without lag resolution, we report pair-level metrics: the lag dimension is collapsed (an edge exists if any lag is active) before computing precision, recall, and F1. We use their published default settings: threshold \(0.5\), 64 training epochs, GRU hidden size 32, and Gumbel temperature annealing from \(\tau=1\) to \(\tau=0.1\). The input window is set to match the true maximum lag (\(L=5\)).

Table 6 shows that CUTS+ pair-level F1 degrades sharply with dimensionality: from \(0.510\) at \(d=20\) to \(0.055\) at \(d=100\) (\(T=1000\)). Table 7 reveals the underlying cause: CUTS+ achieves high recall (\(>83\%\)) by opening most Gumbel-Softmax gates, but precision collapses from \(37\%\) at \(d=20\) to \(2.9\%\) at \(d=100\). Its Gumbel-Softmax gates produce continuous probabilities that cluster in a narrow range (typically \(0.36\)–\(0.71\)), making it difficult to separate true from false edges regardless of threshold choice. By contrast, GRACE’s Hard Concrete gates concentrate near exact zero or one (Appendix 5), yielding a robust binary decision at the \(0.5\) boundary.

10 Runtime↩︎

Table 8 and Figure 4 compare wall-clock runtime at \(T=1000\). CDNOTS-Gated adds \(2\)–\(3.5\times\) overhead over raw CDNOTS while achieving substantially higher F1. The runtime is dominated by the CDNOTS skeleton at high \(d\); the gated training phase itself scales linearly in the number of skeleton edges.

11 Ablation Study↩︎

We ablate three design choices in GRACE on the same SCP benchmarks used in the main experiments (\(d \in \{50, 100\}\), \(T \in \{500, 1000\}\), 10 seeds each). The baseline is GRACE with full CDNOTS skeleton and the adaptive \(\lambda\) formula.

11.0.0.1 Skeleton quality.

Table 9 evaluates three skeleton configurations: (i) No skeleton—all \(d^2(L+1) - d\) candidate edges are passed to the gated model; (ii) Shallow CI (\(\texttt{max\_degree}=1\))—only pairwise and first-order conditional independence tests; (iii) Full CDNOTS (default)—conditioning sets up to the maximum graph degree. Without any skeleton pruning, the gated model faces \(\sim\)​15,000 candidate edges at \(d=50\) for \(\sim\)​100 true edges; no \(\lambda\) setting can recover signal at this ratio (F1 \(<\) 0.02). Shallow CI improves substantially (F1 \(=\) 0.46–0.71) but still underperforms full CDNOTS by 25–50% in F1, reflecting the spurious edges that higher-order conditioning would remove.

11.0.0.2 Lambda sensitivity.

Table 10 compares fixed \(\lambda\) values against the adaptive formula (Eq. 12 ). A fixed \(\lambda=0.01\) performs comparably at \(d=50\), since both yield similar values when the skeleton is sparse. At \(d=100\), the adaptive formula’s \(1/d\) term lowers \(\lambda\) appropriately, maintaining high F1 where a fixed value would overregularize. \(\lambda=0.05\) already degrades at \(d=100\) (F1 drops from 0.92 to 0.49), and \(\lambda \geq 0.1\) aggressively penalizes gates, suppressing most true edges.

11.0.0.3 Lag misspecification robustness.

In practice, the true maximum causal lag \(L\) is unknown and must be specified by the practitioner. We evaluate robustness to lag misspecification by generating SCP graphs with true max lag \(L = 5\) (\(d=50\), \(T=1000\), 10 seeds) and running each method with assumed search lags \(L_{\text{search}} \in \{2, 3, 5, 7, 10\}\). Evaluation uses the ground truth padded (or truncated) to \(L_{\text{search}}\), so that edges beyond the search horizon count as false negatives for all methods equally.

Under underspecification (\(L_{\text{search}} < 5\)), late-lag true edges are outside the search space and count as false negatives for all methods equally. GRACE degrades from \(0.906\) at the correct lag to \(0.793\) at \(L=2\), while CDNOTS drops further to \(0.626\)—a gap of \(0.167\) vs.\(0.113\) points respectively, showing that GRACE is more robust even in the underspecified regime. Under overspecification (\(L_{\text{search}} > 5\)), the advantage is more striking: GRACE maintains \(F_1 = 0.878\)–\(0.891\) at \(L \in \{7, 10\}\) (precision \(\geq 0.908\)), because the Hard Concrete \(L_0\) penalty suppresses the extra false-lag candidates. CDNOTS has no such suppression mechanism and loses \(0.053\) F1 points at \(L=10\) relative to \(L=5\), as spurious cross-lag edges inflate false positives.

12 PCMCI-Gated: Refinement Across Skeleton Sources↩︎

To demonstrate that GRACE’s gated refinement is not tied to a specific skeleton algorithm, we replace the CDNOTS skeleton with PCMCI [7] and apply the same gating pipeline (denoted PCMCI-Gated). PCMCI provides an ideal stress test: it achieves very high recall (\(>90\%\)) but poor precision that degrades sharply with dimensionality (from \({\sim}30\%\) at \(d=20\) to \({\sim}6\%\) at \(d=100\)), producing a skeleton with many false positives for gating to prune.

Tables 12 and 14 show that gated refinement consistently improves PCMCI across all dimensionalities. At \(d=50\)/\(T=1000\), F1 improves from 0.236 to 0.758 (\(3.2\times\)) and SHD drops from 627 to 48 (\(13\times\)); at \(d=100\)/\(T=1000\), F1 improves from 0.117 to 0.602 (\(5.1\times\)) and SHD from 2490 to 150 (\(17\times\)). Table 13 reveals the mechanism: gating converts PCMCI’s extreme recall/low-precision profile (e.g., 97% TPR / 6% precision at \(d=100\)) into a more balanced one (59% TPR / 66% precision).

Comparing with CDNOTS-Gated (Table 1), PCMCI-Gated achieves lower F1 at all configurations—e.g., 0.602 vs. at \(d=100\)/\(T=1000\)—because PCMCI’s skeleton has far more false positives than CDNOTS’s. This confirms that GRACE delivers consistent and substantial gains regardless of the skeleton source—up to \(5\times\) F1 improvement—while the final performance ceiling is set by the skeleton’s recall.

13 Score-Based Skeleton: DYNOTEARS↩︎

The preceding appendix demonstrated that GRACE improves CI-based skeletons (PCMCI). Here we test whether the same gating mechanism can refine a score-based skeleton produced by DYNOTEARS [4], which fits an \(L_1\)-penalized vector autoregression (LASSO VAR) per target variable using cross-validated penalty selection. LASSO VAR produces continuous coefficient magnitudes, and the central challenge is converting them to a binary graph: setting the threshold too low retains many false positives, while setting it too high discards true edges. This is precisely the thresholding problem that GRACE’s Hard Concrete gates solve automatically.

13.0.0.1 Oracle threshold analysis.

To quantify the thresholding difficulty, we sweep all possible coefficient thresholds on each seed and record the value that maximizes F1 given the ground truth (the “oracle” threshold). Across all configurations (\(d \in \{20, 50, 100\}\), \(T \in \{500, 1000\}\), 10 seeds each), the oracle threshold clusters in the range 0.04–0.09 but varies up to \(5\times\) across seeds (e.g., 0.038–0.188 at \(d=20\)). A fixed threshold of \(t=0.05\) recovers \({\sim}97\%\) of the oracle F1, making it a reasonable lower bound—but selecting it required access to the ground truth, which is unavailable in practice.

13.0.0.2 Results.

Table 15 compares four configurations: DYNOTEARS with the default threshold (\(t=0\), all nonzero coefficients), DYNOTEARS with the post-hoc oracle-informed threshold (\(t=0.05\)), DYNOTEARS-Gated (GRACE applied to the \(t=0.05\) skeleton), and CDNOTS-Gated (the default GRACE pipeline).

Several patterns emerge. First, DYNOTEARS with the default threshold (\(t=0\)) achieves F1 of only 0.13–0.24, confirming that raw LASSO coefficients without thresholding are unusable as a causal graph—the vast majority of nonzero coefficients are false positives. Second, the oracle threshold \(t=0.05^\dagger\) dramatically improves DYNOTEARS to F1 0.82–0.91, demonstrating that a good threshold exists but requires ground truth to find. Third, DYNOTEARS-Gated (\(t=0.05^\dagger\)) matches or slightly exceeds CDNOTS-Gated at \(d \leq 50\), showing that GRACE can effectively refine score-based skeletons when the skeleton quality is sufficient. Fourth, CDNOTS-Gated retains its advantage at \(d=100\)/\(T=500\) (0.888 vs.), where the LASSO skeleton becomes sparser and misses more true edges.

The key takeaway is that GRACE’s Hard Concrete gates provide the same benefit to score-based skeletons as to CI-based ones: automatic, threshold-free conversion of continuous edge scores into a binary causal graph. However, the DYNOTEARS-Gated results rely on a post-hoc threshold (\(\dagger\)) to pre-filter the skeleton before gating. By contrast, CDNOTS-Gated requires no such threshold—the CI-based skeleton is already binary—making it the preferred default in practice. This experiment confirms that GRACE is compatible with diverse skeleton sources: it delivers consistent improvements over skeletons from both CI-based methods (CDNOTS, PCMCI) and score-based methods (DYNOTEARS), with the gating mechanism providing the same threshold-free binary decision benefit across all skeleton types.

14 Lorenz-96 Benchmark: Multiplicative Interactions↩︎

The Lorenz-96 model [23] is a chaotic dynamical system defined by the ODE \(\dot{x}_i = (x_{i+1} - x_{i-2}) \cdot x_{i-1} - x_i + F\), where indices wrap cyclically. Each variable \(x_i\) has four causal parents at lag 1: \(x_{i-2}\), \(x_{i-1}\), \(x_i\) (self), and \(x_{i+1}\). Critically, the mechanism is multiplicative—\(x_{i-1}\) multiplies \((x_{i+1} - x_{i-2})\)—creating interaction effects that violate GRACE’s additive parent assumption. This benchmark directly tests whether GRACE remains useful when its structural assumptions are misspecified.

We use the same implementation parameters as CUTS+ [5]: forcing \(F=10\) (chaotic regime), integration time step \(\Delta t=0.1\) via scipy.integrate.odeint, additive Gaussian observation noise \(\sigma=0.1\), and burn-in of 1000 steps. The observed time series are thus noisy discretizations of the underlying chaotic ODE, making the causal discovery task more realistic and challenging.

Tables 16 and 17 show that GRACE substantially improves over CDNOTS on Lorenz-96 at higher dimensionality. At \(N=100\)/\(T=1000\), GRACE achieves F1 \(=\) 0.875 with perfect precision (1.000) compared to CDNOTS F1 \(=\) 0.717 with precision 0.665—a pattern consistent with the SCP results in Section 3. Both methods share the same recall (\({\sim}0.78\)), confirming that GRACE prunes false positives from the skeleton without removing true edges. The improvement grows with both \(N\) and \(T\): at \(N=50\)/\(T=500\) the gap is small (F1 0.764 vs.), but at \(N=100\)/\(T=2000\) it is substantial (0.894 vs.).

These results demonstrate that GRACE’s additive MLP decoders can effectively distinguish true from spurious edges even when the underlying causal mechanism is multiplicative. While the model cannot perfectly represent the true functional form \(x_{i-1} \cdot (x_{i+1} - x_{i-2})\), it learns sufficient signal from true parent variables to keep their gates open, while spurious edges provide no predictive value and are pruned. The observation noise (\(\sigma=0.1\)) further validates robustness: the method operates on noisy discretizations of the chaotic ODE, not clean trajectories.

15 Diverse Graph Topologies↩︎

The main experiments (Section 3) use sparse random SCPs where edges are sampled uniformly at random. A natural concern is whether GRACE generalizes to fundamentally different graph structures. We test on three standard graph families generated via NetworkX [24]:

• Scale-free (Barabási–Albert, \(m=2\)): hub-and-spoke topology with power-law degree distribution, characteristic of gene regulatory networks and metabolic pathways where a few master regulators influence many downstream targets [25].

• Erdős–Rényi (\(p=0.08\), directed): uniformly random edges, producing a denser graph (\({\sim}200\) cross-edges at \(d=50\)) with no preferential structure.

• Small-world (Watts–Strogatz, \(k=4\), \(p=0.3\)): clustered ring lattice with random rewiring, combining high local clustering with short path lengths, as observed in neural connectomes and social networks [26].

For undirected generators (BA, WS), edges are directed from lower to higher index to establish a causal ordering. Lags are drawn uniformly from \(\{1, \ldots, 5\}\), and causal mechanisms use the same coefficient ranges and mixed linear/nonlinear function pool as the SCP benchmarks. Autoregressive terms and spectral radius checking ensure stationarity.

We evaluate at \(d=50\), \(T=1000\) (8–10 seeds per configuration) using pair-level AUROC, which is threshold-free and enables a fair comparison between GRACE’s continuous gate values and CUTS+’s Gumbel-Softmax probabilities.

Table 18 shows that GRACE generalizes well across all three topologies. On small-world and Erdős–Rényi graphs GRACE leads in AUROC (0.976 and 0.839 vs. and 0.803). Scale-free is the one topology where CUTS+ achieves slightly higher AUROC (0.932 vs.), possibly because the hub-and-spoke structure suits CUTS+’s GRU message-passing architecture. The method maintains its advantage across fundamentally different graph structures—including the denser Erdős–Rényi graphs (edge density \({\sim}4\times\) higher than SCP)—demonstrating that gated refinement generalizes beyond the sparse random SCPs used in the main experiments.

16 Elbe River Network: Graph Comparison↩︎

Figure 5 visualizes the causal graphs recovered on the Elbe River’s main branch (Section 3.4). The CDNOTS skeleton (center) recovers all 11 true edges but adds 106 spurious connections, making the graph uninterpretable. GRACE-Bootstrap (right) prunes nearly all false positives while retaining the chain structure.

References↩︎