[Paper Review] On the linearity of large non-linear models: when and why the tangent kernel is constant

Chaoyue Liu, Libin Zhu+1 more|arXiv (Cornell University)|Oct 2, 2020

Stochastic Gradient Optimization Techniques20 references33 citations

TL;DR

The paper explains why wide non-linear neural networks can become linear in parameters as width grows, due to Hessian scaling, and clarifies when tangent kernels stay constant or not.

ABSTRACT

The goal of this work is to shed light on the remarkable phenomenon of transition to linearity of certain neural networks as their width approaches infinity. We show that the transition to linearity of the model and, equivalently, constancy of the (neural) tangent kernel (NTK) result from the scaling properties of the norm of the Hessian matrix of the network as a function of the network width. We present a general framework for understanding the constancy of the tangent kernel via Hessian scaling applicable to the standard classes of neural networks. Our analysis provides a new perspective on the phenomenon of constant tangent kernel, which is different from the widely accepted "lazy training". Furthermore, we show that the transition to linearity is not a general property of wide neural networks and does not hold when the last layer of the network is non-linear. It is also not necessary for successful optimization by gradient descent.

Motivation & Objective

• Clarify the relationship between tangent kernel constancy and model linearity in wide neural networks.
• Identify the scaling properties of the Hessian that drive the transition to linearity as width grows.
• Establish conditions under which the tangent kernel remains constant or fails to be constant.
• Differentiate this phenomenon from the lazy training narrative and explain its limitations.
• Provide theoretical and experimental insights into how network architecture affects NTK behavior.

Proposed method

• Define and analyze the tangent kernel K(w) = ∇w f(w; x)^T ∇w f(w; z) and its relation to linearity.
• Show that f is linear in w if and only if the tangent kernel is constant (Proposition 2.2).
• Derive a sufficient condition for near-constant tangent kernel via small Hessian norm (Proposition 2.3).
• Develop a general Hessian scaling framework that links Hessian norms to infinity-norms of layer-wise derivatives and (2,1,1)-norms of higher-order tensors.
• Prove a bound on the Hessian for deep networks showing it can be controlled by infinity-norms of partial derivatives and by a 1/√m factor from the output layer (Theorem 3.1).
• Demonstrate that non-constancy of the tangent kernel can occur when the last layer is non-linear or when a bottleneck is inserted, and discuss optimization implications.

Experimental results

Research questions

• RQ1 Under what conditions does the tangent kernel stay constant as network width grows?
• RQ2 How does Hessian scaling relate to the width and architecture of the network to produce linearity?
• RQ3 Does the transition to linearity depend on lazy training or model rescaling, and when does it fail to occur?
• RQ4 What architectural features disrupt near-linearity and NTK constancy, and how does this affect gradient-based optimization?

Key findings

• The tangent kernel is constant if and only if the model is linear in its parameters.
• For networks with linear output layers, the Hessian norm scales down with width, leading to vanishing Hessian and constant tangent kernel in the infinite-width limit.
• The Hessian spectral norm is governed by the infinity-norms of layer-wise derivatives, while the gradient and tangent kernel are governed by 2-norms, creating a norm-disparity that drives linearity.
• Constancy of the tangent kernel is not universal for wide networks and can fail when the last layer is non-linear or when a bottleneck is present.
• Even without linearity, gradient descent can efficiently optimize wide networks, indicating that NTK constancy is not a strict prerequisite for successful optimization.