When Do Gaussian Priors Fail? A Latent Geometry Investigation of RL Fine-Tuning
Organizations: Department of Chemical Engineering, Tsinghua University
Abstract
Latent diffusion models almost universally place a Gaussian prior on the VAE latent space. For base generation this default is unproblematic; under reinforcement-learning fine-tuning, however, it carries a structural limitation. The standard VAE's KL term does not merely ignore mode structure - it actively taxes separation and tightness. When data are multimodal, the ELBO optimum compresses distinct modes into an overlapping latent. RL fine-tuning amplifies this overlap: the score network is under-resolved on the saddle region between modes, dynamic-to-variational amplification rises sharply, and the collapse boundary beta_c is pushed far above its variational scale. On the real alanine-dipeptide Boltzmann distribution, a Gaussian-prior VAE collapses at beta_c = 0.224 while a mode-tracked Gaussian-mixture prior collapses at 0.00316, a 71x gap from the same pipeline. We trace this gap to a single geometric mediator, d_min/(2 sigma_0), and verify the causal chain through controlled interventions: moving geometry at fixed prior, and switching prior at matched geometry (full mediation; prior-label coefficient 0.995x, 95% CI [0.79, 1.25]). A three-source decomposition localises the amplification to the learned policy (zero ODE-solver contribution, bounded anchor mismatch); the coverage plateau is an unconditional benefit of mode-separating priors that no amount of RL capacity recovers. The remedy is conditional: a non-Gaussian prior is necessary precisely when the latent geometry would fall into overlap and the task is coverage-sensitive. A machine-checked theory skeleton supports the KL pricing structure - a divergent compactness tax under a Gaussian prior versus a finite assignment fee under a mixture - and the position-versus-plateau distinction delimits what is finite-resource from what requires changing the prior.
Figures & tables
| condition | non-Gaussian prior needed? |
|---|---|
| multimodal data tight latent dimension | yes — Gaussian ELBO packs modes into overlap |
| RL fine-tuning with a single-target reward | yes — RL amplifies whatever overlap exists |
| unimodal data, or a wide latent dimension | no — the Gaussian prior suffices |
| base generation only, no RL fine-tuning | no — overlap does not hurt the base model |
| Quantity | Gaussian prior | GMM prior |
|---|---|---|
| 0.96 | 0.027 | |
| 1.84 | ||
| empirical | 0.224 | |
| 0.12 | 2.24 | |
| 1.61 | 2.06 |
| Quantity | Gaussian | GMM | Ratio |
|---|---|---|---|
| (analytic) | |||
| (empirical) | |||
| dyn / var | — | — |
| Quantity | Gaussian | GMM | Interpretation |
|---|---|---|---|
| (collapse position) | Main Line 1: prior-independent | ||
| Plateau coverage ( ) | Main Line 2: prior matters | ||
| Raw-prior coverage | GMM target captures all modes |
| Seed | plat dec | dyn/var | ||||
|---|---|---|---|---|---|---|
| Gauss | GMM | Gauss | GMM | Gauss | GMM | |
| 0 | 1.50 | 1.16 | 0.63 | 0.70 | 0.83 | 0.67 |
| 1 | 1.83 | 1.17 | 0.58 | 0.73 | 0.88 | 0.82 |
| 2 | 1.55 | 1.90 | 0.53 | 0.77 | 0.67 | 1.18 |
| mean | 1.63 | 1.41 | 0.58 | 0.73 | 0.79 | 0.89 |
| Seed | gap | ||||||
|---|---|---|---|---|---|---|---|
| Gauss | GMM | Gauss | GMM | Gauss | GMM | ratio | |
| 0 | 0.224 | 0.00316 | 0.01 | 0.958 | 0.027 | ||
| 1 | 0.224 | 0.014 | 0.955 | 0.027 | |||
| 2 | 0.014 | 0.01 | 1.14 | 0.030 | |||
| Gaussian prior ( ) | GMM prior ( ) | |||
| hard | soft | hard | soft | |
| 0.333 | 0.560 | 1.000 | 0.847 | |
| 0.667 | 0.652 | 1.000 | 0.929 | |
| 0.667 | 0.717 | 1.000 | 0.935 | |
| 1.000 | 0.850 | 1.000 | 0.934 | |
| 1.000 | 0.892 | 1.000 | 0.937 | |
| Basin | 20 ns frame frac. | 60 ns frame frac. |
|---|---|---|
| C7eq (dominant) | 43.5% | 43.4% |
| (secondary) | 33.8% | 35.2% |
| /C5 (rarest) | 15.7% | 14.4% |
Appendix figures & tables28 assets
Supplementary material from the paper’s appendix.
Appendix
| minimum KDE mass | ||||
|---|---|---|---|---|
| 20 ns ( ) | 3 | 3 | 3 | 3 |
| 60 ns ( ) | 3 | 3 | 3 | 3 |
| peak | KDE mass | frame fraction | ||||
|---|---|---|---|---|---|---|
| Basin | 20 ns | 60 ns | 20 ns | 60 ns | 20 ns | 60 ns |
| C7eq (dominant) | 44.52% | 44.59% | 43.46% | 43.39% | ||
| (secondary) | 38.87% | 40.20% | 33.82% | 35.16% | ||
| /C5 (rarest) | 16.61% | 15.21% | 15.71% | 14.40% | ||
| unassigned ( rad from any peak) | — | — | 7.01% | 7.05% | ||
| MD | seed | prior | cls | plat lat | plat dec | ||||
|---|---|---|---|---|---|---|---|---|---|
| 20 ns | 0 | Gaussian | 0.958 | 1.61 | 0.875 | 0.224 | 1.00 | 1.00 | |
| 20 ns | 0 | GMM | 0.027 | 2.06 | 0.919 | 1.00 | 1.00 | ||
| 20 ns | 1 | Gaussian | 0.955 | 1.56 | 0.875 | 0.224 | 0.014 | 1.00 | 0.67 |
| 20 ns | 1 | GMM | 0.027 | 1.82 | 0.911 | 1.00 | 1.00 | ||
| 20 ns | 2 | Gaussian | 1.140 | 0.99 | 0.898 | 0.014 | 1.00 | 1.00 | |
| 20 ns | 2 | GMM | 0.030 | 1.71 | 0.917 | 1.00 | 1.00 |
| empirical | numeric theory | |||||
|---|---|---|---|---|---|---|
| reverse | forward | fwd/rev | reverse | forward | fwd/rev | |
| 0.1 | 0.122 | — | 0.0493 | 0.0612 | 1.24 | |
| 1 | 0.914 | 0.157 | 0.17 | 0.493 | 0.612 | 1.24 |
| 10 | 11.34 | 31.06 | 2.74 | 4.93 | 6.12 | 1.24 |
| 100 | 85.05 | 299.6 | 3.52 | 49.25 | 61.19 | 1.24 |
| dim | boundary | |||||
|---|---|---|---|---|---|---|
| 1 | empirical reverse | 48 | 1.013 | 0.994 | ||
| 1 | empirical forward | 48 | 0.777 | 0.802 | ||
| 1 | numeric theory | 44 | 1.000 | 0.995 | ||
| 2 | empirical reverse | 18 | 1.016 | 0.992 | ||
| 2 | empirical forward | 18 | 0.679 | 0.545 | ||
| 2 | numeric theory | 18 | 1.000 | 0.999 |
| mode | saddle | valley | MC | sad/mode | |||
|---|---|---|---|---|---|---|---|
| 5 | 0.25 | 10.0 | 0.418 | 251.82 | 145.74 | 0.452 | 603.1 |
| 10 | 0.25 | 10.0 | 1.119 | 143.80 | 61.78 | 1.257 | 128.6 |
| 20 | 0.25 | 10.0 | 0.831 | 225.94 | 130.65 | 0.854 | 271.9 |
| 5 | 0.50 | 5.0 | 0.221 | 15.82 | 11.40 | 0.263 | 71.6 |
| 10 | 0.50 | 5.0 | 0.366 | 8.37 | 8.35 | 0.399 | 22.9 |
| 20 | 0.50 | 5.0 | 0.387 | 13.55 | 7.95 | 0.423 | 35.0 |
| dim | mode | saddle | valley | MC | sad/mode | cover | ||
|---|---|---|---|---|---|---|---|---|
| 2 | 5 | 3 | 0.072 | 0.048 | 0.864 | 0.104 | 0.67 | 1.00 |
| 2 | 5 | 5 | 0.098 | 0.301 | 1.184 | 0.127 | 3.08 | 1.00 |
| 2 | 10 | 3 | 0.107 | 0.055 | 0.672 | 0.118 | 0.51 | 1.00 |
| 2 | 10 | 5 | 0.199 | 0.240 | 0.790 | 0.222 | 1.20 | 1.00 |
| 2 | 20 | 3 | 0.117 | 0.132 | 0.335 | 0.121 | 1.13 | 1.00 |
| 2 | 20 | 5 | 0.165 | 0.358 | 0.582 | 0.191 | 2.18 | 1.00 |
| configuration | cover | |||
|---|---|---|---|---|
| 2D , | 170 | 3087 | 6743 | 1.00 |
| 2D , | 3338 | 186 | 6476 | 1.00 |
| 4D , | 307 | 25 | 9668 | 1.00 |
| 4D , | 79 | 209 | 9712 | 0.85 |
| 5D , | 93 | 5 | 9902 | 0.60 |
| 5D , | 35 | 4 | 9961 | 0.30 |
| configuration | mode | saddle | valley | MC | sad/mode | |
|---|---|---|---|---|---|---|
| 2D , | 0.02 | 0.418 | 251.82 | 145.74 | 0.452 | 603 |
| 0.05 | 0.070 | 123.72 | 127.77 | 0.076 | 1773 | |
| 0.10 | 0.048 | 31.25 | 40.36 | 0.045 | 648 | |
| 5D , (cover 0.15) | 0.02 | 0.114 | 0.314 | 1.670 | 0.147 | 2.8 |
| 0.05 | 0.031 | 0.129 | 1.716 | 0.051 | 4.2 | |
| 0.10 | 0.018 | 0.081 | 1.461 | 0.025 | 4.6 |
| res. | grid/path | |||||||
|---|---|---|---|---|---|---|---|---|
| 5 | 0.25 | 0.20 | 1.23 | 10.15 | 149.1 | 14.7 | 84.0 | 24.8 |
| 10 | 0.25 | 0.40 | 1.07 | 32.15 | 53.8 | 1.7 | 5.0 | 4.5 |
| 20 | 0.25 | 0.25 | 1.00 | 46.52 | 67.9 | 1.5 | 3.7 | 4.9 |
| 5 | 0.50 | 0.40 | 1.10 | 3.94 | 6.1 | 1.5 | 9.8 | 4.0 |
| 10 | 0.50 | 0.20 | 1.08 | 5.04 | 5.8 | 1.1 | 3.5 | 1.7 |
| 20 | 0.50 | 0.35 | 1.01 | 8.74 | 10.4 | 1.2 | 1.5 | 1.6 |
| cell | emp/var | ||||
|---|---|---|---|---|---|
| w0.03 | 0.03 | 2.08 | 3.21 | 1.37 | 0.43 |
| w0.1 | 0.10 | 1.55 | 1.99 | 0.50 | 0.25 |
| w0.3 | 0.30 | 1.14 | 1.17 | 0.85 | 0.73 |
| w1.0 | 1.00 | 1.12 | 1.20 | 0.67 | 0.56 |
| w3.0 | 3.00 | 0.65 | 0.77 | 0.96 | 1.24 |
| w1 h64 | 1.00 | 1.16 | 1.32 | 0.74 | 0.56 |
| prior | placement | emp/var | plateau | |||||
|---|---|---|---|---|---|---|---|---|
| gaussian | — | 1.00 | ||||||
| gmm | 1.00 | |||||||
| vamp | 1.00 | |||||||
| gmvae | – |
| 1 | 2 | 4 | 10 | ||
|---|---|---|---|---|---|
| ( ) | sat | 4.1 | 19.8 | 128 | |
| 0.11 | 0.46 | 6.5 | 14.5 | ||
| ( ) | 1.9 | 18.2 | 78.3 | 494 | |
| 3.5 | 55.9 | ||||
| ( ) | 10.5 | 67.7 | 283 | 1772 | |
| 4.5 | 17.0 | 71.1 | 445 |
| arm | matched Gaussian cell | ||||
|---|---|---|---|---|---|
| width | 1.0 | 0.35 | 2.15 | 0.48 | — (above image edge ) |
| width | 1.0 | 0.50 | 2.02 | 0.67 | ( ) |
| width | 1.0 | 0.70 | 1.93 | 0.93 | |
| width | 1.0 | 1.00 | 1.88 | 1.32 | |
| width | 1.0 | 1.30 | 1.79 | 1.69 | |
| contract | 0.9 | 0.50 | 1.49 | 0.40 | ( ) |