Hamiltonian Mechanics of Feature Learning: Bottleneck Structure in Leaky ResNets

TL;DR

This paper models Leaky ResNets using Hamiltonian mechanics to explain the emergence of bottleneck structures in feature learning, revealing a balance between kinetic and potential energies that governs representation dynamics.

Contribution

It introduces a Hamiltonian framework for analyzing Leaky ResNets, providing new insights into the feature learning process and the formation of bottleneck structures.

Findings

Hamiltonian reformulation highlights key forces in feature learning

Bottleneck structure explained by separation of timescales

Adaptive layer step-size improves training efficiency

Abstract

We study Leaky ResNets, which interpolate between ResNets and Fully-Connected nets 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.