Hamiltonian Mechanics of Feature Learning: Bottleneck Structure in Leaky ResNets

TL;DR

The paper analyzes Leaky ResNets as continuous-depth systems with an effective depth $\tilde{L}$ and studies the emergent feature-learning dynamics via Lagrangian and Hamiltonian formalisms. It introduces representation geodesics $A_p$ and decomposes the optimization into a kinetic term and a Cost of Identity (COI) term that measures representation dimensionality, explaining how large $\tilde{L}$ drives a bottleneck: rapid transitions to a low-dimensional space, slow evolution within it, then a rapid return to high-dimensional outputs. A Hamiltonian formulation with momenta $B_p$ formalizes the conserved energy $\mathcal{H}$ and clarifies the separation of timescales, while a gamma-stabilized version resolves pseudo-inverse instabilities. Practical discretization schemes, including adaptive layer stepping, are proposed to align layer density with the bottleneck, enabling efficient training and revealing a robust bottleneck rank $k^*$. Overall, the work provides a mechanistic view of bottleneck formation in deep nets and a principled path to adaptive training strategies rooted in Hamiltonian dynamics.

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.

Figures (3)

• Figure 1: Leaky ResNet structures: We train adaptive networks with a fixed $L=50$ over a range of effective depths $\tilde{L}$. The true function $f^*:\mathbb{R}^{20}\to\mathbb{R}^{20}$ is the composition of two random FCNNs $g_1,g_2$ mapping from dim. 20 to 5 to 20, the network recovers the true rank of $k^*=5$. (a) Estimates of the Hamiltonian constants for networks trained with different $\tilde{L}$. The Hamiltonian refers to $-\frac{2}{\tilde{L}}\mathcal{H}$ which estimates the true rank $k^*$. The COI refers to $\min_p ||A_p||_{K_p}$. The trend line follows the median estimate for $-\frac{2}{\tilde{L}}\mathcal{H}$ across each network's layers, whereas the error bars signify its minimum and maximum over $p\in[0,1]$. The "stable" Hamiltonians utilize the relaxation from Theorem \ref{['thm:stable_energy_decomposition']}. (b,c,d) Top: The 10 largest singular values of $W_p$ throughout the layers, exhibiting a BN structure. Bottom: the rescaled Hamiltonian, stable Hamiltonian, COI and kinetic energy. The Hamiltonian remains constant throughout the layers, and the stable Hamiltonian approximates it well - except in the first layers, where the kinetic energy (and COI) blows up which is in line with the bound of Equation \ref{['eq:approx_hamiltonians']} in Theorem \ref{['thm:stable_energy_decomposition']}. Inside the bottleneck, the kinetic energy approaches zero and the COI approaches $k^*$.
• Figure 2: Discretization: We train networks with a fixed $\tilde{L}=3$ over a range of depths $L$ and definitions of $\rho_\ell$s. The true function $f^*:\mathbb{R}^{30}\to\mathbb{R}^{30}$ is the composition of three random ResNets $g_1,g_2,g_3$ mapping from dim. 30 to 6 to 3 to 30. (a) Test error as a function of $L$ for different discretization schemes. (b) Weight spectra across layers for adaptive $\rho_\ell$ ($L=18$), grey vertical lines represents the steps $p_\ell$. The Bottleneck structure is more complex for this task, reflecting the more complex true function $f^*=g_3\circ g_2 \circ g_1$: we see a dimension 3 bottleneck around $p=0.7$, but the network remains at an intermediate dimension aroung $p=0.3$, perhaps reflecting the intermediate dimension of $6$ in the true function. (c) 2D projection of the representation paths $A_p$ for $L=18$. Observe how adaptive $\rho_\ell$s appears to better spread out the steps.
• Figure 3: Various properties of the Hamiltonian dynamics of Leaky ResNets which remain bounded

Theorems & Definitions (18)

• Remark
• Proposition 1
• proof
• Proposition 2
• Proposition 3
• Remark
• Theorem 4
• Proposition 5
• Proposition 6: Proposition \ref{['prop:stable_minima_are_positive']} in the main
• proof