Panagiotis (Panos) Toulis
Harvard University,
Department of Statistics
ptoulis at fas dot harvard dot edu

Homepage Implicit SGD Publications Code

What is implicit stochastic gradient descent?

Let $z_t = (x_t, y_t)$ be i.i.d. observations of a model with parameters $\theta_\star$. To estimate $\theta_\star$ through classical stochastic gradient descent (SGD) amounts to the procedure $\theta_t = \theta_{t-1} + \gamma_t C \nabla \ell(\theta_{t-1}; z_t)$, where $\ell$ is the log-likelihood function, $\gamma_t$ is the learning rate of the procedure (e.g., $\gamma_t = \gamma / (\gamma_0 + t))$, and $C$ is an appropriate symmetric psd condition matrix. This approach is well-studied and has been used extensively in applications. However, one major issue is stability: classical SGD is heavily affected by misspecifications of the learning rate $\gamma_t$, and can often explode. Currently, this issue is solved in an ad-hoc way (e.g., projected SGD, gradient clipping, robust transformations, etc.) Implicit SGD is a modification of the classical algorithm that aims to solve this problem in a principled way. Often, we can write the method as follows $\theta_t = \theta_{t-1} + \gamma_t C \nabla \ell(\theta_t; z_t)$, The update in this procedure is implicit because $\theta_t$ appears on both sides of the equation. In general such updates are hard to compute, however for a large family of statistical models (e.g., generalized linear models) the implicit update can be computed very fast (see Algorithm 1, Toulis et.al. (2014)). Running implicit SGD is then equivalent to running $\theta_t = \theta_{t-1} + \gamma_t \xi_t \nabla \ell(\theta_{t-1}^\intercal x_t; y_t)$, where the classical gradient $\ell(\theta_{t-1}^\intercal x_t; y_t)$ is scaled by $\xi_t \in \mathbf{R}$ that satisfies a fixed-point equation: $\mathbf{\xi_t} s_t = \ell'(\theta_{t-1}^\intercal x_t + \gamma_t \xi_t ||x_t||^2s_t; y_t)$, where $s_t = \ell'(\theta_{t-1}^\intercal x_t; y_t)$---$\ell'$ denotes the derivative of $\ell(u; y)$ with respect to the scalar $u$. The scalar $\xi_t$ is one-dimensional and so easy to compute as narrow search bounds are available. We can also show that $\theta_t \to \theta_\star$ and also give asymptotic variance results. In particular, if $t \gamma_t \to \gamma > 0$, then for both standard and implicit SGD, $t \cdot \mathbf{E}(\theta_t - \theta_\star) (\theta_t - \theta_\star)^\intercal \to \gamma^2 (2\gamma C \mathcal{I}(\theta_\star) - \mathbb{I})^{-1} C \mathcal{I}(\theta_\star) C, \quad \quad (1)$ where $\mathcal{I}(\theta) = \mathbf{E} (\nabla \ell(\theta; z) \nabla \ell(\theta; z)^\intercal) = - \mathbf{E}(\nabla^2 \ell(\theta; z))$ is the Fisher information matrix of the model, and $\mathbb{I}$ is the identity matrix -- form (1) is proved here (note: $C$ needs to commute with $\mathcal{I}(\theta_\star)$). Result (1) reveals that SGD procedures incur an information loss. In order to reduce the information loss, the condition matrix needs to account for the Fisher information; generally, $C$ needs to account for the curvature of the regression function of the statistic used in the SGD iterate; in the typical cases shown above, this curvature coincides with the Fisher information. Optimal estimation is obtained only when $C = \mathcal{I}(\theta_\star)^{-1}$, in which case the variance of $\theta_t$ is $(1/t) \mathcal{I}(\theta_\star)^{-1}$ i.e., the variance of the asymptotically optimal MLE. I discuss more details in my tutorial on Stochastic Gradient Descent (under progress).

Why implicit?

Implicit SGD has significant stability advantages over the classical SGD. Intuitively, it is a shrinked version of the typical (explicit) SGD and the shrinkage factor depends on the observed Fisher information. Thus, implicit SGD is significantly more stable than explicit SGD. In the limit, both methods are statistically equivalent, but in small-to-moderate samples, implicit SGD is more stable, and robust to misspecifications of the learning rate, or outliers. Implicit SGD can also be viewed as a stochastic proximal algorithm since it solves $\theta_t = \arg \max_{\theta} \{ \gamma_t\ell(\theta; z_t) - \frac{1}{2} ||\theta-\theta_{t-1}||^2 \}$. This formulation has a Bayesian interpretation in which $\theta_t$ is the posterior mode of a model where the prior distribution of $\theta_t$ is a normal $\mathcal{N}(\theta_{t-1}, \gamma_t \mathbb{I})$ and the data distribution is $\exp(\ell(\theta; z_t))$; note that the prior is shrinking towards $\theta_{t-1}$ as $t\to\infty$. In a more recent paper (Panos Toulis, Edoardo M. Airoldi, "Implicit stochastic approximation" (under review) (  pdf) I show how implicit SGD is a special instance of a more general implicit stochastic approximation procedure, much like how classical SGD is a special instance of stochastic approximation.

Relevant theory and code