PaperlayerView exact revision on arXiv

Source-linked AI summary

Deep Relaxation: partial differential equations for optimizing deep neural networks

Pratik Chaudhari, Adam Oberman, Stanley Osher, Stefano Soatto, Guillaume Carlier

arXiv:1704.04932v2cs.LGmath.APmath.OC

TL;DR

Training deep neural networks requires optimizing large, non-convex losses, motivating alternatives to standard SGD and a clearer account of successful relaxation methods. The paper interprets local-entropy relaxation through viscous Hamilton-Jacobi PDEs and stochastic homogenization, then proves improved expected loss for modified dynamics and derives regularity-based explanations of landscape smoothing.

  • Problem

    Deep-network training uses SGD on non-convex losses, while the mechanisms and theoretical guarantees behind effective relaxation methods remain insufficiently understood.

  • Method

    The paper connects local entropy to viscous Hamilton-Jacobi PDEs, derives its gradient through stochastic homogenization, and analyzes the resulting dynamics using PDE regularity and stochastic control.

  • Results

    The modified local-entropy dynamics provably improve the expected value of the original loss relative to SGD, while PDE estimates quantify smoothing and regularity of the relaxed landscape.

  • Takeaways & Limitations

    PDEs and stochastic analysis provide a unified interpretation of relaxation algorithms and contribute to designing improved optimization methods for neural networks.

  • Takeaways & Limitations

    The paper leaves an alternate Wasserstein variational interpretation of the distributed algorithm for future study.

Abstract

from arXiv · show

In this paper we establish a connection between non-convex optimization methods for training deep neural networks and nonlinear partial differential equations (PDEs). Relaxation techniques arising in statistical physics which have already been used successfully in this context are reinterpreted as solutions of a viscous Hamilton-Jacobi PDE. Using a stochastic control interpretation allows we prove that the modified algorithm performs better in expectation that stochastic gradient descent. Well-known PDE regularity results allow us to analyze the geometry of the relaxed energy landscape, confirming empirical evidence. The PDE is derived from a stochastic homogenization problem, which arises in the implementation of the algorithm. The algorithms scale well in practice and can effectively tackle the high dimensionality of modern neural networks.

1. INTRODUCTION

The paper reframes local-entropy relaxation for deep-network optimization through nonlinear PDEs and stochastic dynamics. This framework yields smoothing, algorithmic equivalences, geometric regularity results, and a proof of improved expected loss over SGD.

  • The paper studies a modification of SGD for non-convex neural-network loss functions, targeting training time and generalization error.
  • Local entropy is identified as the solution of a viscous Hamilton-Jacobi PDE, giving a PDE interpretation of the relaxation.
  • Homogenization of an auxiliary SDE recovers the local-entropy gradient in the fast-dynamics limit, while exponentially convergent local dynamics make smoothing practical in high dimensions.
  • The framework proves that Elastic-SGD is equivalent to regularization by local entropy, because ergodicity connects temporal and spatial averages.
  • Scoping applies more smoothing early and recovers the original loss gradients and minima near termination, with a PDE interpretation through nonlinear forward-backward equations.
  • Stochastic-control and viscosity-solution arguments prove an improvement in expected loss over standard SGD for slightly modified dynamics.

2. BACKGROUND

The background presents deep-network training as large-scale non-convex optimization, where first-order stochastic methods are practical but can suffer from metastability. Smoothing is introduced as a principal strategy for improving optimization on such landscapes.

  • Deep networks are nested compositions with millions of parameters, and supervised training minimizes a generally non-convex empirical loss.
  • SGD updates parameters using randomly sampled data or mini-batches, avoiding computation of the full gradient at every iteration.
  • For non-convex losses, SGD can have an O(k^-1/2) convergence rate under unbiased bounded-variance stochastic gradients.
  • The continuous-time SGD model leads to generator, backward Kolmogorov, and Fokker-Planck equations describing expected values and evolving parameter densities.
  • Although the Gibbs distribution is the unique steady state under mild assumptions, non-convex dynamics may converge exponentially slowly through metastability.
  • Smoothing is presented as an effective way to improve optimization, and the paper develops a unified mathematical framework for comparing smoothing techniques.

3. PDE INTERPRETATION OF LOCAL ENTROPY

Local entropy is connected to viscous and non-viscous Hamilton-Jacobi equations. These PDE formulations provide smoothing, proximal representations, gradient dynamics, and tractable local approximations near quadratic minima.

  • Local entropy is a modified loss function associated with dense solution clusters in shallow neural-network models.
  • The local entropy function fγ(x)=u(x,γ) solves the viscous Hamilton-Jacobi initial-value problem with initial data f(x).
  • The non-viscous Hamilton-Jacobi equation is the zero-viscosity limit of the viscous equation, and its solution is given by the Hopf-Lax inf-convolution formula.
  • The Hopf-Lax gradient can be expressed through a proximal operator and, under λ-convexity, computed by exponentially convergent gradient dynamics.
  • Local entropy computes a non-linear neighborhood average of gradients using the steady-state measure ρ∞(y;x).
  • For quadratic losses, the invariant measure is Gaussian with mean μ=x−Σp and covariance Σ=(Q+γI)^−1, with a Neumann-series approximation that avoids Hessian inversion.

4. DERIVATION OF LOCAL ENTROPY VIA HOMOGENIZATION OF SDES

Homogenization analyzes fast–slow stochastic dynamics and shows that Entropy-SGD converges to gradient descent on local entropy. The same framework establishes Elastic-SGD's equivalence to local entropy and contrasts its smoothing with heat-equation dynamics.

  • 4.3. Elastic-SGD as local entropy: Homogenization rigorously shows that Elastic-SGD is equivalent to local entropy, replacing temporal averages with spatial averages across workers.Elastic-SGD was originally connected to local entropy heuristically; ergodicity supplies the averaging connection.
  • 4.1. Background on homogenization: Homogenization averages fast variables against their invariant measure to derive effective dynamics for slowly evolving variables.The averaged vector field also equals a long-term temporal average by ergodicity.
  • 4.2. Derivation of local entropy via homogenization of SDEs: Theorem 5 states that as ε →0, Entropy-SGD converges to homogenized dynamics determined by the fast-variable solution.The homogenized vector field is obtained from the invariant measure of the fast dynamics.
  • 4.2. Derivation of local entropy via homogenization of SDEs: The homogenized dynamics converge to gradient descent on fγ, the local entropy objective.This connects the auxiliary stochastic dynamics used by Entropy-SGD to optimization of the relaxed loss.
  • 4.4. Heat equation versus the viscous Hamilton-Jacobi equation: Local entropy corresponds to viscous Hamilton-Jacobi smoothing, whereas heat-equation dynamics perform Gaussian averaging of gradients.The paper reports better empirical performance and improved theoretical convergence for local entropy than for heat-equation dynamics; at critical points, the operators coincide.

5. STOCHASTIC CONTROL INTERPRETATION

The viscous Hamilton-Jacobi formulation gives local entropy an optimal-control interpretation. Using stochastic control and viscosity-solution comparison, the paper proves an expected-value improvement over stochastic gradient descent for modified dynamics.

  • 5. Stochastic control interpretation: Local entropy's viscous Hamilton-Jacobi representation permits interpreting gradient descent as an optimal control problem.The paper notes that this interpretation itself has no immediate algorithmic implications.
  • 5.1. Stochastic control formulation: The controlled SDE adds a control term to stochastic gradient dynamics, with a quadratic running cost and terminal cost V(x).The value function minimizes expected cost over admissible controls and paths.
  • 5.1. Stochastic control formulation: Dynamic programming identifies the value function as the viscosity solution of a Hamilton-Jacobi-Bellman PDE.The optimal control is α(x,t) = ∇u(x,t), where u solves the HJB equation with terminal data.
  • 5.1. Stochastic control formulation: The viscous Hamilton-Jacobi equation is recovered by dropping the ∇f term from the controlled dynamics.Retaining that term enables the comparison principle used later in the proof.
  • 5.2. Improvement in the value function: Theorem 12 compares controlled-gradient dynamics with SGD from the same initial data and proves an improvement in the expected value function.The comparison uses stochastic control techniques and the comparison principle for viscosity solutions.

6. REGULARIZATION, WIDENING AND SEMI-CONCAVITY

The PDE formulation smooths the loss landscape and provides estimates showing how local minima widen, with curvature-sensitive behavior and limits over time.

  • The PDE formulation produces a smoother loss function and enables quantitative semiconcavity estimates.These estimates apply to both viscous and non-viscous Hamilton-Jacobi equations.
  • Semiconcavity bounds imply that high-curvature local minima widen faster than flatter ones under viscous-HJ evolution.At a local minimum, no Hessian eigenvalue exceeds the semiconcavity constant.
  • For quadratic initial data, the semiconcavity constant is C(t)=1/(c^-1+t), so larger initial curvature yields faster improvement.For c=10^8, the passage reports C(t=.01)≈10.
  • For short times, Hamilton-Jacobi evolution preserves local minima while widening surrounding convex regions.In one dimension, the convexity interval expands as inflection points move outward.
  • In higher dimensions, convexity is preserved under inf-convolution, local minima persist briefly, and some minima vanish at longer times.A local minimum persists for short times when the minimizer displacement is O(√t).
  • The harmonic mean of Hessian eigenvalues better captures nearly flat directions than the arithmetic mean.A typical network has arithmetic mean 0.0029 and harmonic mean approximately 10^-10.

7. ALGORITHMIC DETAILS

The algorithms discretize auxiliary stochastic dynamics with different noise, averaging, iteration, and momentum choices, while gradually reducing the smoothing parameter.

  • The implementation compares Entropy-SGD, HJ, and related discretizations through auxiliary y dynamics and gradient updates.Euler-Maruyama discretizes the y dynamics, with Gaussian noise in the updates.
  • Mini-batch gradients and extrinsic noise provide two separate noise sources in the practical dynamics.The combined equation includes coefficients for mini-batch and extrinsic noise.
  • Entropy-SGD uses L=20, β^-1_ex=10^8, exponential averaging, and α=0.75 in the reported experiments.The forward-looking average gives more weight to later iterates when invariant-measure convergence is slow.
  • Momentum is implemented with an auxiliary velocity variable and fixed parameter δ=0.9.The momentum modification is applied to the x updates and related algorithmic updates.
  • Scoping starts with larger smoothing and reduces γ toward the original loss during training.The paper states that fγ(x) approaches f(x) as γ→0.

8. EMPIRICAL VALIDATION

Experiments on deep neural networks evaluate whether the PDE-based methods regularize effectively, aid optimization, and improve classification on modern datasets.

  • The empirical validation tests PDE-based methods for regularization, optimization, and classification performance on deep neural networks.

8.1. Setup for deep networks.

The experiments use MNIST and CIFAR-10 classification with cross-entropy loss, compare PDE-based algorithms against SGD, and report variation across random seeds.

  • MNIST contains 70,000 grayscale 28×28 digit images, split into 60,000 training and 10,000 validation images.
  • CIFAR-10 contains 60,000 RGB 32×32 images across 10 object classes, split into 50,000 training and 10,000 validation images.
  • CIFAR-10 is processed with global contrast normalization followed by ZCA whitening, while MNIST receives no preprocessing.
  • All experiments use cross-entropy loss and compare Entropy-SGD, HEAT, HJ, and SGD.SGD is the baseline for performance comparisons.
  • Results report means and standard deviations across 6 random seeds because the algorithms sample mini-batches stochastically.Entropy-SGD adds noise beyond the intrinsic mini-batch noise, whereas HJ does not.
  • SGD uses one back-propagation per update, compared with L=20 for Entropy-SGD and heat-equation methods and L=5 for HJ.

8.2. MNIST.

On MNIST, the methods achieve similar final validation errors, while their convergence rates differ across fully connected and convolutional networks. Entropy-SGD converges fastest for the fully connected network.

  • The fully connected MNIST network uses dense layers with ReLU, batch normalization, and dropout.
  • Final validation errors are similar across algorithms on mnistfc, with Entropy-SGD at 1.08% versus HJ at 1.17%.Because mnistfc is small, the performance difference is minor.
  • LeNet is a convolutional network that performs better than mnistfc on MNIST.LeNet has 131,220 parameters.
  • LeNet reaches approximately 0.50% final validation error for most algorithms, while the heat equation reaches 0.59%.The other algorithms converge in about half the effective epochs required by mnistfc.

8.3. CIFAR.

On CIFAR-10, the All-CNN experiment shows that Hamilton-Jacobi methods outperform the tested alternatives, with HJ achieving the best validation error and training loss while converging faster than SGD.

  • The CIFAR-10 experiment uses the All-CNN-C architecture with batch normalization because fully connected networks typically perform poorly on this more complex dataset.
  • The network downsamples images through strided convolution rather than max-pooling and contains 1.67 million parameters.
  • 7.89% validation error is achieved by HJ in 145 epochs, making it the best result among the tested algorithms.
  • HJ also obtains the lowest training cross-entropy loss at 0.046, while HJ and Entropy-SGD converge faster than SGD.The heat equation performs worse, reaching 9.04% validation error.

9. DISCUSSION

The discussion connects neural-network optimization with nonlinear PDEs, stochastic control, and homogenization. It interprets smoothing as a mechanism for improving non-convex optimization and motivates practical algorithmic refinements.

  • The analysis applies nonlinear PDEs, stochastic optimal control, and stochastic homogenization to two effective neural-network optimization algorithms.
  • Entropy-SGD replaces SGD on f(x) with SGD on H(x,y;γ), and homogenization connects its dynamics to a viscous Hamilton-Jacobi PDE.The paper states that this modification can quantitatively improve the expected original loss.
  • The analysis clarifies hyperparameters, interprets γ as PDE time, and rigorously justifies decreasing γ during training through scoping.
  • Unlike simulated annealing, which modulates noise, these algorithms modulate the smoothness of the underlying loss function for non-convex optimization.
Loading 1704.04932v2…