SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Takahiro Mimori
takahiro.mimori@aoni.waseda.jp
Michiaki Hamada
AIST-Waseda CBBD-OIL
Nippon Medical School
mhamada@waseda.jp
Phylogenetic inference, grounded in molecular evolution models, is essential for understanding the evolutionary relationships in biological data. Accounting for the uncertainty of phylogenetic tree variables, which include tree topologies and evolutionary distances on branches, is crucial for accurately inferring species relationships from molecular data and tasks requiring variable marginalization. Variational Bayesian methods are key to developing scalable, practical models; however, it remains challenging to conduct phylogenetic inference without restricting the combinatorially vast number of possible tree topologies. In this work, we introduce a novel, fully differentiable formulation of phylogenetic inference that leverages a unique representation of topological distributions in continuous geometric spaces. Through practical considerations on design spaces and control variates for gradient estimations, our approach, GeoPhy, enables variational inference without limiting the topological candidates. In experiments using real benchmark datasets, GeoPhy significantly outperformed other approximate Bayesian methods that considered whole topologies.
Phylogenetic inference, the reconstruction of tree-structured evolutionary relationships between biological units, such as genes, cells, individuals, and species ranging from viruses to macro-organisms, is a fundamental problem in biology. As the phylogenetic relationships are often indirectly inferred from molecular observations, including DNA, RNA and protein sequences, Bayesian inference has been an essential tool to quantify the uncertainty of phylogeny. However, due to the complex nature of the phylogenetic tree object, which involve both a discrete topology and dependent continuous variables for branch lengths, the default approach for the phylogenetic inference has typically been an Markov-chain Monte Carlo (MCMC) method [29], enhanced with domain-specific techniques such as a mixed strategy for efficient exploration of tree topologies.
As an alternative approach to the conventional MCMCs, a variational Bayesian approach for phylogenetic inference was proposed in [42] (VBPI), which has subsequently been improved in the expressive powers of topology-dependent branch length distributions [39, 40]. Although these methods have presented accurate joint posterior distributions of topology and branch lengths for real datasets, they required reasonable preselection of candidate tree topologies to avoid a combinatorial explosion in the number of weight parameters beforehand. There have also been proposed variational approaches [24, 16] on top of the combinatorial sequential Monte Carlo method (CSMC[34]), which iteratively updated weighted topologies without the need for the preselection steps. However, the fidelity of the joint posterior distributions was still largely behind that of MCMC and VBPI as reported in [16].
In this work, we propose a simple yet effective scheme for parameterizing a binary tree topological distribution with a transformation of continuous distributions. We further formulate a novel differentiable variational Bayesian approach named GeoPhy<sup>1</sup> to approximate the posterior distribution of the tree topology and branch lengths without preselection of candidate topologies. By constructing practical topological models based on Euclidean and hyperbolic spaces, and employing variance reduction techniques for the stochastic gradient estimates of the variational lower bound, we have developed a robust algorithm that enables stable, gradient-based phylogenetic inference. In our experiments using real biological datasets, we demonstrate that GeoPhy significantly outperforms other approaches, all without topological restrictions, in terms of the fidelity of the marginal log-likelihood (MLL) estimates to gold-standard provided with long-run MCMCs.
- We propose a representation of tree topological distributions based on continuous distributions, thereby we construct a fully differentiable method named GeoPhy for variational Bayesian phylogenetic inference without restricting the support of the topologies.
- By considering variance reduction in stochastic gradient estimates and simple construction of topological distributions on Euclidean and hyperbolic spaces, GeoPhy offers unprecedently close marginal log-likelihood estimates to gold standard MCMC runs, outperforming comparable approaches without topological preselections.
2.1 Phylogenetic models
Let  $\tau$  be represent an unrooted binary tree topology with  $N$  leaf nodes (tips), and let  $B_{\tau}$  denote a set of evolutionary distances defined on each of the branches of  $\tau$ . A phylogenetic tree  $(\tau, B_{\tau})$  represents an evolutionary relationship between  $N$  species, which is inferred from molecular data, such as DNA, RNA or protein sequences obtained for the species. Let  $Y = \{Y_{ij} \in \Omega\}_{1 \leq i \leq N, 1 \leq j \leq M}$  be a set of aligned sequences with length  $M$  from the species, where  $Y_{ij}$  denote a character (base) of the  $i$ -th sequence at  $j$ -th site, and is contained in a set of possible bases  $\Omega$ . For DNA sequences,  $\Omega$  represents a set of 4-bit vectors, where each bit represents 'A', 'T', 'G', or 'C'. A likelihood model of the sequences  $P(Y|\tau, B_{\tau})$  is determined based on evolutionary assumptions. In this study, we follow a common practice for method evaluations [42, 40] as follows:  $Y$  is assumed to be generated from a Markov process along the branches of  $\tau$  in a site-independent manner; The base mutations are assumed to follow the Jukes-Cantor model [12]. The log-likelihood  $\ln P(Y|\tau, B_{\tau})$  can be calculated using Felsenstein's pruning algorithm [5], which is also known as the sum-product algorithm, and differentiable with respect to  $B_{\tau}$ .
2.2 Variational Bayesian phylogenetic inference
The variational inference problem for phylogenetic trees, which seeks to approximate the posterior probability  $P(\tau, B_{\tau}|Y)$ , is formulated as follows:
$$
\min  _ {Q} D _ {\mathrm {K L}} \left(Q (\tau) Q \left(B _ {\tau} | \tau\right) \| P (\tau , B _ {\tau} | Y)\right), \tag {1}
$$
where  $D_{\mathrm{KL}}$ ,  $Q(\tau)$  and  $Q(B_{\tau}|\tau)$  denote the Kullback-Leibler divergence, a variational tree topology distribution, and a variational branch length distribution, respectively. The first variational Bayesian phylogenetic inference method (VBPI) was proposed by Zhang and Matsen IV [42], which has been successively improved for the expressiveness of  $Q(B_{\tau}|\tau)$  [39, 40]. For the expression of variational topology mass function  $Q(\tau)$ , they all rely on a subsplit Bayesian network (SBN) [41], which represents tree topology mass function  $Q(\tau)$  as a product of conditional probabilities of splits (i.e., bipartition of the tip nodes) given their parent splits. However, SBN necessitates a preselection of the likely set of tree topologies, which hinders an end-to-end optimization strategy of the distribution over all the topologies and branch lengths.
<sup>1</sup>Our implementation is found at https://github.com/m1m0r1/geophy
Figure 1: The proposed scheme for constructing a variational distribution  $Q(\tau, B_{\tau})$  of a tree topology and branch lengths by using a distribution  $Q(z)$  defined on the continuous space. Left: The marginal distributions of four tip nodes  $Q(z_1), \ldots, Q(z_4)$ . Middle: Coordinates of the tip nodes  $z = \{z_1, \ldots, z_4\}$  sampled from the distribution  $Q(z)$ , and a tree topology  $\tau(z)$  determined from  $z$ . Right: A set of branch lengths  $B_{\tau}$  (red figures) of the tree  $\tau$  is sampled from a topology dependent distribution  $Q(B_{\tau}|\tau)$ .
2.3 Learnable topological features (LTFs)
Zhang [40] introduced a deterministic embedding method for tree topologies, termed learnable topological features (LTFs), which provide unique signatures for each tree topology. Specifically, given a set of nodes  $V$  and edges  $E$  of a tree topology  $\tau$ , a node embedding function  $f_{\mathrm{emb}}: V \to \mathbb{R}^d$  is determined to minimize the Dirichlet energy defined as follows:
$$
\sum_ {(u, v) \in E} \| f _ {\mathrm {e m b}} (u) - f _ {\mathrm {e m b}} (v) \| ^ {2}. \tag {2}
$$
When feature vectors for the tip nodes are determined, Zhang [40] showed that the optimal  $f_{\mathrm{emb}}$  for the interior nodes can be efficiently computed with two traversals of the tree topology. They also utilized  $f_{\mathrm{emb}}$  to parameterize the conditional distribution of branch length given tree topology  $Q(B_{\tau}|\tau)$  by using graph neural networks with the input node features  $\{f_{\mathrm{emb}}(v)|v\in V\}$ .
2.4 Hyperbolic spaces and probability distributions
Hyperbolic spaces are known to be able to embed hierarchical data with less distortion in considerably fewer dimensions than Euclidean spaces [32]. This property has been exploited to develop various data analysis applications such as link predictions for relational data [26], as well as phylogenetic analyses [21, 19]. For representing the uncertainty of embedded data on hyperbolic spaces, the wrapped normal distribution proposed in [25] has appealing characteristics: it is easy to sample from and evaluate the density of the distribution.
The Lorentz model  $\mathbb{H}^d$  is a representation of the  $d$ -dimensional hyperbolic space, which is a submanifold of  $d + 1$  dimensional Euclidean space. Denote  $\mu \in \mathbb{H}^d$  and  $\Sigma \in \mathbb{R}^{d\times d}$  be a location and scale parameters, respectively, a random variable  $z\in \mathbb{H}^d$  that follows a wrapped normal distribution  $\mathcal{W}\mathcal{N}(\mu ,\Sigma)$  is defined as follows:
$$
u \sim \mathcal {N} (0, \Sigma), \quad z = \exp_ {\mu} \circ \mathrm {P T} _ {\mu^ {o} \rightarrow \mu} (u), \tag {3}
$$
where  $\mu^o$ ,  $\exp_{\mu}: T_{\mu} \mathbb{H}^d \to \mathbb{H}^d$ , and  $\mathrm{PT}_{\mu^o \to \mu}: T_{\mu^o} \mathbb{H}^d \to T_{\mu} \mathbb{H}^d$  denote the origin, the exponential map, and the parallel transport defined on the Lorentz model. Also, the probability density  $\mathcal{W}\mathcal{N}(\mu, \Sigma)$  has a closed form and is differentiable with respect to the parameters  $\mu, \Sigma$ . More details and properties are summarized in Appendix A.
3 Proposed methods
3.1 Geometric representations of tree topology ensembles
Considering the typically infeasible task of parameterizing the probability mass function of unrooted tree topologies  $\mathcal{T}$ , which requires  $(2N - 5)!! - 1$  degrees of freedom, we propose an alternative approach. We suggest constructing the mass function  $Q(\tau)$  through a transformation of a certain probability density  $Q(z)$  over a continuous domain  $\mathcal{Z}$ , as follows:
$$
Q (\tau) := \mathbb {E} _ {Q (z)} [ \mathbb {I} [ \tau = \tau (z) ] ], \tag {4}
$$
where  $\tau : \mathcal{Z} \to \mathcal{T}$  denotes a deterministic link function that maps  $N$  coordinates to the corresponding tree topology (Fig. 1). Note that we have overloaded  $\tau$  to represent both a variable and function for notational simplicity. An example of the representation space  $\mathcal{Z}$  is a product of the Euclidean space  $\mathbb{R}^{N \times d}$  or hyperbolic space  $\mathbb{H}^{N \times d}$ , where  $d$  denotes the dimension of each tip's representation coordinate. For the link function, we can use  $\tau(z) = T_{\mathrm{NJ}} \circ D(z)$ , where  $D: \mathcal{Z} \to \mathbb{R}^{N \times N}$  denotes a function that takes  $N$  coordinates and provides a distance matrix between those based on a geometric measure such as the Euclidean or hyperbolic distance.  $T_{\mathrm{NJ}}: \mathbb{R}^{N \times N} \to \mathcal{T}$  denotes a map that takes this distance matrix and generates an unrooted binary tree topology of their phylogeny, determined using the Neighbor-Joining (NJ) algorithm [31]. While the NJ algorithm offers a rooted binary tree topology accompanied by estimated branch lengths, we only use the topology information and remove the root node from it to obtain the unrooted tree topology  $\tau \in \mathcal{T}$ .
3.2 Derivation of variational lower bound
Given a distribution of tip coordinates  $Q(z)$  and an induced tree topology distribution  $Q(\tau)$  according to equation (4), the variational lower bound (1) is evaluated as follows:
$$
\mathcal {L} [ Q ] = \mathbb {E} _ {Q (z)} \left[ \mathbb {E} _ {Q \left(B _ {\tau} \mid \tau (z)\right)} \left[ \ln \frac {P \left(Y , B _ {\tau} \mid \tau (z)\right)}{Q \left(B _ {\tau} \mid \tau (z)\right)} \right] + \ln \frac {P (\tau (z))}{Q (\tau (z))} \right]. \tag {5}
$$
Thanks to the deterministic mapping  $\tau(z)$ , we can obtain an unbiased estimator of  $\mathcal{L}[Q]$  by sampling from  $Q(z)$  without summing over the combinatorial many topologies  $\mathcal{T}$ . However, even when the density  $Q(z)$  is computable, the evaluation of  $\ln Q(\tau)$  remains still infeasible according to the definition (4). We resolve this issue by introducing the second lower bound with respect to a conditional variational distribution  $R(z|\tau)$  as follows:
$$
\mathcal {L} [ Q, R ] = \mathbb {E} _ {Q (z)} \left[ \mathbb {E} _ {Q \left(B _ {\tau} \mid \tau (z)\right)} \left[ \ln \frac {P \left(Y , B _ {\tau} \mid \tau (z)\right)}{Q \left(B _ {\tau} \mid \tau (z)\right)} \right] + \ln \frac {P (\tau (z)) R (z \mid \tau (z))}{Q (z)} \right]. \tag {6}
$$
This formulation is similar in structure to the variational auto-encoders [14], where  $Q(\tau | z)$  and  $R(z|\tau)$  are viewed as a probabilistic decoder and encoder of data  $\tau$ , respectively. However, unlike typical auto-encoders, we design  $Q(\tau | z)$  to be deterministic, and instead adjust  $Q(z)$  to approximate  $Q(\tau)$ .
Proposition 1. Assuming that  $\operatorname{supp} R(z|\tau) \supseteq \operatorname{supp} Q(z|\tau)$ , the inequality  $\ln P(Y) \geq \mathcal{L}[Q] \geq \mathcal{L}[Q, R]$  holds, where  $\operatorname{supp}$  denotes the support of distribution. The first and the second equality holds when  $Q(\tau, B_{\tau}) = P(\tau, B_{\tau}|Y)$  and  $R(z|\tau) = Q(z|\tau)$ , respectively.
Proof. The first variational lower bound of the marginal log-likelihood is given as follows:
$$
\ln P (Y) \geq \ln P (Y) - D _ {\mathrm {K L}} \left[ Q (\tau , B _ {\tau}) \| P (\tau , B _ {\tau} | Y) \right] = \mathbb {E} _ {Q (\tau , B _ {\tau})} \left[ \ln \frac {P (Y , \tau , B _ {\tau})}{Q (\tau , B _ {\tau})} \right] := \mathcal {L} [ Q ], \tag {7}
$$
where the equality condition of the first inequality holds when  $Q(\tau, B_{\tau}) = P(\tau, B_{\tau}|Y)$ . Since we have defined  $Q(\tau)$  in equation (4), we can further transform the lower bound as  $\mathcal{L}[Q]$
$$
\mathcal {L} [ Q ] = \mathbb {E} _ {Q (z)} \left[ \sum_ {\tau \in \mathcal {T}} \mathbb {I} [ \tau = \tau (z) ] \mathbb {E} _ {Q \left(B _ {\tau} | \tau\right)} \left[ \ln \frac {P (Y , B _ {\tau} | \tau)}{Q \left(B _ {\tau} | \tau\right)} + \ln \frac {P (\tau)}{Q (\tau)} \right] \right], \tag {8}
$$
from which equation (5) immediately follows. Hence the inequality  $\ln P(Y)\geq \mathcal{L}[Q]$  and its equality condition  $Q(\tau ,B_{\tau}) = P(\tau ,B_{\tau}|Y)$  have been proven.
Next, the entropy term  $-\mathbb{E}_{Q(z)}[\ln Q(\tau (z))] = -\mathbb{E}_{Q(\tau)}[\ln Q(\tau)]$  in equation (5) can be transformed to derive further lower bound as follows:
$$
\begin{array}{l} - \mathbb {E} _ {Q (\tau)} [ \ln Q (\tau) ] \geq - \mathbb {E} _ {Q (\tau)} [ \ln Q (\tau) + D _ {\mathrm {K L}} (Q (z | \tau) \| R (z | \tau)) ] \\ = - \mathbb {E} _ {Q (\tau) Q (z | \tau)} \left[ \ln \frac {Q (\tau) Q (z | \tau)}{R (z | \tau)} \right] = - \mathbb {E} _ {Q (z)} \left[ \ln \frac {Q (z)}{R (z | \tau (z))} \right], \\ \end{array}
$$
where the last equality is derived by using the relation  $\mathbb{E}_{Q(\tau)Q(z|\tau)}[\cdot] = \mathbb{E}_{Q(z)}[\sum_{\tau} \mathbb{I}[\tau = \tau(z)]\cdot]$  and  $Q(\tau)Q(z|\tau) = Q(z)\mathbb{I}[\tau = \tau(z)]$ . The equality condition of the first inequality holds when  $R(z|\tau) = Q(z|\tau)$ . Hence, the inequality  $\mathcal{L}[Q] \geq \mathcal{L}[Q,R]$  and the equality condition is proven.
Similar to Burda et al. [1], we can also derive a tractable importance-weighted lower-bound of the model evidence (IW-ELBO), which is used for estimating the marginal log-likelihood (MLL),  $\ln P(Y)$ , or an alternative lower-bound objective for maximization. The details and derivations are described in Appendix B.
3.3 Differentiable phylogenetic inference with GeoPhy
Algorithm 1 GeoPhy algorithm
1:  $\theta, \phi, \psi \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:  $\theta, \phi, \psi \gets$  Update parameters using an SGA algorithm, given gradients  $\widehat{g}_{\theta}, \widehat{g}_{\phi}, \widehat{g}_{\psi}$
8: end while
9: return  $\theta, \phi, \psi$
Here, we introduce parameterized distributions,  $Q_{\theta}(z)$ ,  $Q_{\phi}(B_{\tau}|\tau)$ , and  $R_{\psi}(z|\tau)$ , aiming to optimize the variational objective  $\mathcal{L}[Q_{\theta,\phi},R_{\psi}]$  using stochastic gradient ascent (SGA). The framework, which we term GeoPhy, is summarized in Algorithm 1. While aligning with black-box variational inference [27] and its extension to phylogenetic inference [42], our unique lower bound objective enables us to optimize the distribution over entire phylogenetic trees within continuous geometric spaces.
Gradients of lower bound We derive the gradient terms for stochastic optimization of  $\mathcal{L}$ , omitting parameters  $\theta, \phi, \psi$  for notational simplicity. We assume that  $Q_{\theta}(z)$  and  $Q_{\phi}(B_{\tau}|\tau(z))$  are both reparameterizable; namely, we can sample from these distributions as follows:
$$
z = h _ {\theta} \left(\epsilon_ {z}\right), \quad B _ {\tau} = h _ {\phi} \left(\epsilon_ {B}, \tau (z)\right), \tag {9}
$$
where the terms  $\epsilon_z\sim p_z(\epsilon_z)$  and  $\epsilon_B\sim p_B(\epsilon_B)$  represent sampling from parameter-free distributions  $p_z$  and  $p_B$ , respectively, and the functions  $h_\theta$  and  $h_\phi$  are differentiable with respect to  $\theta$  and  $\phi$ , respectively. Although we cannot take the derivative of  $\tau (z)$  with respect to  $z$ , the gradient of  $\mathcal{L}$  with respect to  $\theta$  is evaluated as follows:
$$
\begin{array}{l} \nabla_ {\theta} \mathcal {L} = \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) ], \tag {10} \\ \end{array}
$$
where  $\mathbb{H}$  denotes the differential entropy. The gradient of  $\mathcal{L}$  with respect to  $\phi$  and  $\psi$  are evaluated as follows:
$$
\nabla_ {\phi} \mathcal {L} = \mathbb {E} _ {Q _ {\theta} (z)} \mathbb {E} _ {p _ {B} (\epsilon)} \left[ \nabla_ {\phi} \ln \frac {P \left(Y , B _ {\tau} = h _ {\phi} (\epsilon , \tau) \mid \tau (z)\right)}{Q _ {\phi} \left(B _ {\tau} = h _ {\phi} (\epsilon , \tau) \mid \tau (z)\right)} \right], \quad \nabla_ {\psi} \mathcal {L} = \mathbb {E} _ {Q _ {\theta} (z)} \left[ \nabla_ {\psi} \ln R _ {\psi} (z \mid \tau (z)) \right], \tag {11}
$$
where we assume a tractable density model for  $R_{\psi}(z|\tau)$ .
Gradient estimators and variance reduction From equation (10), an unbiased estimator of  $\nabla_{\theta}\mathcal{L}$  using  $K$  -sample Monte Carlo samples, can be derived 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 {12}
$$
where we denote  $\epsilon_z^{(k)}\sim p_zz^{(k)} = h_\theta (\epsilon_z^{(k)})$ $\epsilon_B^{(k)}\sim p_B,B_{\tau}^{(k)} = h_{\phi}(\epsilon_B^{(k)},\tau (z^{(k)}))$  , and
$$
f (z, B _ {\tau}) := \ln \frac {P (Y , B _ {\tau} | \tau (z))}{Q _ {\phi} (B _ {\tau} | \tau (z))} + \ln P (\tau (z)) R _ {\psi} (z | \tau (z)). \tag {13}
$$
Note that we explicitly distinguish  $z$  and  $h_{\theta}(\epsilon_z)$  to indicate the target of differentiation with respect to  $\theta$ .
In practice, it is known to be crucial to reducing the variance of the gradient estimators proportional to the score function  $\nabla_{\theta}\ln Q_{\theta}$  to make optimization feasible. This issue is addressed by introducing a control variate  $c(z)$ , which has a zero expectation  $\mathbb{E}_z[c(z)] = 0$  and works to reduce the variance of the term by subtracting it from the gradient estimator. For the case of  $K > 1$ , it is known to be effective to use simple Leave-one-out (LOO) control variates [15, 28] as follows:
$$
\widehat {g} _ {\theta , \mathrm {L O O}} ^ {(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 {14}
$$
where we defined  $\overline{f_k} (z^{(\backslash k)},B_\tau^{(\backslash k)})\coloneqq \frac{1}{K - 1}\sum_{k' = 1,k'\neq k}^{K}f(z^{(k)},B_\tau^{(k)})$ , which is a constant with respect to the  $k$ -the sample.
We can also adopt a gradient estimator called LAX [7], which uses a learnable surrogate function  $s_{\chi}(z)$  differentiable in  $z$  to form control variates as follows:
$$
\widehat {g} _ {\theta , \mathrm {L A X}} = \left(\nabla_ {\theta} \ln Q _ {\theta} (z)\right) \left(f (z, B _ {\tau}) - s _ {\chi} (z)\right) + \nabla_ {\theta} s _ {\chi} \left(h _ {\theta} \left(\epsilon_ {z}\right)\right) - \nabla_ {\theta} \ln Q _ {\theta} \left(h _ {\theta} \left(\epsilon_ {z}\right)\right), \tag {15}
$$
where we write the case with  $K = 1$  for simplicity. To optimize the surrogate function  $s_{\chi}$ , we use the estimator  $\widehat{g}_{\chi} = \nabla_{\chi}\left\langle \widehat{g}_{\theta}^{2}\right\rangle^{2}$ , where  $\langle \cdot \rangle$  denotes the average of the vector elements. We summarize the procedure with LAX estimator in Algorithm 2.
We also consider other options for the gradient estimators, such as the combinations of the LOO and LAX estimators, the gradient of IW-ELBO, and its variance-reduced estimator called VIMCO [22] in our experiments. For the remaining gradients terms  $\nabla_{\phi}\mathcal{L}$  and  $\nabla_{\psi}\mathcal{L}$ , we can simply employ reparameterized estimators as follows:
$$
\widehat {g} _ {\phi} = \nabla_ {\phi} \ln \frac {P (Y , B _ {\tau} = h _ {\phi} (\epsilon_ {B} , \tau) | \tau (z))}{Q _ {\phi} \left(B _ {\tau} = h _ {\phi} (\epsilon_ {B} , \tau) | \tau (z)\right)}, \quad \widehat {g} _ {\psi} = \nabla_ {\psi} \ln R _ {\psi} (z | \tau (z)), \tag {16}
$$
where we denote  $\epsilon_B\sim p_B$  and  $z\sim Q_{\theta}$ . More details of the gradient estimators are summarized in Appendix B.