Bernoulli Differentiable Causal Subgraph Sampling

• The paper introduces a variational Bayesian approach that separates Bernoulli-like edge sampling from order-consistent edge orientation to efficiently generate valid DAGs.
• It employs continuous relaxations using Gumbel-Softmax and a tempered sigmoid mask to ensure acyclicity by construction without explicit penalties.
• The method scales quadratically and outperforms existing approaches on metrics such as AUC-ROC and AUC-PR in high-dimensional causal discovery.

Searching arXiv for the cited paper to ground the article. Bernoulli differentiable causal subgraph sampling, in the sense most closely supported by current causal discovery work, denotes a graph-generation mechanism in which candidate directed edges are sampled in a Bernoulli-like manner and then filtered or oriented by a differentiable latent topological order so that the resulting graph is acyclic by construction. In the formulation introduced in "Scalable Variational Causal Discovery Unconstrained by Acyclicity" (Hoang et al., 2024), the method is best understood as a differentiable Bayesian DAG sampler, and only partly as “Bernoulli differentiable causal subgraph sampling”: it learns a posterior distribution over causal graphs $P(\mathcal G \mid \mathbf X)$ from observational data $\mathbf X$, while separating edge presence from order-consistent edge direction and combining them through a continuous latent representation that generates valid DAGs efficiently.

1. Problem formulation in Bayesian causal discovery

Given observational data $\mathbf X \in \mathbb R^{n\times d}$, the objective is to infer a causal DAG $\mathcal G=\langle \mathbf V,\mathbf E\rangle$ over $d$ variables. The joint distribution factorizes as

$$P(X) = \prod_{i=1}^d P(X_i \mid X_{\mathrm{pa}(i)}).$$

The structural model is an additive noise model,

$$X_i = f_i(X_{\mathrm{pa}(i)}) + \epsilon_i, \qquad \epsilon_i \sim \mathcal N(0,\sigma^2),$$

with adjacency matrix $\mathbf A \in \{0,1\}^{d\times d}$, where $A_{ij}=1$ means $X_i \to X_j$ (Hoang et al., 2024).

The central difficulty is combinatorial and inferential. The space of DAGs is super-exponential in $\mathbf X$0, posterior inference is intractable, and explicit acyclicity constraints are expensive. The Bayesian target is

$\mathbf X$1

Because the denominator is intractable, the method uses variational inference to approximate the posterior.

Within this setting, “Bernoulli differentiable causal subgraph sampling” is not the full problem formulation but rather a partial description of the graph generator: candidate edges are sampled, then pruned or oriented according to a latent order. This suggests that the relevant object is not an arbitrary Bernoulli random graph, but a posterior over DAGs implemented through structured latent variables.

2. Variational parameterization of latent graph-generating variables

The variational family is defined over two latent components:

• $\mathbf X$2: an edge-sampling matrix
• $\mathbf X$3: a priority-score vector that induces a topological order

The factorization is mean-field: $\mathbf X$4

The evidence lower bound is

$\mathbf X$5

Using the factorization, the KL term splits into

$\mathbf X$6

In the final implementation, this becomes

$\mathbf X$7

This variational construction is significant because it relocates the hard combinatorial object—the DAG—into a simpler continuous-domain latent parameterization. A plausible implication is that the method’s tractability derives less from approximating acyclicity itself than from shifting inference to $\mathbf X$8 and $\mathbf X$9, where reparameterized gradient estimators are available.

3. Differentiable construction of a DAG from latent order and edge samples

The core sampling mechanism begins with a priority vector $\mathbf X \in \mathbb R^{n\times d}$0. Sorting $\mathbf X \in \mathbb R^{n\times d}$1 defines an implicit order over nodes. Pairwise order differences are encoded by

$\mathbf X \in \mathbb R^{n\times d}$2

If $\mathbf X \in \mathbb R^{n\times d}$3, then edge $\mathbf X \in \mathbb R^{n\times d}$4 is allowed; otherwise it is suppressed (Hoang et al., 2024).

To convert this order into a differentiable orientation mask, the method adapts the No-Curl idea but replaces the ReLU used for weighted DAGs with a tempered sigmoid suitable for binary adjacency: $\mathbf X \in \mathbb R^{n\times d}$5 For finite temperature $\mathbf X \in \mathbb R^{n\times d}$6, the implementation is

$\mathbf X \in \mathbb R^{n\times d}$7

The interpretation is direct. The quantity $\mathbf X \in \mathbb R^{n\times d}$8 is close to $\mathbf X \in \mathbb R^{n\times d}$9 if $\mathcal G=\langle \mathbf V,\mathbf E\rangle$0, close to $\mathcal G=\langle \mathbf V,\mathbf E\rangle$1 if $\mathcal G=\langle \mathbf V,\mathbf E\rangle$2, and becomes sharper as $\mathcal G=\langle \mathbf V,\mathbf E\rangle$3. The priority scores therefore define a soft, differentiable orientation mask. The final graph is obtained by intersecting edge existence with order consistency.

This mechanism gives the phrase “causal subgraph sampling” its technical meaning: the sampled graph is a subgraph of the order-induced DAG. It is not merely that edges are sampled independently; rather, an order-consistent mask determines which sampled edges survive.

4. Bernoulli-like edge sampling and the continuous–discrete interface

The edge matrix $\mathcal G=\langle \mathbf V,\mathbf E\rangle$4 is sampled in a Bernoulli-like way, but the final graph is not just Bernoulli edge sampling. Each potential edge has probability $\mathcal G=\langle \mathbf V,\mathbf E\rangle$5, and samples are drawn using Gumbel-Softmax to approximate discrete Bernoulli/Categorical edge draws: $\mathcal G=\langle \mathbf V,\mathbf E\rangle$6 A straight-through estimator is then used so that the forward pass is effectively discrete while the backward pass remains continuous (Hoang et al., 2024).

The graph is sampled by four steps:

1. sample edge indicators $\mathcal G=\langle \mathbf V,\mathbf E\rangle$7
2. sample priority scores $\mathcal G=\langle \mathbf V,\mathbf E\rangle$8
3. orient or filter edges using $\mathcal G=\langle \mathbf V,\mathbf E\rangle$9
4. obtain $d$0

This is why the method is better described as Bernoulli-like edge subsampling plus differentiable topological masking rather than pure Bernoulli subgraph sampling.

The continuous and discrete components are explicitly separated. The final adjacency matrix $d$1, candidate edge existence in the intended DAG, and the actual causal graph sampled during inference are discrete. The priority scores $d$2, the relaxed edge variables $d$3, and the tempered sigmoid mask are continuous. Differentiability is maintained through Gumbel-Softmax for edges, the straight-through estimator for discrete-like forward samples, and pathwise gradients through the continuous relaxations.

For $d$4, the variational family is Gaussian,

$d$5

For $d$6, Gumbel-Softmax provides a low-variance pathwise gradient estimator, preferable to REINFORCE-like score-function gradients.

5. Acyclicity by construction

A central property of the method is that acyclicity is not enforced by a separate penalty. Instead, DAG validity is built into the sampling rule itself. The theorem stated in the paper asserts that $d$7 is a DAG iff there exists $d$8 and $d$9 such that

$$P(X) = \prod_{i=1}^d P(X_i \mid X_{\mathrm{pa}(i)}).$$0

with distinct scores $$P(X) = \prod_{i=1}^d P(X_i \mid X_{\mathrm{pa}(i)}).$$1 (Hoang et al., 2024).

The logic is structural: sorting $$P(X) = \prod_{i=1}^d P(X_i \mid X_{\mathrm{pa}(i)}).$$2 yields a topological order, only edges from lower score to higher score are permitted, no directed cycle can then exist, and multiplication by $$P(X) = \prod_{i=1}^d P(X_i \mid X_{\mathrm{pa}(i)}).$$3 can only remove edges. Consequently, the resulting graph remains a DAG.

This distinguishes the method from approaches that rely on explicit acyclicity penalties, smooth DAG constraints, augmented Lagrangian procedures, permutation-matrix decompositions with Gumbel-Sinkhorn, or explicit DAG regularizers. No separate acyclicity function $$P(X) = \prod_{i=1}^d P(X_i \mid X_{\mathrm{pa}(i)}).$$4 is needed. A common misconception is that differentiable causal discovery must repeatedly optimize against an acyclicity surrogate; here the constraint is replaced by a generative parameterization in which only order-consistent edges can be expressed.

6. Algorithmic workflow, scalability, and empirical positioning

The training loop proceeds as follows:

1. initialize variational parameters $$P(X) = \prod_{i=1}^d P(X_i \mid X_{\mathrm{pa}(i)}).$$5 and model parameters $$P(X) = \prod_{i=1}^d P(X_i \mid X_{\mathrm{pa}(i)}).$$6
2. for each minibatch:
   • sample $$P(X) = \prod_{i=1}^d P(X_i \mid X_{\mathrm{pa}(i)}).$$7
   • sample $$P(X) = \prod_{i=1}^d P(X_i \mid X_{\mathrm{pa}(i)}).$$8
   • form

   $$P(X) = \prod_{i=1}^d P(X_i \mid X_{\mathrm{pa}(i)}).$$9

• reconstruct data using the sampled DAG:

  $$X_i = f_i(X_{\mathrm{pa}(i)}) + \epsilon_i, \qquad \epsilon_i \sim \mathcal N(0,\sigma^2),$$0

• maximize the ELBO

1. update $$X_i = f_i(X_{\mathrm{pa}(i)}) + \epsilon_i, \qquad \epsilon_i \sim \mathcal N(0,\sigma^2),$$1 with Adam
2. use early stopping

The reconstruction term is implemented as squared error,

$$X_i = f_i(X_{\mathrm{pa}(i)}) + \epsilon_i, \qquad \epsilon_i \sim \mathcal N(0,\sigma^2),$$2

The computational profile follows from simple vector and matrix operations: sampling $$X_i = f_i(X_{\mathrm{pa}(i)}) + \epsilon_i, \qquad \epsilon_i \sim \mathcal N(0,\sigma^2),$$3 is $$X_i = f_i(X_{\mathrm{pa}(i)}) + \epsilon_i, \qquad \epsilon_i \sim \mathcal N(0,\sigma^2),$$4, sampling $$X_i = f_i(X_{\mathrm{pa}(i)}) + \epsilon_i, \qquad \epsilon_i \sim \mathcal N(0,\sigma^2),$$5 is $$X_i = f_i(X_{\mathrm{pa}(i)}) + \epsilon_i, \qquad \epsilon_i \sim \mathcal N(0,\sigma^2),$$6, forming $$X_i = f_i(X_{\mathrm{pa}(i)}) + \epsilon_i, \qquad \epsilon_i \sim \mathcal N(0,\sigma^2),$$7 is $$X_i = f_i(X_{\mathrm{pa}(i)}) + \epsilon_i, \qquad \epsilon_i \sim \mathcal N(0,\sigma^2),$$8, and final DAG construction is elementwise $$X_i = f_i(X_{\mathrm{pa}(i)}) + \epsilon_i, \qquad \epsilon_i \sim \mathcal N(0,\sigma^2),$$9. Overall DAG sampling is therefore $\mathbf A \in \{0,1\}^{d\times d}$0, in contrast with $\mathbf A \in \{0,1\}^{d\times d}$1-style costs in methods relying on NOTEARS-like DAG constraints, Gumbel-Sinkhorn permutation sampling, or Hungarian-style permutation recovery. This is why the method can sample DAGs with thousands of nodes quickly (Hoang et al., 2024).

Empirically, the paper reports that VCUDA outperforms DiBS, DDS, BaDAG, GraNDAG, Masked-DAG, and others on many metrics. It is especially strong on AUC-ROC, AUC-PR, and MSE of causal mechanism reconstruction; it scales better in high dimensions, including $\mathbf A \in \{0,1\}^{d\times d}$2; and it is faster than methods that repeatedly enforce acyclicity. On the Sachs protein signaling dataset, it achieves best AUC-ROC and second-best AUC-PR. Ablation results show that a lower temperature helps, with $\mathbf A \in \{0,1\}^{d\times d}$3 chosen as a good trade-off.

7. Relation to prior graph samplers and interpretive scope of the term

Relative to prior Bayesian graph samplers, the method occupies a specific position. MCMC methods are characterized as accurate but slow in mixing and convergence. GFlowNets are sequential and often expensive. Gumbel-Sinkhorn permutation models are differentiable but costly, typically $\mathbf A \in \{0,1\}^{d\times d}$4. NOTEARS-based variational methods require explicit acyclicity penalties and extra optimization overhead (Hoang et al., 2024).

The present construction differs by using an implicit order latent variable $\mathbf A \in \{0,1\}^{d\times d}$5, a differentiable orientation mask, direct sampling of valid DAGs, and no separate evaluation of an acyclicity constraint. These features delimit what can and cannot be meant by “Bernoulli differentiable causal subgraph sampling” in this context.

The closest interpretation is that $\mathbf A \in \{0,1\}^{d\times d}$6 acts like a set of independent Bernoulli edge selectors over all candidate directed edges, the resulting sampled edge set is a candidate directed subgraph, and the topological gate removes any edge inconsistent with the sampled order. However, the method is not merely Bernoulli sampling over subgraphs in the usual sense, because Bernoulli sampling alone does not ensure acyclicity, a second latent variable $\mathbf A \in \{0,1\}^{d\times d}$7 determines which Bernoulli edges survive, and the final DAG is the intersection of sampled candidate edges and order-consistent edges.

Accordingly, the technically precise description is not pure Bernoulli subgraph sampling but Bernoulli-like edge subsampling plus differentiable topological masking. This suggests that the phrase “Bernoulli differentiable causal subgraph sampling” is useful only if it is understood as shorthand for a two-part mechanism in which Bernoulli-like edge existence is subordinated to an order-induced DAG construction.