Mean Field Models for Neural Networks in Teacher-student Setting

Lexing Ying, Yuandong Tian

Keywords: mean field model, optimal transport, ResNet

TL;DR: We discuss mean field models for two-layer fully-connected networks and ResNet models and characterize stationary distributions in the teacher-student setting.

Abstract: Mean field models have provided a convenient framework for understanding the training dynamics for certain neural networks in the infinite width limit. The resulting mean field equation characterizes the evolution of the time-dependent empirical distribution of the network parameters. Following this line of work, this paper first focuses on the teacher-student setting. For the two-layer networks, we derive the necessary condition of the stationary distributions of the mean field equation and explain an empirical phenomenon concerning training speed differences using the Wasserstein flow description. Second, we apply this approach to two extended ResNet models and characterize the necessary condition of stationary distributions in the teacher-student setting.