Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

The Expected Loss of Preconditioned Langevin Dynamics Reveals the Hessian Rank

Amitay Bar, Rotem Mulayoff, Tomer Michaeli, Ronen Talmon

TL;DR

This work derives a closed-form expression for the expected loss of preconditioned LD near stationary points of the objective function from the fact that at the vicinity of such points, LD reduces to an Ornstein–Uhlenbeck process, which is amenable to convenient mathematical treatment.

Abstract

Langevin dynamics (LD) is widely used for sampling from distributions and for optimization. In this work, we derive a closed-form expression for the expected loss of preconditioned LD near stationary points of the objective function. We use the fact that at the vicinity of such points, LD reduces to an Ornstein-Uhlenbeck process, which is amenable to convenient mathematical treatment. Our analysis reveals that when the preconditioning matrix satisfies a particular relation with respect to the noise covariance, LD's expected loss becomes proportional to the rank of the objective's Hessian. We illustrate the applicability of this result in the context of neural networks, where the Hessian rank has been shown to capture the complexity of the predictor function but is usually computationally hard to probe. Finally, we use our analysis to compare SGD-like and Adam-like preconditioners and identify the regimes under which each of them leads to a lower expected loss.

The Expected Loss of Preconditioned Langevin Dynamics Reveals the Hessian Rank

TL;DR

This work derives a closed-form expression for the expected loss of preconditioned LD near stationary points of the objective function from the fact that at the vicinity of such points, LD reduces to an Ornstein–Uhlenbeck process, which is amenable to convenient mathematical treatment.

Abstract

Langevin dynamics (LD) is widely used for sampling from distributions and for optimization. In this work, we derive a closed-form expression for the expected loss of preconditioned LD near stationary points of the objective function. We use the fact that at the vicinity of such points, LD reduces to an Ornstein-Uhlenbeck process, which is amenable to convenient mathematical treatment. Our analysis reveals that when the preconditioning matrix satisfies a particular relation with respect to the noise covariance, LD's expected loss becomes proportional to the rank of the objective's Hessian. We illustrate the applicability of this result in the context of neural networks, where the Hessian rank has been shown to capture the complexity of the predictor function but is usually computationally hard to probe. Finally, we use our analysis to compare SGD-like and Adam-like preconditioners and identify the regimes under which each of them leads to a lower expected loss.
Paper Structure (9 sections, 13 theorems, 56 equations, 4 figures, 1 table, 1 algorithm)

This paper contains 9 sections, 13 theorems, 56 equations, 4 figures, 1 table, 1 algorithm.

Table of Contents

  1. Proofs
  2. Proof of Theorem \ref{['thm: E[f(x)] expression']}
  3. Proof of Proposition \ref{['thm: E[f(x)] expression for large t']}
  4. Proof of Proposition \ref{['prop: rank H and loss']}
  5. Proof of Proposition \ref{['prop: Ef(x) > Ef(y) two preconditioning']}
  6. Proof of Proposition \ref{['prop: escape time from saddle point']}
  7. Auxiliary Lemmas
  8. Connection to Hessian Trace Estimation
  9. Additional Experimental Results

Key Result

Theorem 1

The expected loss of a process governed by the SDE in (eq: OU process for dx_t) is given by

Figures (4)

  • Figure 1: The loss (in blue) for a linear NN of depth $5$ with input and output dimensions of $32$ (the number of parameters is $5120$). The theoretical expressions according to Theorem \ref{['thm: E[f(x)] expression']} and Proposition \ref{['thm: E[f(x)] expression for large t']} are in red and black, respectively.
  • Figure 2: Estimated rank of linear networks with different dimensions. Our method is in blue and the U&S method proposed in ubaru2016fast is in red. The error bars represent one standard deviation. We set the same number of iterations for both methods.
  • Figure 3: The normalized cumulative sum of the eigenvalues of the Hessian for DnCNN (blue). The solid black line and the dashed gray line represent the estimated rank of the proposed approach and the U&S method, respectively.
  • Figure 4: The loss for DnCNN with ${\bm{G}}={\bm{I}}$.

Theorems & Definitions (20)

  • Theorem 1: Expected loss over time
  • Proposition 1: Expected loss for large $t$
  • Corollary 1: Expected loss for PD Hessian
  • Corollary 2: Expected loss and Hessian rank
  • Remark 1
  • Proposition 2
  • Corollary 3: Maximal expected loss
  • Proposition 3
  • Definition 1: Escape time
  • Proposition 4: Escaping a saddle point
  • ...and 10 more