Stein's Lemma for the Reparameterization Trick with Exponential Family Mixtures
This work provides a theoretical foundation for gradient estimation in machine learning, particularly for variational inference and reinforcement learning, by addressing limitations in existing methods, though it is incremental as it builds on prior Stein's lemma applications.
The authors tackled the problem of applying Stein's lemma to a broader class of distributions by extending it to exponential-family mixtures, including Gaussian distributions with full covariance structures, which enabled the derivation of new reparameterizable gradient identities for distributions like Student's t-distribution and skew Gaussian.
Stein's method (Stein, 1973; 1981) is a powerful tool for statistical applications and has significantly impacted 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 reparameterization trick to derive gradients of expectations of a large class of functions under weak assumptions. Using this connection, we can derive many new reparameterizable gradient identities that go beyond the reach of existing works. For example, we give gradient identities when the expectation is taken with respect to Student's t-distribution, skew Gaussian, exponentially modified Gaussian, and normal inverse Gaussian.