Disentangling feature and lazy training in deep neural networks

Abstract

Abstract Two distinct limits for deep learning have been derived as the network width h → ∞, depending on how the weights of the last layer scale with h . In the neural tangent Kernel (NTK) limit, the dynamics becomes linear in the weights and is described by a frozen kernel Θ (the NTK). By contrast, in the mean-field limit, the dynamics can be expressed in terms of the distribution of the parameters associated with a neuron, that follows a partial differential equation. In this work we consider deep networks where the weights in the last layer scale as αh −1/2 at initialization. By varying α and h , we probe the crossover between the two limits. We observe two the previously identified regimes of ‘lazy training’ and ‘feature training’. In the lazy-training regime, the dynamics is almost linear and the NTK barely changes after initialization. The feature-training regime includes the mean-field formulation as a limiting case and is characterized by a kernel that evolves in time, and thus learns some features. We perform numerical experiments on MNIST, Fashion-MNIST, EMNIST and CIFAR10 and consider various architectures. We find that: (i) the two regimes are separated by an α * that scales as <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" display="inline" overflow="scroll"> <mml:mfrac> <mml:mrow> <mml:mn>1</mml:mn> </mml:mrow> <mml:mrow> <mml:msqrt> <mml:mrow> <mml:mi>h</mml:mi> </mml:mrow> </mml:msqrt> </mml:mrow> </mml:mfrac> </mml:math> . (ii) Network architecture and data structure play an important role in determining which regime is better: in our tests, fully-connected networks perform generally better in the lazy-training regime, unlike convolutional networks. (iii) In both regimes, the fluctuations δF induced on the learned function by initial conditions decay as <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" display="inline" overflow="scroll"> <mml:mi>δ</mml:mi> <mml:mi>F</mml:mi> <mml:mo>∼</mml:mo> <mml:mn>1</mml:mn> <mml:mo>/</mml:mo> <mml:msqrt> <mml:mrow> <mml:mi>h</mml:mi> </mml:mrow> </mml:msqrt> </mml:math> , leading to a performance that increases with h . The same improvement can also be obtained at an intermediate width by ensemble-averaging several networks that are trained independently. (iv) In the feature-training regime we identify a time scale <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" display="inline" overflow="scroll"> <mml:msub> <mml:mrow> <mml:mi>t</mml:mi> </mml:mrow> <mml:mrow> <mml:mn>1</mml:mn> </mml:mrow> </mml:msub> <mml:mo>∼</mml:mo> <mml:msqrt> <mml:mrow> <mml:mi>h</mml:mi> </mml:mrow> </mml:msqrt> <mml:mi>α</mml:mi> </mml:math> , such that for t ≪ t 1 the dynamics is linear. At t ∼ t 1 , the output has grown by a magnitude <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" display="inline" overflow="scroll"> <mml:msqrt> <mml:mrow> <mml:mi>h</mml:mi> </mml:mrow> </mml:msqrt> </mml:math> and the changes of the tangent kernel | |ΔΘ| | become significant. Ultimately, it follows <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" display="inline" overflow="scroll"> <mml:mo stretchy="false">|</mml:mo> <mml:mspace width="-0.17em"/> <mml:mo stretchy="false">|</mml:mo> <mml:mi mathvariant="normal">Δ</mml:mi> <mml:mi mathvariant="normal">Θ</mml:mi> <mml:mo stretchy="false">|</mml:mo> <mml:mspace width="-0.17em"/> <mml:mo stretchy="false">|</mml:mo> <mml:mo>∼</mml:mo> <mml:msup> <mml:mrow> <mml:mrow> <mml:mo stretchy="false">(</mml:mo> <mml:mrow> <mml:msqrt> <mml:mrow> <mml:mi>h</mml:mi> </mml:mrow> </mml:msqrt> <mml:mi>α</mml:mi> </mml:mrow> <mml:mo stretchy="false">)</mml:mo> </mml:mrow> </mml:mrow> <mml:mrow> <mml:mo>−</mml:mo> <mml:mi>a</mml:mi> </mml:mrow> </mml:msup> </mml:math> for ReLU and Softplus activation functions, with a &lt; 2 and a → 2 as depth grows. We provide scaling arguments supporting these findings.

Topics & Concepts