SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics

TL;DR

Introduces the 'leap' measure to quantify hierarchical complexity, analyzing SGD training time on low-dimensional functions with neural networks.

cs.LG 🔴 Advanced 2023-02-22 50 views
Emmanuel Abbe Enric Boix-Adsera Theodor Misiakiewicz
deep learning neural networks complexity theory optimization dynamics hierarchical functions

Key Findings

Methodology

This paper introduces the 'leap' metric to quantify the hierarchical structure of target functions, defined via the support set order. Analyzing two-layer neural networks trained with SGD on Gaussian and Boolean data, the authors connect the learning dynamics to saddle-to-saddle transitions, revealing a sequential support learning process. By expanding Hermite and Fourier-Walsh bases, and controlling the full complexity beyond mean-field and gradient flow approximations, the study demonstrates that the training time scales exponentially with the leap size. The approach combines theoretical analysis with empirical validation, establishing a link between the hierarchical structure of functions and the efficiency of SGD.

Key Results

  • On Gaussian data, functions with support size 1 (leap-1) are learned in Θ(d) steps, leap-4 functions in Θ(d^3), and leap-k functions in Θ(d^{max(k-1,1)}), confirming the leap metric's predictive power.
  • Experimental results show the network learns support sequentially, with saddle-to-saddle dynamics, matching theoretical time estimates.
  • The analysis extends previous work limited to leap-1, providing a comprehensive understanding of the entire learning trajectory, surpassing mean-field and continuous-time approximations.

Significance

This work uncovers the intrinsic relationship between the hierarchical complexity of target functions and the training time of neural networks. By formalizing the 'leap' measure, it offers a quantitative framework to understand how neural networks exploit low-dimensional structures efficiently. The results demonstrate that SGD can achieve near-optimal complexity bounds, matching information-theoretic lower bounds from CSQ models. Such insights deepen our understanding of feature learning, support the design of more efficient algorithms, and clarify the role of hierarchical structures in high-dimensional data processing. The findings have broad implications for theoretical foundations and practical training strategies in deep learning, especially for functions with layered support structures.

Technical Contribution

The paper pioneers the formalization of the 'leap' metric to characterize the hierarchical complexity of target functions, and rigorously connects it to the training dynamics of two-layer neural networks. It overcomes limitations of mean-field and gradient flow analyses by establishing full control over the training complexity, leveraging saddle-to-saddle support learning mechanisms. The analysis demonstrates that the entire learning trajectory is governed by the sequential acquisition of support, with the complexity scaling exponentially with leap size. It also proves that SGD achieves the same complexity bounds as optimal CSQ algorithms, establishing a theoretical equivalence. These contributions provide new mathematical tools for analyzing non-convex learning dynamics and extend the understanding of hierarchical feature learning.

Novelty

This work introduces the 'leap' metric as a novel measure of hierarchical support structure, extending prior models limited to leap-1 staircase functions. It integrates saddle-to-saddle dynamics into the analysis, offering a detailed description of the sequential support learning process. Unlike previous studies constrained by mean-field or infinite-width assumptions, this paper controls the full complexity of finite neural networks, providing a comprehensive picture of the entire training process. The approach bridges the gap between statistical and computational complexity, aligning SGD dynamics with information-theoretic lower bounds, marking a significant advancement in deep learning theory.

Limitations

  • The analysis is primarily focused on two-layer neural networks with specific initialization and smooth activations, limiting immediate generalization to deeper architectures.
  • Assumes idealized conditions for SGD operation, such as specific step sizes and technical assumptions, which may differ from practical training scenarios.
  • The framework is currently tailored to Gaussian and Boolean low-dimensional functions; extension to more complex distributions remains open.

Future Work

Future research will extend the analysis to deeper networks, explore the effects of different initializations and optimization algorithms, and examine more general data distributions. Investigating how these hierarchical learning dynamics manifest in real-world datasets and large-scale models is also a key direction. Additionally, integrating these theoretical insights into practical training protocols could lead to more efficient algorithms that exploit hierarchical structures more effectively.

AI Executive Summary

Deep learning has revolutionized high-dimensional data analysis by enabling neural networks to learn hierarchical features efficiently. However, understanding the fundamental limits of this process remains a challenge. This paper introduces the 'leap' metric, a measure of the hierarchical complexity of target functions, capturing how support sets grow during learning. Analyzing two-layer neural networks trained with SGD on Gaussian and Boolean data, the authors reveal a saddle-to-saddle dynamic where the network sequentially acquires support features. The core finding is that the training time scales exponentially with the leap size, confirming that functions with higher hierarchy require significantly more steps. For instance, leap-1 functions are learned in Θ(d) steps, while leap-4 functions need Θ(d^3). These results are validated through experiments showing the network's support learning process aligns with theoretical predictions. Importantly, the complexity bounds match the lower bounds established by CSQ models, indicating that SGD achieves near-optimal efficiency for low-dimensional functions. This work advances the theoretical understanding of hierarchical feature learning, bridging the gap between statistical and computational complexity. It suggests that neural networks inherently implement an adaptive curriculum, learning simple features first and progressing to more complex ones. Future directions include extending the analysis to deeper architectures, diverse data distributions, and practical training regimes, promising to deepen our grasp of deep learning's capabilities and limitations.

Deep Dive

Abstract

We investigate the time complexity of SGD learning on fully-connected neural networks with isotropic data. We put forward a complexity measure -- the leap -- which measures how "hierarchical" target functions are. For $d$-dimensional uniform Boolean or isotropic Gaussian data, our main conjecture states that the time complexity to learn a function $f$ with low-dimensional support is $\tildeΘ(d^{\max(\mathrm{Leap}(f),2)})$. We prove a version of this conjecture for a class of functions on Gaussian isotropic data and 2-layer neural networks, under additional technical assumptions on how SGD is run. We show that the training sequentially learns the function support with a saddle-to-saddle dynamic. Our result departs from [Abbe et al. 2022] by going beyond leap 1 (merged-staircase functions), and by going beyond the mean-field and gradient flow approximations that prohibit the full complexity control obtained here. Finally, we note that this gives an SGD complexity for the full training trajectory that matches that of Correlational Statistical Query (CSQ) lower-bounds.

cs.LG stat.ML