SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Given location and scale parameters denoted as  $\mu \in \mathbb{H}^d$  and  $\Sigma \in \mathbb{R}^{d\times d}$ , respectively, the procedure for sampling from a wrapped normal distribution  $z\sim \mathcal{W}\mathcal{N}(\mu ,\Sigma)$  defined over  $\mathbb{H}^d$  is given as follows:
$$
z = \exp_ {\mu} \circ \mathrm {P T} _ {\mu^ {o} \rightarrow \mu} (u), \quad u _ {1: d} \sim \mathcal {N} (0, \Sigma), \tag {22}
$$
where  $\mu^o = (1,0,\ldots ,0)^\top$  denotes the origin of the  $\mathbb{H}^d$ . Note that we set  $u_{0} = 0$ , and  $u\coloneqq u_{0:d}\in T_{\mu^o}\mathbb{H}^d$  represents a tangent vector at the origin  $T_{\mu^o}\mathbb{H}^d$ .
From the sampling definition in equation (22), the probability density function  $\mathcal{W}\mathcal{N}(z;\mu ,\Sigma)$  can be derived as follows:
$$
\log \mathcal {W N} (z; \mu , \Sigma) = \log \mathcal {N} (u; 0, \Sigma) - (d - 1) \ln \left(\frac {\sinh \| u \| _ {L}}{\| u \| _ {L}}\right), \tag {23}
$$
where  $u$  is defined as  $u = \mathrm{PT}_{\mu \to \mu^{\circ}}\circ \log_{\mu}(z)$ . For detailed derivation, we refer to Appendix A of [25].
A.2 GNN-based parameterization for variational branch length distributions
In this work, we employ a variational branch length distribution  $Q_{\phi}(B_{\tau}|\tau)$  parameterized with a graph neural network (GNN) as described in [40]. In concrete, each of the branch lengths follows an independent lognormal distribution, where its location and scale parameters are predicted with a GNN that takes the tree topology  $\tau$  and the learnable topological features (LTFs) of the topology  $\tau$ , which are computed with a method described in [40]. Below, we summarize an architecture that we use in this study.
Branch length parameterizations Let  $V_{\tau}$  and  $E_{\tau}$  respectively represent the sets of nodes and branch edges for a given unrooted binary tree topology  $\tau$ . The input to the GNN consists of node features represented by LTFs denoted as  $\{h_v^{(0)}\}_{v \in V_{\tau}}$ . These features undergo transformation  $L$  times as follows:
$$
\left\{h _ {v} ^ {(L)} \right\} _ {v \in V _ {\tau}} = \operatorname {G N N} \left(\left\{h _ {v} ^ {(0)} \right\} _ {v \in V _ {\tau}}\right) = g ^ {(L)} \circ \dots \circ g ^ {(1)} \left(\left\{h ^ {(0)} \right\} _ {v \in V _ {\tau}}\right), \tag {24}
$$
where we set  $L = 2$ . The function  $g^{(\ell)}$  represents a GNN layer. For this function, we utilize edge convolutional layers, which will be described in more detail in the following paragraph.
Next, the last node features  $h^{(L)}$  are transformed to output parameters of edge length as follows:
$$
\widetilde {h} _ {v} = \mathrm {M L P} _ {V} \left(h _ {v} ^ {(L)}\right), \quad (\forall v \in V _ {\tau}) \tag {25}
$$
$$
\widetilde {m} _ {(v, u)} = \operatorname {M A X} \left(\widetilde {h} _ {v}, \widetilde {h} _ {u}\right), \quad (\forall (v, u) \in E _ {\tau}) \tag {26}
$$
$$
\mu_ {(v, u)}, \log \sigma_ {(v, u)} = \mathrm {M L P} _ {E} (\widetilde {m} _ {(v, u)}) \quad (\forall (v, u) \in E _ {\tau}) \tag {27}
$$
where  $\mathrm{MLP}_N$ , MAX, and  $\mathrm{MLP}_E$  denotes a multi-layer perceptron for node features with two hidden layers, the element-wise max operation, and a multi-layer perceptron with a hidden layer that outputs the location and scale parameter  $(\mu_e,\sigma_e)$  of the lognormal distributions for each edge  $e\in \mathbb{E}_{\tau}$ . For each of the hidden layers employed in  $\mathrm{MLP}_N$  and  $\mathrm{MLP}_E$ , we set its width to 100 and apply the ELU activation function after the linear transformation of input values.
Edge convolutional layers In a previous study [40], a GNN with edge convolutional layers, referred to as EDGE, demonstrated strong performance when predicting the posterior tree distributions. In EDGE, the function  $g^{(\ell)}$  transforms node features  $\{h_v^{(\ell)}\}_{v\in V_\tau}$  according to the following scheme:
$$
\left\{h _ {v} ^ {(\ell + 1)} \right\} _ {v \in V _ {\tau}} = g ^ {(\ell)} \left(\left\{h _ {v} ^ {(\ell)} \right\} _ {v \in V _ {\tau}}\right), \tag {28}
$$
where  $g^{(\ell)}$  is comprised of the edge convolutional operation with the exponential linear unit (ELU) activation function. Specifically, the transformation with the layer  $g^{(\ell)}$  is computed as follows:
$$
e _ {u \rightarrow v} ^ {(\ell)} = \operatorname {M L P} ^ {(\ell)} \left(h _ {v} ^ {(\ell)} \| h _ {u} ^ {(\ell)} - h _ {v} ^ {(\ell)}\right), \quad \forall u \in N _ {\tau} (v) \tag {29}
$$
$$
h _ {v} ^ {\prime (\ell + 1)} = \underset {u \in N _ {\tau} (v)} {\operatorname {A G G}} ^ {(\ell)} e _ {u \rightarrow v} ^ {(\ell)}, \tag {30}
$$
$$
h _ {v} ^ {(\ell + 1)} = \operatorname {E L U} \left(h _ {v} ^ {\prime (\ell + 1)}\right), \tag {31}
$$
where  $N_{\tau}(v)$  represents a set of neighboring nodes connected to node  $v$  in the tree topology  $\tau$ , i.e. refers to the concatenation operation of elements,  $\mathrm{MLP}^{(\ell)}$  denotes a full connection layer and the exponential linear unit (ELU) activation unit, and  $\mathrm{AGG}^{(\ell)}$  represents an aggregation operation that takes the maximum value of neighboring edge features  $e_{u \to v}^{(\ell)}$ ,  $\forall u \in N(v)$  for each element.
B Variational Lower Bounds and Gradient Estimators
B.1 Variational lower bound
In Proposition 1, we present that the following functional is a lower bound of the marginal log-likelihood  $\ln P(Y)$ .
$$
\begin{array}{l} \mathcal {L} [ Q, R ] := \mathbb {E} _ {Q (z, B _ {\tau})} \left[ \ln F ^ {\prime} (z, B _ {\tau}) \right] = \mathbb {E} _ {Q (z)} \left[ \mathbb {E} _ {Q \left(B _ {\tau} | z\right)} \left[ \ln F (z, B _ {\tau}) \right] - \ln Q (z) \right] (32) \\ \leq \ln P (Y), (33) \\ \end{array}
$$
where  $F$  and  $F^{\prime}$  are respectively defined as follows:
$$
F (z, B _ {\tau}) := \frac {P (Y , B _ {\tau} | \tau (z))}{Q \left(B _ {\tau} | \tau (z)\right)} P (\tau (z)) R (z | \tau (z)), \quad F ^ {\prime} (z, B _ {\tau}) := \frac {F (z , B _ {\tau})}{Q (z)}. \tag {34}
$$
B.2 Gradient estimators for variational lower bound
The gradient of  $\mathcal{L}[Q_{\theta, \phi}, R_{\psi}]$  with respect to  $\theta$  is given by
$$
\begin{array}{l} \nabla_ {\theta} \mathcal {L} = \nabla_ {\theta} \mathbb {E} _ {Q _ {\theta} (z)} \left[ \mathbb {E} _ {Q _ {\phi} \left(B _ {\tau} \mid z\right)} [ \ln F (z, B _ {\tau}) ] - \ln Q _ {\theta} (z) \right] (35) \\ = \mathbb {E} _ {Q _ {\theta} (z)} \left[ (\nabla_ {\theta} \ln Q _ {\theta} (z)) \left(\mathbb {E} _ {Q _ {\phi} (B _ {\tau} | \tau (z))} \left[ \ln \frac {P (Y , B _ {\tau} | \tau (z))}{Q _ {\phi} (B _ {\tau} | \tau (z))} \right] + \ln P (\tau (z)) R _ {\psi} (z | \tau (z))\right) \right] \\ + \nabla_ {\theta} \mathbb {H} [ Q _ {\theta} (z) ], (36) \\ \end{array}
$$
where  $\mathbb{H}$  denotes the differential entropy. We assume that  $Q_{\phi}(B_{\tau}|\tau)$  is reparameterizable as in [40]: namely,  $B_{\tau}$  can be sampled through  $B_{\tau} = h_{\phi}(\epsilon_B,\tau)$ , where  $\epsilon_B\sim p_B(\epsilon_B)$ , where  $p_B(\epsilon)$  and  $h_\phi$  denote a parameter-free base distribution and a differentiable function with  $\phi$ , respectively. Consequently, the gradient of  $\mathcal{L}$  with respect to  $\phi$  is evaluated as follows:
$$
\begin{array}{l} \nabla_ {\phi} \mathcal {L} = \nabla_ {\phi} \mathbb {E} _ {Q _ {\theta} (z)} \mathbb {E} _ {Q _ {\phi} \left(B _ {\tau} | z\right)} [ \ln F (z, B _ {\tau}) ] (37) \\ = \mathbb {E} _ {Q _ {\theta} (z)} \mathbb {E} _ {p _ {B} \left(\epsilon_ {B}\right)} \left[ \nabla_ {\phi} \ln \frac {P \left(Y , B _ {\tau} = h _ {\phi} \left(\epsilon_ {B} , \tau\right) \mid \tau (z)\right)}{Q _ {\phi} \left(B _ {\tau} = h _ {\phi} \left(\epsilon_ {B} , \tau\right) \mid \tau (z)\right)} \right]. (38) \\ \end{array}
$$
Lastly, the gradient of  $\mathcal{L}$  with respect to  $\psi$  can be evaluated with a tractable density model  $R_{\psi}(z|\tau)$  as follows:
$$
\nabla_ {\psi} \mathcal {L} = \nabla_ {\psi} \mathbb {E} _ {Q _ {\theta} (z)} \mathbb {E} _ {Q _ {\phi} \left(B _ {\tau} \mid z\right)} [ \ln F (z, B _ {\tau}) ] = \mathbb {E} _ {Q _ {\theta} (z)} \left[ \nabla_ {\psi} \ln R _ {\psi} (z \mid \tau (z)) \right]. \tag {39}
$$
Given samples  $\epsilon_z\sim p_z$  and  $\epsilon_B\sim p_B$ , we can compute  $z = h_{\theta}(\epsilon_z)$ ,  $\tau (z)$ , and  $B_{\tau} = h_{\phi}(\epsilon_B,\tau (z))$ . Then, the below equations are estimators of gradients  $\nabla_{\theta}\mathcal{L}$ ,  $\nabla_{\phi}\mathcal{L}$ , and  $\nabla_{\psi}\mathcal{L}$ , respectively:
$$
\widehat {g} _ {\theta} = \nabla_ {\theta} \ln Q _ {\theta} (z) \cdot \ln F (z, B _ {\tau}) - \nabla_ {\theta} \ln Q _ {\theta} \left(h _ {\theta} \left(\epsilon_ {z}\right)\right), \tag {40}
$$
$$
\widehat {g} _ {\phi} = \nabla_ {\phi} \ln F (z, h _ {\phi} (\epsilon_ {B}, \tau (z))) = \nabla_ {\phi} \ln \frac {P (Y , h _ {\phi} (\epsilon_ {B} , \tau (z)) | \tau (z))}{Q _ {\phi} (h _ {\phi} (\epsilon_ {B} , \tau (z)) | \tau (z))}, \tag {41}
$$
$$
\widehat {g} _ {\psi} = \nabla_ {\psi} \ln F (z, h _ {\phi} (\epsilon_ {B}, \tau (z))) = \nabla_ {\psi} \ln R _ {\psi} (z | \tau (z)). \tag {42}
$$
The gradients can be computed through the auto-gradient of the following target:
$$
\hat {\mathcal {L}} ^ {\prime} = \ln Q _ {\theta} (z) \cdot \operatorname {d e t a c h} [ f (z, B _ {\tau}) ] + f (z, h _ {\phi} (\epsilon_ {B}, \tau (z))) - \ln Q _ {\theta} (h _ {\theta} (\epsilon_ {z})), \tag {43}
$$
where we denote  $f(z, B_{\tau}) = \ln F(z, B_{\tau})$ , and detach[.] refers to an operation that blocks backpropagation through its argument. For clarity in terms of differentiability with respect to the parameters, we distinguish between expressions  $(z, B_{\tau})$  and  $(h_{\theta}(\epsilon_z), h_{\phi}(\epsilon_B, \tau(z)))$ .
B.3 Multi-sample gradient estimators
Given a  $K$  set of Monte Carlo (MC) samples from  $Q_{\theta, \phi}(z, B_{\tau})$ , i.e.  $\{\epsilon_Z^{(k)}, z^{(k)} = h_\theta(\epsilon_Z^{(k)})\}_{k=1}^K$  and  $\{\epsilon_B^{(k)}, B_\tau^{(k)} = h_\phi(\epsilon_B^{(k)}, \tau(z^{(k)}))\}_{k=1}^K$ , we can simply estimate  $\nabla L_\theta[Q_{\theta, \phi}, R_\psi]$  as follows:
$$
\widehat {g} _ {\theta} ^ {(K)} = \frac {1}{K} \sum_ {k = 1} ^ {K} \left(\nabla_ {\theta} \ln Q _ {\theta} \left(z ^ {(k)}\right) \cdot f \left(z ^ {(k)}, B _ {\tau} ^ {(k)}\right) - \nabla_ {\theta} \ln Q _ {\theta} \left(h _ {\theta} \left(\epsilon_ {z} ^ {(k)}\right)\right)\right). \tag {44}
$$
As a simple extension of equation (43), the gradients are obtained through an auto-gradient computation of the following target:
$$
\widehat {\mathcal {L} ^ {\prime}} ^ {(K)} = \frac {1}{K} \sum_ {k = 1} ^ {K} \left(\ln Q _ {\theta} \left(z ^ {(k)}\right) \cdot \operatorname {d e t a c h} \left[ f \left(z ^ {(k)}, B _ {\tau} ^ {(k)}\right) \right] + f \left(z ^ {(k)}, h _ {\phi} \left(\epsilon_ {B} ^ {(k)}, \tau \left(z ^ {(k)}\right)\right)\right) - \ln Q _ {\theta} \left(h _ {\theta} \left(\epsilon_ {z} ^ {(k)}\right)\right)\right), \tag {45}
$$
B.4 Leave-one-out (LOO) control variates for variance reduction
For the term of  $K$ -sample gradient estimator  $\widehat{g}_{\theta}^{(K)}$  proportional to the score function  $\nabla_{\theta}\ln Q_{\theta}$ , a leave-one-out (LOO) variance reduction is known to be effective [15, 28], which is denoted as follows:
$$
\widehat {g} _ {\mathrm {L O O}, \theta} ^ {(K)} = \frac {1}{K} \sum_ {k = 1} ^ {K} \left[ \nabla_ {\theta} \ln Q _ {\theta} (z ^ {(k)}) \cdot \left(f \left(z ^ {(k)}, B _ {\tau} ^ {(k)}\right) - \overline {{f _ {k}}} \left(z ^ {(\backslash k)}, B _ {\tau} ^ {(\backslash k)}\right)\right) - \nabla_ {\theta} \ln Q _ {\theta} \left(h _ {\theta} \left(\epsilon_ {z} ^ {(k)}\right)\right) \right], \tag {46}
$$
where  $\overline{f_k}$  denotes:
$$
\overline {{f _ {k}}} \left(z ^ {(\backslash k)}, B _ {\tau} ^ {(\backslash k)}\right) := \frac {1}{K - 1} \sum_ {k ^ {\prime} = 1, k ^ {\prime} \neq k} ^ {K} f \left(z ^ {\left(k ^ {\prime}\right)}, B _ {\tau} ^ {\left(k ^ {\prime}\right)}\right). \tag {47}
$$
To employ the LOO gradient estimator for  $\theta$ , the target of auto-gradient computation in equation (45) needs to be adjusted as follows:
$$
\begin{array}{l} \widehat {\mathcal {L}} _ {\mathrm {L O O}} ^ {(K)} = \frac {1}{K} \sum_ {k = 1} ^ {K} \left(\ln Q _ {\theta} (z ^ {(k)}) \cdot \det a c h [ f (z ^ {(k)}, B _ {\tau} ^ {(k)}) - \overline {{f _ {k}}} (z ^ {(\backslash k)}, B _ {\tau} ^ {(\backslash k)}) ] \right. \\ \left. + f \left(z ^ {(k)}, h _ {\phi} \left(\epsilon_ {B} ^ {(k)}, \tau \left(z ^ {(k)}\right)\right)\right) - \ln Q _ {\theta} \left(h _ {\theta} \left(\epsilon_ {z} ^ {(k)}\right)\right)\right), \tag {48} \\ \end{array}
$$
B.5 LAX estimators for adaptive variance reduction
Algorithm 2 GeoPhy algorithm with LAX
1:  $\theta, \phi, \psi, \chi \gets$  Initialize variational parameters
2: while not converged do
3:  $\epsilon_z^{(1:K)},\epsilon_B^{(1:K)}\gets$  Random samples from distributions  $p_z,p_B$
4:  $z^{(1:K)}\gets h_{\theta}(\epsilon_z^{(1:K)})$
5:  $B_{\tau}^{(1;K)} \gets h_{\phi}(\epsilon_B^{(1;K)}, \tau(z^{(1;K)}))$
6:  $\widehat{g}_{\theta}, \widehat{g}_{\phi}, \widehat{g}_{\psi} \gets$  Estimate the gradients  $\nabla_{\theta} \mathcal{L}, \nabla_{\phi} \mathcal{L}, \nabla_{\psi} \mathcal{L}$
7:  $\widehat{g}_{\chi}\gets \nabla_{\chi}\left\langle \widehat{g}_{\theta}^{2}\right\rangle$
8:  $\theta, \phi, \psi, \chi \gets$  Update parameters using an SGA algorithm, given gradients  $\widehat{g}_{\theta}, \widehat{g}_{\phi}, \widehat{g}_{\psi}, \widehat{g}_{\chi}$
9: end while
10: return  $\theta, \phi, \psi$
The LAX estimator [7] is a stochastic gradient estimator based on a surrogate function, which can be adaptively learned to reduce the variance regarding the term  $\nabla_{\theta}\ln Q_{\theta}(z)$ . In our case, the LAX estimator is given as follows:
$$
\widehat {g} _ {\mathrm {L A X}, \theta} := \nabla_ {\theta} \ln Q _ {\theta} (z) \cdot \left(f (z, B _ {\tau}) - s _ {\chi} (z)\right) + \nabla_ {\theta} s _ {\chi} \left(h _ {\theta} \left(\epsilon_ {z}\right)\right). \tag {49}
$$
We summarize the modified procedure of GeoPhy with LAX in Algorithm 2. As we assume  $Q_{\theta}(z)$  is differentiable with respect to  $z$ , we can also use a modified estimator as follows:
$$
\widehat {g} _ {\mathrm {L A X}, \theta} := \nabla_ {\theta} \ln Q _ {\theta} (z) \cdot (f (z, B _ {\tau}) - s _ {\chi} (z)) + \nabla_ {\theta} s _ {\chi} (h _ {\theta} (\epsilon_ {z})) - \nabla_ {\theta} \ln Q _ {\theta} (h _ {\theta} (\epsilon_ {z})). \tag {50}
$$
Since it is favorable to reduce the variance of  $\widehat{g}_{\mathrm{LAX},\theta}$ , we optimize  $\chi$  to minimize the following objective as proposed in [7]:
$$
\left\langle \mathbb {V} _ {Q _ {\theta} (z)} [ \widehat {g} _ {\theta} ] \right\rangle := \frac {1}{n _ {\theta}} \sum_ {i = 1} ^ {n _ {\theta}} \mathbb {V} _ {Q _ {\theta} (z)} [ \widehat {g} _ {\theta_ {i}} ] = \frac {1}{n _ {\theta}} \sum_ {i = 1} ^ {n _ {\theta}} \left(\mathbb {E} _ {Q _ {\theta} (z)} [ \widehat {g} _ {\theta_ {i}} ^ {2} ] - \mathbb {E} _ {Q _ {\theta} (z)} [ \widehat {g} _ {\theta_ {i}} ] ^ {2}\right), \tag {51}
$$
where  $n_{\theta}$  denotes the dimension of  $\theta$ . As the gradient in equation (49) is given as an unbiased estimator of  $\nabla_{\theta} \mathcal{L}$ , which is not dependent on  $\chi$ , we can use the relation  $\nabla_{\chi} \mathbb{E}_{Q_{\theta}(z)}[\widehat{g}_{\mathrm{LAX}, \theta_i}] = 0$ . Therefore, the unbiased estimator of the gradient  $\nabla_{\chi} \left\langle \mathbb{V}_{Q_{\theta}(z)}[\widehat{g}_{\mathrm{LAX}, \theta}]\right\rangle$  is given as follows:
$$
\widehat {g} _ {\chi} = \frac {1}{n _ {\theta}} \sum_ {i = 1} ^ {n _ {\theta}} \nabla_ {\chi} \widehat {g} _ {\mathrm {L A X}, \theta_ {i}} ^ {2}. \tag {52}
$$
As we require the gradient of  $\nabla_{\theta}\mathcal{L}$  with respect to  $\chi$  for the optimization, we use different objectives for auto-gradient computation with respect to  $\theta$  and the other parameters  $\phi$  and  $\psi$  as follows:
$$
\widehat {\mathcal {L}} _ {\mathrm {L A X}, \theta} ^ {\prime} = \ln Q _ {\theta} (z) \cdot (\operatorname {d e t a c h} [ f (z, B _ {\tau}) ] - s _ {\chi} (z)) + s _ {\chi} \left(h _ {\theta} \left(\epsilon_ {z}\right)\right) - \ln Q _ {\theta} \left(h _ {\theta} \left(\epsilon_ {z}\right)\right), \tag {53}
$$
$$
\widehat {\mathcal {L}} _ {\phi , \psi} ^ {\prime} = f (z, h _ {\phi} (\epsilon_ {B}, \tau (z))). \tag {54}
$$
B.6 LAX estimators with multiple MC-samples
For the cases with  $K$  MC-samples, we use LAX estimators by differentiating the following objectives:
$$
\widehat {\mathcal {L}} _ {\mathrm {L A X}, \theta} ^ {(K)} = \frac {1}{K} \sum_ {k = 1} ^ {K} \left(\ln Q _ {\theta} \left(z ^ {(k)}\right) \cdot \left(\det a c h [ f \left(z ^ {(k)}, B _ {\tau} ^ {(k)}\right) ] - s _ {\chi} \left(z ^ {(k)}\right)\right) + s _ {\chi} \left(h _ {\theta} \left(\epsilon_ {z} ^ {(k)}\right)\right) - \ln Q _ {\theta} \left(h _ {\theta} \left(\epsilon_ {z} ^ {(k)}\right)\right)\right), \tag {55}
$$
$$
\widehat {\mathcal {L}} _ {\phi , \psi} ^ {(K)} = \frac {1}{K} \sum_ {k = 1} ^ {K} f \left(z ^ {(k)}, h _ {\phi} \left(\epsilon_ {B} ^ {(k)}, \tau \left(z ^ {(k)}\right)\right)\right). \tag {56}
$$
When we combine LAX estimators with LOO control variates, the target for auto-gradient computation changes to the following:
$$
\begin{array}{l} \hat {\mathcal {L}} _ {\mathrm {L O O} + \mathrm {L A X}, \theta} ^ {(K)} = \frac {1}{K} \sum_ {k = 1} ^ {K} \left(\ln Q _ {\theta} (z ^ {(k)}) \cdot \left(\det a c h [ f (z ^ {(k)}, B _ {\tau} ^ {(k)}) - \overline {{f _ {k}}} (z ^ {(\backslash k)}, B _ {\tau} ^ {(\backslash k)}) ] - s _ {\chi} (z ^ {(k)})\right)\right) \\ \left. + s _ {\chi} \left(h _ {\theta} \left(\epsilon_ {z} ^ {(k)}\right)\right) - \ln Q _ {\theta} \left(h _ {\theta} \left(\epsilon_ {z} ^ {(k)}\right)\right)\right). \tag {57} \\ \end{array}
$$
We note that  $\widehat{\mathcal{L}}_{\phi, \psi}^{(K)}$  is not affected by the introduction of LOO control variates.
B.7 Derivation of importance-weighted evidence lower bound (IW-ELBO)
An importance-weighted evidence lower bound (IW-ELBO) [1], is a tighter lower bound of the log-likelihood  $\ln P(Y)$  than ELBO. For our model, a conventional  $K$ -sample IW-ELBO is given as follows:
$$
\mathcal {L} _ {\mathrm {I W}} ^ {(K)} [ Q ] := \mathbb {E} _ {Q (z ^ {(1)}, B _ {\tau} ^ {(1)}) \dots Q (z ^ {(K)}, B _ {\tau} ^ {(K)})} \left[ \ln \frac {1}{K} \sum_ {k = 1} ^ {K} \frac {P \left(Y , B _ {\tau} ^ {(k)} , \tau \left(z ^ {(k)}\right)\right)}{Q \left(B _ {\tau} ^ {(k)} , \tau \left(z ^ {(k)}\right)\right)} \right]. \tag {58}
$$
The fact that  $\mathcal{L}_{\mathrm{IW}}^{(K)}[Q]$  is the lower bound of  $\ln P(Y)$  is directly followed from Theorem 1 in [1]. However, as our model cannot directly evaluate the mass function  $Q(\tau)$ , we must resort to considering the second lower bound, similar to the case of  $K = 1$  as depicted in Proposition 1. We define the  $K$ -sample tractable IW-ELBO as follows:
$$
\mathcal {L} _ {\mathrm {I W}} ^ {(K)} [ Q, R ] := \mathbb {E} _ {Q (z ^ {(1)}, B _ {\tau} ^ {(1)}) \dots Q (z ^ {(K)}, B _ {\tau} ^ {(K)})} \left[ \ln \frac {1}{K} \sum_ {k = 1} ^ {K} F ^ {\prime} \left(z ^ {(k)}, B _ {\tau} ^ {(k)}\right) \right], \tag {59}
$$
where  $F'$  is defined in equation (34). We will prove in Theorem 1 that  $\mathcal{L}_{\mathrm{IW}}^{(K)}[Q,R]$  serves as a lower bound of the  $\ln P(Y)$ . Although this inequality holds when  $K = 1$ , as shown by  $\ln P(Y) \geq \mathcal{L}[Q] \geq \mathcal{L}[Q,R]$  in Proposition 1, the relationship is less obvious when  $K > 1$ . Before delving into that, we prepare the following proposition.
Proposition 2. Given  $Q(z,\tau)$  as defined in equation 4 and an arbitrary conditional distribution  $R(z|\tau)$ , it follows that
$$
\mathbb {E} _ {R (z | \tau)} [ \mathbb {I} [ \tau = \tau (z) ] ] \leq 1, \tag {60}
$$
where setting  $R(z|\tau) = Q(z|\tau)$  is a sufficient condition for the equality to hold.
Proof. The inequality immediately follows from the definition as follows:
$$
\mathbb {E} _ {R (z | \tau)} [ \mathbb {I} [ \tau = \tau (z) ] ] \leq \mathbb {E} _ {R (z | \tau)} [ 1 ] = 1. \tag {61}
$$
Next, when we set  $R(z|\tau) = Q(z|\tau)$ , the condition for equality is satisfied as follows:
$$
\mathbb {E} _ {Q (z \mid \tau)} [ \mathbb {I} [ \tau = \tau (z) ] ] = \frac {\mathbb {E} _ {Q (z)} [ \mathbb {I} [ \tau = \tau (z) ] ^ {2} ]}{Q (\tau)} = \frac {Q (\tau)}{Q (\tau)} = 1, \tag {62}
$$
where we have used the definition of  $Q(\tau) \coloneqq \mathbb{E}_{Q(z)}[\mathbb{I}[\tau = \tau(z)]]$  from equation (4) and the resulting relation  $Q(z|\tau)Q(\tau) = \mathbb{I}[\tau = \tau(z)]Q(z)$ .
Theorem 1. Given  $Q(z, \tau)$  as defined in equation 4 and an arbitrary conditional distribution  $R(z|\tau)$  that satisfies  $\mathrm{supp}R(z|\tau) \supseteq \mathrm{supp}Q(z|\tau)$ , for any natural number  $K > 1$ , the following relation holds:
$$
\ln P (Y) \geq \mathcal {L} _ {\mathrm {I W}} ^ {(K)} [ Q, R ] \geq \mathcal {L} _ {\mathrm {I W}} ^ {(K - 1)} [ Q, R ]. \tag {63}
$$
Additionally, if  $F'(z, B_{\tau})$  is bounded and  $\forall \tau, \mathbb{E}_{R(z|\tau)}[\mathbb{I}[\tau = \tau(z)]] = 1$ , then  $\mathcal{L}_{\mathrm{IW}}^{(K)}[Q,R]$  approaches  $\ln P(Y)$  as  $K \to \infty$ .
Proof. We first show that for any natural number  $K > M$ ,
$$
\mathcal {L} _ {\mathrm {I W}} ^ {(K)} [ Q, R ] \geq \mathcal {L} _ {\mathrm {I W}} ^ {(M)} [ Q, R ]. \tag {64}
$$
For simplicity, we denote  $Q(z^{(k)}, B_{\tau}^{(k)})$  and  $F'(z^{(k)}, B_{\tau}^{(k)})$  as  $Q_{k}$  and  $F_{k}'$ , respectively, in the following discussion. Let  $U_{M}^{K}$  represent a uniform distribution over a subset with  $M$  distinct indices chosen from the  $K$  indices  $\{1, \ldots, K\}$ . Similar to the approach used in [1], we will utilize the following relationship:
$$
\frac {1}{K} \sum_ {k = 1} ^ {K} F _ {k} ^ {\prime} = \mathbb {E} _ {\left\{i _ {1}, \dots , i _ {M} \right\}} \sim U _ {M} ^ {K} \left[ \frac {1}{M} \sum_ {m = 1} ^ {M} F _ {i _ {m}} ^ {\prime} \right]. \tag {65}
$$
Now, the inequality (64) is derived as follows:
$$
\begin{array}{l} \mathbb {E} _ {Q _ {1} \dots Q _ {K}} \left[ \ln \left(\frac {1}{K} \sum_ {k = 1} ^ {K} F _ {k} ^ {\prime}\right) \right] = \mathbb {E} _ {Q _ {1} \dots Q _ {K}} \left[ \ln \mathbb {E} _ {\{i _ {1}, \dots , i _ {m} \} \sim U _ {M} ^ {K}} \left[ \left(\frac {1}{M} \sum_ {m = 1} ^ {M} F _ {i _ {m}} ^ {\prime}\right) \right] \right] (66) \\ \geq \mathbb {E} _ {Q _ {1} \dots Q _ {K}} \left[ \mathbb {E} _ {\{i _ {1}, \dots , i _ {m} \} \sim U _ {M} ^ {K}} \left[ \ln \left(\frac {1}{M} \sum_ {m = 1} ^ {M} F _ {i _ {m}} ^ {\prime}\right) \right] \right] (67) \\ = \mathbb {E} _ {Q _ {1} \dots Q _ {M}} \left[ \ln \left(\frac {1}{M} \sum_ {m = 1} ^ {M} F _ {m} ^ {\prime}\right) \right], (68) \\ \end{array}
$$
where we have also used Jensen's inequality.
Next, we show that  $\ln P(Y) \geq \mathcal{L}_{\mathrm{IW}}^{(K)}[Q,R]$ . We again use Jensen's inequality as follows:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {I W}} ^ {(K)} [ Q, R ] = \mathbb {E} _ {Q _ {1} \dots Q _ {K}} \left[ \ln \frac {1}{K} \sum_ {k = 1} ^ {K} F _ {k} ^ {\prime} \right] (69) \\ \leq \ln \mathbb {E} _ {Q _ {1} \dots Q _ {K}} \left[ \frac {1}{K} \sum_ {k = 1} ^ {K} F _ {k} ^ {\prime} \right] = \ln \mathbb {E} _ {Q (z, B _ {\tau})} \left[ F ^ {\prime} (z, B _ {\tau}) \right]. (70) \\ \end{array}
$$
The last term is further transformed as follows: