Hamiltonian Mechanics of Feature Learning: Bottleneck Structure in Leaky ResNets

Arthur Jacot, Alexandre Kaiser

Keywords: Feature Learning, Bottleneck Structure, NeuralODE, Hamiltonian mechanics

TL;DR: The path of representations from input to output, optimizes a kinetic energy that favors `short' paths, and a potential energy that favors low dimensional representations.

Abstract:

We study Leaky ResNets, which interpolate between ResNets ($\tilde{L}=0$)
and Fully-Connected nets ($\tilde{L}\to\infty$) depending on an 'effective
depth' hyper-parameter $\tilde{L}$. In the infinite depth limit,
we study 'representation geodesics' $A_{p}$: continuous paths in
representation space (similar to NeuralODEs) from input $p=0$ to
output $p=1$ that minimize the parameter norm of the network. We
give a Lagrangian and Hamiltonian reformulation, which highlight the
importance of two terms: a kinetic energy which favors small layer
derivatives $\partial_{p}A_{p}$ and a potential energy that favors
low-dimensional representations, as measured by the 'Cost of Identity'.
The balance between these two forces offers an intuitive understanding
of feature learning in ResNets. We leverage this intuition to explain
the emergence of a bottleneck structure, as observed in previous work:
for large $\tilde{L}$ the potential energy dominates and leads to
a separation of timescales, where the representation jumps rapidly
from the high dimensional inputs to a low-dimensional representation,
move slowly inside the space of low-dimensional representations, before
jumping back to the potentially high-dimensional outputs. Inspired
by this phenomenon, we train with an adaptive layer step-size
to adapt to the separation of timescales.

Primary Area: Learning theory

Submission Number: 19362