Sparse Additive Model Pruning for Order-Based Causal Structure Learning

AI summary

Overview

Research area: Causal structure learning (causal discovery) from observational data, specifically the "pruning" stage of order-based algorithms; the methods draw on sparse additive models, randomized tree embeddings, and group lasso regression.

Technical level: Intermediate. Readers should be comfortable with causal DAGs, additive models, regularization (lasso/group lasso), and standard evaluation metrics for graph recovery.

Scope: The paper proposes SARTRE, a pruning method that learns sparse additive models via randomized tree embedding plus group lasso, replacing the hypothesis-testing-based CAM-pruning step in order-based causal structure learning.

What This Paper Is About

Order-based causal structure learning splits the problem into two stages: first estimate a topological order of the variables, then remove the spurious edges from the fully-connected DAG implied by that order. The paper targets the second stage, where the dominant technique (CAM-pruning) must repeatedly fit generalized additive models and run hypothesis tests for every variable and candidate parent. The authors aim to replace that with a single sparse-model fit per variable that prunes edges directly, avoiding both the repeated GAM fitting cost and the accuracy loss from multiple testing.

Key Contributions

1. A new sparse additive model learning framework (SARTRE). The shape function for each variable is a linear combination of weighted indicator functions over intervals, so that a zero weight vector for a candidate parent means that parent can be dropped. Intervals are generated by randomized tree embedding and coefficients are learned by group lasso regression.

2. An efficient pruning algorithm for order-based causal structure learning. Given an estimated topological order, SARTRE-pruning removes redundant edges from the fully-connected DAG without hypothesis testing. It is described as combinable with any ordering algorithm, such as SCORE.

3. A representation result. Proposition 1 shows that the paper's piece-wise constant shape function can approximate any continuous function on an interval arbitrarily well, indicating the model retains rich representational capacity despite its simple form.

4. Empirical validation on synthetic and real data. Experiments on nonlinear additive-noise-model synthetic datasets (Erdős–Rényi and scale-free graphs) plus the Sachs and fMRI datasets show significant speedups over existing pruning methods while keeping comparable or superior accuracy.

Main Findings

• Baseline comparison by number of variables (n = 2000, d from 10 to 50). SARTRE was significantly faster than the baselines as d approached 50. It maintained comparable SHD and SID on ER4 and SF4, and achieved the best SHD and SID on ER1 and SF1.

• Sample-size comparison (d = 20, n from 1000 to 5000, ER1 and ER4). The running time of DAS approached that of SCORE as n increased, while SARTRE remained faster than both. SCORE attained better SHD and SID than SARTRE on ER4, and the authors suggest CAM-pruning's hypothesis testing remains accurate when n is large relative to d and the DAG is dense. SARTRE outperformed SCORE on ER1 and DAS on both ER1 and ER4.

• High-dimensional cases (d ∈ {64, 128, 256, 512}, n = 2000). With the ordering step replaced by the ground-truth topological order, SARTRE was faster than DAS while attaining better SHD and SID in all cases.

• Real and semi-real datasets (bootstrap sampling with n = 2000, 10 trials). On Sachs, SARTRE recorded SHD 22.7, SID 58.0, and running time 3.62 s, versus SCORE at 43.2, 102.4, 14.7 s and DAS at 27.6, 70.3, 6.42 s. On fMRI, SARTRE recorded SHD 12.9, SID 60.0, and running time 3.18 s, versus SCORE at 19.6, 71.8, 11.4 s and DAS at 11.6, 58.6, 4.75 s. SARTRE was faster than both baselines on both datasets without significantly degrading SHD and SID.

• Accuracy pattern. Across settings, SARTRE was more accurate than the baselines when the underlying causal graph was sparse, and remained competitive on dense graphs.

• Hyperparameters used. λ = 0.1 (best in the authors' sensitivity analyses), 5 trees per ensemble, maximum leaf size 8, giving at most 40 intervals per variable; each experiment was repeated 10 times.

Methodology in Plain English

The pipeline starts from an estimated topological order (from SCORE in the experiments). For each variable X_j, the method builds an ensemble of completely randomized decision trees that take only X_j as input. Each leaf of each tree corresponds to an interval of X_j, so collecting leaves yields a set of intervals R_j and a binary "indicator" embedding vector φ_j(X_j) that marks which interval a value falls into. Because these trees are built by randomly choosing split points and never look at a target variable, this embedding step is unsupervised and can be computed once and reused for every variable that has X_j among its candidate parents.

With the embeddings fixed, regressing X_i on its candidate parents becomes a linear model over concatenated indicator vectors. The coefficients are grouped by candidate parent, and a group lasso penalty is applied so that entire groups can be driven to exactly zero. If every coefficient for a candidate parent j is zero, the edge from X_j to X_i is deleted from the fully-connected DAG. This avoids the hypothesis tests CAM-pruning relies on; the authors solve the group lasso problem with existing algorithms such as block-coordinate descent or dual extrapolation. A short proof shows the resulting piece-wise constant shape functions are universal approximators for continuous functions on a bounded interval.

Why This Matters

Impact on research. The pruning stage is described as often being the bottleneck of the whole order-based pipeline, and CAM-pruning's reliance on repeated GAM fitting plus multiple testing limits both scalability and accuracy in high dimensions. SARTRE reframes pruning as a single sparse variable-selection problem with a reusable, unsupervised embedding, and it can be plugged into any ordering algorithm (CAM, SCORE, CaPS), so it isolates an improvement in one modular component.

Real-world applications:

• Biology and biomedicine — recovering gene regulatory or protein-signaling structure from observational measurements, as in the Sachs dataset used here.
• Neuroscience — estimating directed connectivity between brain regions, as in the fMRI dataset used here.
• Economics and the social sciences — identifying which variables drive others in observational economic or social data, where interventions are infeasible.
• Decision-making under intervention — causal DAGs allow prediction of intervention effects, which the authors note is crucial for decision-making across applications.

Industry relevance. The work comes from Fujitsu's Artificial Intelligence Laboratory, and the practical selling point is scale: the method was validated at d up to 512, where the ordering step was bypassed with the ground-truth order, and it was faster than the scalable alternative DAS in every high-dimensional case tested. For teams that already run score-based ordering algorithms at scale, SARTRE is a drop-in replacement for the pruning phase.

Future Directions

• Automatic hyperparameter tuning. The authors note they used the same λ and interval count l_j across all settings and that developing a data-driven method for tuning these values would be beneficial.
• Theoretical guarantees for pruning quality. Proposition 1 covers representational ability, but the authors state that guaranteeing the correctness of pruning remains an open problem, including determining a sufficient number of intervals l_j.
• Latent confounders. The paper frames investigation under latent confounders as important future work, since the current model assumes all relevant variables are observed.
• Broader settings and tasks. The authors note SARTRE is also a nonlinear variable selection method and could be applied beyond causal structure learning, though they leave this to future work.

Target Audience

Researchers and practitioners working on causal discovery, especially those using or extending order-based methods such as SCORE, DAS, or CAM; methodologists interested in sparse additive models, tree embeddings, and group lasso for variable selection; and applied scientists in biology, neuroscience, or economics who need scalable structure recovery from high-dimensional observational data. A working knowledge of DAGs, additive models, and regularization is assumed, so the paper is best suited to readers at an intermediate level or above.

Authors’ abstract

Causal structure learning, also known as causal discovery, aims to estimate causal relationships between variables as a form of a causal directed acyclic graph (DAG) from observational data. One of the major frameworks is the order-based approach that first estimates a topological order of the underlying DAG and then prunes spurious edges from the fully-connected DAG induced by the estimated topological order. Previous studies often focus on the former ordering step because it can dramatically reduce the search space of DAGs. In practice, the latter pruning step is equally crucial for ensuring both computational efficiency and estimation accuracy. Most existing methods employ a pruning technique based on generalized additive models and hypothesis testing, commonly known as CAM-pruning. However, this approach can be a computational bottleneck as it requires repeatedly fitting additive models for all variables. Furthermore, it may harm estimation quality due to multiple testing. To address these issues, we introduce a new pruning method based on sparse additive models, which enables direct pruning of redundant edges without relying on hypothesis testing. We propose an efficient algorithm for learning sparse additive models by combining the randomized tree embedding technique with group-wise sparse regression. Experimental results on both synthetic and real datasets demonstrated that our method is significantly faster than existing pruning methods while maintaining comparable or superior accuracy.