Source-linked AI summary
Spherical Latent Spaces for Stable Variational Autoencoders
Jiacheng Xu, Greg Durrett
TL;DR
Text VAEs can collapse when KL regularization makes the posterior match the prior, causing the latent variable to be ignored. The paper uses vMF distributions on the unit hypersphere with fixed concentration to simplify optimization and avoid collapse. Across language and document modeling, vMF achieves better likelihoods and uses latent representations more effectively, while fixed concentration reduces flexibility and requires tuning.
Problem
Gaussian text VAEs can reach a KL-collapse optimum where the posterior matches the prior and the latent variable is unused.
Method
The paper uses vMF priors and posteriors on the unit hypersphere, fixing concentration κ so the KL term is controlled as a hyperparameter.
Results
vMF models are more robust than Gaussian models, achieve better likelihoods across recurrent language and document modeling, and use latent variables more effectively.
Takeaways & Limitations
vMF representations make more effective use of latent space and are more sensitive to word-order information than Gaussian representations.
Takeaways & Limitations
Fixing κ reduces per-example flexibility and introduces a parameter that must be tuned by the system designer.
Abstract
from arXiv · showhide
A hallmark of variational autoencoders (VAEs) for text processing is their combination of powerful encoder-decoder models, such as LSTMs, with simple latent distributions, typically multivariate Gaussians. These models pose a difficult optimization problem: there is an especially bad local optimum where the variational posterior always equals the prior and the model does not use the latent variable at all, a kind of "collapse" which is encouraged by the KL divergence term of the objective. In this work, we experiment with another choice of latent distribution, namely the von Mises-Fisher (vMF) distribution, which places mass on the surface of the unit hypersphere. With this choice of prior and posterior, the KL divergence term now only depends on the variance of the vMF distribution, giving us the ability to treat it as a fixed hyperparameter. We show that doing so not only averts the KL collapse, but consistently gives better likelihoods than Gaussians across a range of modeling conditions, including recurrent language modeling and bag-of-words document modeling. An analysis of the properties of our vMF representations shows that they learn richer and more nuanced structures in their latent representations than their Gaussian counterparts.
1 Introduction
The paper replaces Gaussian latent distributions with vMF distributions on the unit hypersphere to address KL collapse, and evaluates the approach across language and document modeling.
- Gaussian VAEs can suffer KL collapse, where the posterior matches the prior and the latent structure becomes unused.
- The proposed vMF latent space uses a mean direction and concentration parameter on the unit hypersphere instead of a Gaussian distribution.
- Fixing κ makes the KL divergence depend only on concentration, structurally preventing collapse and simplifying optimization.
- Across recurrent language modeling and bag-of-words document modeling, vMF models are more robust and achieve better held-out likelihoods than Gaussian models.
- The analysis finds that vMF codes capture ordering information more effectively than Gaussian codes.
2 Variational Autoencoders for Text
VAEs train an encoder and generative model through an ELBO, but Gaussian latent spaces can collapse when KL regularization overpowers latent-code usage. The paper addresses this by changing the latent distribution and fixing its concentration.
- Variational Autoencoders: A VAE models data with a continuous latent variable z and optimizes an evidence lower bound combining KL divergence and expected reconstruction error.
- Variational Autoencoders: The encoder qφ(z|x) maps an input sentence x to a dense latent code z that the generative model uses for reconstruction.
- Case Study: NVRNN: In an NVRNN, the decoder conditions on z at every generation step while recurrently modeling the word sequence.
- Posterior Collapse: KL regularization can trap training at qφ(z|x) = pθ(z), making z act as useless noise that the model ignores.
- Posterior Collapse: Annealing the KL weight can remain insufficient with strong decoders, especially when z has little impact on pθ(x|z).
- Posterior Collapse: On Penn Treebank, a 1-layer vMF model achieves NLL 117 and KL 18.6, outperforming comparable Gaussian settings while relying more on the latent variable.
- Proposed Solution: Changing to vMF distributions places representations on the unit hypersphere and permits fixing κ to avert KL collapse across two modeling paradigms.
3 von Mises-Fisher VAE
The vMF VAE places latent representations on a hypersphere and fixes posterior concentration, making KL independent of the mean direction and preventing KL collapse.
- The vMF distribution lies on the (d −1)-dimensional sphere in R^d and is governed by a unit direction vector µ and concentration κ ≥ 0.
- The model uses a uniform vMF prior with κ = 0 and a variational posterior whose mean direction µ is encoded from the input while κ remains fixed.
- Because the KL divergence depends only on κ and not µ, fixing κ makes the KL term constant and renders KL collapse impossible.
- Unlike Gaussian VAEs, vMF optimization allows µ to vary freely because fixed κ removes pressure toward a single prior-centered distribution.
- Higher κ concentrates samples more closely around µ, increasing both cosine similarity and KL cost; maintaining dispersion in higher dimensions requires larger κ.
- vMF sampling uses a change magnitude and a tangent-space unit vector, allowing gradients of the sample with respect to µ.
4 Experiments on Language Modeling
The NVRNN experiments compare vMF and Gaussian VAEs across datasets, decoder settings, and hyperparameters. vMF achieves stronger language-modeling results while using nonzero KL values and offers more robust control of the KL–reconstruction trade-off.
- Experimental settings: The Inputless setting requires the decoder to predict the sequence without the previous word, making latent representations more useful than in the Standard setting.In the Standard setting, the decoder receives the latent code and the previous ground-truth word at each step.
- Results: vMF VAE gives a performance boost over Gaussian VAE across PTB and Yelp in both Standard and Inputless settings.It works without KL-weight annealing and can use nonzero KL values even in the more powerful Standard setting.
- Evaluation: Table 3 reports test NLL and perplexity, with KL divergence shown in parentheses alongside total NLL.The table covers PTB and Yelp under four experimental settings.
- Hyperparameters: Permanent Gaussian KL annealing with weights 0.2 or 0.5 mitigates KL collapse but produces worse overall performance than vMF.The comparison indicates that the issue is not only the numerical KL–NLL trade-off but also how latent representations are structured.
- Hyperparameters: κ = 80 gives the best perplexity for vMF VAE, while reconstruction error remains around 4.5 because of task difficulty and limited LSTM-decoder capacity.Performance is not very sensitive to κ overall, and reasonable values transfer across similar tasks.
5 Experiments on Document Modeling
The document-modeling experiments evaluate vMF and Gaussian latent distributions in NVDM on two news corpora. vMF outperforms Gaussian despite the absence of KL collapse, with performance depending on concentration, latent dimension, and dataset size.
- Model: NVDM represents documents as fixed-size word-presence vectors and uses two-layer MLPs for both the encoder and decoder.The decoder models a multinomial distribution over vocabulary words.
- Evaluation: Table 4 reports test-set perplexities and compares v-NVDM with fDARN and Gaussian-based NVDM.The caption states that v-NVDM outperforms prior models by a substantial margin.
- Results: vMF achieves better document-modeling performance than Gaussian even though NVDM does not exhibit KL collapse.NVDM relies fully on the latent code to predict the word distribution.
- Hyperparameters: Larger latent dimensions require larger κ values and correspondingly larger KL terms to maintain the same dispersion in variational-posterior samples.The relationship is illustrated in Figure 3.
- Dataset effects: For 20NG, smaller latent dimensions perform better because the corpus is much smaller than RCV1.The reported settings use κ = 100, with the remaining configuration continuing in the source passage.
6 What do our VAEs encode?
The probing experiments show that vMF latent codes capture information beyond bag-of-words, including word-order information, and use the latent space more effectively than Gaussian codes.
- Bag-of-words information: The NVRNN-BoW variant gives the decoder an additional bag-of-words representation, enabling tests of whether latent codes encode information beyond word choice.The decoder conditions on the sentence’s average word embedding as a bag-of-words vector.
- Latent-code utilization: The NVRNN-BoW experiment shows better vMF perplexity despite Gaussian KL collapse, while the latent code remains useful for modeling information not supplied by the decoder.The setup makes bag-of-words information independently available to test what else the latent code captures.
- Bag-of-words information: vMF latent codes reconstruct bag-of-words more accurately than bag-of-words reconstructs the codes, indicating that they capture information beyond word choice.The asymmetry in reconstruction accuracy supports richer latent representations.
- Bag-of-words information: When the decoder already receives bag-of-words information, vMF achieves lower cosine similarity than G-VAE for predicting the latent code from bag-of-words: 0.23 versus 0.32.The lower similarity indicates less redundant information in the latent code.
- Word-order information: vMF representations change faster than G-VAE representations when adjacent words are increasingly swapped, indicating greater sensitivity to word order.This pattern appears in both NVRNN and NVRNN-BoW settings.
7 Controlling Variance with κ
Fixing the vMF concentration parameter κ improves optimization stability and offers control over latent variance, although κ must be tuned and fixed values reduce flexibility.
- Fixed κ: Fixing κ makes optimization easier by preventing local optima associated with unused latent variables, but it prevents per-example concentration and adds a tuning parameter.The trade-off is explicitly between optimization stability and model flexibility.
- Latent variance and structure: The Gaussian latent representations change little under adjacent-word swaps, whereas vMF representations are more sensitive to these ordering changes.The comparison concerns cosine similarity between original and swapped-sentence latent codes.
- Choosing κ: A wide range of κ values works well across the evaluated tasks, and reasonable ranges transfer between comparable tasks.The tasks differ in how strongly reconstruction depends on the latent representation.
- Learned κ: Learning κ per instance is possible because the vMF KL divergence is differentiable with respect to κ, but this reintroduces KL collapse.The KL term encourages κ toward low values, undermining latent usage.
- Learned κ: Clipping κ for numerical stability still leads models toward the smallest values and substantially worse performance than fixed-κ models.This result supports fixing the KL-related concentration during optimization.
8 Related Work
Related work applies deep generative models and VAEs across NLP tasks, studies alternatives for optimizing the VAE objective, and explores non-Gaussian priors with different trade-offs.
- Applications of VAEs in NLP: Deep generative models have been applied to NLP tasks including document modeling, language modeling, dialogue generation, text generation, translation, and style transfer.The cited applications span both generation and structured language tasks.
- VAE objective: Prior work addresses VAE objective optimization through information-theoretic analysis and reweighting schemes with theoretical and empirical justification.These approaches target issues arising during optimization of the VAE objective.
- Choices of priors: Alternative VAE priors include piecewise constant and vMF distributions, with prior studies examining multimodality, continuous interpolation, sentence editing, image modeling, and link prediction.The prior choices involve different modeling properties and application settings.
9 Conclusion
The paper concludes that vMF VAEs address text-VAE optimization issues by controlling latent capacity and outperform Gaussian VAEs across settings, while encoding more word-order and nonredundant information.
- Conclusion: vMF VAEs resolve optimization issues by explicitly controlling the balance between decoder capacity and latent-representation utilization.The conclusion attributes this control to the choice of latent distribution.
- Conclusion: Across a range of settings, vMF VAEs achieve better performance than Gaussian VAEs.The conclusion summarizes the paper’s experimental comparison.
- Conclusion: Further analysis finds that vMF VAEs are more sensitive to word order and use the latent code space more effectively.These findings characterize the structure and utilization of the learned representations.