Stein's Lemma for the Reparameterization Trick with Exponential Family Mixtures

Abstract

Stein's method (Stein, 1973;1981) is a powerful tool for statistical applications, and has had a significant impact in machine learning. Stein's lemma plays an essential role in Stein's method. Previous applications of Stein's lemma either required strong technical assumptions or were limited to Gaussian distributions with restricted covariance structures.In this work, we extend Stein's lemma to exponential-family mixture distributions including Gaussian distributions with full covariance structures. Our generalization enables us to establish a connection between Stein's lemma and the reparamterization trick to derive gradients of expectations of a large class of functions under weak assumptions. Using this connection, we can derive many new reparameterizable gradientidentities that goes beyond the reach of existing works. For example, we give gradient identities when expectation is taken with respect to Student's t-distribution, skew Gaussian, exponentially modified Gaussian, and normal inverse Gaussian.

Tractable structured natural gradient descent using local parameterizations

“…and second, instead of using ( 13) with Monte-Carlo (MC), we will use Stein's identity (Opper & Archambeau, 2009;Lin et al, 2019b):…”

“…is the multivariate trigamma function. Moreover, we can use re-parameterizable gradients (Figurnov et al, 2018;Lin et al, 2019b) for G S −1 t and g n due to the Bartlett decomposition (Smith et al, 1972) (see Appendix E.1 for details).…”

“…Our approach makes it easy to use different distributions and derivation possible for a variety of parameterizations and maps, unlike the method originally described by Glasmachers et al (2010). Moreover, we have Stein's first-order identity (Lin et al, 2019b; known as the re-parameterization trick for Gaussian, g Σ = 1 2 E q S(w − µ)∇ T w (w) , to avoid using the Hessian ∇ 2 w (w). We also obtain a NGD update in a unconstrained space for univariate EFs in Appendix G.…”

“…Since Gamma distribution is implicitly re-parametrizable, we can also compute the gradient ∇ n L 1 thanks to the implicit reparametrization trick (Figurnov et al, 2018;Lin et al, 2019b) for Gamma distribution.…”