VAE-Inf: A statistically interpretable generative paradigm for imbalanced classification

Abstract: Imbalanced classification remains a pervasive challenge in machine learning, particularly when minority samples are too scarce to provide a robust discriminative boundary. In such extreme scenarios, conventional models often suffer from unstable decision boundaries and a lack of reliable error control. To bridge the gap between generative modeling and discriminative classification, we propose a two-stage framework \textbf{VAE-Inf} that integrates deep representation learning with statistically interpretable hypothesis testing. In the first stage, we adopt a one-class modeling perspective by training a variational autoencoder (VAE) exclusively on majority-class data to capture the underlying reference distribution. The resulting latent posteriors are aggregated via a Wasserstein barycenter to construct a global Gaussian reference model, providing a geometrically principled baseline for the majority class. In the second stage, we transform this generative foundation into a discriminative classifier by fine-tuning the encoder with limited minority samples. This is achieved through a novel distribution-aware loss that enforces probabilistic separation between classes based on variance-normalized projection statistics. For inference, we introduce a projection-based score that admits a natural hypothesis testing interpretation, allowing for a distribution-free calibration procedure. This approach yields exact finite-sample control of the Type-I error (false positive rate) without relying on restrictive parametric assumptions. Extensive experiments on diverse real-world benchmarks demonstrate that our framework achieves competitive performance against other approaches. The codes are available upon request.

• The paper introduces a two-stage framework that combines VAE-based majority reference modeling with distribution-aware fine-tuning to control Type-I error.
• It leverages geometric aggregation and hypothesis testing to achieve high AUC-PR and F1 scores, demonstrating superior performance over traditional deep anomaly detection methods.
• The approach offers finite-sample error control and operational guarantees, positioning it as a robust solution for high-stakes applications such as fraud detection and rare disease screening.

Statistically Interpretable Generative Modeling for Imbalanced Classification: An Analysis of VAE-Inf

Introduction

The paper "VAE-Inf: A statistically interpretable generative paradigm for imbalanced classification" (2604.25334) introduces a novel two-stage framework leveraging deep generative modeling for robust classification under severe class imbalance. Addressing a fundamental challenge in ML where minority class scarcity undermines discriminative boundary formation, the approach combines deep probabilistic representation learning with statistically principled error control. The methodology hinges on the extraction of a reference latent distribution for the majority class via Variational Autoencoding, followed by discriminative refinement incorporating limited minority samples. The theoretical backbone involves hypothesis testing and variance-normalized anomaly scoring, providing finite-sample guarantees on Type-I error via exchangeability-driven calibration.

Figure 1: Overview of the proposed two-stage VAE-Inf architecture for imbalanced classification, detailing majority reference learning, discriminative fine-tuning, and calibrated inference.

Methodological Framework

Stage 1: Majority Reference Modeling via VAE

The initial phase consists of training a VAE exclusively on majority-class data to capture the dominant latent structure. Encoder posteriors, parameterized as diagonal Gaussians, are aggregated using the Wasserstein barycenter, enabling closed-form estimates of global majority mean and variance in latent space. This geometric aggregation ensures statistically coherent modeling of the majority class, independent of minority sample anomalies or outliers.

Stage 2: Distribution-Aware Discriminative Fine-Tuning

With the global reference established, the encoder is fine-tuned using limited minority-class samples. The core innovation is a distribution-aware regularization loss, leveraging projection-based statistics—specifically, variance-normalized squared deviations along random directions. This loss is designed to both tighten majority embeddings within a high-probability latent region and force minority embeddings to fall outside the acceptance region. It exploits the chi-square distribution of projected statistics, permitting selection of a threshold $\alpha$ directly corresponding to desired statistical confidence levels.

Calibration and Hypothesis Testing

Inference proceeds through projection-based scoring of latent representations. Aggregation over multiple random directions macroscopically detects deviations from majority normality. Decision thresholds are calibrated empirically on held-out majority validation scores, yielding finite-sample control of Type-I error with minimal assumptions (exchangeability), thereby distinguishing the approach from parametric or heuristic thresholding paradigms.

Empirical Performance and Numerical Results

Benchmarking across a broad spectrum of real-world tabular, image, and high-dimensional biomedical datasets, VAE-Inf consistently achieves superior AUC-PR and F1 relative to deep anomaly detection baselines, particularly as minority-class proportion diminishes (<1%).

Figure 2: Type-I and Type-II error curves across threshold variations for the TCGA dataset; validation and test results show minimal deviation, confirming the stability of empirical calibration.

Highlighting quantitative results:

• On Credit Card Fraud Detection ($\rho=0.17\%$), VAE-Inf attains AUC-PR of 85.61% and F1 of 83.57%, outperforming DeepSAD and PReNet.
• On TCGA Pan-Cancer (33 one-vs-rest tasks, $20{,}531$ features), VAE-Inf yields average AUC-PR of 95.58% and F1 93.52%, with high stability across rare-cancer experiments.

A pronounced difference emerges with increasing imbalance: projection-based scores retain discriminative power while deep discriminative baselines degrade appreciably, highlighting the robustness of reference modeling plus distribution-aware fine-tuning.

Error Control and Calibration

Type-I and Type-II errors are further characterized as functions of the inference threshold. The proposed calibration mechanism generalizes stably from validation to test splits (mean absolute deviation for Type-I error $\sim$0.0003–0.0031), supporting explicit operational specification of error tolerance.

Hyperparameter Sensitivity and Ablation Analysis

Performance sensitivity to the hyperparameters $\alpha$ (reference margin) and $\beta$ (minority penalty) is empirically evaluated. Optimal discrimination is achieved by balancing majority tolerance and minority separation—too restrictive or too permissive regularization degrades AUC-PR. Ablation reveals that Stage-2 fine-tuning is essential for discriminative improvements: Stage-1 (reference-only) yields low AUC-PR (<5%) and F1 (<10%), both surging after Stage-2 regularization.

Figure 3: Sensitivity of AUC-PR to $\beta$ for various $\alpha$ values on Credit Card data; optimal performance at $(\alpha=16, \beta=2)$.

Practical and Theoretical Implications

The VAE-Inf framework’s statistical interpretability and finite-sample Type-I guarantees position it as a strong candidate for deployment in safety-critical settings characterized by extreme rarity (e.g., rare disease screening, fraud detection, outlier analysis in network traffic). Methodologically, it suggests a promising paradigm where deep generative modeling is systematically integrated with distribution-free inference mechanisms, departing from ad hoc thresholding and asymptotic approximations.

Potential directions for future research include:

• Extension to multi-class imbalance and complex hierarchical anomaly detection.
• Adaptation of projection sets and prior model architectures for domain-specific customization.
• Application in real-time systems requiring fine-grained operational guarantees on false positive rates.

Conclusion

The paper establishes VAE-Inf as a statistically principled, generative discriminative framework for imbalanced classification. By constructing a robust majority-class reference via VAE and exploiting a distribution-aware regularization loss for discriminative refinement, it achieves high sensitivity and specificity under severe minority scarcity. Empirical calibration confers finite-sample error control, enabling transparent inference mechanisms deployable in high-stakes settings. The approach advances the operational utility of deep generative models in imbalanced domains and provides fertile ground for further theoretical and applied research in robust anomaly-sensitive classification.