Latent Optimal Paths by Gumbel Propagation for Variational Bayesian Dynamic Programming

In the stochastic approach, every possible path $\mathbf{y}\in\mathcal{Y}$ on $\mathbf{W}$ follows a Gibbs distribution given a DAG $\mathcal{R}$, edge weights $\mathbf{W}$ and temperature parameter $\alpha$ defined by Definition 4.1.

Despite the intractable form of the denominator in Equation 9, we provide the ingredients necessary for approximate Bayesian inference for latent distribution (unobserved) $\mathcal{D}$. In particular, we can efficiently compute the normalized likelihood (Corollary 4.6), sample (Corollary 4.7), and compute the KL divergence within $\mathcal{D}(\mathcal{R},\cdot,\alpha)$ (Lemma 4.10) in linear time (Corollary 4.11).

The Gumbel propagation offers an equivalent formulation of Definition 4.1 that lends itself to dynamic programming by the properties in Equation 7 and Equation 5 as per the following result. Proofs per lemma of this subsection are in Appendix A.

Let the definitions of $\Omega_{\mathbf{y}}$ and $G_{\mathbf{y}}$ extend to all $\mathbf{y}\in\bigcup_{u=1}^{N}\mathcal{Y}(1,u)$, which is the set of all partial paths. We define for each node $v\in\mathcal{V}$ the real-valued random variable

The motivation for constructing $Q_{v}$ is $Q_{v}$ allows us to set up a recursion on the entire DAG $\mathcal{R}$. We now state the first main result in Lemma 4.4.

Then we state an alternative normalized likelihood of a sampled path $\mathbf{y}$ by Corollary 4.7 as Corollary 4.6 according to a transition matrix defined in Lemma 4.5. The transition matrix in Lemma 4.5 can be computed according to the location parameter $\mu$ defined in Lemma 4.4 directly. Proofs per lemma of this subsection are in Appendix B.

Given $\mathcal{D}(\mathcal{R},\mathbf{W},\alpha)$ and $\mathcal{D}(\mathcal{R},\mathbf{W}^{(r)},\alpha)$, where $\mathbf{W}^{(r)}$ are edge weights have different values to $\mathbf{W}$. We give a tractable closed-form of the KL divergence within the distribution family of $\mathcal{D}(\mathcal{R},\cdot,\alpha)$ in Lemma 4.10. Proofs per lemma of this subsection are in Appendix C.

The quantity in the above definition may be computed using two dynamic programming passes, one topologically ordered and the other reverse topologically ordered, by applying the following

Given a DAG $\mathcal{R}$ with edge $\mathcal{E}$ and nodes $\mathcal{V}$, the distribution of the posterior encoder is denoted as

where NN${}_{\mathbf{W}}(\cdot;\phi)$ is a neural network to learn the edge weights $\mathbf{W}$ of the DAG $\mathcal{R}$ and $\alpha$ is a hyper-parameter. The latent optimal path $\mathbf{y}$ with $\{0,1\}$ can be sampled reversely according to Lemma 4.5 and Corollary 4.7. As a result, $\mathbf{y}$ forms a sparse matrix with dimensions identical to those of the weight matrix $\mathbf{W}$. As shown in Figure 1, the sparsity and structure of $\mathbf{y}$ are inherently achieved through its construction using the DAG $\mathcal{R}$¹¹1In the case where the structure of DAGs $\mathcal{R}$ is unknown. Since we are learning the DAG weights $\mathbf{W}$ and zero weights are equivalent to removing an edge, in some sense the BDP algorithm can in principle learn an approximate DAG structure by imposing small edge weight values. However, in this study, we target to verify our method on applications with a clear prior knowledge of structure DAGs..

Denote the distribution of the conditional prior as

We have provided a closed-form KL divergence within the distribution family $\mathcal{D}(\mathcal{R},\cdot,\alpha)$ in Lemma 4.10 which may be convenient for unconditional generation by pre-setting the prior distribution statistics directly or has tractable $\mathbf{W}^{(0)}$.

In most conditional generation tasks, the non-accessible $\mathbf{x}$ during the inference phase in real applications leads to problems on the prior encoder when forming the edge weights $\mathbf{W}^{(0)}$ given information on $\mathbf{c}$ only, especially in the case that $\mathbf{x}$ has varying lengths $t$. To address this issue, we give a flexible solution for inferring feature information of $\mathbf{x}$ given $\mathbf{c}$ to form the edge weights $\mathbf{W}^{(0)}$ in the conditional prior. Inspired by Ma et al. (2019), we make use of a flow-based model ²²2The flow-based conditional prior strategy is proposed to solve the problem of forming edge weights $\mathbf{W}^{(0)}$ in case of $\mathbf{W}^{(0)}$ is intractable if $\mathbf{x}$ is non-accessible with varying lengths $t$. In practical applications, this strategy is not mandatory if the task has a known and fix prior $\mathcal{D}(\mathcal{R},\mathbf{W}^{0},\alpha)$ or tractable $\mathbf{W}^{(0)}$ (i.e., Lemma 4.10). as the conditional prior to infer information about $\mathbf{x}$ condition on $\mathbf{c}$ and further obtain edge weights $\mathbf{W}^{(0)}$.

Assume there exists a series of invertible transformations of random variables $\mathbf{x}$, such that

where $f=f_{1}\circ\cdots\circ f_{K}$ and $v\sim N(0,1)$ ($\theta$ is omitted for brevity), the KL divergence term in Equation 28 can be written as

The backward pass is to infer the KL divergence during training. The forward pass is to infer feature information of $\mathbf{x}$ given $\mathbf{c}$ and form edge weight $\mathbf{W}^{(0)}$ during inference.

Based on the idea of Mohamed et al. (2020), the gradient of the ELBO (28) with respect to $\theta$ is straightforward, however, the gradient with respect to $\phi$ of the reconstruction error part in the ELBO is non-trivial. We make use of the REINFORCE estimator

where $\mathbf{\tilde{y}}$ is an exact sample via Corollary 4.7 from the posterior $q(\cdot|\mathbf{x},\mathbf{c};\phi)$, and we recall that $\log q(\mathbf{\tilde{y}}|\mathbf{x},\mathbf{c};\phi)$ may be computed using the efficient and exact closed-form of Equation 20, and automatically differentiated.

Alternatively, we provide hints of Gumbel softmax trick to avoid using the REINFORCE estimator in Appendix D for interested readers. In this study, we focus on evaluating the closed form of latent sparse optimal paths. Compared to the reparameterization trick discussed in Appendix D, the log-derivative method (i.e., the REINFORCE estimator) has the advantage of being the most straight-forward, is formally unbiased, does not require the temperature parameter, and allows fast evaluation with the closed form expression of sparse paths $\mathbf{y}$ in Equation 20.

We conduct an experiment to demonstrate how BDP in Section 4 finds the optimal paths on toy DAGs, in which the DAG structures and edge values are randomly generated. In Figure 2, we plot approximated density plots of path distributions with 1, 50, 100, and 250 samples by BDP compared with the corresponding ground truth path density plot. As BDP samples increase, the distribution of path samples will eventually converge to the real path distribution.

We then studied how the value of the temperature parameter $\alpha$ affects the path distribution. The larger $\alpha$ is, the sharper the distribution is, leading to an accurate optimal path result with less sample time. Conversely, for too small $\alpha$, the BDP needs more samples to obtain the optimal path. We conduct another experiment to connect this finding with BDP-VAE in Section 6.5.

We apply the BDP-VAE framework with the computational graph of MA to perform an end-to-end TTS model on the RyanSpeech (Zandie et al., 2021) dataset. RyanSpeech contains 11279 audio clips (10 hours) of a professional male voice actor’s speech recorded at 44.1kHz. We randomly split 2000 clips for validation and 9297 clips for training.

The task of TTS involves an unobserved monotonic hard alignment that maps phonemes to time intervals, since the order of the phonemes should be preserved while the duration of each may vary in speech. Thus, TTS models generate speech according to corresponding phonemes in which the duration of each phoneme is discrete, structural, and unobserved (Mehrish et al., 2023). To show the BDP-VAE framework can be easily adapted into down-stream tasks, we redesigned the popular non-end-to-end TTS model FastSpeech2 (Ren et al., 2020) into the BDP-VAE framework, namely BDPVAE-TTS, which can capture unobserved hard monotonic dependencies between phonemes and utterances jointly on both training and inference.

We verify model performance by an objective metric, the Mel cepstral distortion (MCD) (Kubichek, 1993; Chen et al., 2022), between ground truths and synthesized outputs. We record inference speed by real-time factor (RTF) per generated spectrogram frame. We randomly pick 70 sentences from the test set, the numerical results are shown in Table 1. Our method outperforms the baseline on both MCD and RTF and achieves end-to-end training. This shows the success of BDP-VAE in adapting a non-end-to-end model that relies on external DP aligners into an end-to-end pipeline.

We make a comparison with other end-to-end TTS models (Shen et al., 2018a; Kim et al., 2020; Lu et al., 2021) which target capturing phoneme dependencies. Shen et al. (2018a) models the unobserved temporal dependencies by an auto-regressive architecture with attention. Kim et al. (2020) integrate a monotonic alignment search in parallel in a Glow model (Kingma & Dhariwal, 2018) to obtain hard monotonic alignment. Lu et al. (2021) utilizes the latent Gaussian and captures soft monotonic alignment in the decoder by causality-masked self-attention. Among them, BDPVAE-TTS performs discrete monotonic alignment on both training and inference that ensures model consistency for training and inference. BDPVAE-TTS gets a better MCD and RTF than Shen et al. (2018a) and Kim et al. (2020), but gets higher MCD and RTF than Lu et al. (2021).

For the RTF, since BDPVAE-TTS involves a linear time consumption DP algorithm which causes higher inference time than VAENAR-TTS. We have discussed time complexity for capturing phoneme-to-frame dependencies during training for each end-to-end TTS model in this experiment as below. However, the linear time consumption (Corollary 4.11) is the best we can do for solving a DP problem.

For the MCD, even though BDPVAE-TTS obtains a higher MCD than VAENAR-TTS, BDPVAE-TTS captures discrete monotonic aligned paths in its latent space on training and inference phase which ensures model’s train and test consistency and improves the model’s interpretability. VAENAR-TTS compresses the learned feature of spectrograms with conditions into Gaussian latent variables and uses these in its decoder for reconstruction. During the decoding process, these latent variables are used alongside phonemes to reconstruct the outputs. In contrast, the latent variables in BDPVAE-TTS are designed to capture phoneme-duration dependencies. This focus provides the decoder with a nuanced understanding of how the phoneme contributes to the overall speech utterance, which can be interpreted by the inference fundamental frequency (F0) in Figure 3. We provide detailed additional interpretations in Appendix H.

We extend the BDP-VAE with MA for end-to-end SVS on the popcs dataset (Liu et al., 2022). We perform this experiment not to compare with other models as in Section 6.2, but to demonstrate the utility of our method in a related task. In SVS, the longer phoneme duration also provides an opportunity to visualize the monotonically aligned optimal path clearly. The popcs contains 117 Chinese Mandarin pop songs (5 hours) collected from a qualified female vocalist. We randomly split 50 clips for inference and the rest for training. During the inference phase, the conditional inputs are the fundamental frequency and lyrics of the song clips.

Similar to TTS task, SVS task also involves unobserved structural dependencies. BDP-VAE framework could help the SVS task to capture the discrete structural dependencies in parallel, leading to an end to end framework and a better model interpretability.

Two inference results are visualized in Figure 4 where red rectangles indicate that the temporal structure between the generated Mel-spectrogram and ground truth is almost identical. As expected, the decoder of BDP-VAE synthesizes spectrograms according to the extended conditions (i.e., the phonemes) mapping by the dependency of the latent monotonic path from the prior encoder. Therefore, the decoder synthesizes singing voice according to latent path-aligned phoneme conditions.

To verify the behaviour and generalization of the latent optimal paths in BDP-VAE, we obtain latent optimal paths under the computational graph of DTW on the TIMIT speech corpus dataset (Garofolo et al., 1992) which includes manually time-aligned phonetic and word transcriptions. TIMIT contains English speech of 630 speakers which utterances are recorded as 16-bit 16kHz speech waveform files. We randomly split 50 clips for testing and 580 clips for training.

To the best of our knowledge, there are no existing similar works that enable VAEs to obtain hard latent optimal paths given any defined DAG structures. Thus, inspired by the experimental designs in  Mensch & Blondel (2018) and  Van Den Oord et al. (2017), we set a latent distribution $\mathcal{D}(\mathcal{R}_{\text{DTW}},\mathbf{W}=\{w\in\mathbf{W}\sim U(0,1)\},1)$ as the baseline due to lack of direct comparison. The random baseline still follows the DTW constraint with uniform edge weights, which can verify the behaviour of latent space in BPD-VAE.

We take a summation along the spectrogram dimension to obtain phoneme duration and compute the mean absolute error (MAE) between the ground truth and the phoneme duration. We input phoneme tokens and spectrogram lengths as conditions on inference to obtain latent optimal paths of the prior encoder. We repeat the inference 5 times and take the average and standard deviation of MAEs as our metric of evaluation. As in Table 2, our MAE is lower than the baseline indicating that the BDP-VAE captures meaningful information to obtain stochastic optimal paths in the latent space and is not merely guessing randomly. Figure 5 visualizes ground-truth alignments, latent optimal paths from BDP-VAE, and optimal paths from the baseline latent distribution on two audio clips, which clearly shows the latent optimal path from BDP-VAE under DTW gets a close alignment to the GT. This indicates BDP-VAE has the ability to learn informative edge weights $\mathbf{W}$ and capture sparse optimal paths in latent space.

To study the sensitivity of the temperature parameter $\alpha$, we extend the experiment in Section 6.4. Table 3 shows the MAE value on the training and inference phase with different $\alpha$ settings. As discussed in Section 6.1, the $\alpha$ affects the sharpness of the path distribution. When $\alpha$ is smaller, the distribution is more stochastic. As $\alpha$ increases, the distribution becomes sharper. However, when the $\alpha$ value becomes large, the latent path alignment performance decreases, which is consistent with the contribution of temperature in Platt (2000). In BDP-VAE, the model with too large $\alpha$ may not capture enough variety of data, conversely, for too small $\alpha$ the model becomes too spread out and lacks meaningful structure, with the distribution approaching uniform as $\alpha\rightarrow 0$. In real applications, the value of $\alpha$ in BDP-VAE should be located in a reasonable range which could be found by tuning or setting a learnable parameter.