How to Marginalize in Causal Structure Learning?

AI summary

Overview

Research area: Causal discovery / Bayesian network structure learning, with a focus on tractable probabilistic models (probabilistic circuits).

Technical level: Advanced. The paper assumes familiarity with Bayesian networks, Bayesian scoring functions (BGe, BDeu), dynamic programming over parent sets, and probabilistic circuits with smoothness and decomposability properties.

Scope in one sentence: The paper proposes replacing the exponential-time dynamic programming step used to marginalize Bayesian scores in structure learning with a learned probabilistic circuit that supports exact marginalization over all possible parent sets.

arXiv: 2511.14001v1 [cs.LG], 18 Nov 2025. Authors: William Zhao, Guy Van den Broeck, Benjie Wang (University of California, Los Angeles).

What This Paper Is About

Bayesian structure learners infer a posterior distribution over the directed acyclic graphs (DAGs) that could underlie a Bayesian network, and doing so requires repeatedly marginalizing Bayesian scores over all possible parent sets of each node. The standard approach computes these marginals exactly with a dynamic programming table, but that computation takes time and memory exponential in the number of candidate parents, so existing methods artificially cap each node to a small set of "most likely" parents. The paper's goal is to remove that cap by learning a probabilistic circuit per node that approximates the score function and can be marginalized exactly and cheaply.

Key Contributions

1. A circuit-based replacement for DP marginalization. Instead of computing marginals from a lookup table, the authors train one probabilistic circuit per node over the parent-indicator vector, so that marginal (and "marginal/zero") queries required by structure learners can be answered exactly in time linear in circuit size, with no exponential precomputation.

2. A new two-phase circuit learning routine. Phase one trains on complete, unmarginalized parent sets scored by the Bayesian scorer. Phase two finetunes on a mixture of complete parent sets and marginalized queries, labeled using an exact DP marginalizer that is still restricted to a limited candidate parent set. The number of marginalized variables, k, is increased iteratively up to a limit L.

3. An unnormalized circuit architecture tailored to Bayesian scores. The circuit follows a RAT-SPN-like layout (leaves arranged as an (M, N) matrix, alternating product layers that halve the rows, and sum layers that mix within rows), but removes the usual softmax normalization on sum weights because Bayesian scores are unnormalized. Parameters are stored in the log domain.

4. Empirical validation inside TRUST. The method is substituted for the dynamic-programming marginalization inside the TRUST Bayesian structure learner and compared against DP with and without candidate-set restrictions, producing a structure learner built entirely from probabilistic circuits.

Main Findings

• Exact DP is the bottleneck to remove: The DP algorithm computes all marginal probability masses in O(3^N) time and O(3^N) memory, and becomes computationally infeasible once N exceeds roughly 16, which is why candidate parent sets are limited in practice.

• d = 16 with unrestricted DP: When the DP considers all 15 other nodes (so its marginals are exact), the approximate regression circuit performs comparably to the exact DP, except on E-SHD.

• d = 20 with restricted DP: When the DP is restricted to considering only 8 nodes, as is normal practice, the circuit method performs considerably better on all metrics except E-SHD.

• E-SHD behaves differently and the authors explain why: The circuit's posterior holds a higher E-SHD than exact computation at both d = 16 and d = 20 because the method tends to sample graphs with more edges. The AUROC scores indicate that the circuit-based posterior is nonetheless capable of determining which edges are more likely.

• Log-domain MSE is the loss of choice: Training minimizes (log p(S) − B(S))², which was empirically the best choice compared to alternatives such as KL divergence. MSE has also been used in similar marginal-focused models (Liu et al. 2024).

• Sampling and optimization design choices matter: Complete parent vectors are sampled with weight proportional to 2^(M−T), where T is the number of indicators set to 1, because such a vector contributes to that many marginal probability masses. Phase one uses a relatively high learning rate of 10⁻¹, reduced when training loss plateaus, to counteract vanishing gradients caused by unnormalized log-domain weights. A limit L is also placed on finetuning iterations, since too many cause the circuit to forget previously learned marginals.

• Initialization: The paper describes initializing log-domain parameters according to m log(U(0,1)) with a multiplier m < 0, set to around −10 to break gradient symmetry, while the reported experimental baseline circuit uses an initialization multiplier m = 12.

Methodology in Plain English

A Bayesian structure learner has to answer the same kind of question over and over: "what is the total posterior probability mass of all parent sets that exclude these nodes and include those?" Computing this exactly requires tabulating every combination of "parent," "not a parent," or "either" across all candidate parents, which explodes quickly. The authors' idea is to train a neural-style probabilistic model, a probabilistic circuit, to imitate the Bayesian scoring function for each node. Because their circuit is built from product nodes with disjoint scopes and sum nodes with matching scopes, and because it is decomposable and smooth, it can compute exact marginals over its own learned distribution in one forward pass.

Training happens in two stages. First the circuit learns to reproduce raw scores on complete parent sets, sampled with a weighting that favors sets contributing to many marginal masses. Then it is finetuned on a blend: half the data are complete parent sets scored by the Bayesian scorer, and half are marginal queries whose answers are produced by the exact dynamic programming marginalizer restricted to a small candidate parent set. The number of marginalized variables grows one step at a time, because each (k+1,0) query decomposes into a (k,0) plus a (k,1) query, so solving smaller cases builds toward larger ones. The loss is mean squared error in the log domain. At the end, the circuit replaces the DP table inside TRUST.

Why This Matters

The paper targets a hard limit baked into most Bayesian structure learners: they must give up support over some parent sets to stay computationally feasible, which restricts the structures they can ever recover. By shifting marginalization onto a tractable circuit, the authors keep exact marginalization while expanding support to all parent sets, producing a more accurate posterior.

Real-world applications (as cited in the paper):

• Healthcare modeling with Bayesian networks.
• Industrial fault diagnosis.
• Bayesian model averaging of causal effects, which is one reason posterior distributions over structures are useful in the first place.
• General settings where interpretability and transparency relative to black-box models are valued, since BNs expose their graph structure directly.

Industry relevance: Causal discovery pipelines in medicine, manufacturing, and any domain that needs a readable cause-and-effect model depend on the posterior over graphs being trustworthy. Methods that silently truncate the search space can bias downstream causal effect estimates. This work suggests that a learned, tractable surrogate can replace a hard computational shortcut without sacrificing the exactness of the marginal queries, and it demonstrates this inside TRUST, a method that already relies on circuits for the posterior over orderings.

Future Directions

• Scaling to higher dimensions. The authors identify scaling TRUST to higher dimensions as an important focus, since their method removes the candidate-parent restriction that previously blocked such scaling.
• Denser graph structures. Limited testing suggests the candidate-set restriction hampers performance more significantly when learning denser structures, such as an Erdős–Rényi graph with an average of 4 edges per node; experimenting with denser graphs may reveal larger improvements than on sparser ones.
• Compatibility with other learners and scores. The marginalization circuit is currently limited in applicability to TRUST or linear Gaussian Bayesian networks. Adapting it to different BN mechanisms and scores would make it compatible with methods such as ArCO-GP (Toth et al. 2025), which also restrict candidate parents.
• Closing the E-SHD gap. The method's tendency to sample graphs with more edges raises E-SHD relative to exact computation, an issue the AUROC results suggest is not about edge ranking but about edge density.

Target Audience

Researchers and graduate students working on causal discovery, Bayesian structure learning, and tractable probabilistic modeling, particularly those already familiar with probabilistic circuits and Bayesian scoring functions such as BGe or BDeu. It is also relevant to practitioners building structure learning or causal inference systems who need to understand what candidate-parent restrictions cost in accuracy, and to anyone interested in replacing brute-force dynamic programming with learned surrogates that retain exact inference guarantees.

Authors’ abstract

Bayesian networks (BNs) are a widely used class of probabilistic graphical models employed in numerous application domains. However, inferring the network's graphical structure from data remains challenging. Bayesian structure learners approach this problem by inferring a posterior distribution over the possible directed acyclic graphs underlying the BN. The inference process often requires marginalizing over probability distributions, which is typically done using dynamic programming methods that restrict the set of possible parents for each node. Instead, we present a novel method that utilizes tractable probabilistic circuits to circumvent this restriction. This method utilizes a new learning routine that trains these circuits on both the original distribution and marginal queries. The architecture of probabilistic circuits then inherently allows for fast and exact marginalization on the learned distribution. We then show empirically that utilizing our method to answer marginals allows Bayesian structure learners to improve their performance compared to current methods.