A Geometric Interpretation of Stochastic Gradient Descent Using Diffusion Metrics

Abstract

This paper is a step towards developing a geometric understanding of a popular algorithm for training deep neural networks named stochastic gradient descent (SGD). We built upon a recent result which observed that the noise in SGD while training typical networks is highly non-isotropic. That motivated a deterministic model in which the trajectories of our dynamical systems are described via geodesics of a family of metrics arising from a certain diffusion matrix; namely, the covariance of the stochastic gradients in SGD. Our model is analogous to models in general relativity: the role of the electromagnetic field in the latter is played by the gradient of the loss function of a deep network in the former.

1. Introduction

Deep neural networks are high-dimensional machine learning models that have demonstrated impressive performance on a number of challenging tasks in computer vision, natural language processing and reinforcement learning [1,2]. These models are typically trained to minimize the misprediction error compared to large amounts of human-annotated data. A large number of diverse models with varying properties are prevalent in these application domains. Despite this diversity, stochastic gradient descent (SGD) is the gold standard for training deep neural networks. It has been shown to obtain good generalization performance; i.e., to train a model that performs well on new data, across a wide range of applications. In spite of this popularity and efficacy, a precise understanding of SGD for deep learning remains elusive.

This paper develops a geometric understanding of stochastic gradient descent. We build upon the work of [3], wherein the authors model the dynamics of SGD as a stochastic differential equation with state-dependent Gaussian noise. We interpret the covariance of this noise, called the diffusion matrix henceforth, as a metric on the parameter space. Our result provides a deterministic Equation (10) that can be compared to SGD near equilibrium points Equation (1). We write the diffusion matrix $D\left(x\right)$ in the form Equation (3) to show how it fundamentally captures the anisotropy of the dynamical system underlying SGD. This clarifies how $D\left(x\right)$ is one of the key factors that differentiates steady-state solutions of SGD from those of ordinary gradient descent GD (see comparison in [3,4,5]). Using the diffusion matrix, we then define a family of metrics on the parameter space that we call diffusion metrics. We then take the Einstein equation describing the geodesic on a Riemannian manifold, for the motion of a particle subject to a gravitational and electromagnetic field. We replace the electromagnetic force by the ordinary gradient while gravity is taken into account using the diffusion metric itself. After some mild hypotheses on the architecture of the neural network, we obtain the result that geodesics with respect to this equation correspond precisely to the evolution of a dynamical system, which is not subject to Euclidean gradient descent but to relativistic gradient descent (RGD) with respect to the family of diffusion metrics.

In the end, we obtain Equation (10), which is along the same vein as natural gradient descent in [6], but whose significance is much deeper in the context of SGD, since it stems from the anisotropy of the gradients with respect to the various parameters. This anisotropy encodes the difference between the dynamics of GD and those of SGD. We also compare our result with the ones in [3] and show them to be perfectly compatible. In the end we provide Appendix A for the reader convenience, with a quick review of some key facts of Riemannian geometry.

2. Continuous-Time SGD and the Diffusion Matrix

Stochastic gradient descent performs an update of the weights x of a neural network, replacing the ordinary gradient of the loss function $f=\frac{1}{N}{\sum }_{i=1}^N{f}_i$ with ${\nabla }_{\mathcal{B}}f$:

$$dx=-{\nabla }_{\mathcal{B}}fdt,{\nabla }_{\mathcal{B}}f=\frac{1}{|\mathcal{B}|}\sum _{i\in \mathcal{B}}\nabla f_i,$$

(1)

where $dx$ represents the continuous version of the weight update at step j: $x_{j+1}=x_j-\eta {\nabla }_{\mathcal{B}}f\left(x_j\right)$, with the learning rate $\eta$ incorporated into the expression of ${\nabla }_{\mathcal{B}}f$, and $\mathcal{B}$ is the mini-batch. In the expression of the loss function $f=\frac{1}{N}{\sum }_{i=1}^N{f}_i$, $f_i$ is the loss relative to the i-th element in our dataset $\mathsf{\Sigma }$ of size $|\mathsf{\Sigma }|=N$. We assume that weights belong to a compact subset $\mathsf{\Omega }\subset {\mathbb{R}}^d$ and that the $f_{i}s$ satisfy suitable regularity conditions (see [3] Section 2 for more details).

We define the diffusion matrix as the product of the size of the mini-batch $|\mathcal{B}|$ and the variance of ${\nabla }_{\mathcal{B}}f$, viewed as a random variable, $\varphi :\mathsf{\Sigma }⟶{\mathbb{R}}^d$, $\varphi \left(z_i\right)=\nabla f_i$:

$$D\left(x\right)=\mathbb{E}\left[(\varphi -\mathbb{E}\left[\varphi \right]){(\varphi -\mathbb{E}\left[\varphi \right])}^t\right]$$

(2)

Notice that $D\left(x\right)\in {\mathbb{R}}^{d\times d}$ and does not depend on the size of the mini-batch; it only depends on the weights x, loss function f and the dataset $\mathsf{\Sigma }$. With a direct calculation one can show that:

$$\begin{array}{cc}\hfill D& =\frac{1}{N}{\sum }_k(\nabla f_k){(\nabla f_k)}^t-(\nabla f){(\nabla f)}^t=\frac{1}{N^2}\left(\langle {\partial }_r\widehat{f},{\partial }_s\widehat{f}\rangle \right)\hfill \end{array}$$

(3)

where:

$$\widehat{f}=(f_1-f_2,f_1-f_3,\cdots ,f_{N-1}-f_N)\in {\mathbb{R}}^{N(N-1)/2}$$

and $\langle ,\rangle$ is the euclidean scalar product. In fact:

$$\begin{array}{cc}\hfill D_{rs}& =\frac{1}{N}{\sum }_{k=1}^N{\partial }_r{f}_k{\partial }_s{f}_k-\frac{1}{N^2}\sum _{i,j=1}^N{\partial }_r{f}_i{\partial }_s{f}_j\hfill \\ & =\frac{1}{N^2}[N({\partial }_r{f}_1{\partial }_s{f}_1+\cdots +{\partial }_r{f}_N{\partial }_s{f}_N)+\hfill \\ & -({\partial }_r{f}_1{\partial }_s{f}_1+{\partial }_r{f}_1{\partial }_s{f}_2+\cdots +{\partial }_r{f}_N{\partial }_s{f}_N)]=\hfill \\ & =\frac{1}{N^2}[(N-1){\partial }_r{f}_1{\partial }_s{f}_1-{\partial }_r{f}_1{\partial }_s{f}_2-\cdots -{\partial }_r{f}_1{\partial }_s{f}_N+\hfill \\ & -{\partial }_r{f}_2{\partial }_s{f}_1+(N-1){\partial }_r{f}_2{\partial }_s{f}_2-\cdots -{\partial }_r{f}_2{\partial }_s{f}_N+\cdots \hfill \\ & -{\partial }_r{f}_N{\partial }_s{f}_1-{\partial }_r{f}_N{\partial }_s{f}_2+\cdots +(N-1){\partial }_r{f}_N{\partial }_s{f}_N]\hfill \end{array}$$

which gives:

$$\begin{array}{cc}\hfill D_{rs}& =\frac{1}{N^2}[({\partial }_r{f}_1-{\partial }_r{f}_2)({\partial }_s{f}_1-{\partial }_s{f}_2)+({\partial }_r{f}_1-{\partial }_r{f}_3)({\partial }_s{f}_1-{\partial }_s{f}_3)+\cdots \hfill \\ & +({\partial }_r{f}_1-{\partial }_r{f}_N)({\partial }_s{f}_1-{\partial }_s{f}_N)+({\partial }_r{f}_2-{\partial }_r{f}_3)({\partial }_s{f}_2-{\partial }_s{f}_3)+\cdots \hfill \\ & +({\partial }_r{f}_{N-1}-{\partial }_r{f}_N)({\partial }_s{f}_{N-1}-{\partial }_s{f}_N)=\frac{1}{N^2}\left(\langle {\partial }_r\widehat{f},{\partial }_s\widehat{f}\rangle \right).\hfill \end{array}$$

The diffusion matrix measures effectively the anisotropy of our data: $D=0$ if and only if ${\partial }_r\left(f_i\right)={\partial }_r\left(f_j\right)$ for all $r=1,\cdots ,d$ and $i,j=1,\cdots ,N$. In other words, the diffusion matrix measures how the loss of each datum depends, at first order, on the weights in a different way with respect to the loss of each of the other data-points. So, it tells us how much we should expect the SGD dynamics to differ from the GD one. Notice that the expression (Equation (3)) gives us immediately a bound on the rank of D; namely, $\mathrm{rk}\left(D\right)\le N-1$.

The

Table 1. Values for N and d for various architectures on CIFAR and SVHN datasets (see [8]).

Architecture	$\mathit{d}=|\mathbf{Weights}|$	$\mathit{N}=|\mathbf{Data}|$, CIFAR	$\mathit{N}=|\mathbf{Data}|$, SVHN
ResNet	1.7 M	60 K	600 K
Wide ResNet	11 M	60 K	600 K
DenseNet (k = 12)	1 M	60 K	600 K
DenseNet (k = 24)	27.2 M	60 K	600 K

suggests that in many algorithms currently available, the diffusion matrix has low rank; hence, it is singular; this just by comparing the size d of D and its rank which bound by N. This fact turns out to be very important in the construction of the diffusion metrics, as we will see below.

3. Diffusion Metrics and General Relativity

The evolution of a dynamical system in general relativity takes place along the geodesics according to the metric imposed on the Minkowski space by the presence of gravitational masses. The equation for such geodesics, once Einstein’s equation is solved, is:

$$\frac{d^2{x}^{\mu }}{d{t}^2}+{\mathsf{\Gamma }}_{\rho \sigma }^{\mu }\frac{d{x}^{\rho }}{dt}\frac{d{x}^{\sigma }}{dt}=\frac{q}{m}{F}_{\nu }^{\mu }\frac{d{x}^{\nu }}{dt},$$

(4)

where ${\mathsf{\Gamma }}_{\rho \sigma }^{\mu }$ are the Christoffel symbols for the Levi–Civita connection:

$${\mathsf{\Gamma }}_{uv}^w=\frac{1}{2}{g}^{wz}\left({\partial }_u{g}_{vz}-{\partial }_z{g}_{uv}+{\partial }_v{g}_{uz}\right)$$

(5)

and $\frac{q}{m}{F}_{\nu }^{\mu }$ is a term regarding an external force; e.g., one coming from an electromagnetic field.

If we take the time derivative of the differential equation underlying the ordinary (i.e., non stochastic) gradient descent:

$$\frac{d^2{x}^{\mu }}{d{t}^2}=-\frac{d}{dt}{\partial }_{\mu }f$$

and compare with Equation (4), observe that $-\frac{d}{dt}{\partial }_{\mu }f$ effectively replaces the force term $(q/m)F_{\nu }^{\mu }\frac{d{x}^{\nu }}{dt}$. Hence, the geodesic equation Equation (4) models the ordinary GD equation if we take a constant metric and replace the force term with the gradient of the loss; furthermore, this corresponds to the condition $D=0$ in SGD dynamics Equation (1).

This suggests that one may define a metric dependent on the diffusion matrix; this metric should become constant when $D=0$. As a side remark, notice that since D is singular in many important practical applications (see Table 1), it is not reasonable to use it to define the metric itself. On the other hand, since D measures the anisotropy of the weight space, it is reasonable to employ it to perturb the euclidean metric. So the stochastic nature of the dynamical system ruled by the SGD is replaced by a perturbation of the dynamics for the ordinary gradient descent. As an analogy, the presence of (small) masses in space, in the weak field approximation (see [7]) of general relativity, generates gravity; this motivates our (small) deformation of the euclidean metric using the diffusion matrix.

At each point $x\in \mathsf{\Omega }$, we define a metric called diffusion metric as

$$g\left(x\right)=\mathrm{id}+\epsilon \left(x\right)D\left(x\right)$$

(6)

with $\epsilon \left(x\right)<1/M_x$, where $M_x=\max \{{\lambda }_k:{\lambda }_k\mathrm{is}\mathrm{an}\mathrm{eigenvalue}\mathrm{of}D\left(x\right)\}$. This ensures that g is non-singular at each $x\in \mathsf{\Omega }$. We have, thus, defined a family of metrics that depend on the real parameter $\epsilon$. We expect this model to approximate the solution to the Fokker–Planck equation when the parameter ${\beta }^{-1}$, whose interpretation is related with the temperature, is very small (see [3] for the notation and more details).

Notice that our heuristic hypothesis on $\epsilon$ allows us to make the so called weak field approximation (see [7]):

$$g^{-1}=\mathrm{id}-\epsilon D\left(x\right).$$

Hence, we have the following expression for the Christoffel’s symbols (in this approximation we discard $\mathcal{O}\left({\epsilon }^2\right)$):

$${\mathsf{\Gamma }}_{uv}^w=\frac{1}{2}\sum _z({\delta }_{wz}-{\epsilon }_{wz}^d){\epsilon }^{}\left({\partial }_u{d}_{vz}-{\partial }_z{d}_{uv}+{\partial }_v{d}_{uz}\right)=\frac{\epsilon }{2}\left.({\partial }_u{d}_{vw}-{\partial }_w{d}_{uv}+{\partial }_v{d}_{uw}\right)$$

(7)

where d_ij are the coefficients of D, and δ_wz is the Kronecker delta.

Let us now compute the Christoffel’s symbols and then substitute them into the geodesic equation given by Equation (4).

$$\begin{array}{cc}\hfill {\mathsf{\Gamma }}_{ij}^k& =\frac{\epsilon }{2N^2}[\langle {\partial }_i{\partial }_j\widehat{f},{\partial }_k\widehat{f}\rangle +\langle {\partial }_j\widehat{f},{\partial }_i{\partial }_k\widehat{f}\rangle -\langle {\partial }_i{\partial }_k\widehat{f},{\partial }_j\widehat{f}\rangle -\langle {\partial }_i\widehat{f},{\partial }_k{\partial }_j\widehat{f}\rangle \hfill \\ & +\langle {\partial }_j{\partial }_k\widehat{f},{\partial }_i\widehat{f}\rangle +\langle {\partial }_k\widehat{f},{\partial }_i{\partial }_j\widehat{f}\rangle ]=\frac{\epsilon }{N^2}\langle {\partial }_i{\partial }_j\widehat{f},{\partial }_k\widehat{f}\rangle .\hfill \end{array}$$

Let us substitute in Equation (4) (writing the sum now):

$$\frac{d^2{x}^k}{d{t}^2}+\frac{\epsilon }{N^2}\sum _{i,j}\langle {\partial }_i{\partial }_j\widehat{f},{\partial }_k\widehat{f}\rangle \frac{d{x}^i}{dt}\frac{d{x}^j}{dt}=\frac{d}{dt}{\partial }_{k}f.$$

(8)

Let us concentrate on the expression:

$$\begin{array}{cc}\hfill \frac{\epsilon }{N^2}{\sum }_{i,j}\langle {\partial }_i{\partial }_j\widehat{f},{\partial }_k\widehat{f}\rangle \frac{d{x}^i}{dt}\frac{d{x}^j}{dt}& =\frac{\epsilon }{N^2}{\sum }_{i,j,\alpha }{\partial }_i{\partial }_j{\widehat{f}}_{\alpha }{\partial }_k{\widehat{f}}_{\alpha }\frac{d{x}^i}{dt}\frac{d{x}^j}{dt}\hfill \\ & =\frac{\epsilon }{N^2}{\sum }_{\alpha }\frac{d^2{\widehat{f}}_{\alpha }}{d{t}^2}{\partial }_k{\widehat{f}}_{\alpha }.\hfill \end{array}$$

Now, we take the integral in dt (we compute the parts twice):

$$\begin{array}{cc}\hfill \frac{\epsilon }{N^2}\int {\sum }_{\alpha }\frac{d^2{\widehat{f}}_{\alpha }}{d{t}^2}{\partial }_k{\widehat{f}}_{\alpha }dt& =\frac{\epsilon }{N^2}{\sum }_{\alpha }\left[\frac{d{\widehat{f}}_{\alpha }}{dt}{\partial }_k{\widehat{f}}_{\alpha }-\int \frac{d{\widehat{f}}_{\alpha }}{dt}\frac{{\partial }_k{\widehat{f}}_{\alpha }}{dt}dt\right]\hfill \\ & =\frac{\epsilon }{N^2}{\sum }_{\alpha }\left[\frac{d{\widehat{f}}_{\alpha }}{dt}{\partial }_k{\widehat{f}}_{\alpha }-{\widehat{f}}_{\alpha }\frac{d}{dt}{\partial }_k{\widehat{f}}_{\alpha }+\int {\widehat{f}}_{\alpha }\frac{d^2}{d{t}^2}{\partial }_k{\widehat{f}}_{\alpha }dt\right].\hfill \end{array}$$

Notice that, in many practical applications we have:

$$\frac{d^2}{d{t}^2}{\partial }_k{\widehat{f}}_{\alpha }=0$$

(9)

because ${\partial }_i{\partial }_j{\partial }_k\widehat{f}=0.$

Hence

$$\begin{array}{cc}\hfill \frac{\epsilon }{N^2}\int {\sum }_{\alpha }\frac{d^2{\widehat{f}}_{\alpha }}{d{t}^2}{\partial }_k{\widehat{f}}_{\alpha }dt& =\frac{\epsilon }{N^2}{\sum }_{\alpha ,\ell }\left[{\partial }_{\ell }{\widehat{f}}_{\alpha }\frac{d{x}^{\ell }}{dt}{\partial }_k{\widehat{f}}_{\alpha }-{\widehat{f}}_{\alpha }\frac{d}{dt}{\partial }_k{\widehat{f}}_{\alpha }\right]\hfill \\ & =\epsilon {\sum }_{\ell }{d}_{k,\ell }\frac{d{x}^{\ell }}{dt}-\frac{\epsilon }{N^2}{\sum }_{\alpha }{\widehat{f}}_{\alpha }\frac{d}{dt}{\partial }_k{\widehat{f}}_{\alpha }.\hfill \end{array}$$

We now substitute the obtained expression into the Equation (8), taking the integral in dt:

$$\frac{d{x}^k}{dt}+\epsilon \sum _{\ell }{d}_{k,\ell }\frac{d{x}^{\ell }}{dt}-\frac{\epsilon }{N^2}\sum _{\alpha }{\widehat{f}}_{\alpha }\frac{d}{dt}{\partial }_k{\widehat{f}}_{\alpha }=-{\partial }_{k}f.$$

We may assume $\frac{\epsilon }{N^2}$ very small, as motivated by Equation (1); hence, discard this term. Writing the equation into vector form, we have:

$$\frac{dx}{dt}+\epsilon D\frac{dx}{dt}=-\nabla f.$$

By the weak field approximation ${(I+\epsilon D)}^{-1}\cong (I-\epsilon D)$, we can write:

$$\frac{dx}{dt}=-(I-{\epsilon }^D)\nabla f=-{\nabla }_{D}f,$$

(10)

where ${\nabla }_{D}f$ is the gradient computed according to the diffusion metric Equation (6).

We can summarize our result as follows: the SGD equation Equation (1) can be replaced, provided the approximation Equation (9) holds, by the deterministic equation Equation (10), where the dynamical system evolves with respect to the gradient computed according to the diffusion metric Equation (6).

We now want to compare our result Equation (10) with [3] Section 1, in order to understand how the steady state solutions of Equation (10) compare to the SGD steady state solutions described by Equation (1). In [3], the authors regard SGD as minimizing the function Φ instead of our loss f. Let us focus on (8) in [3], where the relation between f and Φ is discussed. In our case, we take ∇Φ = ∇_Df so that Equation (11) becomes

$$-\nabla f\left(x\right)+\tilde{D}\left(x\right)\nabla \mathsf{\Phi }\left(x\right)=0,$$

(11)

where $\tilde{D}\left(x\right)=(I+\epsilon D\left(x\right))$ is the diffusion metric. If the term D(x) in (8) in [3] is spelled out as our $\tilde{D}\left(x\right)$, we can write such equation as:

$$j\left(x\right)=-\nabla f\left(x\right)+\tilde{D}\left(x\right)\nabla \mathsf{\Phi }\left(x\right)-{\beta }^{-1}\nabla ·\tilde{D}\left(x\right).$$

(12)

Notice that according to our approximation Equation (9), $\nabla ·\tilde{D}\left(x\right)=0$. Hence, Equation (12) (that is (8) in [3]) is perfectly compatible with our treatment, and furthermore, the assumption 4 in [3] is fully justified by the fact that j(x) = 0.

4. Conclusions

The general relativity (GR) model helps to provide a deterministic approach the evolution of the dynamical system described by SGD, through the use of the diffusion metric accounting for the anisotropy of the system. Our results are compatible with [3]; moreover, they give GD for an isotropic system. We plan to explore deep neural networks mixing our GR model with energy models in a forthcoming paper (see also [9]). This will provide new geometric insight to the theory.