Understanding the dynamics of the frequency bias in neural networks

TL;DR

This paper develops a PDE model to analyze the frequency bias in neural networks, showing how initialization distributions can control this bias, validated through experiments on Fourier Features and multi-layer NNs.

Contribution

It introduces a PDE framework for understanding frequency dynamics in neural networks and demonstrates how initialization can mitigate frequency bias.

Findings

PDE accurately models frequency learning dynamics

Initialization distributions influence frequency bias control

Method extends from 2-layer to multi-layer NNs

Abstract

Recent works have shown that traditional Neural Network (NN) architectures display a marked frequency bias in the learning process. Namely, the NN first learns the low-frequency features before learning the high-frequency ones. In this study, we rigorously develop a partial differential equation (PDE) that unravels the frequency dynamics of the error for a 2-layer NN in the Neural Tangent Kernel regime. Furthermore, using this insight, we explicitly demonstrate how an appropriate choice of distributions for the initialization weights can eliminate or control the frequency bias. We focus our study on the Fourier Features model, an NN where the first layer has sine and cosine activation functions, with frequencies sampled from a prescribed distribution. In this setup, we experimentally validate our theoretical results and compare the NN dynamics to the solution of the PDE using the finite element method. Finally, we empirically show that the same principle extends to multi-layer NNs.