SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $q(\mathbf{c}_0 = 0 \mid \mathbf{z}_0) + q(\mathbf{c}_0 = 1 \mid \mathbf{z}_0) = 1$  and KL ( $q(\mathbf{z}_0 \mid \mathbf{x}) \| p(\mathbf{z}_0)$ ) is the KL between two Gaussians, which can be computed analytically.
A.1.2 KL Decisions
The second term (25) corresponds to the KL of the decisions. We can pull out the product from the log, yielding
$$
\begin{array}{l} \mathrm {K L} _ {\text {d e c i s i o n s}} = \sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} q (\boldsymbol {z} _ {0} \mid \boldsymbol {x}) \prod_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} q \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}\right) q (\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}) \\ \times \log \left(\prod_ {j \in \mathcal {P} _ {l} \backslash \{0 \}} \frac {q \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {x}\right)}{p \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {z} _ {p a (j)}\right)}\right) \tag {31} \\ \end{array}
$$
$$
= \sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} \sum_ {j \in \mathcal {P} _ {l} \backslash \{0 \}} q (\boldsymbol {z} _ {0} \mid \boldsymbol {x}) \prod_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} q \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right) \log \left(\frac {q \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {x}\right)}{p \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {z} _ {p a (j)}\right)}\right) \tag {32}
$$
Let us define as  $\mathcal{P}_{l\in j}$  all paths that go through node  $j$ , as  $\mathcal{P}_{\leq j}$  (denoted as  $\mathcal{P}_j$  in the main text for brevity) the unique path that ends in the node  $j$ , and as  $\mathcal{P}_{>j}$  all the possible paths that start from the node  $j$  and continue to a leaf  $l\in \mathbb{L}$ . Similarly, let us define as  $\pmb{z}_{\leq j}$  all the latent embeddings that are contained in the path from the root to node  $j$  and as  $\pmb{z}_{>j}$  all the latent embeddings of the nodes  $i > j$  that can be reached from node  $j$ .
To factorize the above equation, we first change from a pathwise view to a nodewise view. Instead of summing over all possible leaves in the tree  $(\sum_{l\in \mathbb{L}})$  and then over each contained node  $(\sum_{j\in \mathcal{P}_l\backslash \{0\}})$ , we sum over all nodes  $(\sum_{j\in \mathbb{V}\backslash \{0\}})$  and then over each path that leads through the selected node  $(\sum_{\mathcal{P}_{l\in j}})$ .
$$
\begin{array}{l} \sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} \sum_ {j \in \mathcal {P} _ {l} \backslash \{0 \}} q (\boldsymbol {z} _ {0} \mid \boldsymbol {x}) \prod_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} q (c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}) q (\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}) \log \left(\frac {q (c _ {p a (j) \rightarrow j} \mid \boldsymbol {x})}{p (c _ {p a (j) \rightarrow j} \mid \boldsymbol {z} _ {p a (j)})}\right) \\ = \sum_ {j \in \mathbb {V} \backslash \{0 \}} \sum_ {\mathcal {P} _ {l \in j}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} q (\boldsymbol {z} _ {0} \mid \boldsymbol {x}) \prod_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} q (c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}) q (\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}) \log \left(\frac {q (c _ {p a (j) \rightarrow j} \mid \boldsymbol {x})}{p (c _ {p a (j) \rightarrow j} \mid \boldsymbol {z} _ {p a (j)})}\right) \\ \end{array}
$$
The above can be proved with the following Lemma, where we rewrite  $\sum_{\mathcal{P}_{l\in j}} = \sum_{l\in \mathbb{L}}\mathbb{1}[j\in \mathcal{P}_l]$ .
Lemma A.1. Given a binary tree  $\mathcal{T}$  as defined in Section 2.1, composed of a set of nodes  $\mathbb{V} = \{0,\dots ,V\}$  and leaves  $\mathbb{L}\subset \mathbb{V}$ , where  $\mathcal{P}_l$  is the decision path from root 0 to leaf  $l$ , and  $\mathbf{z}_{\mathcal{P}_l} = \{\mathbf{z}_i\mid i\in \mathcal{P}_l\}$  the set of latent variables selected by the path  $\mathcal{P}_l$ . Then it holds
$$
\sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} \sum_ {j \in \mathcal {P} _ {l} \backslash \{0 \}} f (j, l, \boldsymbol {z} _ {\mathcal {P} _ {l}}) = \sum_ {j \in \mathbb {V} \backslash \{0 \}} \sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} \mathbb {1} [ j \in \mathcal {P} _ {l} ] f (j, l, \boldsymbol {z} _ {\mathcal {P} _ {l}}), \tag {34}
$$
Proof. The proof is as follows:
$$
\begin{array}{l} \sum_ {j \in \mathbb {V} \backslash \{0 \}} \sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} \mathbb {1} [ j \in \mathcal {P} _ {l} ] f (j, l, \boldsymbol {z} _ {\mathcal {P} _ {l}}) = \sum_ {j \in \mathbb {V} \backslash \{0 \}} \sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} f (j, l, \boldsymbol {z} _ {\mathcal {P} _ {l}}) \sum_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} \mathbb {1} [ i = j ] (35) \\ = \sum_ {l \in \mathbb {L}} \sum_ {j \in \mathbb {V} \backslash \{0 \}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} \sum_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} f (j, l, \boldsymbol {z} _ {\mathcal {P} _ {l}}) \mathbb {1} [ i = j ] (36) \\ = \sum_ {l \in \mathbb {L}} \sum_ {j \in \mathbb {V} \backslash \{0 \}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} \sum_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} f (i, l, \boldsymbol {z} _ {\mathcal {P} _ {l}}) \mathbb {1} [ i = j ] (37) \\ = \sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} \sum_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} f (i, l, \boldsymbol {z} _ {\mathcal {P} _ {l}}) \sum_ {j \in \mathbb {V} \backslash \{0 \}} \mathbb {1} [ i = j ] (38) \\ = \sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} \sum_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} f (i, l, \boldsymbol {z} _ {\mathcal {P} _ {l}}) (39) \\ = \sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} \sum_ {j \in \mathcal {P} _ {l} \backslash \{0 \}} f (j, l, \boldsymbol {z} _ {\mathcal {P} _ {l}}). (40) \\ \end{array}
$$
Having proven the equality, we can continue with the KL of the decisions as follows:
$$
\begin{array}{l} \mathrm {K L} _ {\text {d e c i s i o n s}} = \\ \sum_ {j \in \mathbb {V} \backslash \{0 \}} \sum_ {\mathcal {P} _ {l} \in j} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} q (\boldsymbol {z} _ {0} \mid \boldsymbol {x}) \prod_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} q \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right) \log \left(\frac {q \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {x}\right)}{p \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {z} _ {p a (j)}\right)}\right) (41) \\ = \sum_ {j \in \mathbb {V} \backslash \{0 \}} \sum_ {\mathcal {P} _ {l \in j}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} [ q (\boldsymbol {z} _ {0} \mid \boldsymbol {x}) \prod_ {i \in \mathcal {P} _ {\leq j} \backslash \{0 \}} q (c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}) q (\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}) \log \left(\frac {q (c _ {p a (j) \rightarrow j} \mid \boldsymbol {x})}{p (c _ {p a (j) \rightarrow j} \mid \boldsymbol {z} _ {p a (j)})}\right) (42) \\ \times \prod_ {k \in \mathcal {P} _ {> j}} q \left(c _ {p a (k) \rightarrow k} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {k} \mid \boldsymbol {z} _ {p a (k)}\right) ] \\ = \sum_ {j \in \mathbb {V} \backslash \{0 \}} \sum_ {\mathcal {P} _ {> j}} \int_ {\boldsymbol {z} _ {\leq j}, \boldsymbol {z} _ {> j}} [ q (\boldsymbol {z} _ {0} \mid \boldsymbol {x}) \prod_ {i \in \mathcal {P} _ {\leq j} \backslash \{0 \}} q \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right) \log \left(\frac {q \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {x}\right)}{p \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {z} _ {p a (j)}\right)}\right) (43) \\ \times \prod_ {k \in \mathcal {P} > j} q \left(c _ {p a (k) \rightarrow k} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {k} \mid \boldsymbol {z} _ {p a (k)}\right) ] \\ = \sum_ {j \in \mathbb {V} \backslash \{0 \}} \int_ {\boldsymbol {z} _ {\leq j}} q (\boldsymbol {z} _ {0} \mid \boldsymbol {x}) \prod_ {i \in \mathcal {P} _ {\leq j} \backslash \{0 \}} q \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right) \log \left(\frac {q \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {x}\right)}{p \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {z} _ {p a (j)}\right)}\right) (44) \\ \times \sum_ {\mathcal {P} _ {> j}} \int_ {\mathbf {z} _ {> j}} \left[ \prod_ {k \in \mathcal {P} _ {> j}} q \left(c _ {p a (k) \rightarrow k} \mid \mathbf {x}\right) q \left(\mathbf {z} _ {k} \mid \mathbf {z} _ {p a (k)}\right)\right] \\ \end{array}
$$
From Eq. 41 to Eq. 42, we split the inner product into the nodes of the paths  $\mathcal{P}_{l\in j}$  that are before and after the node  $j$
From Eq. 42 to Eq. 43, we observe that the sum over all paths going through  $j$  can be reduced to the sum over all paths starting from  $j$ , because there is only one path to  $j$ , which is specified in the product that comes after.
From Eq. 43 to Eq. 44, we observe that the sum over paths starting from  $j$  and integral over  $z_{>j}$  concern only the terms of the second line. Observe that the term on the second line of Eq. 44
integrates out to 1 and we get
$$
\begin{array}{l} \mathrm {K L} _ {\text {d e c i s i o n s}} = \\ \sum_ {j \in \mathbb {V} \backslash \{0 \}} \int_ {\boldsymbol {z} \leq j} q \left(\boldsymbol {z} _ {0} \mid \boldsymbol {x}\right) \prod_ {i \in \mathcal {P} _ {\leq j} \backslash \{0 \}} q \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right) \log \left(\frac {q \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {x}\right)}{p \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {z} _ {p a (j)}\right)}\right) (45) \\ = \sum_ {j \in \mathbb {V} \backslash \{0 \}} \int_ {\boldsymbol {z} _ {<   j}} \int_ {\boldsymbol {z} _ {j}} q (\boldsymbol {z} _ {0} \mid \boldsymbol {x}) \prod_ {i \in \mathcal {P} _ {\leq j} \backslash \{0 \}} q \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right) \log \left(\frac {q \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {x}\right)}{p \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {z} _ {p a (j)}\right)}\right) (46) \\ = \sum_ {j \in \mathbb {V} \backslash \{0 \}} \int_ {\boldsymbol {z} _ {<   j}} q \left(\boldsymbol {z} _ {0} \mid \boldsymbol {x}\right) \prod_ {i \in \mathcal {P} _ {<   j} \backslash \{0 \}} q \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right) \log \left(\frac {q \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {x}\right)}{p \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {z} _ {p a (j)}\right)}\right) (47) \\ \times \int_ {\boldsymbol {z} _ {j}} q \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {j} \mid \boldsymbol {z} _ {p a (j)}\right) \\ = \sum_ {j \in \mathbb {V} \backslash \{0 \}} \int_ {\boldsymbol {z} _ {<   j}} q \left(\boldsymbol {z} _ {0} \mid \boldsymbol {x}\right) \prod_ {i \in \mathcal {P} _ {<   j} \backslash \{0 \}} q \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right) q \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {x}\right) (48) \\ \times \log \left(\frac {q (c _ {p a (j) \rightarrow j} \mid \boldsymbol {x})}{p (c _ {p a (j) \rightarrow j} \mid \boldsymbol {z} _ {p a (j)})}\right). \\ \end{array}
$$
From Eq. 46 to Eq. 47, we single out the term in the product that corresponds to  $j = i$ , which is the only term that depends on  $\int_{z_i}$ .
From Eq. 47 to Eq. 48, we observe that in the singled-out term,  $\int_{z_j} q(z_j \mid \mathbf{z}_{pa(j)}) = 1$ , which leaves only  $q(c_{pa(j) \to j} \mid \mathbf{x})$ .
This equation can be rewritten in a more interpretable way. Let us define the probability of reaching node  $j$  and observing  $\mathbf{z}_j$  as
$$
P (j; \boldsymbol {z}, \boldsymbol {c}) = q (\boldsymbol {z} _ {0} \mid \boldsymbol {x}) \prod_ {i \in \mathcal {P} _ {\leq j} \backslash \{0 \}} q \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right). \tag {49}
$$
Then the KL term of the decisions can be simplified as
$$
\begin{array}{l} \mathrm {K L} _ {\text {d e c i s i o n s}} = \sum_ {j \in \mathbb {V} \backslash \{0 \}} \int_ {\boldsymbol {z} <   j} P (p a (j); \boldsymbol {z}, \boldsymbol {c}) q \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {x}\right) \log \left(\frac {q \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {x}\right)}{p \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {z} _ {p a (j)}\right)}\right) (50) \\ = \sum_ {i \in \mathbb {V} \backslash \mathbb {L}} \sum_ {k \in \{0, 1 \}} \int_ {\boldsymbol {z} <   i} P (i; \boldsymbol {z}, \boldsymbol {c}) q \left(c _ {i} = k \mid \boldsymbol {x}\right) \log \left(\frac {q \left(c _ {i} = k \mid \boldsymbol {x}\right)}{p \left(c _ {i} = k \mid \boldsymbol {z} _ {i}\right)}\right). (51) \\ \end{array}
$$
This term requires Monte Carlo sampling for the expectations over the latent variables  $\mathbf{z}$ , while we can analytically compute the sum over all decisions  $\mathcal{P}_l$ .
$$
\begin{array}{l} \mathrm {K L} _ {\text {d e c i s i o n s}} = \sum_ {i \in \mathbb {V} \backslash \mathbb {L}} \int_ {\boldsymbol {z} <   i} P (i; \boldsymbol {z}, \boldsymbol {c}) (52) \\ \times \left[ q (c _ {i} = 0) \mid \boldsymbol {x}) \log \left(\frac {q (\mathbf {c} _ {i} = 0 \mid \boldsymbol {x})}{p (\mathbf {c} _ {i} = 0 \mid \boldsymbol {z} _ {i})}\right) + q (\mathbf {c} _ {i} = 1 \mid \boldsymbol {x}) \log \left(\frac {q (\mathbf {c} _ {i} = 1 \mid \boldsymbol {x})}{p (\mathbf {c} _ {i} = 1 \mid \boldsymbol {z} _ {i})}\right) \right] \\ \approx \frac {1}{M} \sum_ {m = 1} ^ {M} \sum_ {i \in \mathbb {V} \backslash \mathbb {L}} P (i; \boldsymbol {z} ^ {(m)}, \boldsymbol {c}) (53) \\ \times \left[ q (c _ {i} = 0) \mid \boldsymbol {x}) \log \left(\frac {q (\mathbf {c} _ {i} = 0 \mid \boldsymbol {x})}{p (\mathbf {c} _ {i} = 0 \mid \boldsymbol {z} _ {i} ^ {(m)})}\right) + q (\mathbf {c} _ {i} = 1 \mid \boldsymbol {x}) \log \left(\frac {q (\mathbf {c} _ {i} = 1 \mid \boldsymbol {x})}{p (\mathbf {c} _ {i} = 1 \mid \boldsymbol {z} _ {i} ^ {(m)})}\right) \right], \\ \end{array}
$$
where  $P(i; \boldsymbol{z}^{(m)}, \boldsymbol{c}) = \prod_{j \in \mathcal{P}_{\leq i}} q(c_{pa(j) \to j} \mid \boldsymbol{x})$  is the probability of reaching node  $i$ , defined as  $P(i; \boldsymbol{c})$  in Eq. 16/19/20 for simplicity.
A.1.3 KL Nodes
Finally, we can analyze the last term of the KL term, which corresponds to the KL of the nodes (26). The reasoning is similar to the equations above and we will use the same notation. The KL of the nodes can be written as
$$
\begin{array}{l} \mathrm {K L} _ {\text {n o d e s}} = \sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} q \left(\boldsymbol {z} _ {0} \mid \boldsymbol {x}\right) \prod_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} q \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right) (54) \\ \times \log \left(\frac {\prod_ {j \in \mathcal {P} _ {l} \backslash \{0 \}} q \left(\boldsymbol {z} _ {j} \mid \boldsymbol {z} _ {p a (j)}\right)}{\prod_ {k \in \mathcal {P} _ {l} \backslash \{0 \}} p \left(\boldsymbol {z} _ {k} \mid \boldsymbol {z} _ {p a (k)}\right)}\right) \\ = \sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} \sum_ {j \in \mathcal {P} _ {l} \backslash \{0 \}} q (\boldsymbol {z} _ {0} \mid \boldsymbol {x}) \prod_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} q \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right) \log \left(\frac {q \left(\boldsymbol {z} _ {j} \mid \boldsymbol {z} _ {p a (j)}\right)}{p \left(\boldsymbol {z} _ {j} \mid \boldsymbol {z} _ {p a (j)}\right)}\right) (55) \\ \end{array}
$$
We now change from a pathwise view to a nodewise view.
where  $P(pa(j); \boldsymbol{z}, \boldsymbol{c})$  is defined in Eq. 49 and where  $P(i; \boldsymbol{z}^{(m)}, \boldsymbol{c}) = P(i; \boldsymbol{c}) = \prod_{j \in \mathcal{P}_{\leq i}} q(c_{pa(j) \to j} | \boldsymbol{x})$  is the probability of reaching node  $i$ .
A.1.4 KL terms
Using the above factorization, the KL term of the ELBO can be written as
$$
\begin{array}{l} \operatorname {K L} \left(q \left(\boldsymbol {z}, \mathcal {P} _ {l} \mid \boldsymbol {x}\right) \| p \left(\boldsymbol {z}, \mathcal {P} _ {l}\right)\right) \approx \operatorname {K L} \left(q \left(\boldsymbol {z} _ {0} \mid \boldsymbol {x}\right) \| p \left(\boldsymbol {z} _ {0}\right)\right) \\ + \frac {1}{M} \sum_ {m = 1} ^ {M} \sum_ {i \in \mathbb {V} \backslash \mathbb {L}} P (i; \boldsymbol {z} ^ {(m)}, \boldsymbol {c}) \sum_ {c _ {i} \in \{0, 1 \}} q \left(c _ {i} \mid \boldsymbol {x}\right) \log \left(\frac {q \left(c _ {i} \mid \boldsymbol {x}\right)}{p \left(c _ {i} \mid \boldsymbol {z} _ {i} ^ {(m)}\right)}\right) \tag {64} \\ + \frac {1}{M} \sum_ {m = 1} ^ {M} \sum_ {i \in \mathbb {V} \backslash \{0 \}} P (i; \boldsymbol {z} ^ {(m)}, \boldsymbol {c}) \operatorname {K L} (q (\boldsymbol {z} _ {i} ^ {(m)} \mid p a (\boldsymbol {z} _ {i} ^ {(m)})) \| p (\boldsymbol {z} _ {i} ^ {(m)} \mid p a (\boldsymbol {z} _ {i} ^ {(m)}))), \\ \end{array}
$$
where  $P(i;\pmb{z}^{(m)},\pmb {c}) = P(i;\pmb {c}) = \prod_{j\in \mathcal{P}_{\leq i}}q(c_{pa(j)\to j}\mid \pmb {x}).$
A.1.5 Reconstruction Loss
Finally, to compute the full ELBO, the KL terms are added to the reconstruction loss defined in (16). Here, assumptions about the distribution of the inputs are required. For the grayscale datasets such as MNIST, Fashion-MNIST, and Omniglot, as well as the one-hot-encoded 20Newsgroup, we assume that the inputs are Bernoulli distributed, such that the resulting reconstruction loss is the binary cross entropy. On the other hand, for the colored datasets CIFAR-10, CIFAR-100, and CelebA, we assume that their pixel values are normally distributed, which leads to the mean squared error as loss function, where we assume that  $\sigma = 1$ .
A.2 Computational Complexity
All terms of the Evidence Lower Bound, Equation 12, can be computed efficiently and the computational complexity of a single joint update of the parameters is  $O(MBVDC_p)$ , where  $M$  is the number of MC samples,  $B$  is the batch size,  $V$  is the number of nodes in the tree,  $D$  is the maximum depth of the tree, and  $C_p$  is the cost to compute the KL between Gaussians. It should be noted that the computational complexity is, in practice, reduced to  $O(LBVC_p)$ , as the term  $P(i; c)$  can be computed dynamically from parent nodes.
In addition to the review presented in Section 3, which encompasses a broad range of related work in the field of deep latent variable models, we provide a review of relevant work in the domains of hierarchical clustering and decision trees. By doing so, we hope to shed further light on the current state-of-the-art approaches and contribute to a deeper understanding of the challenges and opportunities that lie in the intersection of hierarchical clustering, decision trees, and latent variable models.
B.1 Hierarchical Clustering
Hierarchical clustering algorithms have long been employed in the field of data mining and machine learning to extract hierarchical structures from data (Sneath, 1957; Ward, 1963; Murtagh & Contreras, 2012). Agglomerative clustering is among the earliest and most well-known hierarchical clustering algorithms. These methods start with each data point as an individual cluster and then iteratively merge the closest pairs of clusters, according to a predefined distance metric, until a stopping criterion is met. While single-linkage and complete-linkage agglomeration clustering are widely used as baselines, we observe better performance when using the bottom-up strategy proposed by Ward (1963). Ward's minimum variance criterion minimizes the total within-cluster variance (Murtagh & Legendre, 2014), thus providing balanced and compact clusters. In contrast, the Bayesian Hierarchical Clustering (BHC) proposed by Heller & Ghahramani (2005) takes a different approach by employing hypothesis testing to determine when to merge the clusters. The divisive clustering algorithms, on the other hand, provide a different strategy to hierarchical clustering. Unlike agglomerative methods, divisive clustering starts with all data points in a single cluster and recursively splits clusters into smaller ones. The proposed TreeVAE is an example of a divisive clustering method. Among a variety of proposed methods, the Bisecting-K-means algorithm (Steinbach et al., 2000; Nistér & Stewénius,
2006) is widely used for its simplicity; it applies k-means with two clusters recursively. Similarly, Williams (1999) learn a hierarchical probabilistic Gaussian mixture model. Neal (2003) introduce a hierarchical generalization of Dirichlet process mixture models.
More recent approaches include PERCH (Kobren et al., 2017b), which is a non-greedy, incremental algorithm that scales to both the number of data points and the number of clusters, GHC (Monath et al., 2019), which leverages continuous representations of trees in a hyperbolic space and optimizes a differentiable cost function, and Mathieu et al. (2019), which similarly use a Poincaré ball model of hyperbolic geometry as a latent space to learn continuous hierarchical representations with a VAE. Furthermore, RSSCOMP (You et al., 2015), which explores a subspace clustering method based on an orthogonal matching pursuit. Finally, Deep ECT (Mautz et al., 2020) proposes a divisive hierarchical embedded clustering method, which jointly optimizes an autoencoder that compresses the data into an embedded space and a hierarchical clustering layer on top of it. Density-based clustering algorithms, such as DBSCAN, belong to a distinct category of clustering techniques. They aim to identify regions in a dataset where points are densely concentrated and classify outliers as noise (Ester et al., 1996). R. J. Campello et al. (2015) build on this idea to learn a hierarchy based on the distances between datapoints where distance is roughly determined by the density. However, one limitation of density-based methods lie in their performance when confronted with complex datasets requiring high-dimensional representations. In such cases, estimating density requires an exponentially growing number of data points, which leads to scalability issues that TreeVAE does not have.
B.2 Decision Trees