Implicit and explicit reparameterizations of the Folded Normal and Rice distributions

Abstract

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}.