🤖 machine learning

Optimization-Free Topological Sort for Causal Discovery via the Schur Complement of Score Jacobians

This paper introduces the Score-Schur Topological Sort (SSTS) algorithm, which bypasses non-convex structural optimization by extracting causal order directly from the Schur complement of score Jacobians, thereby reframing scalable causal discovery as a statistical estimation problem capable of handling high-dimensional non-linear graphs.

Original authors: Rui Wu, Hong Xie

Published 2026-04-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Rui Wu, Hong Xie

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to figure out the family tree of a large, chaotic family reunion based only on a group photo. You don't know who is the parent, who is the child, or who is just a cousin. In the world of data science, this is called Causal Discovery: figuring out "what causes what" from a pile of observations.

For a long time, solving this puzzle was like trying to find the perfect arrangement of 1,000 people in a line by shuffling them around blindly, checking every single possible order. This is slow, prone to getting stuck in "local optima" (thinking you found the best line when you actually just found a good one), and breaks down when the family gets too big.

This paper introduces a new way to solve the puzzle called SSTS (Score-Schur Topological Sort). Here is how it works, using simple analogies:

1. The Old Way: The Exhaustive Shuffler

Previous methods tried to learn the family tree and the rules of the family at the same time. They used a complex, non-linear "penalty" system to force the rules to make sense (no loops, everyone has a parent).

The Problem: It's like trying to solve a Rubik's cube while simultaneously painting the stickers. The math gets messy, the computer gets stuck in local loops, and it takes forever for large families.

2. The New Way: The "Score" Detective (SSTS)

The authors propose a decoupled approach. They split the job into two distinct stages, like a two-step investigation.

Step 1: The "Generative Model" (The Artist)

First, they train a computer program (a neural network) just to understand the data. Think of this as an artist who studies the photo and learns to draw a perfect copy of the crowd.

The Magic: This artist doesn't care about the family tree yet. They just learn the "shape" of the data.
The Score: Once trained, this artist can calculate a "score" for every person in the photo. This score tells you how likely it is for that person to be in that exact spot.

Step 2: The "Algebraic Sort" (The Architect)

This is the paper's big breakthrough. Instead of shuffling people around, the authors realized that the mathematical shape of the artist's "score" contains a hidden map of the family tree.

The Metaphor: Imagine the family tree is a building. The "leaf nodes" (the youngest generation with no children) are the roof tiles. The authors found that if you look at the "energy" of the roof tiles in the artist's score, they stand out clearly.
The Schur Complement: This is a fancy math term for a specific way of "peeling off" layers of an onion. Once the algorithm identifies the "roof tiles" (the leaves), it uses a mathematical trick (the Schur complement) to mathematically remove them from the picture.
The Result: By peeling off the leaves one by one (or in groups), the algorithm reveals the order of the family from youngest to oldest without ever having to guess or shuffle. It turns a messy guessing game into a clean, deterministic calculation.

Why is this a big deal?

Speed and Scale: The old way was like trying to count every grain of sand on a beach to find a specific shell. The new way is like using a metal detector. The authors tested this on graphs with 1,000 variables (a very large family). The old methods would crash or take days; this new method did it in seconds.
No More "Stuck" Moments: Because they removed the messy "shuffling" optimization, the algorithm doesn't get stuck in local traps. It follows a straight mathematical path.
The "Expectation Gap": The paper admits that for very complex, non-linear families (where the rules change depending on the situation), the math isn't perfectly exact. It's like a slightly blurry photo. However, they created a "Block" version that groups people together to minimize this blur, keeping the error very low.

The Bottom Line

The paper claims that by separating the "learning the data" part from the "finding the order" part, and by using a specific mathematical trick (Schur complement) on the data's "score," we can discover cause-and-effect relationships much faster and more reliably than before.

They successfully moved the problem from a hard optimization puzzle (trying to find the best path through a maze) to a statistical estimation challenge (measuring the height of the walls to see where the exit is).

What they did NOT claim:

They did not claim this works for every type of data (it struggles if the noise is very weird or if the relationships are post-nonlinear).
They did not claim this is a medical diagnosis tool or a clinical application.
They did not claim it solves the problem of "hidden confounders" (unseen variables) perfectly, though they tested it on real-world biological data with some success.

In short: They found a way to turn a chaotic, slow guessing game into a fast, clean math problem.

1. Problem Statement

Continuous causal discovery aims to recover the underlying Directed Acyclic Graph (DAG) from observational data. Current state-of-the-art methods (e.g., NOTEARS, DAGMA) rely on continuous constrained optimization, coupling representation learning with structural optimization via non-convex acyclicity penalties (e.g., $tr(e^{W \circ W}) = d$ ).

Key Limitations:

Local Optima: The non-convex nature of acyclicity penalties causes gradient-based solvers to get trapped in local minima, especially in high-dimensional regimes.
Scalability: These methods suffer from $O(d^3)$ computational costs per gradient step due to repeated acyclicity checks, limiting scalability to high-dimensional datasets ( $d \gg 100$ ).
Retraining Overhead: Existing score-based approaches often require recursive retraining of generative models on marginalized subsets to identify node order, creating severe computational bottlenecks.

2. Methodology: Score-Schur Topological Sort (SSTS)

The authors propose a decoupled paradigm that shifts the bottleneck from non-convex optimization to statistical score estimation. The core idea is that the causal hierarchy leaves a geometric signature in the Score-Jacobian Information Matrix (SJIM).

A. Theoretical Foundation

Score Function & SJIM: Let $s(x) = \nabla_x \log p(x)$ be the score function. The SJIM, $I$ , is defined as the expected negative Hessian of the log-density: $I = \mathbb{E}[-\nabla^2_x \log p(x)]$ .
Leaf Identifiability (Theorem 3.1): In linear Gaussian Additive Noise Models (ANMs), the diagonal entries of the SJIM encode the topological hierarchy. Specifically, a node $l$ is a topological leaf (sink) if and only if its diagonal entry $I_{ll}$ is minimal ( $1/\sigma^2$ ).
Exact Marginalization via Schur Complement (Theorem 3.2): Once a leaf is identified, removing it from the graph is mathematically equivalent to computing the Schur complement of the SJIM with respect to that leaf. This allows for recursive elimination without retraining the generative model.
Non-Linear Extension & Expectation Gap: In non-linear ANMs, the expectation operator does not commute with matrix inversion ( $\text{Schur}(\mathbb{E}[H]) \neq \mathbb{E}[\text{Schur}(H)]$ ). The authors derive an Expectation Gap ( $\Delta$ ) proportional to the covariance of the gradient of the noise mechanism. This gap introduces structural error in non-linear settings.

B. Algorithm: Block-SSTS

To address the expectation gap and finite-sample variance, the authors introduce Block-SSTS:

Generative Modeling: Train an unconstrained score network (e.g., MLP) to estimate $s_\theta(x)$ .
Sparse Jacobian Prior: Apply a Group Lasso ( $\ell_{1,2}$ ) penalty to the input layer weights to induce sparsity in the empirical Jacobian, reducing estimation variance in high dimensions.
Block Marginalization: Instead of sequential single-node removal, the algorithm groups parallel leaf nodes (those with similar diagonal energies within a tolerance $\gamma$ $γ$ ) and performs a block Schur complement inversion.
- This compresses the extraction depth, truncating the accumulation of the non-linear expectation gap.
- It avoids the $O(Nd^3)$ cost of exact sample-wise covariance patching.
Edge Pruning: Once the topological order is extracted, the DAG is recovered via standard sparse regression (Lasso) for each node against its permitted ancestors.

3. Key Contributions

Algebraic Mapping of Acyclicity: Proved that iterative leaf-node marginalization is exactly equivalent to the Schur complement of the SJIM in linear Gaussian ANMs. For non-linear ANMs, they characterized the approximation error (Expectation Gap).
Decoupled Extraction: Eliminated the non-convex optimization bottleneck entirely. Causal discovery is reframed as a statistical estimation problem followed by deterministic matrix operations ( $O(d^3)$ total, but only once).
Scalability to $d=1000$ : Introduced Block-SSTS to handle high-dimensional non-linear graphs, achieving structural recovery up to 1000 variables, a regime where continuous optimization methods fail.
Metric Correction: Argued against using rank correlation (Kendall's $\tau$ ) for evaluating parallel DAGs, proposing Edge Violations (EV) as the rigorous metric invariant to topological ambiguity.

4. Experimental Results

Linear Regimes: SSTS achieves zero causal edge violations for dimensions up to $d=1000$ on linear ANMs, validating the theoretical exactness.
Non-Linear Regimes:
- Efficiency: At $d=50$ , SSTS reduces execution time from ~2700s (DAGMA) to 3.66s while achieving a significantly lower Structural Hamming Distance (SHD: 2.6 vs 93.0).
- Scalability: At $d=1000$ , SSTS completes extraction in under 1 second, whereas continuous methods are computationally intractable.
- Trade-off: While SSTS introduces a small structural error due to the non-linear expectation gap, it remains bounded by the finite-sample estimation variance of the score network.
Real-World Data: On the Sachs protein-signaling network ( $d=12$ ), SSTS outperformed DAGMA and Scalable SCORE in both accuracy (SHD 13.8 vs 15.0/16.2) and speed (6.62s vs 182s/56s).
Ablation Studies:
- Sparsity: Group Lasso priors are critical for stability in data-starved regimes ( $N=500, d=100$ ).
- Block Tolerance: Moderate block tolerance ( $\gamma \approx 0.05$ ) optimally truncates error accumulation without violating causal hierarchy.

5. Significance and Conclusion

This work fundamentally reframes causal discovery from a constrained optimization problem to a statistical estimation challenge. By leveraging the geometric properties of the score function and the algebraic properties of the Schur complement, SSTS bypasses the local optima and scalability issues inherent in continuous acyclicity penalties.

Implications:

Theoretical: Establishes a direct link between score matching, information geometry, and topological sorting.
Practical: Enables causal structure learning on high-dimensional, non-linear datasets (up to $d=1000$ ) that were previously inaccessible to continuous optimization methods.
Limitations: The method's fidelity in non-linear regimes is ultimately bounded by the accuracy of the score function estimation (finite-sample variance) and the magnitude of the non-linear expectation gap, rather than optimization convergence.

In summary, SSTS offers a deterministic, optimization-free pathway to causal discovery, proving that once the non-convex bottleneck is removed, the primary challenge becomes the statistical estimation of the score geometry.