Heavy-Tailed Class-Conditional Priors for Long-Tailed Generative Modeling

1. Problem Statement

The paper addresses the critical issue of class imbalance (long-tailed distributions) in generative modeling, specifically within Variational Autoencoders (VAEs).

• The Core Issue: Standard VAEs trained on imbalanced data tend to allocate latent probability mass proportionally to class frequency. This results in a geometric bias where majority classes occupy larger regions of the latent space, while minority (tail) classes are confined to smaller regions with reduced representational capacity.
• Limitations of Existing Solutions:
  • Gaussian Priors: Standard VAEs use isotropic Gaussian priors, which fail to capture heavy-tailed data structures and exacerbate mode collapse for rare classes.
  • Global Heavy-Tailed Priors (e.g., t3VAE): While models like t3VAE introduce Student's $t$-distributions to handle outliers, they still rely on a single global prior. Consequently, the latent mass remains aligned with empirical class frequencies, failing to fully correct the geometric bias for minority classes.
  • GANs and Diffusion: While alternatives like GANs and Diffusion models exist, they suffer from training instability or require complex re-weighting strategies that do not explicitly decouple representational capacity from empirical frequency in the latent geometry.

2. Methodology: C-t3VAE

The authors propose Conditional-t3VAE (C-t3VAE), a generative model that combines the robustness of heavy-tailed distributions with a class-conditional structure to enforce uniform prior mass across classes.

A. Model Architecture

• Class-Conditional Joint Prior: Unlike t3VAE, which uses a global prior $p(z)$, C-t3VAE defines a joint prior over the latent variable $z$ and output $x$ conditioned on the class label $y$:
  $$p_\theta(x, z|y) \propto \left[ 1 + \frac{(z - \mu_y)^\top \Sigma_y^{-1} (z - \mu_y) + (x - \mu_\theta(z))^\top \Sigma_x^{-1} (x - \mu_\theta(z))}{\nu} \right]^{-\frac{\nu+m+n}{2}}$$
  Here, $\mu_y$ represents a learnable mean vector for each class in the latent space.
• Heavy-Tailed Latent Space: The latent prior $p(z|y)$ is a multivariate Student's $t$-distribution. This allows the model to capture intra-class variations (heavy tails) while the class-conditional structure ensures each class has its own distinct region in the latent space.
• Posterior: The approximate posterior $q_\phi(z|x)$ is also modeled as a Student's $t$-distribution to maintain consistency with the heavy-tailed assumption.

B. Training Objective

• $\gamma$-Power Divergence: Since the Kullback-Leibler (KL) divergence between Student's $t$-distributions lacks a closed-form solution, the authors utilize the $\gamma$-power divergence ($D_\gamma$).
• Closed-Form Loss: They derive a closed-form objective function for the class-conditional setting. The loss function $L(\gamma, y)$ for a specific class $y$ includes:
  1. Reconstruction Term: Minimizes input reconstruction error.
  2. Alignment Term: Encourages the inferred posterior mean to align with the class-specific prior mean $\mu_y$.
  3. Regularization: Includes trace and determinant terms to regularize the posterior covariance, preventing singularity and encouraging compactness.
• Uniform Weighting: The total loss is the sum of class-wise losses with uniform weighting, explicitly removing the scaling effect of empirical class frequency during training.

C. Sampling Strategy

To generate samples, the authors propose an equal-weight latent mixture:
$$p^\star_\nu(z) = \sum_{y=1}^K \alpha_y \cdot t_m(\mu_y, \tau^2 I, \nu + n), \quad \text{where } \alpha_y = \frac{1}{K}$$

• Equal Weights ($\alpha_y$): By setting $\alpha_y = 1/K$, the model ensures that every class is sampled with equal probability, regardless of its frequency in the training set. This directly counteracts the frequency-aligned bias.
• Theoretical Variance ($\tau$): The authors derive a closed-form expression for the sampling variance $\tau^2$ based on the $\gamma$-divergence geometry, providing a principled initialization rather than a heuristic.

3. Key Contributions

1. C-t3VAE Model: Introduction of a class-conditional VAE using Student's $t$-distributions for both the latent prior and the decoder, specifically designed to mitigate geometric bias in long-tailed settings.
2. Theoretical Derivation: Derivation of a closed-form training objective using $\gamma$-power divergence and an analytically derived optimal sampling variance $\tau^2$.
3. Equal-Weight Sampling: A generative sampling scheme that decouples representational capacity from empirical frequency, ensuring uniform prior mass across all classes.
4. Imbalance Threshold Discovery: Empirical identification of a critical imbalance ratio threshold ($\rho \approx 5$). Below this threshold, Gaussian-based models remain competitive; above it, heavy-tailed class-conditional priors become essential.

4. Experimental Results

The model was evaluated on SVHN-LT, CIFAR100-LT, and CelebA under varying degrees of class imbalance ($\rho \in \{1, 10, 50, 100\}$).

• FID Performance:
  • C-t3VAE consistently outperformed Gaussian-based VAEs (VAE, C-VAE) and the global heavy-tailed t3VAE in severely imbalanced settings ($\rho \ge 10$).
  • On SVHN-LT ($\rho=100$), C-t3VAE achieved an FID of 44.02 (with $\beta$-optimization) compared to 47.11 for $\beta$-VAE and 51.62 for $\beta$-t3VAE.
  • On CelebA (Mustache attribute, $\rho=25$), C-t3VAE reduced FID by up to 15 points compared to C-VAE.
• Per-Class Metrics (Precision, Recall, F1):
  • Recall: C-t3VAE significantly improved Recall for tail classes, indicating better mode coverage and reduced mode collapse.
  • Precision: While Precision was slightly lower than Gaussian models in some cases (due to the trade-off between coverage and sharpness), the F1 score was consistently higher for tail classes.
• Threshold Analysis:
  • For mild imbalance ($\rho < 5$), Gaussian priors (C-VAE) performed comparably or slightly better.
  • For $\rho \ge 5$, the performance gap widened significantly in favor of C-t3VAE, confirming the necessity of heavy-tailed, class-conditional priors for long-tail generation.

5. Significance and Conclusion

This work provides a geometric solution to the class imbalance problem in generative modeling. By shifting from a global prior to a class-conditional heavy-tailed prior, the authors successfully decouple the model's representational capacity from the empirical data distribution.

• Theoretical Insight: The paper demonstrates that simply using heavy tails is insufficient; the allocation of prior mass must also be controlled to prevent majority-class dominance.
• Practical Impact: C-t3VAE offers a stable, single-step sampling alternative to Diffusion models and GANs for long-tailed data, achieving superior generation quality for rare classes without the training instability associated with adversarial methods.
• Future Directions: The authors suggest extending this framework to multi-label settings and developing adaptive methods for the sampling variance $\tau$ to further reduce the need for dataset-specific tuning.