P E T E R: Post-Training Robustification of Probabilistic Circuits
Adrian Ciotinga1 1
Yeming Dai1
School of Computing and Augmented Intelligence, Arizona State University University
arXiv:2607.07671v1 [cs.LG] 8 Jul 2026
Abstract Probabilistic circuits (PCs) can model complex joint distributions while supporting exact and efficient computation of many inference queries. However, standard likelihood-based PC learning is vulnerable to overfitting and fragile generalization when confronted with data noise, small sample sizes, or distribution shifts. This can be mitigated using distributionally-robust optimization which consider worst-case distributions within a Wasserstein ball of the empirical distribution, but current methods are limited to training a model from scratch in this framework. Instead, we propose P E T E R: a novel, data-free post-training framework designed to robustify pre-trained PCs against distribution shifts without retraining from scratch. Empirical evaluations across multiple density estimation benchmarks demonstrate that P E T E R effectively robustifies baseline models against both random and adversarial perturbations, achieving competitive or superior performance to data-dependent robust learning baselines.
1
YooJung Choi1
INTRODUCTION
Distributionally robust optimization (DRO) has emerged as a principled framework for learning estimators that remain reliable under data noise, limited sample sizes, and distribution shift, by hedging against an uncertainty set of plausible distributions rather than committing to the empirical distribution alone [Delage and Ye, 2010, Bertsimas and Nohadani, 2019]. A common approach is to define this uncertainty set as a Wasserstein ball around the empirical distribution, yielding a minimax formulation in which the estimator is optimized against the worst-case distribution, according to the appropriate performance measure, within some radius ϵ. Although widely studied in the context of
deep generative models [Liu et al., 2023, Kuhn et al., 2019, Bauso et al., 2017], its application to tractable probabilistic models (TPMs) remains scarce; to the best of our knowledge, Peddi et al. [2022] is the sole work addressing this intersection. While effective, existing approaches to robust learning typically focus on training a model from scratch on a dataset, tying robustness to the availability of training data and requiring that any existing models be completely replaced rather than improved. We argue that TPMs, and probabilistic circuits (PCs) in particular, are uniquely positioned for an alternative, more flexible framework of post-training robustification. Rather than learning a new model from scratch or fine-tuning it with additional data as classic DRO settings, post-training robustification takes a pre-trained model and grants it the same distributionally-robust guarantees within an ϵ-Wasserstein ball without the use of any training data or complete retraining of the model. As we show in this paper, circuit structural properties enabling tractable inference allow us to not only compute the distance between two distributions, but also to explicitly represent and optimize against the worst-case adversarial distribution efficiently. Based on this insight, we propose P E T E R, a post-training robustification framework for PCs that hedges against worst-case perturbations within a Wasserstein ball, bounded using the recently introduced Circuit-Wasserstein distance [Ciotinga and Choi, 2025]. Our method reduces the optimization problem to an unconstrained problem and solves it via the gradient ascent-descent algorithm, leveraging tractable gradient computation for PCs. We empirically demonstrate the efficacy of our approach in comparison to an existing robust MLE baseline for PCs that requires a training dataset, as—to the best of our knowledge—no existing methods are fully data-free.
2
PRELIMINARIES
We use capital letters (X) to denote random variables and lowercase letters (x) to denote their assignments. Boldface denotes a set of random variables and their assignments
respectively (e.g., X and x). Probabilistic Circuits A probabilistic circuit (PC) C over a set of variables X is a parameterized, rooted directed acyclic graph (DAG) with three types of nodes: sum, product, and input nodes. Each input node n is associated with function fn that encodes a probability distribution over a variable Xi ∈ X, also called its scope sc(n). The set of child nodes for an internal node (sum or product) S n is denoted ch(n), and its scope is given by sc(n) = c∈ch(n) sc(c). Each sum node n has normalized parameters θn,c for each child node c. For an assignment x, let xn denote its projection onto the scope sc(n) of node n. Then a PC rooted at node n recursively defines a probability distribution pn (x) as follows: if n is an input node, fQn (x) if n is a product node, pn (x) = c∈ch(n) pc (xc ) P θ p (x ) if n is a sum node. c∈ch(n) n,c c c Probabilistic circuits admit exact and efficient computation of many probabilistic inference queries, enabled by enforcing certain structural constraints. In particular, throughout this paper we assume two properties, smoothness and decomposability, which enable tractable (polytime) computation of marginal and conditional queries. We further state when PCs are structured-decomposable or compatible; see Appendix A for the formal definitions of these properties. Robust Maximum-Likelihood Estimation In many machine learning paradigms, one assumes access to samples drawn from an underlying but unknown distribution P(X). In the limit, the empirical distribution P̂ (X) approximates P(X); however, assuming that P̂ ≈ P is not always principled. If the observed data is contaminated by noise, limited to a small sample size, or subject to distribution shifts, optimizing directly over the empirical distribution can lead to severe overfitting and fragile generalization. To mitigate this, one can adopt a distributionally robust optimization framework [Delage and Ye, 2010, Bertsimas and Nohadani, 2019] that considers an uncertainty set of candidate distributions around P̂ . A popular choice to define this neighborhood is the Wasserstein distance, which safeguards the learned estimator P by optimizing against the worst-case distribution Q within an ϵ-Wasserstein ball around P̂ , yielding the following minimax formulation: max min EQ [log P (X)] P
Rn
Rn
Intuitively, guaranteeing the performance of estimator P against distributions within Wasserstein distance ϵ of P̂ also guarantees the performance of P against similarly bounded corruptions, adversarial perturbations, or sampling noise.
3
ROBUST POST-TRAINING
We introduce robust post-training, a modification to robust MLE, which aims to take a pre-trained generative model encoding a baseline distribution and update it to be robust to distribution shifts. Specifically, we guarantee strong performance across an uncertainty set containing all distributions within a Wasserstein distance of ϵ around the baseline distribution. Similar to robust MLE, this also safeguards our model against corrupted data or adversarial perturbations of bounded magnitude. Note that the framework does not assume any access to data—including the model’s original training data, making the task generally harder than robust MLE but also more flexible and widely applicable. To solve this for PCs, we propose P E T E R: a Post-Training Robustification framework that takes a pre-trained PC and updates its parameters to be robust to distribution shifts within a Wasserstein ball. Our approach is enabled by tractability of probabilistic circuits. Specifically, while computing the Wasserstein distance is #P-hard even for product distributions [Taşkesen et al., 2022], the recently introduced Circuit-Wasserstein distance CW [Ciotinga and Choi, 2025] can be computed exactly and efficiently for structured-decomposable PCs. We introduce a differentiable approach to computing CW, which we use to formulate the constrained optimization problem within an ϵ-CW ball as an unconstrained gradient descent-ascent problem. 3.1
THE OPTIMIZATION PROBLEM
Given a pre-trained PC P̂ , our objective is to learn a robust circuit Pθ parameterized by θ that has high (log)-likelihood w.r.t. any distribution within an ϵ-CW ball around P̂ :
Q
s.t. W(P̂ , Q) ≤ ϵ The Wasserstein Distance Let P and Q be probability measures on metric space Rn with distance metric d(·, ·). The d-Wasserstein distance between P and Q is defined as: Wd (P, Q) =
where Γ(P, Q) denotes the set of all couplings, i.e. joint distributions whose marginal distributions match P and Q. That is, the following holds for all γ ∈ Γ(P, Q): Z Z P (x) = γ(x, y)dy, Q(y) = γ(x, y)dx.
inf γ∈Γ(P,Q)
Eγ(x,y) [d(x, y)]
max min EQϕ [log Pθ (X)] Pθ
Qϕ
(1)
s.t. CW(P̂ , Qϕ ) ≤ ϵ The Circuit-Wasserstein distance is a natural choice of a bounding metric because of its efficient and exact algorithm when applied to compatible circuits. Note that P̂ being structured-decomposable implies that both Pθ and Qϕ are
compatible with it, as all three circuits share the same structure.1 CW acts as an upper-bound on the true Wasserstein distance and allows us to explicitly parameterize the worstcase distribution Qϕ (X) using a PC. 3.2
GRADIENT-BASED ROBUST POST-TRAINING
To solve the constrained optimization problem for robust post-training, we first rewrite it as an unconstrained problem. Proposition 1. The optimal solution to the following unconstrained Lagrangian optimization problem is a feasible and optimal solution to the constrained problem in Equation 1.
Algorithm 1 P E T E R(P̂ , ϵ): returns parameters θ that make P̂ robust to distributions within CW-ball of radius ϵ 1: Pθ , Qϕ ← P̂ 2: repeat
▷ Initialize Pθ , Qϕ i.i.d.
x1 , x2 , ..., xn ∼ Q ▷ Sample from Qϕ Pϕ (X) ▷ Estimate θ ← θ + ηθ ∇θ ( n1 i log Pθ (xi )]) for k steps do ϕ ← ϕ−ηϕ ∇ϕ [log EQϕ [Pθ (X)]+λCW(P̂ , Qϕ )] 7: λ ← max(0, λ + ηλ [CW(P̂ , Qϕ ) − ϵ]) 8: end for 9: until convergence; 3: 4: 5: 6:
max min max EQϕ [log Pθ (X)] + λ(CW(P̂ , Qϕ ) − ϵ).
Letting n and m be corresponding nodes in Pθ and Qϕ respectively, we have three cases:
Saddle-point optimization problems arise frequently in machine learning—e.g., GAN training, multi-agent reinforcement learning, etc. [Goodfellow et al., 2014, Albrecht et al., 2024]—and can be solved using the gradient ascent-descent algorithm. Our problem admits the following update rules with learning rates ηθ , ηϕ , ηλ :
• Sum Nodes: The distance is formulated as an optimal transport problem between the children of n and m, subject to the constraints imposed by the mixture weights θ and ϕ. By applying the Envelope Theorem to the dual formulation of this linear program, we can compute the partial derivatives of the distance with respect to both the children’s distances and the mixture (n,m) (n,m) weights: ∂ CW∂θ = x∗i and ∂ CW∂ϕ = yj∗ . i j
Pθ
Qϕ λ≥0
θ ← θ + ηθ ∇θ EQϕ [log Pθ (X)] ϕ ← ϕ − ηϕ [∇ϕ EQϕ [log Pθ (X)] + λ∇ϕ CW(P̂ , Qϕ )]
• Product Nodes: The distance is defined as the unweighted sum of distances between corresponding children, leading to an identity gradient.
λ ← max(0, λ + ηλ [CW(P̂ , Qϕ ) − ϵ]) If Pθ and Qϕ are compatible and deterministic, the necessary gradients can be computed exactly and efficiently in a single forward and backward pass; however, without determinism, computing EQϕ [log Pθ (X)] is #Phard [Vergari et al., 2021]. To circumvent this, we estimate EQϕ [log Pθ (X)] by sampling from Qϕ when computing gradients w.r.t. θ, and apply Jensen’s inequality to bound EQϕ [log Pθ (X)] ≤ log EQϕ [Pθ (X)] and update ϕ to instead minimize this upper-bound—which we can compute exactly thanks to the compatibility of Pθ and Qϕ . Algorithm 1 shows the pseudocode of our post-training method. 3.3
DIFFERENTIABLE COMPUTATION OF CW
While Ciotinga and Choi [2025] introduced the recursive algorithm for CW, computing its gradients is nontrivial because the standard forward pass relies on linear program solvers, which are not inherently compatible with backpropagation. To address this, we propose a differentiable framework for CW that allows for computing gradients ∇θ CW(Pθ , Qϕ ) and ∇ϕ CW(Pθ , Qϕ ). CW is computed recursively by traversing circuits Pθ and
Qϕ , where the nature of the computation—and thus the gradient flow—at each step is dependent on the node type. 1
We note that same structure is sufficient, not necessary, for efficient CW computation; see Ciotinga and Choi [2025] for details.
• Input Nodes: We assume the distance between base distributions is differentiable, allowing for efficient computation of ∇θ W(n, m) and ∇ϕ W(n, m). This holds for standard families such as Gaussian or categorical distributions. A rigorous derivation of above can be found in Appendix B.
4
EXPERIMENTS
In this section, we empirically evaluate P E T E R on seven of the density estimation benchmark datasets [Lowd and Davis, 2010, Van Haaren and Davis, 2012, Bekker et al., 2015, Larochelle and Murray, 2011], a set of binary datasets ranging from 16 to 1000 variables. We use Hidden Chow-Liu Tree (HCLT) [Liu and Van den Broeck, 2021] implemented with PyJuice [Liu et al., 2024] to learn the PC structure for each dataset and estimate the parameters using P E T E R and two other baselines for comparison. Baselines We compare against the robust MLE method (RL-TPM) proposed by Peddi et al. [2022] and standard maximum-likelihood estimation via expectationmaximization (MLE-PC) [Desana and Schnörr, 2016]. RLTPM, designed for binary datasets, poses robust MLE as a minimax game between the PC estimator maximizing likelihood and an adversary corrupting bits with a fixed budget to
T
Ta
Tr
Dataset
h
MLE-PC
RL-TPM
PETER
MLE-PC
RL-TPM
PETER
MLE-PC
RL-TPM
PETER
NLTCS
1 3 5
-6.09
-9.36 -10.34 -10.86
-6.79 -7.79 -9.90
-11.26 -18.43 -23.02
-10.91 -11.14 -10.61
-9.31 -11.56 -11.75
-8.41 ±0.03 -11.46 ±0.05 -13.24 ±0.09
-10.13 ±0.03 -11.22 ±0.02 -11.44 ±0.02
-8.04 ±0.02 -9.79 ±0.02 -10.76 ±0.01
1 MSNBC 3 5
-6.43
-9.14 -8.47 -9.92
-7.05 -8.69 -10.76
-12.01 -19.88 -25.39
-10.33 -11.52 -11.85
-8.45 -10.96 -11.95
-8.31 ±0.01 -11.29 ±0.01 -13.52 ±0.02
-9.88 ±0.01 -10.17 ±0.01 -10.97 ±0.00
-8.13 ±0.00 -9.99 ±0.00 -11.29 ±0.00
Plants
1 3 5
-14.46
-21.01 -22.30 -23.64
-14.98 -16.16 -17.23
-23.39 -38.81 -52.40
-22.59 -23.22 -27.54
-19.62 -26.69 -33.30
-18.78 ±0.03 -26.26 ±0.09 -32.36 ±0.11
-23.65 ±0.04 -28.76 ±0.05 -32.63 ±0.04
-18.03 ±0.02 -23.31 ±0.05 -27.49 ±0.07
Netflix
1 3 5
-57.59
-58.82 -59.27 -60.10
-57.66 -58.09 -59.21
-61.13 -67.12 -72.17
-60.67 -63.08 -65.32
-60.61 -65.06 -66.48
-58.27 ±0.02 -59.56 ±0.03 -60.78 ±0.06
-59.37 ±0.01 -60.55 ±0.02 -61.88 ±0.04
-58.32 ±0.02 -59.69 ±0.03 -61.28 ±0.04
DNA
1 3 5
-79.97
-82.69 -84.31 -85.85
-81.05 -82.62 -84.23
-87.68 -102.88 -117.82
-86.67 -93.38 -98.24
-85.12 -90.70 -96.08
-83.44 ±0.09 -90.33 ±0.15 -97.00 ±0.23
-84.65 ±0.05 -88.80 ±0.05 -92.00 ±0.08
-83.42 ±0.05 -88.07 ±0.07 -91.18 ±0.10
Movie
1 3 5
-52.81
-55.02 -57.06 -59.34
-54.07 -53.38 -56.09
-61.93 -79.51 -96.75
-61.21 -69.61 -78.24
-58.59 -67.62 -79.51
-58.53 ±0.09 -69.77 ±0.20 -80.66 ±0.16
-59.77 ±0.07 -68.82 ±0.12 -77.02 ±0.11
-58.99 ±0.08 -68.91 ±0.17 -75.07 ±0.16
BBC
1 3 5
-250.20
-250.72 -251.45 -252.52
-250.30 -250.55 -250.58
-258.72 -275.07 -290.24
-254.99 -262.65 -269.24
-256.75 -266.11 -274.58
-253.04 ±0.12 -253.43 ±0.10 -253.10 ±0.11 -258.74 ±0.15 -259.25 ±0.12 -258.82 ±0.13 -264.30 ±0.30 -265.01 ±0.26 -264.16 ±0.29
Table 1: Test-set log-likelihood. ϵ ∈ {1, 3, 5}: corruption budget for test-set generation. T : original test set; Ta : adversariallyperturbed test set; Tr : average over 10 randomly-perturbed test sets. Boldface indicates best approach, underline second best. minimize the estimator’s likelihood. Its corruption budget ϵ corresponds to the maximum Hamming distance a data point may be moved, subsequently bounding the Wasserstein distance between the original and corrupted datasets by ϵ. Lastly, we use MLE-PC as our base distribution P̂ . Evaluation Sets For each dataset, we construct three test sets to evaluate each method: (1) the original test set T ; (2) an adversarially-perturbed test set Ta , where each data point receives ϵ corruptions applied greedily to maximize the drop in likelihood of the MLE-PC; and (3) Tr generated by randomly corrupting each data point ϵ times. Results Table 1 summarizes the results. As we expect, the distributionally robust methods achieve slightly worse likelihood than MLE-PC given the unperturbed test set T ; nevertheless, we observe that P E T E R consistently outperforms RL-TPM on T . P E T E R also consistently outperforms RL-TPM on the randomly-perturbed test set Tr , which aims to capture an “average-case” corruption rather than a single worst case w.r.t. MLE-PC. Lastly, P E T E R outperforms RLTPM on the majority of the adversarially-perturbed datasets. We note that instances where P E T E R is outperformed tend to be at higher corruption levels, likely due to the CWdistance being a looser upper bound on the W-distance than the Hamming distance used by RL-TPM is; this results in the true Wasserstein ball that P E T E R optimizes for being
smaller than the one RL-TPM does, so the true worst-case distribution—typically lying on the edge of the ball—may be outside of the ϵ-CW ball. Despite the strict ϵ-CW constraint, P E T E R remains highly stable during training, as discussed further in Appendix C.
5
CONCLUSION
We introduced P E T E R, a data-free post-training framework that robustifies probabilistic circuits against distribution shifts within a Wasserstein ball without requiring access to any training data or retraining a circuit from scratch. By leveraging the Circuit-Wasserstein distance and deriving gradients for its computation, we formulated robust posttraining within a Wasserstein ball as an unconstrained gradient descent-ascent problem that can be solved efficiently for structured-decomposable PCs. Empirically, P E T E R consistently outperforms the data-dependent baseline across both unperturbed and perturbed test sets, particularly at lower corruption budgets, while maintaining stable training despite the non-smooth nature of Circuit-Wasserstein gradients. These results suggest that post-training robustification offers a practical and flexible alternative to training robust models from scratch and motivate further development of optimal transport-based losses for PCs. Future work includes extending P E T E R to decomposable but non-structured PCs.
Author Contributions Briefly list author contributions. This is a nice way of making clear who did what and to give proper credit. This section is optional. H. Q. Bovik conceived the idea and wrote the paper. Coauthor One created the code. Coauthor Two created the figures. Acknowledgements Briefly acknowledge people and organizations here. All acknowledgements go in this section. References Stefano V. Albrecht, Filippos Christianos, and Lukas Schäfer. Multi-Agent Reinforcement Learning: Foundations and Modern Approaches. MIT Press, 2024. URL https://www.marl-book.com. Dario Bauso, Jian Gao, and Hamidou Tembine. Distributionally robust games: f-divergence and learning. In Proceedings of the 11th EAI International Conference on Performance Evaluation Methodologies and Tools, pages 148–155, 2017. Jessa Bekker, Jesse Davis, Arthur Choi, Adnan Darwiche, and Guy Van den Broeck. Tractable learning for complex probability queries. Advances in Neural Information Processing Systems, 28, 2015. Dimitris Bertsimas and Omid Nohadani. Robust maximum likelihood estimation. INFORMS Journal on Computing, 31(3):445–458, 2019. Adrian Ciotinga and YooJung Choi. Optimal transport for probabilistic circuits. In The 41st Conference on Uncertainty in Artificial Intelligence, 2025. URL https: //openreview.net/forum?id=lNau87jgPx. Erick Delage and Yinyu Ye. Distributionally robust optimization under moment uncertainty with application to data-driven problems. Operations research, 58(3):595– 612, 2010. Mattia Desana and Christoph Schnörr. Expectation maximization for sum-product networks as exponential family mixture models, 04 2016. Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K. Weinberger, editors, Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014. URL https://proceedings.neurips.
cc/paper_files/paper/2014/file/ f033ed80deb0234979a61f95710dbe25-Paper. pdf. Daniel Kuhn, Peyman Mohajerin Esfahani, Viet Anh Nguyen, and Soroosh Shafieezadeh-Abadeh. Wasserstein distributionally robust optimization: Theory and applications in machine learning. In Operations research & management science in the age of analytics, pages 130– 166. Informs, 2019. Hugo Larochelle and Iain Murray. The neural autoregressive distribution estimator. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 29–37. JMLR Workshop and Conference Proceedings, 2011. Anji Liu and Guy Van den Broeck. Tractable regularization of probabilistic circuits. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 3558–3570. Curran Associates, Inc., 2021. URL https://proceedings.neurips. cc/paper_files/paper/2021/file/ 1d0832c4969f6a4cc8e8a8fffe083efb-Paper. pdf. Anji Liu, Kareem Ahmed, and Guy Van den Broeck. Scaling tractable probabilistic circuits: A systems perspective, 2024. URL https://arxiv.org/abs/2406. 00766. Jiashuo Liu, Zheyan Shen, Peng Cui, Linjun Zhou, Kun Kuang, and Bo Li. Distributionally robust learning with stable adversarial training. IEEE Trans. on Knowl. and Data Eng., 35(11):11288–11300, November 2023. ISSN 1041-4347. doi: 10.1109/TKDE.2022.3224056. URL https://doi.org/10.1109/TKDE.2022. 3224056. Daniel Lowd and Jesse Davis. Learning markov network structure with decision trees. In 2010 IEEE International Conference on Data Mining, pages 334–343. IEEE, 2010. Rohith Peddi, Tahrima Rahman, and Vibhav Gogate. Robust learning of tractable probabilistic models. In James Cussens and Kun Zhang, editors, Proceedings of the Thirty-Eighth Conference on Uncertainty in Artificial Intelligence, volume 180 of Proceedings of Machine Learning Research, pages 1572–1581. PMLR, 01–05 Aug 2022. URL https://proceedings.mlr.press/ v180/peddi22a.html. Bahar Taşkesen, Soroosh Shafieezadeh-Abadeh, Daniel Kuhn, and Karthik Natarajan. Discrete optimal transport with independent marginals is #p-hard, 2022. URL https://arxiv.org/abs/2203.01161.
Jan Van Haaren and Jesse Davis. Markov network structure learning: A randomized feature generation approach. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 26, pages 1148–1154, 2012. Antonio Vergari, YooJung Choi, Anji Liu, Stefano Teso, and Guy den Broeck. A Compositional Atlas of Tractable Circuit Operations for Probabilistic Inference. In Advances in Neural Information Processing Systems, volume 34, pages 13189–13201, 2021.
P E T E R: Post-Training Robustification of Probabilistic Circuits (Supplementary Material)
A
PROBABILISTIC CIRCUITS BACKGROUND 1 1 Adrian Ciotinga
YooJung Choi1
Yeming Dai
This section formally establishes the core structural properties of probabilistic circuits that dictate their tractability for 1 School of Computing various probabilistic inference queries. and Augmented Intelligence, Arizona State University University Definition 1. A probabilistic circuit C is smooth if every sum node n ∈ C has the same scope as its children: ∀ni ∈ ch(n), sc(ni ) = sc(n). Definition 2. A probabilistic T circuit C is decomposable if the children of every product node n ∈ C have disjoint scope: ∀ ni ̸= nj ∈ ch(n), sc(ni ) sc(nj ) = ∅. Definition 3. A probabilistic circuit C is structured-decomposable if it is decomposable and any two product nodes with the same scope partition their scope identically: ∀ n, m ∈ P , sc(n) = sc(m) =⇒ {sc(ni ) | ni ∈ ch(n)} = {sc(mi ) | mi ∈ ch(m)}. Definition 4. Two probabilistic circuits C1 and C2 are compatible if they are both structured-decomposable and any product node in C1 partitions its scope identically to any product node in C2 with the same scope: ∀ n ∈ C1 , m ∈ C2 , sc(n) = sc(m) =⇒ {sc(ni ) | ni ∈ ch(n)} = {sc(mi ) | mi ∈ ch(m)}. Nodes n and m in compatible PCs C1 and C2 that share the same type (e.g., sum, product, or input) and same scope are called corresponding nodes.
B
DIFFERENTIABLE COMPUTATION OF CW
Ciotinga and Choi [2025] detailed the recursive algorithm for computing the Circuit-Wasserstein distance, but the use of linear program solvers during the forward pass prevents the use of classic automatic differentiation techniques to compute gradients. Here, we derive the gradients for each step of the Circuit-Wasserstein algorithm. CW looks at pairs of corresponding nodes n and m in Pθ and Qϕ respectively; let ni (ext. mj ) denote the i’th child of n (ext. m). There are three cases: n and m are Sum Nodes Let θi and ϕj denote the weights of the i’th and j’th children of n and m respectively. The Circuit-Wasserstein distance between corresponding sum nodes is the solution to the discrete optimal transport problem where the distance between child ni and mj is their recursively-computed Circuit-Wasserstein distance. This is formulated as a linear programming problem: P min i,j CW(ni , mj )wi,j w P s.t. ∀i, j wi,j = θi CW(n, m) = (2) P ∀j, i wi,j = ϕj wi,j ≥ 0 CW(n,m) ∗ With solution w∗ and applying the Envelope Theorem, we see that ∂∂CW (ni ,mj ) = wi,j . Now, denoting the dual variables for the θi and ϕj constraints as xi and yj respectively, the above linear programming problem has the equivalent dual form:
( CW(n, m) =
max
P
s.t.
∀i, j, xi + yj ≤ CW(ni , mj )
x,y
i θi xi +
P
j ϕj yj
(3)
(n,m) (n,m) Another application of the Envelope Theorem at the optimal point yields ∂ CW∂θ = x∗i and ∂ CW∂ϕ = yj∗ respectively. i j Calculating ∇θ CW(n, m) and ∇ϕ CW(n, m) in this manner yields a sub-gradient—as the solution to a linear programming problem is non-smooth at its optimal point—leading to the computed gradient vectors being non-smooth step functions. This would typically necessitate the use of e.g., a quadratic regularizer to smooth out the gradients, but we forego this for differentiating Circuit-Wasserstein due to the unique optimization landscape of the problem effectively nullifying this discontinuity.
n and m are Product Nodes Since circuits Pθ and Qϕ are compatible, there is a bijective map between children ni and mi of n and m according to their scopes. The Circuit-Wasserstein distance between n amd m is simply an unweighted sum CW(n,m) of the distances between corresponding children. Consequently, ∂∂CW (ni ,mi ) = 1. n and m are Input Nodes The Circuit-Wasserstein distance between input nodes is equal to the Wasserstein distance between their encoded distributions, which can be computed efficiently for many continuous and discrete families of distributions. By the assumption of tractability for input node distributions commonly made for probabilistic circuits, we assume that ∇θ W(n, m) and ∇ϕ W(n, m) can be computed efficiently. This assumption holds for many classes of distributions, included (but not limited to) categorical and Gaussian distributions.
C
FEASIBILITY AND STABILITY OF PETER
Proposition 1. The optimal solution to the following unconstrained Lagrangian optimization problem is a feasible and optimal solution to the constrained problem in Equation 1. max min max EQϕ [log Pθ (X)] + λ(CW(P̂ , Qϕ ) − ϵ). Pθ
Qϕ λ≥0
(4)
Proof. Let θ∗ , ϕ∗ , λ∗ denote an optimal solution to 4. If CW(P̂ , Qϕ∗ ) > ϵ, then λ → ∞ results in the innermost maximization going to positive infinity (+∞). Since the outer optimization with respect to Qϕ seeks to minimize this objective, it will strictly avoid any Qϕ that causes the inner maximum to diverge, provided there is at least one feasible solution (for instance, setting Qϕ = P̂ yields CW(P̂ , P̂ ) = 0 ≤ ϵ, which results in a finite value). Consequently, to prevent the objective from diverging to +∞, the optimal distribution must satisfy CW(P̂ , Qϕ∗ ) ≤ ϵ. This ensures that the optimal solution Qϕ∗ strictly resides within the ϵ-CW ball. Thus, the optimal solution to the unconstrained Lagrangian optimization problem is guaranteed to be feasible with respect to the original constraint. Stability of CW Gradients As noted in Section 3.3, the gradients computed for the linear programs in CW computation are actually discontinuous sub-gradients that behave like step functions. This would typically make for highly instable training—manifesting as the optimizer’s inability to maintain a feasible Qϕ —but we observed that training within a CW-ball is actually highly stable, with our adversarial distribution Qϕ maintaining negligible constraint violation with nearly zero hyperparameter tuning. We attribute this to the fact that a single sum node has multiple corresponding sum nodes in the other circuit. Thus, its parameter gradients are the sum of multiple sub-gradients with offset discontinuities, mitigating the instability.