RAVE ramps the KL regularization weight over a warmup schedule to avoid posterior collapse

The BetaWarmupCallback ramps the KL weight (beta_factor) from an initial value to a target value over a number of training steps. This mirrors the beta-VAE technique: starting with low beta keeps the latent space unregularized early, letting the model first minimize reconstruction error. As beta increases, the KL term forces the latent distribution toward N(0,1), producing a smooth, traversable space. Warmup can follow a linear or log-linear schedule. Starting beta too high causes the model to collapse the latent space prematurely (posterior collapse), producing poor reconstruction; leaving it too low yields a structured but potentially non-smooth latent space.

Examples

The callback computes the current beta and assigns pl_module.beta_factor = math.exp(beta) on a log-linear ramp over warmup_len steps.

Assessment

Why does starting with a high KL weight cause posterior collapse in a VAE? What symptom would you see in the RAVE latent PCA fidelity curves?