Reproduce “On the Superlinear Relationship between SGD Noise Covariance and Loss Landscape Curvature” on OpenResearch

About this paper

Yikuan Zhang, Ning Yang, Yuhai Tu

This research introduces a general theoretical framework based on Activity-Weight Duality to characterize the relationship between SGD noise covariance and the loss landscape curvature, finding a superlinear dependence for cross-entropy loss that resolves previous dimensional inconsistencies. The study empirically shows $C_{ii} \propto H_{ii}^\gamma$ with $\gamma \approx 1$ for MSE and $\gamma \in [1.08, 1.51]$ for cross-entropy, providing theoretical bounds of $1 \le \gamma \le 2$.

Key insights

• The SGD noise covariance (C) and the loss Hessian (H) exhibit strong, approximate alignment, as indicated by a characteristic 'arrowhead' structure when C is represented in the H eigenbasis.
• The relationship between C and H is generally superlinear ($C \propto H^\gamma$) with a theoretically derived lower bound of $\gamma \ge 1$ and an upper bound of $\gamma \le 2$.
• The superlinear scaling for cross-entropy loss ($\gamma > 1$) stems from an entanglement between the magnitude of dominant per-sample Hessian eigenvalues and their alignment with global Hessian eigenvectors, a geometric property not observed in MSE loss (where $\gamma \approx 1$).

Summary by alphaXiv.