The Riemannian Langevin equation and conic programs

TL;DR

This paper introduces the Riemannian Langevin equation (RLE) as a generalization of stochastic gradient methods to Riemannian manifolds, providing explicit formulas for Brownian motion on cones, advancing understanding of stochastic processes in geometric spaces.

Contribution

It formulates the Riemannian Langevin equation and derives explicit formulas for Brownian motion on fundamental cones, expanding stochastic analysis on manifolds.

Findings

Explicit formulas for Brownian motion on cones

Generalization of Langevin dynamics to Riemannian manifolds

Framework for analyzing stochastic processes on geometric spaces

Abstract

Diffusion limits provide a framework for the asymptotic analysis of stochastic gradient descent (SGD) schemes used in machine learning. We consider an alternative framework, the Riemannian Langevin equation (RLE), that generalizes the classical paradigm of equilibration in R^n to a Riemannian manifold (M^n, g). The most subtle part of this equation is the description of Brownian motion on (M^n, g). Explicit formulas are presented for some fundamental cones.

Equations

$$x^{k+1}=x^k-{\gamma }_k\nabla {\epsilon }_{i_k}(x^k),$$

$$d{X}_t=-\nabla E(X_t)dt+\sqrt{\frac{2}{\beta }}d{B}_t,$$

$$d{X}_t=-\mathrm{grad}E(X_t)dt+d{B}_t^{g,\beta }.$$

$$\mathcal{L}f=-\mathrm{grad}E(f)+\frac{1}{\beta }\Delta f,\mathrm{where}\Delta f=\frac{1}{\sqrt{\det g}}{\partial }_i(\sqrt{\det g}{g}^{ij}{\partial }_{j}f),$$

$${\partial }_t\rho =\mathrm{div}\left(\rho \mathrm{grad}F\right),F=E+\frac{1}{\beta }\log \rho .$$

$$\rho (x)=\frac{1}{Z_{\beta }}{e}^{-\beta E(x)},Z_{\beta }={\int }_{\mathcal{M}}{e}^{-\beta E(x)}\sqrt{\det g}dx.$$

$$F=\frac{1}{2}\log \det D^{2}F,x\in K^o.$$

$$x(\theta )={\mathrm{argmin}}_{s\in C}\{F(s)+\theta c^{T}s\}.$$

$$\frac{dx}{d\theta }=-{\mathrm{grad}}_g(c^{T}x),g=D^{2}F,x(0)=\mathrm{argmin}F.$$

$$g(Z_1,Z_2)=\mathrm{Tr}(A_N^{-1}(Z_1)Z_2),A_N(Z)=\frac{1}{N}\sum _{i=1}^N(W{W}^T{)}^{\frac{N-i}{N}}Z(W^{T}W{)}^{\frac{i}{N}}.$$

$$f({\mathbf{X}}_t)=f({\mathbf{X}}_0)+\frac{1}{\beta }{\int }_0^t\Delta f({\mathbf{X}}_s)\mathrm{d}s+\mathrm{a}\mathrm{local}\mathrm{martingale}.$$

$$[f({\mathbf{B}}^{g,\beta }){]}_t=\frac{1}{\beta }{\int }_0^t|\nabla f({\mathbf{B}}_s^{g,\beta }){|}_g^2\mathrm{d}s$$

$$[f_1({\mathbf{B}}^{g,\beta }),f_2({\mathbf{B}}^{g,\beta }){]}_t=\frac{1}{\beta }{\int }_0^t(\nabla f_1\cdot \nabla f_2)({\mathbf{B}}_s^{g,\beta })\mathrm{d}s$$

$$g_{ij}=\frac{{\partial }^{2}F}{\partial x^i\partial x^j}.$$

$$\mathrm{grad}F=-x,|\nabla F{|}^2=n\mathrm{and}\Delta F=0.$$

$$x^i{\partial }_{i}F(x)=-n.$$

$$(\mathrm{grad}F{)}^k=g^{jk}{\partial }_{j}F=-x^k.$$

$$|\nabla F{|}^2=(g^{ij}{\partial }_{i}F){\partial }_{j}F=-x^j{\partial }_{j}F=n.$$

$$\Delta F=\mathrm{div}(\mathrm{grad}F)$$

$$=-{\partial }_i(x^i)-x^i{\partial }_i(\frac{1}{2}\log \det g)=-n-x^i{\partial }_{i}F=0,$$

$$F(x)=-\sum _{i=1}^n\log x^i,g=\sum _{i=1}^n\frac{\mathrm{d}{x}^i\mathrm{d}{x}^i}{(x^i{)}^2}.$$

$$\log x^i({\mathbf{B}}_t^{g,\beta })-\log x^i({\mathbf{B}}_0^{g,\beta })=\frac{1}{\sqrt{\beta }}{B}_t^{i}i=1,...,n$$

$$F(x)=-\sum _{i=1}^n\log \frac{\sin (\pi x^i)}{\pi },g=\sum _{i=1}^n{\pi }^2\frac{\mathrm{d}{x}^i\mathrm{d}{x}^i}{{\sin }^2(\pi x^i)}.$$

$$\log (\tan (\frac{\pi x^i({\mathbf{B}}_t^{g,\beta })}{2}))-\log (\tan (\frac{\pi x^i({\mathbf{B}}_0^{g,\beta })}{2}))=\frac{1}{\sqrt{\beta }}{B}_t^i$$

$$K_{n+1}=\left\{x\in {\mathbb{R}}^{n+1}|(x^0)\ge \sqrt{\sum _{i=1}^n(x^i{)}^2}\right\},$$

$$F(x)$$

$$g_{ij}$$

$$g^{ij}$$

$$\sqrt{\det (g_{ij})}$$

$$f_b(x)=\frac{\sqrt{n+1}}{2}\log \frac{(b^{T}x{)}^2}{x^{T}Ax},b\in L_{n+1}^{+}.$$

$$L_{n+1}^{+}=\left\{b\in {\mathbb{R}}^{n+1}|b^0>0,b^{T}Bb=0\right\}.$$

$$f^0=\frac{1}{\sqrt{n+1}}F,\mathrm{and}{f}^i=f_{b_i}i=1,\dots ,n.$$

Full text

1 Introduction

1.1 Stochastic gradient descent

Stochastic gradient descent (SGD) schemes in machine learning typically arise as follows. An empirical loss function $E:\mathbb{R}^{n}\to\mathbb{R}$, for a training parameter $x$, is defined through a finite sum $E(x)=\frac{1}{N}\sum_{i=1}^{N}\varepsilon_{i}(x)$, where $\varepsilon_{i}(x)=\varepsilon(x;y_{i},z_{i})$ denotes a loss function evaluated on a finite set of training data $\{(y_{i},z_{i})\}_{i=1}^{N}$. The loss function is minimized using the stochastic gradient descent scheme

[TABLE]

where $i_{k}$ is chosen randomly from the set $\{1,\ldots,N\}$ and $\gamma_{k}$ is a time step.

Several variants of SGD have been explored since the classic work of Robbins and Monro [13]. What is different in modern maching learning is the large size of $n$ and $N$ and the protocol for the learning rate $\gamma$. Diffusion limits of SGD replace the discrete iteration above with stochastic differential equations (SDE); these SDE depend on the manner in which $\gamma_{k}\to 0$, $n\to\infty$ and $N\to\infty$. Some examples of this approach are the stochastic modified equation (SME) proposed in [10], the variational analysis using Kullback-Leibler divergence proposed in [11], and homogenized SGD (HSGD) defined in [12].

1.2 Riemannian Langevin equation

SDE limits of SGD schemes begin with an algorithm and study its scaling limits. The approach in this paper is different. We begin with diffusions that extend the classical Langevin equation to a Riemannian setting. The relation to optimization lies in the nature of the underlying Riemannian geometry. Let us first explain the model; we then explain why it is a natural extension of ideas used in classical and modern optimization theory.

Recall that the Langevin equation associated to the loss function $E:\mathbb{R}^{n}\to\mathbb{R}$, at inverse temperature $\beta>0$, is formulated mathematically as the Itô SDE

[TABLE]

where $\{B_{t}\}_{t\geq 0}$ denotes standard Brownian motion on $\mathbb{R}^{n}$. Given an $n$-dimensional Riemannian manifold $(\mathcal{M},g)$ with metric $g$ and a loss function $E:\mathcal{M}\to\mathbb{R}$ we consider the Riemannian Langevin equation (RLE)

[TABLE]

Both the gradient and the Brownian motion at inverse temperature $\beta$ are now computed with respect to the Riemannian metric $g$. In particular, the Brownian motion $B_{t}^{g,\beta}$ on $(\mathcal{M},g)$ must be defined carefully as discussed below. Let $f\in C^{\infty}(\mathcal{M})$ be a test function. The infinitesimal generator $L$ of the diffusion (3) is

[TABLE]

is the Laplace-Beltrami operator, $\mathrm{grad}$ denotes the gradient with respect to $g$, and the volume form is computed in coordinates using $\det g=\det(g_{ij})$.

The Fokker-Planck equation, $\partial_{t}\rho=L^{*}\rho$, where the dual is with respect to the volume form of $g$, takes the form

[TABLE]

The free energy, $F$, is constant in equilibrium and we find the Gibbs density

[TABLE]

RLE is a method a method to study the Gibbs measure associated to $F$, whereas SGD schemes seek the minimum of $F$. However, these techniques are closely related. When $\beta\to\infty$ and $F$ has a unique global minimum at $x_{*}\in\mathcal{M}$, the Gibbs measure concentrates at $x_{*}$ as $\beta\to\infty$ with rigorous asymptotics provided by large deviations theory. A subtle feature of the metrics arising in optimization is that the volume $\int_{M}\sqrt{\det g}\,dx$ may be infinite.

1.3 Riemannian geometries in optimization

The framework of RLE provides a natural geometric unity between conic programs and deep learning. What changes is the underlying Riemannian manifold $(\mathcal{M},g)$. Let us explain this idea through examples.

Bayer and Lagarias systematized the Riemannian geometry discovered by Karmarkar for interior-point methods [3, 9]. We focus on the canonical barrier [7]. Associated to every regular convex cone $K\subset\mathbb{R}^{n}$ is a unique convex function $F$ defined in the interior $K^{o}$ of $K$ such that $F(x)\to+\infty$ as $x\to\partial K$. This function is the Cheng-Yau solution to the Monge-Ampère equation

[TABLE]

Given a barrier and a vector $c\in\mathbb{R}^{n}$, the conic program $\min_{x\in C}c^{T}x$ is solved by taking the $\theta\to\infty$ limit of the central path

[TABLE]

Further, $x(\theta)$ above is the solution to the Riemannian gradient flow

[TABLE]

That is, the Hessian of the barrier $F$ provides the underlying Riemannian metric. The canonical barrier has several striking geometric properties [7].

Riemannian metrics have also been extensively used in geometric deep learning [4]. A model problem that allows a comparison between deep learning and classical optimization is the deep linear network [1, 2, 5]. The training space for a network of depth $N$ is the product space of $d\times d$ matrices $\mathbb{M}_{d}^{N}$. Given $\mathbf{W}=(W_{N},W_{N-1},\ldots,W_{1})$ the observable is the product $V=W_{N}W_{N-1}\cdots W_{1}$. Learning problems like matrix completion may be modeled as a Euclidean gradient descent for a cost function $L(\mathbf{W}):=E(V)$. Then for suitable initial conditions, the Euclidean gradient flow, $\dot{W}_{i}=-\nabla_{W_{i}}L(W)$, $1\leq i\leq N$ corresponds to the Riemannian gradient flow, $\dot{V}=-\mathrm{grad}_{g}E(V)$, where the metric $g$ acts by

[TABLE]

In order to explore the nature of the Riemannian Langevin equation in optimization, we must understand Brownian motion on Riemannian manifolds like those above. This is a problem of some depth. We illustrate this by computing explicit expressions for Brownian motion in some fundamental cones, using expressions for the barrier from [6].

2 Brownian motion and conic programs

Manifold-valued Brownian motion may be defined in several ways [8]. We use the following definition in this note: an $\mathcal{M}$-valued semimartingale $\boldsymbol{X}_{t}$ is called a Brownian motion on $\mathcal{M}$, with temperature $T=\frac{1}{\beta}$, if for any $f\in C^{\infty}(\mathcal{M})$,

[TABLE]

We denote Brownian motion on $(\mathcal{M},g)$ at temperature $T=\frac{1}{\beta}$ by $\boldsymbol{B}^{g,\beta}_{t}$.

Proposition 1

The quadratic variation process of $f(\boldsymbol{B}^{g,\beta}_{t})$ for $f\in C^{\infty}(\mathcal{M})$ is

[TABLE]

where $|\nabla f|^{2}:=g^{ij}\partial_{i}f\partial_{j}f$.

Corollary 1

The covariation process of $f_{1}(\boldsymbol{B}^{g,\beta}_{t})$ and $f_{2}(\boldsymbol{B}^{g,\beta}_{t})$ for $f_{1},f_{2}\in C^{\infty}(\mathcal{M})$ is

[TABLE]

where $\nabla f_{1}\cdot\nabla f_{2}:=g^{ij}\partial_{i}f_{1}\partial_{j}f_{2}$.

Proposition 1 allows us to analyze Brownian motion through a careful choice of coordinate functions. We will choose $f$ such that $\Delta f=0$ and $|\nabla f|^{2}=1$, so that $f(\boldsymbol{B}^{g,\beta}_{t})-f(\boldsymbol{B}^{g,\beta}_{0})$ has the same law as an $\mathbb{R}$-valued standard Brownian motion.

Let us now assume $\mathcal{M}$ is a regular convex cone $K\subset\mathbb{R}^{n}$, let $F$ denote its canonical barrier, and equip $K$ with the Hessian metric

[TABLE]

Theorem 2.1

Consider Brownian motion $\boldsymbol{B}^{g,\beta}_{t}$ on $(K,g)$. The process $\frac{\sqrt{\beta}}{\sqrt{n}}(F(\boldsymbol{B}^{g,\beta}_{t})-F(\boldsymbol{B}^{g,\beta}_{0}))$ has the same law as a standard Brownian motion on the line.

Proof

We will use the logarithmic homogeneity of $F$ and the Monge-Ampère equation to establish the identities

[TABLE]

Theorem 2.1 follows immediately from these identities.

The first identity uses logarithmic homogeneity. Since $F(\lambda x)=F(x)-n\log\lambda$ for $\lambda\in\mathbb{R}^{+}$, $x\in K$, we may differentiate with respect to $\lambda$ and set $\lambda=1$ to find

[TABLE]

Next, the differential of $F$ with respect to $x^{j}$ is $\partial_{j}F(x)=-x^{i}\partial_{i}\partial_{j}F(x)=-g_{ij}x^{i}$ by equation (14). Thus, each component of the gradient of $F$ is

[TABLE]

This proves the first identity in equation (15). It immediately follows that

[TABLE]

Finally, we show that $\Delta F=0$ as follows

[TABLE]

where we have used the Monge-Ampére equation (7) and equation (18).

The above theorem sheds new light on the mysterious reappearance of the Cheng-Yau metric in optimization theory. Let us understand it better with examples.

3 Brownian motion examples

3.1 Positive orthant

Denoted by $\mathbb{R}^{n}_{+}=\{x\in\mathbb{R}^{n}|x^{i}>0,i=1,...,n\}$ the positive orthant. The canonical barrier and its Hessian metric are

[TABLE]

Then for each choice of coordinate, we have the identity in law

[TABLE]

where $\{B^{i}_{t}\}_{i=1}^{n}$ are $n$ independent standard Brownian motion on $\mathbb{R}$.

3.2 Cube

Next we consider a convex set, the cube $B_{n}=(0,1)^{n}$. We find that

[TABLE]

Similar calculations yield the identity in law

[TABLE]

where $\{B^{i}_{t}\}_{i=1}^{n}$ are $n$ independent standard Brownian motions on $\mathbb{R}$.

3.3 Lorentz cone

A deeper example is provided by the Lorentz cone

[TABLE]

where $x=(x^{0},\ldots,x^{n})$. The canonical barrier $F$ on $K_{n+1}$ is given by

[TABLE]

The metric $g$, its inverse, and volume form are as follows

[TABLE]

where $B=A^{-1}$ is the inverse of $A$. In our case we have $B=A$, but these matrices are conceptually distinct.

We characterize Brownian motion on $K_{n+1}$ using the auxiliary functions

[TABLE]

Here we have introduced the light-cone

[TABLE]

Pick $n$ vectors $\{b_{i}\}_{i=1}^{n}\subset L_{n+1}^{+}$ and define the $n+1$ functions

[TABLE]

We choose a drift and covariance tensor as follows

[TABLE]

[TABLE]

Finally, set $\Sigma^{00}=1$ and $\Sigma^{ij}=0$ when exactly one of the indices is zero.

Theorem 3.1

Denote by $\boldsymbol{B}_{t}$ Brownian motion on $(K_{n+1},g)$ with $\beta=2$. The stochastic processes $f^{i}_{t}:=f^{i}(\boldsymbol{B}_{t})$ satisfy the Itô SDE

[TABLE]

In particular, each $f^{i}_{t}$ is itself identical in law with a Brownian motion with constant drift.

Proof

We only need to check that

[TABLE]

First, when $i=j=0$, this is just the claim of Theorem 2.1. When exactly one of the indices is zero, we use equation (15) and the fact that $\frac{(b^{T}x)^{2}}{x^{T}Ax}$ is a homogeneous polynomial of order [math] to obtain

[TABLE]

Finally, consider the case when both $i$ and $j$ are space-like. We start with the following property: for $b,b^{\prime}\in L_{n+1}^{+}$,

[TABLE]

Using the fact that $\log(b^{T}x)=\frac{1}{\sqrt{n+1}}(f_{b}-f^{0})$ and $\nabla f_{b}\cdot\nabla f^{0}=0$, we have

[TABLE]

In particular, we find that $|\nabla f_{b}|^{2}=1$ because $b^{T}Bb=0$ when $b\in L_{n+1}^{+}$.

The proof of the first identity in equation (31) is a computation:

[TABLE]

Thus, finally we have

[TABLE]

4 Acknowledgements

This work was supported by NSF grant DMS-2107205.