A Prior-Regularized Framework for Knowledge-Guided Differentiable Causal Discovery

Technical Summary: A Prior-Regularized Framework for Knowledge-Guided Differentiable Causal Discovery

Problem Statement
Differentiable causal discovery methods, such as NOTEARS, DAGMA, and GOLEM, have reformulated the combinatorial acyclicity constraint into smooth functions, enabling structure learning via standard gradient-based optimizers. However, these methods treat all candidate edges as equally unknown, ignoring decades of curated domain knowledge (e.g., transcription factor-target interactions or protein-protein associations). This symmetric treatment creates two critical issues:

1. Inefficiency: Optimizers expend capacity rediscovering known edges rather than identifying novel relationships.
2. Instability in Small-Sample Regimes: In high-dimensional, small-sample settings ($n < 2d$), common in transcriptomic studies, baseline differentiable methods often achieve near-random performance (e.g., F1 scores of 0.06–0.07 at $d=100, n=200$), rendering them ineffective precisely when data is scarce.

Existing approaches to incorporate domain knowledge are often restricted to specific settings (e.g., Bayesian priors or specific algorithms) and lack a general-purpose framework for injecting arbitrary pairwise regulatory priors into any differentiable causal discovery backbone.

Methodology
The paper proposes a universal Prior-Regularized Framework that augments existing differentiable causal discovery objectives with a penalty term derived from curated databases.

• Prior Construction: A prior matrix $W_{prior} \in \mathbb{R}^{d \times d}$ is constructed from curated databases (e.g., TRRUST v2 for transcription factors, STRING v12 for protein interactions). If a regulatory interaction $(g_i \to g_j)$ exists in the database, $(W_{prior})_{ij} = \mu$ (default $\mu=0.5$); otherwise, it is 0. For synthetic benchmarks, partial coverage is simulated by encoding a fraction $\alpha$ of true edges.
• Objective Function: The framework modifies the standard loss functions of NOTEARS, DAGMA, and GOLEM by adding a Frobenius norm penalty:
  $$L_{prior}(W) = L_{base}(W) + \lambda_p \|W - W_{prior}\|_F^2$$
  where $\lambda_p$ controls the strength of the prior guidance.
  • NOTEARS: Augments the augmented Lagrangian objective with the prior penalty.
  • DAGMA: Adds the penalty to the log-determinant constraint objective.
  • GOLEM: Incorporates the penalty into the Gaussian likelihood-based loss.
• Optimization: The framework uses the Adam optimizer with specific learning rates and iteration schedules. For DAGMA, Cholesky factorization is performed periodically to mitigate $O(d^3)$ computational costs.

Key Contributions

1. Universal Framework: A database-agnostic and backbone-agnostic computational pipeline that allows any differentiable causal discovery method to leverage curated pairwise knowledge.
2. Implementation: A fully reproducible implementation (approx. 600 lines of Python/PyTorch) with checkpoint-based fault tolerance, supporting synthetic benchmarks and real-data validation.
3. Systematic Benchmarking: A comprehensive evaluation across six dimensions ($d \in \{30, \dots, 200\}$) and four prior coverage rates ($\alpha \in \{0.1, \dots, 0.7\}$), totaling 1,440 independent runs.

Results

• Synthetic Benchmarks:
  • Prior-regularized NOTEARS improved F1 scores in 15 of 24 configurations.
  • The most significant gains occurred in small dimensions ($d \le 50$) where baseline performance is weakest. At $d=30$ with $\alpha=0.7$, F1 improved from 0.046 (baseline) to 0.090 (a 98% relative gain).
  • The mean $\Delta$F1 across all configurations was +0.013.
  • DAGMA showed no systematic benefit in small-$n$ regimes (often converging to zero edges), while GOLEM showed mixed but positive effects at higher dimensions ($d \ge 80$).
• Real-Data Validation (TCGA):
  • Applied to 33 TCGA cancer types ($d=200$).
  • Coverage Analysis: The STRING database (filtered score > 900) provided 100% coverage of cancer gene sets (mean 138 entries per set), whereas TRRUST covered only 21% (mean 0.5 entries).
  • Statistical Significance: While 11 of 33 cancer types showed positive $\Delta$F1, the Wilcoxon signed-rank test on the full set yielded $p=0.061$, failing to reach conventional significance ($\alpha=0.05$). The author attributes this gap to measurement noise and unmeasured confounders in real data, contrasting with the clean synthetic environment.
  • Biological Relevance: In cancer types with positive prior effects (e.g., UCEC), prior-regularized models discovered edges with significantly higher Gene Ontology (GO) enrichment (18-fold increase in cancer-relevant hits) compared to baselines.
  • DepMap Screening: Genome-wide CRISPR co-essentiality screening ($d=1000$) revealed no significant pairwise correlations, reinforcing the difficulty of causal discovery at scale without curated priors.
• Sensitivity: The framework is robust for $\lambda_p \in [0.5, 1.0]$ in synthetic data and $\lambda_p \in [0.1, 2.0]$ in real data. The author recommends $\lambda_p = 0.5$ as a practical default.

Significance and Claims
The paper positions this work as a Methodology article, emphasizing its contribution to computational regularization frameworks rather than novel biological discoveries.

• Bridging the Gap: The framework enables causal discovery to leverage existing biological knowledge rather than attempting to rediscover it from scratch.
• Targeted Utility: The method yields the largest improvements in the $n < 2d$ regime, where baseline methods are least reliable, offering a path from "near-random" to "moderately predictive" performance for transcriptomic studies.
• Scalability and Reproducibility: The approach is fully reproducible, database-agnostic (compatible with TRRUST, STRING, or any pairwise knowledge source), and backbone-agnostic.
• Future Outlook: The author notes that the framework's effectiveness is currently limited by the incompleteness of curated databases. As databases expand (e.g., via projects like Biomni), the prior-regularized approach is expected to become increasingly powerful, transforming decades of biological research into actionable constraints for machine learning models.

The author concludes that while the method does not solve the fundamental challenges of real-world noise and nonlinearity, it provides a foundational, generalizable tool for knowledge-guided causal structure learning.