[Paper Review] Generalization bound of globally optimal non-convex neural network training: Transportation map estimation by infinite dimensional Langevin dynamics

Taiji Suzuki|arXiv (Cornell University)|Jul 11, 2020

Neural Networks and Applications4 citations

TL;DR

This paper proposes a novel framework for analyzing deep neural network training by formulating weight optimization as transportation map estimation via infinite-dimensional Langevin dynamics, enabling global convergence and fast generalization error bounds for both finite and infinite width networks. It achieves a fast learning rate, including exponential convergence for classification and minimax optimal rates for regression.

ABSTRACT

We introduce a new theoretical framework to analyze deep learning optimization with connection to its generalization error. Existing frameworks such as mean field theory and neural tangent kernel theory for neural network optimization analysis typically require taking limit of infinite width of the network to show its global convergence. This potentially makes it difficult to directly deal with finite width network; especially in the neural tangent kernel regime, we cannot reveal favorable properties of neural networks beyond kernel methods. To realize more natural analysis, we consider a completely different approach in which we formulate the parameter training as a transportation map estimation and show its global convergence via the theory of the infinite dimensional Langevin dynamics. This enables us to analyze narrow and wide networks in a unifying manner. Moreover, we give generalization gap and excess risk bounds for the solution obtained by the dynamics. The excess risk bound achieves the so-called fast learning rate. In particular, we show an exponential convergence for a classification problem and a minimax optimal rate for a regression problem.

Motivation & Objective

• Address the gap in theoretical analysis of deep learning generalization by unifying finite-width and infinite-width network analysis.
• Overcome limitations of existing frameworks like mean field theory and neural tangent kernel (NTK) that require infinite width for global convergence.
• Provide a generalization error bound that achieves a fast learning rate, avoiding the limitations of kernel methods in the NTK regime.
• Enable analysis of non-convex optimization in deep networks without relying on over-parameterization or dimension-dependent convergence rates.
• Bridge the theoretical understanding of optimization and generalization in deep learning by connecting to nonparametric Bayesian estimation.

Proposed method

• Formulate deep learning training as a transportation map estimation problem in the parameter space, treating network weights as elements of a reproducing kernel Hilbert space (RKHS).
• Apply infinite-dimensional Langevin dynamics to solve the transportation map estimation, leveraging stochastic dynamics to ensure global convergence.
• Use the theory of McKean–Vlasov dynamics and ergodicity to establish global convergence independent of network width.
• Relate the solution to a nonparametric Bayesian Gaussian process estimator to derive generalization error bounds.
• Employ a student-teacher setup and RKHS-based regularization to bound bias and variance terms in the generalization error.
• Derive convergence rates by analyzing the interplay between the regularization parameter, sample size, and eigenvalue decay of the kernel operator.

Experimental results

Research questions

• RQ1 Can global convergence of non-convex deep learning optimization be established without requiring infinite width?
• RQ2 Can generalization error bounds be derived that achieve fast learning rates in a unified framework for both narrow and wide networks?
• RQ3 How does the proposed Langevin dynamics-based approach compare to NTK and mean field theories in terms of generalization and convergence?
• RQ4 What is the relationship between the solution of infinite-dimensional Langevin dynamics and nonparametric Bayesian estimation?
• RQ5 Can the framework achieve minimax optimal rates for regression and exponential convergence for classification?

Key findings

• The proposed framework achieves global convergence for both finite-width and infinite-width neural networks using infinite-dimensional Langevin dynamics, with a convergence rate independent of network size.
• For regression, the excess risk bound achieves the minimax optimal rate, confirming statistical efficiency.
• For classification, the generalization error converges exponentially fast under the student-teacher setup, indicating strong generalization performance.
• The generalization error bound achieves a fast learning rate, with the excess risk decaying as $ \max\{\lambda^{-\tilde{\alpha}}\beta^{-1}, \lambda^{\theta}, n^{-\frac{1}{2-s}}\} $, which reflects favorable dependence on sample size and regularization.
• The framework unifies the analysis of narrow and wide networks by avoiding the need for infinite-width limits, unlike mean field and NTK theories.
• The solution of the Langevin dynamics is shown to be closely related to a nonparametric Bayesian Gaussian process estimator, enabling tight generalization bounds.