SlowGuess/NAACL_2021_OCR huggingface dataset:
APO-VAE: Text Generation in Hyperbolic Space
Shuyang Dai $^{1*}$  Zhe Gan $^{2}$  Yu Cheng $^{2}$  Chenyang Tao $^{1}$  Lawrence Carin $^{1}$  Jingjing Liu $^{2}$
$^{1}$ Duke University  $^{2}$ Microsoft Corporation
{shuyang.dai, chenyang.tao, lcarin}  $@$  duke.edu
{zhe.gan, yu.cheng, jingjl} @microsoft.com
Natural language often exhibits inherent hierarchical structure ingrained with complex syntax and semantics. However, most state-of-the-art deep generative models learn embeddings only in Euclidean vector space, without accounting for this structural property of language. We investigate text generation in a hyperbolic latent space to learn continuous hierarchical representations. An Adversarial Poincaré Variational Autoencoder (APo-VAE) is presented, where both the prior and variational posterior of latent variables are defined over a Poincaré ball via wrapped normal distributions. By adopting the primal-dual formulation of Kullback-Leibler divergence, an adversarial learning procedure is introduced to empower robust model training. Extensive experiments in language modeling, unaligned style transfer, and dialog-response generation demonstrate the effectiveness of the proposed APo-VAE model over VAEs in Euclidean latent space, thanks to its superb capabilities in capturing latent language hierarchies in hyperbolic space.
The Variational Autoencoder (VAE) (Kingma and Welling, 2013; Rezende et al., 2014) is a generative model widely applied to language-generation tasks, which propagates latent codes drawn from a simple prior to manifest data samples through a decoder. The generative model is augmented by an inference network, which feeds observed data samples through an encoder to yield a distribution on the corresponding latent codes. Since natural language often manifests a latent hierarchical structure, it is desirable for the latent code in a VAE to reflect such inherent language structure, so that the generated text can be more natural and expressive. An example of language structure is illustrated in Figure 1, where sentences are organized into a tree structure.
Figure 1: Illustration of the latent hierarchy in natural language. Each tree node is a latent code of its corresponding sentence.
The root node corresponds to simple sentences (e.g., "Yes"), while nodes on outer leaves represent sentences with more complex syntactic structure and richer, more specific semantic meaning (e.g., "The food in the restaurant is awesome")<sup>1</sup>.
In existing VAE-based generative models, such structures are not explicitly considered. The latent code often employs a simple Gaussian prior, and the posterior is approximated as a Gaussian with diagonal covariance matrix. Such embeddings assume Euclidean structure, which is inadequate in capturing geometric structure illustrated in Figure 1. While some variants have been proposed to enrich the prior distributions (Xu and Durrett, 2018; Wang et al., 2019a,b; Shi et al., 2019), there is no evidence that structural information in language can be recovered effectively by the model.
Hyperbolic geometry has recently emerged as an effective method for representation learning from data with hierarchical structure (Mathieu et al.,
Another possible way to organize sentences is a hierarchy of topics, e.g., a parent node can be a sentence on "sports", while its children are sentences on "basketball" and "skiing".
*Work was done when the author interned at Microsoft.
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 416-431 June 6-11, 2021. ©2021 Association for Computational Linguistics
2019; Nickel and Kiela, 2017). Informally, hyperbolic space can be considered as a continuous map of trees. For example, a Poincaré disk (a hyperbolic space with two dimensions) can represent any tree with arbitrary low distortion (De Sa et al., 2018; Sarkar, 2011). In Euclidean space, however, it is difficult to learn such structural representation even with infinite dimensions (Linial et al., 1995).
Motivated by these observations, we propose Adversarial Poincaré Variational Autoencoder (APoVAE), a text embedding and generation model based on hyperbolic representations, where the latent code is encouraged to capture the underlying tree-like structure in language. Such latent structure provides more control of the generated sentences, i.e., an increase of sentence complexity and diversity can be achieved along some trajectory from a root to its children. In practice, we define both the prior and the variational posterior of the latent code over a Poincaré ball, via the use of a wrapped normal distribution (Nagano et al., 2019). To obtain more stable model training and learn more flexible representation of the latent code, we exploit the primal-dual formulation of Kullback-Leibler (KL) divergence (Dai et al., 2018) based on the Fenchel duality (Rockafellar et al., 1966), to adversarially optimize the variational bound. Unlike the primal form that relies on Monte Carlo approximation (Mathieu et al., 2019), our dual formulation bypasses the need for tractable posterior likelihoods via the introduction of an auxiliary dual function.
We apply the proposed approach to language modeling, unaligned style transfer and dialog-response generation. For language modeling, in order to enhance the distribution complexity of the prior, we use an additional "variational mixture of posteriors" prior (VampPrior) design (Tomczak and Welling, 2018) for the wrapped normal distribution. Specifically, VampPrior uses a mixture distribution with components from variational posteriors, coupling the parameters of the prior and variational posterior. For unaligned style transfer, we add a sentiment classifier to our model, and disentangle content and sentiment information by using adversarial training (Zhao et al., 2017a). For dialog-response generation, a conditional model variant of APo-VAE is designed to take into account the dialog context.
Experiments also show that the proposed model addresses posterior collapse (Bowman et al., 2016),
a major obstacle preventing efficient learning of a VAE on text data. In posterior collapse, the encoder learns an approximate posterior similar to the prior, and the decoder tends to ignore the latent code for generation. Experiments show that our proposed model can effectively avoid posterior collapse. We hypothesize that this is due to the use of a more informative prior in hyperbolic space that enhances the complexity of the latent representation, which aligns well with previous work (Tomczak and Welling, 2018; Wang et al., 2019a) that advocates a better prior design.
Our main contributions are summarized as follows. (i) We present Adversarial Poincaré Variational Autoencoder (APo-VAE), a novel approach to text embedding and generation based on hyperbolic latent representations. (ii) In addition to the use of a wrapped normal distribution, an adversarial learning procedure and a VampPrior design are incorporated for robust model training. (iii) Experiments on language modeling, unaligned style transfer, and dialog-response generation benchmarks demonstrate the superiority of the proposed approach compared to Euclidean VAEs, as it benefits from capturing informative latent hierarchies in natural language.
2.1 Variational Autoencoder
Let  $\mathbf{X} = \{\pmb{x}_i\}_{i=1}^N$  be a dataset of sentences, where each  $\pmb{x}_i = [x_{i,1}, \dots, x_{i,T_i}]$  is a sequence of tokens of length  $T_i$ . Our goal is to learn  $p_\theta(\pmb{x})$  that best models the observed sentences so that the expected log-likelihood is maximized, i.e.,  $\mathcal{L}(\pmb{\theta}) = \frac{1}{N} \sum_i \log p_\theta(\pmb{x}_i)$ .
The variational autoencoder (VAE) (Kingma and Welling, 2013; Chen et al., 2018b) considers a latent-variable model  $p_{\theta}(\pmb{x}, \pmb{z})$  to represent sentences, with an auxiliary encoder that draws samples of latent code  $\pmb{z}$  from the conditional density  $q_{\phi}(\pmb{z}|\pmb{x})$ , known as the approximate posterior. Given a latent code  $\pmb{z}$ , the decoder samples a sentence from the conditional density  $p_{\theta}(\pmb{x}|\pmb{z}) = \prod_t p(\pmb{x}_t|\pmb{x}_{<t}, \pmb{z})$ , where the "decoding" pass takes an auto-regressive form. Together with prior  $p(\pmb{z})$ , the model is given by the joint  $p_{\theta}(\pmb{x}, \pmb{z}) = p_{\theta}(\pmb{x}|\pmb{z})p(\pmb{z})$ . The VAE leverages the approximate posterior to derive an evidence lower bound (ELBO) to the (intractable) marginal log-likelihood
$\log p_{\pmb{\theta}}(\pmb{x}) = \log \int p_{\pmb{\theta}}(\pmb{x},\pmb{z})\mathrm{d}\pmb{z}$
$$
\mathcal {L} (\boldsymbol {x}; \boldsymbol {\theta}, \phi) = \mathbb {E} _ {\boldsymbol {z} \sim q _ {\phi} (\boldsymbol {z} | \boldsymbol {x})} \left[ \log \frac {p _ {\boldsymbol {\theta}} (\boldsymbol {x} , \boldsymbol {z})}{q _ {\phi} (\boldsymbol {z} | \boldsymbol {x})} \right], \tag {1}
$$
where  $(\theta, \phi)$  are jointly optimized during training, and the gap is given by the decomposition
$$
\log p _ {\boldsymbol {\theta}} (\boldsymbol {x}) = \mathcal {L} (\boldsymbol {x}; \boldsymbol {\theta}, \phi) + \mathbb {D} _ {\mathrm {K L}} \left(p _ {\boldsymbol {\theta}} (\boldsymbol {z} | \boldsymbol {x}) \| q _ {\phi} (\boldsymbol {z} | \boldsymbol {x})\right), \tag {2}
$$
where  $\mathbb{D}_{\mathrm{KL}}$  denotes Kullback-Leibler divergence. Alternatively, the ELBO can be written as:
$$
\begin{array}{l} \mathcal {L} (\boldsymbol {x}; \boldsymbol {\theta}, \phi) = \mathbb {E} _ {\boldsymbol {z} \sim q _ {\phi} (\boldsymbol {z} | \boldsymbol {x})} \left[ \log p _ {\boldsymbol {\theta}} (\boldsymbol {x} | \boldsymbol {z}) \right] \\ - \mathbb {D} _ {\mathrm {K L}} (q _ {\phi} (\boldsymbol {z} | \boldsymbol {x}) \| p (\boldsymbol {z})) \tag {3} \\ \end{array}
$$
where the first conditional likelihood and second KL terms respectively characterize reconstruction and generalization capabilities. Intuitively, a good model is expected to strike a balance between good reconstruction and generalization. In most cases, both the prior and variational posterior are assumed to be Gaussian for computational convenience. However, such over-simplified assumptions may not be ideal for capturing the intrinsic characteristics of data that have unique geometrical structure, such as natural language.
Riemannian manifolds can provide a more powerful and meaningful embedding space for complex data with highly non-Euclidean structure, that cannot be effectively captured in a vectorial form (e.g., social networks, biology and computer graphics). Of particular interest is the hyperbolic space (Ganea et al., 2018), where  $(i)$  the relatively simple geometry allows tractable computations, and  $(ii)$  the exponential growth of distance in finite dimensions naturally embeds rich hierarchical structure in a compact form.
Riemannian Geometry. An  $n$ -dimensional Riemannian manifold  $\mathcal{M}^n$  is a set of points locally similar to a linear space  $\mathbb{R}^n$ . At each point  $\pmb{x}$  of the manifold  $\mathcal{M}^n$ , we can define a real vector space  $\mathcal{T}_x\mathcal{M}^n$  that is tangent to  $\pmb{x}$ , along with an associated metric tensor  $g_{\pmb{x}}(\cdot, \cdot): \mathcal{T}_{\pmb{x}}\mathcal{M}^{n} \times \mathcal{T}_{\pmb{x}}\mathcal{M}^{n} \to \mathbb{R}$  which is an inner product on  $\mathcal{T}_x\mathcal{M}^n$ . Intuitively, a Riemannian manifold behaves like a vector space only in its infinitesimal neighborhood, allowing the generalization of common notation like angle, straight line and distance to a smooth manifold. For each tangent space  $\mathcal{T}_x\mathcal{M}^n$ , there exists a specific
one-to-one map  $\exp_{\pmb{x}}(\pmb{v}): \mathcal{T}_{\pmb{x}}\mathcal{M}^{n} \to \mathcal{M}^{n}$  from an  $\epsilon$ -ball at the origin of  $\mathcal{T}_{\pmb{x}}\mathcal{M}^{n}$  to a neighborhood of  $\pmb{x}$  on  $\mathcal{M}^{n}$ , called the exponential map. We refer to the inverse of an exponential map as the logarithm map, denoted  $\log_{\pmb{x}}(\pmb{y}): \mathcal{M}^{n} \to \mathcal{T}_{\pmb{x}}\mathcal{M}^{n}$ . In addition, a parallel transport  $P_{\pmb{x} \rightarrow \pmb{x}^{\prime}}: \mathcal{T}_{\pmb{x}}\mathcal{M}^{n} \to \mathcal{T}_{\pmb{x}^{\prime}}\mathcal{M}^{n}$  intuitively transports tangent vectors along a "straight" line between  $\pmb{x}$  and  $\pmb{x}^{\prime}$ , so that they remain "parallel." This is the basic machinery that allows us to generalize distributions and computations in the hyperbolic space, as detailed in later sections.
Poincaré Ball Model. Hyperbolic geometry is one type of non-Euclidean geometry with a constant negative curvature. As a classical example of hyperbolic space, an  $n$ -dimensional Poincaré ball, with curvature parameter  $c \geq 0$  (i.e., radius  $\frac{1}{\sqrt{c}}$ ), can be denoted as  $\mathbb{B}_c^n := \left\{z \in \mathbb{R}^n \mid c \| z \|^2 < 1\right\}$  with its metric tensor given by  $g_z^c = \lambda_z^2 g^E$ , where  $\lambda_z = \frac{2}{1 - c \| z \|^2}$  and  $g^E$  denotes the regular Euclidean metric tensor. Intuitively, as  $z$  moves closer to the boundary  $\frac{1}{\sqrt{c}}$ , the hyperbolic distance between  $z$  and a nearby  $z'$  diverges at a rate of  $\frac{1}{1 - c \| z \|^2} \to \infty$ . This implies significant representation capacity, as very dissimilar objects can be encoded on a compact domain. Note that as  $c \to 0$ , the model recovers the Euclidean space  $\mathbb{R}^n$ , i.e., the lack of hierarchy. In comparison, a larger  $c$  implies a stronger hierarchical organization.
Mathematical Operations. We review the closed-form mathematical operations that enable differentiable training for hyperbolic space models, namely the hyperbolic algebra (vector addition) and tangent space computations (exponential/logarithm map and parallel transport). The hyperbolic algebra is formulated under the framework of gyrovector spaces (Ungar, 2008), with the addition of two points  $z, z' \in \mathbb{B}_c^n$  given by the Möbius addition:
$$
\begin{array}{l} z \oplus_ {c} z ^ {\prime} := \tag {4} \\ \frac {(1 + 2 c \langle z , z ^ {\prime} \rangle + c \| z ^ {\prime} \| ^ {2}) z + (1 - c \| z \| ^ {2}) z ^ {\prime}}{1 + 2 c \langle z , z ^ {\prime} \rangle + c ^ {2} \| z \| ^ {2} \| z ^ {\prime} \| ^ {2}}. \\ \end{array}
$$
For any point  $\pmb{\mu} \in \mathbb{B}_c^n$ , the exponential map and the logarithmic map are given for  $\pmb{u} \neq \pmb{0}$  and  $\pmb{y} \neq \pmb{\mu}$
2The fact that APo-VAE outperforms standard VAE evidences the existence of the hierarchical organization in NLP data.
Figure 2: Model framework of the proposed APO-VAE (red is the prior and blue is the posterior).  $\pmb{x} = [x_{1},\dots,x_{T}]$  is text sequential data, and  $s_k = [s_{k,1},\dots,s_{k,T}]$  is the pseudo-input. The posterior (blue) is obtained by (7), and VampPrior (red) is achieved by (12).  $\nu_{\psi}(\pmb {x},\pmb {z})$  is the dual function.
$$
\exp_ {\boldsymbol {\mu}} ^ {c} (\boldsymbol {u}) := \boldsymbol {\mu} \oplus_ {c} \left(\tanh \left(\sqrt {c} \frac {\lambda_ {\boldsymbol {\mu}} ^ {c} \| \boldsymbol {u} \|}{2}\right) \frac {\boldsymbol {u}}{\sqrt {c} \| \boldsymbol {u} \|}\right),
$$
$$
\log_ {\boldsymbol {\mu}} ^ {c} (\boldsymbol {y}) := \frac {2}{\sqrt {c} \lambda_ {\boldsymbol {\mu}} ^ {c}} \tanh  ^ {- 1} \left(\sqrt {c} \| \kappa_ {\boldsymbol {\mu}, \boldsymbol {y}} \|\right) \frac {\kappa_ {\boldsymbol {\mu} , \boldsymbol {y}}}{\| \kappa_ {\boldsymbol {\mu} , \boldsymbol {y}} \|}, \tag {5}
$$
where  $\kappa_{\mu, y} \coloneqq (-\mu) \oplus_c y$ . Note that the Poincaré ball model is geodesically complete in the sense that  $\exp_{\mu}^c$  is well-defined on the full tangent space  $\mathcal{T}_{\mu}\mathbb{B}_c^n$ . The parallel transport map from a vector  $\pmb{v} \in \mathcal{T}_0\mathbb{B}_c^n$  to another tangent space  $\mathcal{T}_{\mu}\mathbb{B}_c^n$  is given by
$$
P _ {\mathbf {0} \rightarrow \boldsymbol {\mu}} ^ {c} (\boldsymbol {v}) = \log_ {\boldsymbol {\mu}} ^ {c} (\boldsymbol {\mu} \oplus_ {c} \exp_ {\mathbf {0}} ^ {c} (\boldsymbol {v})) = \frac {\lambda_ {\mathbf {0}} ^ {c}}{\lambda_ {\boldsymbol {\mu}} ^ {c}} \boldsymbol {v}. \tag {6}
$$
3 Adversarial Poincaré VAE
We first introduce our hyperbolic encoder and decoder, and how to apply reparametrization. We then provide detailed descriptions on model implementation, explaining how the primal-dual form of KL divergence can help stabilize training. Finally, we describe how to adopt VampPrior (Tomczak and Welling, 2018) to enhance performance. A summary of our model scheme is provided in Figure 2.
3.1 Flexible Wrapped Distribution Encoder
We begin by generalizing the standard normal distribution to a Poincaré ball (Ganea et al., 2018). While there are a few competing definitions of the hyperbolic normal, we choose the wrapped normal as our prior and variational posterior, largely due to its flexibility for more expressive generalization. A wrapped normal distribution  $\mathcal{N}_{\mathbb{B}_c^n}(\boldsymbol{\mu}, \boldsymbol{\Sigma})$  is defined as follows: (i) sample vector  $\mathbf{v}$  from  $\mathcal{N}(\mathbf{0}, \boldsymbol{\Sigma})$ , (ii) parallel transport  $\mathbf{v}$  to  $\mathbf{u} := P_{\mathbf{0} \to \boldsymbol{\mu}}^c(\mathbf{v})$ , and (iii) using exponential map to project  $\mathbf{u}$  back to
$z := \exp_{\mu}^{c}(\boldsymbol{u})$ . Putting these together, a latent sample has the following reparametrizable form:
$$
\boldsymbol {z} = \exp_ {\boldsymbol {\mu}} ^ {c} \left(\frac {\lambda_ {\boldsymbol {0}} ^ {c}}{\lambda_ {\boldsymbol {\mu}} ^ {c}} \boldsymbol {v}\right), \boldsymbol {v} \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {\Sigma}). \tag {7}
$$
For approximate posteriors,  $(\mu, \Sigma)$  depends on  $\pmb{x}$ . We further generalize the (restrictive) hyperbolic wrapped normal by acknowledging that under the implicit VAE (Fang et al., 2019) framework, one does not need the approximate posterior  $q_{\phi}(z|x)$  to be analytically tractable. This allows us to replace the tangent space sampling step  $\pmb{v} \sim \mathcal{N}(\mathbf{0}, \pmb{\Sigma})$  in (7) with a more flexible implicit distribution from which we draw samples as  $\pmb{v} \coloneqq G(\pmb{x}, \pmb{\xi}; \phi_1)$  for  $\pmb{\xi} \sim \mathcal{N}(\mathbf{0}, \pmb{I})$ . Note that now  $\pmb{\mu} \coloneqq F(\pmb{x}; \phi_2)$  can be regarded as a deterministic displacement vector that anchors embeddings to the correct semantic neighborhood, allowing the stochastic  $\pmb{v}$  to only focus on modeling the local uncertainty of the semantic embedding. The synergy between the deterministic and stochastic parts enables efficient representation learning relative to existing alternatives. For simplicity, we denote the encoder neural network as EncNet $_{\phi}$ , which contains  $G$  and  $F$ , with parameters  $\phi = \{\phi_1, \phi_2\}$ .
3.2 Poincaré Decoder
To build a geometry-aware decoder for a hyperbolic latent code, we follow Ganea et al. (2018), and use a generalized linear function analogously defined in the hyperbolic space. A Euclidean linear function takes the form  $f(z) = \langle a, z - b \rangle = \mathrm{sign}(\langle a, z - b \rangle) \| a \| d^{E}(z, H_{a,b})$ , where  $a$  is the coefficient,  $b$  is the intercept,  $H_{a,b}$  is a hyperplane passing through  $b$  with  $a$  as the normal direction, and  $d^{E}(z, H)$  is the distance between  $z$  and hyperplane  $H$ . The
counterpart in Poincaré ball analogously writes
$$
f _ {\boldsymbol {a}, \boldsymbol {b}} ^ {c} (\boldsymbol {z}) = \operatorname {s i g n} \left(\left\langle \boldsymbol {a}, \log_ {\boldsymbol {b}} ^ {c} (\boldsymbol {z}) \right\rangle_ {\boldsymbol {b}}\right) \| \boldsymbol {a} \| _ {\boldsymbol {b}} d _ {c} ^ {\mathbb {B}} (\boldsymbol {z}, H _ {\boldsymbol {a}, \boldsymbol {b}} ^ {c}), \tag {8}
$$
where  $H_{\pmb{a},\pmb{b}}^{c} = \{\pmb {z}\in \mathbb{B}_{c}^{n}|\langle \pmb {a},\log_{\pmb{b}}^{c}(\pmb {z})\rangle_{\pmb{b}} = 0\}$  and  $d_c^{\mathbb{B}}(z,H_{a,b}^c) = \frac{1}{\sqrt{c}}\sinh^{-1}\left(\frac{2\sqrt{c}|\langle\kappa_{b,z},a\rangle|}{(1 - c\|\kappa_{b,z}\|^2)\|a\|}\right)$  are the the gyroplane and the distance between  $\textbf{z}$  and the gyroplane, respectively. Specifically, we use the hyperbolic linear function in (8) to extract features from the Poincaré embedding  $\textbf{z}$ . The feature  $f_{a,b}^{c}(z)$  will be the input to the RNN decoder. We denote the combined network of  $f_{a,b}^{c}$  and the RNN decoder as DecNet, where parameters  $\theta$  contain  $\pmb{a}$  and  $\pmb{b}$ .
3.3 Implementing APo-VAE
While it is straightforward to compute the ELBO (3) via Monte Carlo estimates using the explicit wrapped normal density (Mathieu et al., 2019), we empirically observe that:  $(i)$  the normal assumption restricts the expressiveness of the model, and  $(ii)$  the wrapped normal likelihood makes the training unstable. Therefore, we appeal to a primal-dual view of VAE training to overcome such difficulties (Rockafellar et al., 1966; Dai et al., 2018; Tao et al., 2019; Fang et al., 2019). Specifically, the KL term in (3) can be reformulated as:
$$
\mathbb {D} _ {\mathrm {K L}} \left(q _ {\phi} (\boldsymbol {z} | \boldsymbol {x}) \| p (\boldsymbol {z})\right) = \max  _ {\psi} \tag {9}
$$
$$
\left\{\mathbb {E} _ {\boldsymbol {z} \sim q _ {\phi} (\boldsymbol {z} | \boldsymbol {x})} \nu_ {\boldsymbol {\psi}} (\boldsymbol {x}, \boldsymbol {z}) - \mathbb {E} _ {\boldsymbol {z} \sim p (\boldsymbol {z})} \exp \nu_ {\boldsymbol {\psi}} (\boldsymbol {x}, \boldsymbol {z}) \right\},
$$
where  $\nu_{\psi}(\pmb{x},\pmb{z})$  is the (auxiliary) dual function (i.e., a neural network) with parameters  $\psi$ . The primal-dual view of the KL term enhances the approximation ability, while also being tractable computationally. Meanwhile, since the density function in the original KL term in (3) is replaced by the dual function  $\nu_{\psi}(\pmb{x},\pmb{z})$ , we can avoid direct computation with respect to the probability density function of the wrapped normal distribution.
To train our proposed APo-VAE with the primal-dual form of the VAE objective, we follow the training schemes of coupled variational Bayes (CVB) (Dai et al., 2018) and implicit VAE (Fang et al., 2019), which optimize the objective adversarially. Specifically, we update  $\psi$  in the dual function  $\nu_{\psi}(\pmb{x},\pmb{z})$  to maximize:
$$
\begin{array}{l} \mathcal {L} _ {1} = \mathbb {E} _ {\boldsymbol {x} \sim \boldsymbol {X}} \big [ \mathbb {E} _ {z \sim q _ {\phi} (\boldsymbol {z} | \boldsymbol {x})} \nu_ {\psi} (\boldsymbol {x}, \boldsymbol {z}) \\ \left. - \mathbb {E} _ {\boldsymbol {z} \sim p (\boldsymbol {z})} \exp \nu_ {\boldsymbol {\psi}} (\boldsymbol {x}, \boldsymbol {z}) \right], \quad (1 0) \\ \end{array}
$$
Algorithm 1 Training procedure of APo-VAE.
1: Input: Data samples  $\mathbf{X} = \{\pmb{x}_i\}_{i=1}^N$ , Poincaré curvature  $c$ , and number of pseudo-input  $K$ .
2: Initialize  $\theta, \phi, \psi$ , and  $\delta$ .
3: for iter from 1 to max_iter do
4: Sample a mini-batch  $\{\pmb{x}_m\}_{m=1}^M$  from  $\pmb{X}$  of size  $M$ .
5: # Sampling in the Hyperbolic Space.
6: Obtain  $\pmb{\mu}_{m}$  and  $\pmb{v}_{m}$  from EncNet  $\phi (\pmb {x}_m)$
7: Move  $\pmb{v}_m$  to  $\pmb{u}_m = P_{\pmb{0} \to \pmb{\mu}_m}^c(\pmb{v}_m)$  by (6).
8: Map  $\pmb{u}_m$  to  $\pmb{z}_m = \exp_{\pmb{\mu}_m}^c (\pmb{u}_m)$  by (5).
9: # Update the dual function and the pseudo-input.
10: Sample  $\tilde{z}_m$  by (12).
11: Update  $\psi$  and  $\delta$  by gradient ascent on (10)
12: # Update the encoder and decoder networks.
13: Update  $\theta$  and  $\phi$  by gradient ascent on (11).
14: end for
where  $\mathbb{E}_{\boldsymbol{x} \sim \boldsymbol{X}}[\cdot]$  denotes the expectation over empirical distribution on observations. Accordingly, parameters  $\theta$  and  $\phi$  are updated to maximize:
$$
\begin{array}{l} \mathcal {L} _ {2} = \mathbb {E} _ {\boldsymbol {x} \sim \boldsymbol {X}} \mathbb {E} _ {\boldsymbol {z} \sim q _ {\phi} (\boldsymbol {z} | \boldsymbol {x})} [ \log p _ {\boldsymbol {\theta}} (\boldsymbol {x} | \boldsymbol {z}) \\ \left. - \nu_ {\psi} (\boldsymbol {x}, \boldsymbol {z}) \right]. \tag {11} \\ \end{array}
$$