Convergence of Adam in Deep ReLU Networks via Directional Complexity and Kakeya Bounds

TL;DR

This paper provides the first theoretical convergence and generalization bounds for Adam optimizer in Deep ReLU networks, leveraging stratified Morse theory and Kakeya bounds to address non-smooth, non-convex landscapes.

Contribution

It introduces novel theoretical analysis techniques for Adam in non-smooth deep networks, including convergence to global optima and improved generalization bounds.

Findings

First ((\sqrt{d_{ ext{eff}}/n})) generalization bound for Adam in Deep ReLU networks.

Global-optimal convergence of Adam in non-smooth, non-convex ReLU landscapes without PL or convexity assumptions.

Number of region crossings reduces from exponential to near-linear in effective dimension.

Abstract

First-order adaptive optimization methods like Adam are the default choices for training modern deep neural networks. Despite their empirical success, the theoretical understanding of these methods in non-smooth settings, particularly in Deep ReLU networks, remains limited. ReLU activations create exponentially many region boundaries where standard smoothness assumptions break down. \textbf{We derive the first \(\tilde{O}\!\bigl(\sqrt{d_{\mathrm{eff}}/n}\bigr)\) generalization bound for Adam in Deep ReLU networks and the first global-optimal convergence for Adam in the non smooth, non convex relu landscape without a global PL or convexity assumption.} Our analysis is based on stratified Morse theory and novel results in Kakeya sets. We develop a multi-layer refinement framework that progressively tightens bounds on region crossings. We prove that the number of region crossings collapses from exponential to near-linear in the effective dimension. Using a Kakeya based method, we give a tighter generalization bound than PAC-Bayes approaches and showcase convergence using a mild uniform low barrier assumption.