Search papers, labs, and topics across Lattice.
This paper introduces an analytical formalism to decompose the Hessian of neural networks with DAG architectures into blocks representing inter-layer interactions, separating the Gauss-Newton component from the tensor component. The analysis reveals that for ReLU activations, the tensor component of the input Hessian vanishes, while the full parametric Hessian contains residual terms beyond the GGN approximation. The authors introduce diagnostic metrics like inter-layer resonance and geometric coupling, estimated stochastically, to reveal structural curvature interactions, demonstrating their utility on MLPs and ResNet-18.
Forget monolithic Hessians: this work unveils a new lens into neural network training by decomposing the Hessian into layer-wise interactions, revealing curvature structures previously hidden.
Modern automatic differentiation frameworks (JAX, PyTorch) return the Hessian of the loss function as a monolithic tensor, without exposing the internal structure of inter-layer interactions. This paper presents an analytical formalism that explicitly decomposes the full Hessian into blocks indexed by the DAG of an arbitrary architecture. The canonical decomposition $H = H^{GN} + H^T$ separates the Gauss--Newton component (convex part) from the tensor component (residual curvature responsible for saddle points). For piecewise-linear activations (ReLU), the tensor component of the input Hessian vanishes ($H^{T}_{v,w}\!\equiv\!0$ a.e., $H^f_{v,w}\!=\!H^{GN}_{v,w}\!\succeq\!0$); the full parametric Hessian contains residual terms that do not reduce to the GGN. Building on this decomposition, we introduce diagnostic metrics (inter-layer resonance~$\mathcal{R}$, geometric coupling~$\mathcal{C}$, stable rank~$\mathcal{D}$, GN-Gap) that are estimated stochastically in $O(P)$ time and reveal structural curvature interactions between layers. The theoretical analysis explains exponential decay of resonance in vanilla networks and its preservation under skip connections; empirical validation spans fully connected MLPs (Exp.\,1--5) and convolutional architectures (ResNet-18, ${\sim}11$M~parameters, Exp.\,6). When the architecture reduces to a single node, all definitions collapse to the standard Hessian $\nabla^2_θ\mathcal{L}(θ)\in\mathbb{R}^{p\times p}$.