> For the complete documentation index, see [llms.txt](https://chenchao-zhao.gitbook.io/machine-learning-demystified/llms.txt). Markdown versions of documentation pages are available by appending `.md` to page URLs; this page is available as [Markdown](https://chenchao-zhao.gitbook.io/machine-learning-demystified/blogs/stein-variational-gradient-descent.md). # Stein Variational Gradient Descent A new way to learn ## ## Introduction Gradient descent has become the basic algorithm for the training of almost all deep learning models. Stein variational gradient descent was proposed as a "natural counterpart of gradient descent for optimization." In our previous blog post *Differentiable Bayesian Structure Learning*, we briefly mentioned the core engine of the Bayesian algorithm was the Stein variational gradient descent. In this article, we will expand this topic and articulate the motivation, fundamental logic, mathematical derivation of this novel optimization method. The original paper can be found in the Reference section while in this article we will derive the algorithm in a much more intuitive way. ## Notations The usual inner product of two functions $$f$$ and $$g$$ is defined as $$ \langle f, g \rangle = \int dx ; f(x)g(x) $$ Let $$f$$ and $$g$$ be functions in [Reproducing kernel Hilbert space](https://en.wikipedia.org/wiki/Reproducing_kernel_Hilbert_space) (RKHS) $$\mathcal H$$, then the inner product is denoted $$\langle f, g \rangle\_\mathcal H$$ . Let $${f\_\mu}*{\mu=1}^d$$ and $${g*\nu}\_{\nu=1}^d$$ be functions in RKHS $$\mathcal H^d$$, then the inner product $$ \langle \mathbf f, \mathbf g \rangle\_{\mathcal H^d} = \delta\_{\mu\nu}\langle f\_\mu, g\_\nu \rangle\_{\mathcal H^d}. $$ ## Stein identity and discrepancy Observe that using Stokes' Theorem $$ \int\_{M} f\mathrm dp + p\mathrm df = \int\_{M} \mathrm d (pf) = \int\_{\partial M} pf $$ if $$pf \rightarrow 0$$ at boundary $$\partial M$$ and both $p$ and $$f$$ are smooth functions, then we have $$ \int\_M f\mathrm d p + p \mathrm df = 0. $$ If $$p$$ is a probability density over $$M$$, then we have the Stein **identity** $$ \int\_M fp; \mathrm d\log p + p, \mathrm d f = \mathbb E\_p \left\[ f, \mathrm d \log p + \mathrm d f \right] = 0 $$ for any test function $$f$$ that satisfies the requirements. Now we replace the sampling distribution by $$q$$ $$ \mathbb E\_p \[f \mathrm d\log p + \mathrm d f] \rightarrow \mathbb E\_q \[f \mathrm d\log p + \mathrm d f]\equiv S\_{q,p}f $$ we get the Stein discrepancy which vanishes when $$q$$ is $$p$$. Thus, we obtain a measure of "distance" between $$q$$ and $$p$$ with properly chosen test function $$f$$. ## Variational inference The goal of variational inference is to approximate a target distribution $$p$$ with a tractable ansatz distribution $$q$$ by minimizing the [Kullback–Leibler divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence) $$ D\_{\rm KL} (q||p) = \mathbb E\_q \log \frac{q}{p} = (-\mathbb E\_q\log p) - (-\mathbb E\_q \log q) $$ which is in the form of free energy $$F = U - TS$$ with temperature equals unity. Thus, minimization of KL-divergence is equivalent to striking a balance between minimizing $$q$$-average of energy ($$-\log p$$) and maximizing the entropy of $$q$$. KL-divergence is non-negative and attains minimum zero when $$q=p$$. We wish to $$\min\_q D\_{\rm KL}(q||p)$$ subject to the constraint that $q$ is a probability distribution $$\int dx, q(x) = \langle q, 1 \rangle = 1$$. Thus, the total objective is $$ \mathcal L\[q] = \langle q, \log q - \log p \rangle - \lambda \langle q, 1 \rangle $$ where $$\lambda$$ is a Lagrange multiplier. Take functional derivative $$ \frac{\delta \mathcal L}{\delta q} = \log q - \log p + 1 - \lambda = 0 $$ we get $$ q = p \exp (\lambda -1). $$ Since $$p$$ is a distribution, then we get $$\lambda=1$$. Thus, we showed $$q=p$$ is the solution of the optimization problem. However, for a real-world distribution $$p$$, which may be arbitrarily complicated, it is impossible to obtain an exact equality but a best approximation given the functional form of ansatz $$q$$. The ansatz distribution can be manually constructed based on the knowledge of the target distribution, e.g. mean-field approximation using exponential family distributions. In this article, we will discuss a non-parametric approach using particles of $$q$$. ## Coordinate flow Since both $$q$$ and $$p$$ are smooth functions and more importantly probability distributions, we can adiabatically deform $$q$$ into $$p$$ by shifting the coordinate $$ x^\mu \mapsto x^\mu + v^\mu(x) \delta t $$ where for simplicity we assume $$x \in \mathbb R^d$$. The task of seeking such a transformation is equivalent to searching for a proper velocity field $$v(x)$$. The total mass of $$q$$ is conserved and we have a conserved current of $$q$$-charge $$ j\_q^\mu (x) = q(x) v^\mu(x) $$ and $$ \dot q (t) = - \partial\_\mu j^\mu *q = - v^\mu \partial*\mu q - q \partial\_\mu v^\mu . $$ We have the following equivalent optimization problems $$ \min\_{v} D\_{\rm KL} (q\_{t+\delta t}||p\_t)\Leftrightarrow \min\_{v} D\_{\rm KL}(q\_t||p\_{t - \delta t}) $$ which is the equivalence of active and passive perspectives of coordinate transformations. In the latter case, we have the velocity in reverse direction $$ \dot p (t) = - \partial\_\mu j\_p^\mu = v^\mu \partial\_\mu p + p \partial\_\mu v^\mu . $$ where $$j\_p$$ is the conserved current of $$p$$-charge. We have the following expansion $$ p(x, t-\delta t) = p(x, t) - \dot p (x, t) \delta t + \mathcal O(\delta t^2) $$ $$ \log p(x, t-\delta t) = \log p(x, t) - \frac{d}{dt}\log p (x, t) , \delta t + \mathcal O(\delta t^2) . $$ The increment of the objective $$ \mathcal L\[v] = D\_{\rm KL}(q\_t||p\_{t - \delta t}) = \int dx, q\_t(x) \log \frac{q\_t(x)}{p\_{t-\delta t}(x)} $$ in first order of $$\delta t$$ $$ \delta \mathcal L\[v] = - \delta t \int dx, q\_t(x) \frac{\dot p\_t(x)}{p\_t(x)} . $$ Replace $$\dot p$$ using the continuity equation, we get $$ \frac{ \delta \mathcal L\[v] }{\delta t} = - \int dx, q(x) (v^\mu(x) \partial\_\mu p(x) + \partial\_\mu v^\mu) = - \mathbb E\_q(v^\mu \partial\_\mu p + \partial\_\mu v^\mu). $$ In other words, the gradient of $$D\_{\rm KL}(q\_t||p\_{t - \delta t})$$ is the negative Stein discrepancy of $$(q,p)$$ using test function $$v$$, i.e. $$S\_{q,p} v$$. ## Method of steepest descent Now we got the gradient descent of our variational inference objective, but we wish to get the steepest descent by searching for a proper velocity field. We further assume the velocity field is an element of $$d$$-dimensional RKHS $$\mathcal H^d$$. Then we have the reproducing property $$ v = \langle K, v \rangle\_{\mathcal H^d} = \langle v, K \rangle\_{\mathcal H^d} $$ where $$K(\cdot, \cdot)$$ is the kernel function of $$\mathcal H^d$$, e.g. Gaussian RBF kernel. Furthermore, the linear operator $$S\_{q,p}$$ can be shifted to the $$K$$, $$ S\_{q,p}v = \langle v, S\_{qp}K \rangle\_{\mathcal H^d} $$ which is the "dot product" of $$v$$ and $$S\_{qp} K$$ in RKHS. The solution of the optimization $$ \max\_{v, \Vert v \Vert\_{\mathcal H^d} \le 1} S\_{qp} v = \langle v, S\_{qp}K \rangle\_{\mathcal H^d} $$ is simply $$v^\* = S\_{qp}K/\Vert S\_{qp}K \Vert\_{\mathcal H^d}$$. Also note that the velocity will vanish when $$p=q$$. ## The Stein variational inference algorithm Using method of steepest descent, we obtained the optimal flow field $$v^\*$$. Next, we just need to go with flow $$ x \mapsto x + v^\* \delta t = x + \delta t', S\_{qp}K $$ and incrementally update $$q(x)$$. The algorithm is as follows: * Sample $$m$$ particles of initial $$q$$ * Approximate $$S\_{qp}K = \mathbb E\_q \[K^\mu \partial\_\mu \log p + \partial\_\mu K^\mu]$$ using sample mean of $$m$$ particles * The inverse of the norm of $$v^\*$$ is absorbed into learning rate $$\delta t'$$ * Update the coordinates of the particles using the calculated flow * Repeat the process ## Reference 1. [\[1608.04471\] Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://arxiv.org/abs/1608.04471)