Implicit and explicit reparameterizations of the Folded Normal and Rice distributions
Authors: Luis Aldama, Kevin Dalton, Minhuan Li, Doeke Hekstra
Organizations: The Harvard Biophysics Graduate Program, Harvard University, Cambridge, Massachusetts, United States · Department of Molecular and Cellular Biology, and John A. Paulson School of Engineering and Applied Sciences, Harvard University, Cambridge, Massachusetts, United States · LCLS Data Systems, SLAC National Accelerator Laboratory, Menlo Park, CA 94025. · Center for Computational Mathematics, Flatiron Institute, New York, NY, USA
Variational inference is a flexible tool for modeling experimental data and has proven useful in crystallography. Pathwise gradient estimators provide low-variance gradients for reparameterized variational inference, but are not readily available for some distributions central to crystallography. This work derives and implements implicit and explicit reparameterized gradient estimators for the folded Normal and Rice distributions in \prog{PyTorch}, and compares their variance analytically and empirically. In each case the implicit estimator achieves lower variance for the location derivative, particularly at small location-to-scale ratios, while the scale-derivative variances are nearly equal. For both distributions, implicit estimator is proved to be the conditional expectation of the explicit estimator given the sampled magnitude, so its variance is never larger. As a practical application, French-Wilson scaling is formulated as a variational inference problem, using the folded Normal and Rice distributions as variational families for centric and acentric reflections, respectively. Both estimators recover the exact posterior to better than 1% on the amplitude, and converge at the same rate. The implementations are available in the open-source package \prog{rs-distributions}.
Figures & tables
Figure 1: Theoretical and empirical variance of the FoldedNormal pathwise estimators as a function of Monte Carlo sample size N , at μ=1,σ=10 (small ∣μ∣/σ ). Left, the derivative with respect to μ ; right, the derivative with respect to σ . The two empirical curves match their respective theoretical predictions: \Var[G^μE]=4σ2/N for the explicit estimator and equation 20 for the implicit estimator (left), where the expected conditional variance is the analytical gap from Theorem 1 . For the σ derivative (right) both empirical curves coincide with the theoretical prediction of equation 13 . At this parameter setting, the implicit μ derivative estimator achieves variance roughly 50-fold smaller than the explicit estimator.
Figure 2: Theoretical (top) and empirical (bottom) variance of the explicit ( G^μE ) and implicit ( G^μI ) FoldedNormal derivative estimators of ∂μ\E[Y2] as a function of μ and σ . The colorbar denotes the variance value for each (μ,σ) pair. The empirical variance was computed across 10,000 simulations using 100 Monte Carlo samples each. The theoretical variance was calculated using equations 13 and 20 , for the explicit and implicit estimators, respectively.
Figure 3: Variance ratio \Var[G^I]/\Var[G^E] of the implicit FoldedNormal estimator relative to the explicit estimator, as a function of the location-to-scale ratio ∣μ∣/σ , for the test function f(Y)=Y2 . Because the explicit variance is the upper bound, the ratio lies in [0,1] and is independent of the number of Monte Carlo samples. The location-derivative ratio falls to 0 at small ∣μ∣/σ , while the scale-derivative ratio dips only to 0.90 near ∣μ∣/σ≈0.8 .
Figure 4: Theoretical and empirical variance of the \Rice pathwise estimators as a function of Monte Carlo sample size N , at ν=1,σ=10 (small ν/σ ). Left, the derivative with respect to ν ; right, the derivative with respect to σ . The two empirical curves match their respective theoretical predictions: \Var[G^νE]=4σ2/N for the explicit estimator and equation 34 for the implicit estimator (left), where the expected conditional variance is the analytical gap from Theorem 2 ; for the σ derivative (right) the theoretical curves coincide with the empirical curves. For the ν derivative at this parameter setting, the implicit estimator achieves variance roughly 100-fold smaller than the explicit estimator.
Figure 5: Theoretical (top) and empirical (bottom) variance of the explicit ( G^νE ) and implicit ( G^νI ) \Rice derivative estimators of ∂ν\E[R2] as a function of ν and σ . The colorbar denotes the variance value for each (ν,σ) pair. The empirical variance was computed across 10,000 simulations using 100 Monte Carlo samples each. The theoretical variance was calculated using equations 26 and 34 , for the explicit and implicit estimators, respectively.
Figure 6: Variance ratio \Var[G^I]/\Var[G^E] of the implicit \Rice estimator relative to the explicit estimator, as a function of the location-to-scale ratio ν/σ , for the test function f(R)=R2 . As in the FoldedNormal case, the ratio lies in [0,1] and is independent of the number of Monte Carlo samples. The location-derivative ratio falls to 0 at small ν/σ , while the scale-derivative ratio dips only to 0.90 near ν/σ≈1.2 .
Figure 7: Recovery of the reference posterior, optimized with Adam, as a function of the number of Monte Carlo samples N used per gradient. Panels show the amplitude F and the posterior width σF , plotted separately for centric and acentric reflections. Each point is a median over ten replicates, and the bands span the interquartile range (too narrow to visualize). The same comparison for all five optimization algorithms is given in Fig. S3.
Figure 8: Difference in amplitude mean absolute percentage error (MAPE) against the reference posterior (implicit minus explicit), as a function of number of Monte Carlo samples used at each training step. Left, the centric reflections at or below zero intensity, where the two gradient variances differ; right, intensities with Iobs>0 , where the implicit and explicit gradient variances are approximately the same. Positive values favor the explicit estimator at this number of optimization steps. Points are means over the ten per-replicate paired differences, and the bands are one standard error of the mean.
F error (%)
σF error (%)
Method
centric
acentric
centric
acentric
Explicit
0.22
0.11
1.76
1.54
Implicit
0.28
0.15
2.06
1.82
cctbx
0.06
0.11
0.07
0.39
Table 1: The mean absolute percentage error of each method against the exact posterior, evaluated on an experimental hen egg-white lysozyme dataset (PDB: 9B7C).
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Optimizer
Explicit η
Implicit η
SGD
10−4
10−4
SGD + momentum
3×10−5
3×10−5
RMSprop
3×10−3
3×10−3
Adam
3×10−3
10−2
AdaBelief
3×10−3
10−2
Appendix
Table S1: Learning rates used for the different optimizers, one per gradient estimator.
Figure S1: Theoretical (top) and empirical (bottom) variance of the explicit ( G^σE ) and implicit ( G^σI ) FoldedNormal derivative estimators of ∂σ\E[Y2] as a function of μ and σ . The colorbar denotes the variance value for each (μ,σ) pair. The empirical variance was computed across 10,000 simulations using 100 Monte Carlo samples each. The theoretical variance was calculated using equations 13 and 20 of the main text, for the explicit and implicit estimators, respectively.
Figure S2: Theoretical (top) and empirical (bottom) variance of the explicit ( G^σE ) and implicit ( G^σI ) \Rice derivative estimators of ∂σ\E[R2] as a function of ν and σ . The colorbar denotes the variance value for each (ν,σ) pair. The empirical variance was computed across 10,000 simulations using 100 Monte Carlo samples each. The theoretical variance was calculated using equations 26 and 34 of the main text, for the explicit and implicit estimators, respectively.
Figure S3: Recovery of the reference posterior as a function of the number of Monte Carlo samples N used per gradient, for each of the five optimization algorithms tested. The first and second row show the amplitude F for centric and acentric reflections, respectively. The third and fourth rows show the posterior width σF for centric and for acentric reflections, respectively. The two amplitude rows share a vertical scale, and so do the two posterior-width rows. Lines are medians over ten replicates and bands span the interquartile range.
F error (%)
σF error (%)
r for F
Algorithm
Estimator
centric
acentric
centric
acentric
centric
acentric
SGD
Explicit
3.91
2.33
2.64
2.83
0.99417
0.99539
Implicit
3.91
2.33
2.64
2.83
0.99417
0.99539
SGD + mom.
Explicit
2.10
0.91
2.85
2.60
0.99896
0.99924
Implicit
2.10
0.91
2.85
2.60
0.99896
0.99924
RMSprop
Explicit
0.69
0.33
2.42
1.67
0.99975
0.99992
Appendix
Table S2: Recovery of the reference posterior at N=512 , for every optimization algorithm and both estimators. Errors are mean absolute percentage errors and correlations ( r ) are Pearson coefficients against the reference amplitude. Each value is the median over ten replicates.
Figure S4: Ratio of implicit to explicit gradient variance, plotted against the observed signal-to-noise (SNR) of the reflection. The ratio is measured at the initialization θ0 and at the converged parameters θ^ of an Adam fit after 25,000 steps at N=512 . The top row shows the variance ratio for the location-derivative, and the bottom row shows the variance ratio for the scale-derivative. The ratio is shown separately for the centric reflections (left column) and the acentric reflections (right column). Each point is a median over reflections and ten replicates, and the bands span the interquartile range within each I/σI bin. The I/σI axis is divided into 18 equal-width bins between −3 and 9 .
Figure S5: Ratio of implicit to explicit gradient variance at the initialization θ0 parameters, and at each algorithm’s own converged parameters after 25,000 steps at N=512 . The top row is the centric location derivative and the bottom row is the centric scale derivative. The initialization parameters are the same in all cases, so the θ0 curve is identical across each row. Points are medians over reflections and ten replicates and bands span the interquartile range within each bin. The SNR axis is divided into 18 equal-width bins between −3 and 9 . The three lowest bins hold nine reflections each, and the fullest hold about three hundred.
Figure S6: Centric amplitude error and objective against training step, one column per optimization algorithm, at N=512 Monte Carlo samples per gradient. Top, mean absolute percentage error against the reference posterior; bottom, the negative ELBO. Lines are medians over ten replicates and bands span the interquartile range. Every algorithm drives the objective to within 0.5% of the same value, so the fits agree closely on the optimum they approach. Stochastic gradient descent’s objective is still descending at the end of training, so it has not converged.
In variational inference (VI), the practitioner approximates a high-dimensional distribution π with a simple surrogate one, often a (product) Gaussian distribution. However, in many cases of practical interest, Gaussian distributions might not capture the correct radial profile of π, resulting in poor coverage. In this work, we approach the VI problem from the perspective of optimizing over these radial profiles. Our algorithm radVI is a cheap, effective add-on to many existing VI schemes, such as Gaussian (mean-field) VI and Laplace approximation. We provide theoretical convergence guarantees for our algorithm, owing to recent developments in optimization over the Wasserstein space--the space of probability distributions endowed with the Wasserstein distance--and new regularity properties of radial transport maps in the style of Caffarelli (2000).
For approximating a target distribution given only its unnormalized log-density, stochastic gradient-based variational inference (VI) algorithms are a popular approach. For example, Wasserstein VI (WVI) and black-box VI (BBVI) perform gradient descent in measure space (Bures-Wasserstein space) and parameter space, respectively. Previously, for the Gaussian variational family, convergence guarantees for WVI have shown superiority over existing results for black-box VI with the reparametrization gradient, suggesting the measure space approach might provide some unique benefits. In this work, however, we close this gap by obtaining identical state-of-the-art iteration complexity guarantees for both. In particular, we identify that WVI's superiority stems from the specific gradient estimator it uses, which BBVI can also leverage with minor modifications. The estimator in question is usually associated with Price's theorem and utilizes second-order information (Hessians) of the target log-density. We will refer to this as Price's gradient. On the flip side, WVI can be made more widely applicable by using the reparametrization gradient, which requires only gradients of the log-density. We empirically demonstrate that the use of Price's gradient is the major source of performance improvement.
Kyurae Kim, Qiang Fu, Yi-An Ma +2
University of Pennsylvania, Philadelphia, U.S. · Yale University, New Haven, U.S. · University of California San Diego, San Diego, U.S. +1
Variational inference (VI) is a central tool in modern machine learning, used to approximate an intractable target density by optimising over a tractable family of distributions. As the variational family cannot typically represent the target exactly, guarantees on the quality of the resulting approximation are crucial for understanding which of its properties VI can faithfully capture. Recent work has identified instances in which symmetries of the target and the variational family enable the recovery of certain statistics, even under model misspecification. However, these guarantees are inherently problem-specific and offer little insight into the fundamental mechanism by which symmetry forces statistic recovery. In this paper, we overcome this limitation by developing a general theory of symmetry-induced statistic recovery in variational inference. First, we characterise when variational minimisers inherit the symmetries of the target and establish conditions under which these pin down identifiable statistics. Second, we unify existing results by showing that previously known statistic recovery guarantees in location-scale families arise as special cases of our theory. Third, we apply our framework to distributions on the sphere to obtain novel guarantees for directional statistics in von Mises-Fisher families. Together, these results provide a modular blueprint for deriving new recovery guarantees for VI in a broad range of symmetry settings.
Daniel Marks, Dario Paccagnan, Mark van der Wilk
Department of Computer Science University of Oxford · Department of Computing Imperial College London