Estimating Staged Event Tree Models via Hierarchical Clustering on the Simplex

Muhammad Shoaib; Eva Riccomagno; Manuele Leonelli; Gherardo Varando

stat.MLMar 16, 2026

Estimating Staged Event Tree Models via Hierarchical Clustering on the Simplex

Authors: Muhammad Shoaib, Eva Riccomagno, Manuele Leonelli, Gherardo Varando

Organizations: University of Genova, Department of Mathematics, Genova, 16146, Italy · IE University, School of Science and Technology, Madrid, Spain · Universitat de València, Department of Statistics and Operational Research, València, , Spain

Abstract

Staged tree models enhance Bayesian networks by incorporating context-specific dependencies through a stage-based structure. In this study, we present a new framework for estimating staged trees using hierarchical clustering on the probability simplex, utilizing simplex basesd divergences. We conduct a thorough evaluation of several distance and divergence metrics including Total Variation, Hellinger, Fisher, and Kaniadakis; alongside various linkage methods such as Ward.D2, average, complete, and McQuitty. We conducted the simulation experiments that reveals Total Variation, especially when combined with Ward.D2 linkage, consistently produces staged trees with better model fit, structure recovery, and computational efficiency. We assess performance by utilizing relative Bayesian Information Criterion (BIC), and Hamming distance. Our findings indicate that although Backward Hill Climbing (BHC) delivers competitive outcomes, it incurs a significantly higher computational cost. On the other, Total Variation divergence with Ward.D2 linkage, achieves similar performance while providing significantly better computational efficiency, making it a more viable option for large-scale or time sensitive tasks.

Figures & tables

Panel of Figure 1 — Figure 1: Event tree (left) and staged event tree (right) for $(X_{1},X_{2},X_{3})$ with $\mathcal{X}_{1}={a,b,c}$ , $\mathcal{X}_{2}={0,1}$ , and $\mathcal{X}_{3}={1,2,3}$ . Left: edges labeled as $X_{i}=\cdot$ . Right: edges labeled as $X_{i}=\cdot$ with stage transition probabilities in parentheses.

$X_{1}$	$X_{2}$	$X_{3}$	$P(X_{1},X_{2},X_{3})$	$X_{1}$	$X_{2}$	$X_{3}$	$P(X_{1},X_{2},X_{3})$
$a$	$0$	$1$	$0.3\times 0.6\times 0.7=0.1260$	$b$	$1$	$1$	$0.5\times 0.4\times 0.4=0.0800$
$a$	$0$	$2$	$0.3\times 0.6\times 0.2=0.0360$	$b$	$1$	$2$	$0.5\times 0.4\times 0.4=0.0800$
$a$	$0$	$3$	$0.3\times 0.6\times 0.1=0.0180$	$b$	$1$	$3$	$0.5\times 0.4\times 0.2=0.0400$
$a$	$1$	$1$	$0.3\times 0.4\times 0.7=0.0840$	$c$	$0$	$1$	$0.2\times 0.6\times 0.4=0.0480$
$a$	$1$	$2$	$0.3\times 0.4\times 0.2=0.0240$	$c$	$0$	$2$	$0.2\times 0.6\times 0.4=0.0480$
$a$	$1$	$3$	$0.3\times 0.4\times 0.1=0.0120$	$c$	$0$	$3$	$0.2\times 0.6\times 0.2=0.0240$

Table 1: Joint distribution induced by the staged tree in Figure 1 (right).

Distance	Formula	Range
Total Variation distance	$d_{TV}(p,q)=\frac{1}{2}\sum_{i=1}^{s}\|p_{i}-q_{i}\|$	$[0,1]$
Hellinger distance	$d_{H}(p,q)=\frac{1}{\sqrt{2}}\sqrt{\sum_{i=1}^{s}\left(\sqrt{p_{i}}-\sqrt{q_{i}}\right)^{2}}$	$[0,1]$

Table 2: Distances used in Section 4 .

Divergence	Formula	Range
Fisher divergence	$\displaystyle d_{F}(p\\|q)=4\,\arccos^{2}\!\left(\sum_{i=1}^{s}\sqrt{p_{i}\,q_{i}}\right)$	$\left[0,\frac{\pi}{2}\right]$
Jensen–Shannon divergence	$\displaystyle d_{\mathrm{JS}}(p\\|q)=\tfrac{1}{2}\sum_{i=1}^{s}\!\left[p_{i}\log\!\frac{2p_{i}}{p_{i}+q_{i}}+q_{i}\log\!\frac{2q_{i}}{p_{i}+q_{i}}\right]$	$[0,\ln 2]$
Kaniadakis divergence	$\displaystyle{d}_{\kappa}(p\,\\|\,q)=\sum_{x\in\Omega}\left(\log_{\kappa}(p(x))-\log_{\kappa}(q(x))\right)\cdot\frac{A(p(x))}{\sum_{x\in\Omega}A(p(x))}$	$[0,+\infty)$
Total KL divergence	$\displaystyle d_{T}(p\\|q)=\sum_{i=1}^{s}\left(p_{i}\log\!\frac{p_{i}}{q_{i}}+q_{i}\log\!\frac{q_{i}}{p_{i}}\right)$	$[0,+\infty)$

Table 3: Divergences used in Section 4 .

Clustering method	Formula
Average linkage	$\displaystyle D_{A}(\mathcal{C}_{1},\mathcal{C}_{2})=\frac{1}{\|\mathcal{C}_{1}\|\|\mathcal{C}_{2}\|}\sum_{p\in\mathcal{C}_{1}}\sum_{q\in\mathcal{C}_{2}}d(p,q)$
Complete linkage	$\displaystyle D_{C}(\mathcal{C}_{1},\mathcal{C}_{2})=\max_{p\in\mathcal{C}_{1},\;q\in\mathcal{C}_{2}}d(p,q)$
McQuitty (WPGMA)	$\displaystyle D_{M}(\mathcal{C}_{3},\mathcal{C}_{1}\!\cup\!\mathcal{C}_{2})=\tfrac{1}{2}\!\left[D(\mathcal{C}_{3},\mathcal{C}_{1})+D(\mathcal{C}_{3},\mathcal{C}_{2})\right]$
Ward’s minimum variance (Ward.D2)	$\displaystyle D_{W}(\mathcal{C}_{1},\mathcal{C}_{2})=\frac{\|\mathcal{C}_{1}\|\|\mathcal{C}_{2}\|}{\|\mathcal{C}_{1}\|+\|\mathcal{C}_{2}\|}\,\\|\bar{p}_{1}-\bar{p}_{2}\\|^{2}$

Table 4: Hierarchical clustering linkages used in Section 4 . Here

\mathcal{C}_{1},\mathcal{C}_{2},\mathcal{C}_{3}

denote clusters of probability vectors

p\in\Delta_{s-1}

;

d(p,q)

is a dissimilarity between individual distributions;

\bar{p}_{i}=|\mathcal{C}_{i}|^{-1}\sum_{p\in\mathcal{C}_{i}}p

is the centroid of cluster

\mathcal{C}_{i}

; and

\|\cdot\|

is the Euclidean norm.

Figure 2: Comparison of

\Delta_{BIC}

and

\Delta_{HD}

across divergence-based staged event tree classifiers for

p=11

, under the following configurations:

k_{0},k=NA,q=0.9.

Figure 3: Comparison of $\Delta_{\mathrm{BIC}}$ and $\Delta_{\mathrm{HD}}$ and time across divergence-based staged event tree classifiers for sample sizes $512$ with $k_{0},k=NA,q=0.9$

Figure 4: Accuracy and f1 for all the staged event tree classifiers utilizing various divergence measures

k=2

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning
$T=(V,E)$	Event tree with vertices $V$ and edges $E$
$\mathring{V}$	Set of non-leaf (internal) vertices
$\kappa:\mathring{V}\to\mathcal{C}$	Stage assignment function
$\eta:E\to\mathcal{C}\times\bigcup_{i}\mathcal{X}_{i}$	Edge-labeling map
$\mathcal{S}_{i}$	Set of situations at level $i-1$
$\hat{\theta}_{v}$	Empirical conditional probability vector at vertex $v$

Appendix

Table 5: Summary of key mathematical notation used in the learning algorithm 1

Dataset	# Observations	# Variables	# Atomic Events	Imbalance
Asym	1,000	4	13	0.7
chestSim500	500	8	27	0.546
energy1	768	9	1,728	1.000
energy2	768	9	1,728	1.000
fallEld	5,000	4	64	0.888
fertility	100	10	15,552	0.529

Appendix

Table 6: Summary of the 14 datasets used in the experimental study, including number of observations, number of variables, number of atomic events, and imbalance measure.

Figure 5: Comparison of

\Delta_{BIC}

and

\Delta_{HD}

across divergence-based staged event tree classifiers for

p=11

, under the following configurations:

k_{0}=2,k=NA,q=NA

Figure 6: Comparison of $\Delta_{\mathrm{BIC}}$ and $\Delta_{\mathrm{HD}}$ and time across divergence-based staged event tree classifiers for sample sizes $512$ with $k_{0}=2,k=NA,q=NA$

Bayesian Networks

Bayesian networks represent joint probability distributions as graphs of variables and conditional dependencies, enabling inference under uncertainty.

See all 47 papers

Algorithmically Aligned Neural Agglomerative Tree Construction
Oct 5, 2026Robert R Nerem, Pranav Singh, Cheyenne Ward +1Clustering Hierarchical Clustering
University of California San Diego · University of Washington
Decomposition for Bayesian Networks: Local and Parallel Inference
Jul 6, 2026Pei Heng, Xinyi Hu, Yi SunParameter Estimation Bayesian Networks
Institute of Statistics and Data Science, Xinjiang University of Finance and Economics, Urumqi, China
Simultaneous Latent Budget Trees for Stratified Classification
Jun 11, 2026Cristian Buoncompagni, Stefano Pellegrino, Giulia Vannucci +2Latent Variable Models Interpretable ML
Department of Physics, University of Naples Federico II, via Cintia 21, Napoli, 80125, Italy. · Department of Electrical Engineering and Information Technologies, University of Naples Federico II, via Claudio 21, Napoli, 80125, Italy. · Department of Neuroscience, Reproductive Sciences and Odontostomatology, University of Naples Federico II, via Sergio Pansini 5, Napoli, 80131, Italy.

Oct 5, 2026J/K move · Enter open · S save

Algorithmically Aligned Neural Agglomerative Tree Construction

Robert R Nerem, Pranav Singh, Cheyenne Ward, Yusu Wang

University of California San Diego · University of Washington

Clustering Hierarchical Clustering Phylogenetic Inference

Open paper AskPDF arXiv

Linkage algorithms for hierarchical clustering (HC) are a powerful and efficient framework for constructing clustering trees, yet it is often unclear which merge rule best suits a given dataset or task. In contrast, neural approaches can learn from data, but often fail to retain the efficiency and size generalization of classical algorithms. We introduce NN-linkage, a neural network (NN) model that can learn task-specific and locally dependent merge rules while retaining the recursive structure and efficient inference of classical linkage algorithms. In particular, our model is algorithmically aligned with the Lance-Williams (LW) recurrence, a parameterized framework for defining a broad, continuous family of linkage rules for agglomerative HC. Classical methods such as single linkage (SL), complete linkage (CL), and average linkage arise as discrete choices within this broader family. We show that NN-linkage is a universal approximator for continuous linkage functions, including LW recurrences, and, when paired with a transformer encoding, can also approximate globally dependent rules such as robust single-linkage. We further show that NN-linkage can exactly implement any symmetric constant-coefficient LW recurrence across all input sizes. On the empirical front, we evaluate NN-linkage in real-world applications, clock-tree routing and phylogenetic reconstruction, using both synthetic and real datasets, demonstrating its effectiveness over both classical algorithms and other neural approaches. By learning merge rules directly from target trees, NN-linkage extends efficient HC to scientific and engineering objectives not adequately captured by existing hand-designed linkage rules.

Abstract

Bayesian Networks

Explore similar work

Algorithmically Aligned Neural Agglomerative Tree Construction

Decomposition for Bayesian Networks: Local and Parallel Inference

Simultaneous Latent Budget Trees for Stratified Classification