Understanding the training of infinitely deep and wide ResNets with Conditional Optimal Transport (2403.12887v1)

Abstract: We study the convergence of gradient flow for the training of deep neural networks. If Residual Neural Networks are a popular example of very deep architectures, their training constitutes a challenging optimization problem due notably to the non-convexity and the non-coercivity of the objective. Yet, in applications, those tasks are successfully solved by simple optimization algorithms such as gradient descent. To better understand this phenomenon, we focus here on a ``mean-field'' model of infinitely deep and arbitrarily wide ResNet, parameterized by probability measures over the product set of layers and parameters and with constant marginal on the set of layers. Indeed, in the case of shallow neural networks, mean field models have proven to benefit from simplified loss-landscapes and good theoretical guarantees when trained with gradient flow for the Wasserstein metric on the set of probability measures. Motivated by this approach, we propose to train our model with gradient flow w.r.t. the conditional Optimal Transport distance: a restriction of the classical Wasserstein distance which enforces our marginal condition. Relying on the theory of gradient flows in metric spaces we first show the well-posedness of the gradient flow equation and its consistency with the training of ResNets at finite width. Performing a local Polyak-\L{}ojasiewicz analysis, we then show convergence of the gradient flow for well-chosen initializations: if the number of features is finite but sufficiently large and the risk is sufficiently small at initialization, the gradient flow converges towards a global minimizer. This is the first result of this type for infinitely deep and arbitrarily wide ResNets.

• The paper introduces a mean-field model for infinite-depth ResNets and analyzes gradient flow convergence using Conditional Optimal Transport.
• It leverages a metric space formulation and a local Polyak-Łojasiewicz condition to prove the existence, uniqueness, and stability of gradient flows.
• The research offers practical insights on activation choices and initialization strategies that ensure proper kernel conditioning and effective loss minimization.

Convergence Analysis of Infinitely Deep ResNets through Conditional Optimal Transport and Polyak-Łojasiewicz Conditions

Introduction

The effective training of Residual Neural Networks (ResNets) has posed a significant challenge due to the optimization difficulties arising from their depth. The incorporation of skip connections, facilitating deeper architectures, highlights the necessity of underpinning the success of these networks with a solid theoretical foundation. This research focuses on a "mean-field" model representation of infinitely deep and arbitrarily wide ResNet architectures, addressing the convergence of the gradient flow in training such networks. A novel application of Conditional Optimal Transport (COT) is proposed to model the parameter space, aligning with the practical layer-wise $L^2$ training approach. The analysis contributes to recognizing conditions under which gradient flow convergence towards a minimizer of the training loss is guaranteed.

Metric Structure and Gradient Flow Dynamics

Our initial investigation furnishes the parameter set $\Pp^Leb_2([0, 1] \times \Om)$ with a metric, enforcing the marginal condition via the Conditional Optimal Transport (COT) distance. This metric space's completeness and its dynamic formulation underlie the conditional continuity equation inherent to absolutely continuous curves. Broadly, this facilitates the characterization of the gradient flow for the loss function within a conditional metric space framework.

Gradient Flow as Curves of Maximal Slope and Well-posedness

The equivalence between our gradient flow definition and the concept of curves of maximal slope in metric spaces confirms the existence, uniqueness, and stability of gradient flow solutions. This result, underpinning well-posedness, leverages a local Polyak-Łojasiewicz (P-Ł) condition, elucidating the conditions under which the loss decreases at a constant rate along the flow.

Convergence Conditions and Polyak-Łojasiewicz Property

The crux of achieving convergence in the gradient flow dynamic rests on a local P-Ł property. This research delineates sufficient conditions where this property holds, especially in reference to the architecture of Single Hidden Layer (SHL) perceptrons. Key to this analysis is the kernel's conditioning formed from the training data and network parameterization, linking the numerical aspect of training with the analytical framework.

Practical Implications and Activation Choices

For practical applicability, the paper dives into specific instances such as Identity (or FixUp) initialization and positively homogeneous activations including $ReLU$. The strict positivity of the kernel matrix, crucial for the P-Ł property, warrants the conditioning for varied activation choices and initialization strategies. Discussions extend to implications for network training, identifying how data separation and initialization modulate the convergence guarantee.

Modifications for Assurance of Convergence

Adjustments like lifting and rescaling the embedding space are proposed as mechanisms to ensure convergence criteria are met. This facet of the analysis emphasizes adaptability in model architecture to align with theoretical assurances of gradient flow convergence towards loss minimization.

Conclusion

Historically, the training efficacy of ResNets is well-noted, though underlying theoretical explanations have been disjointed. This research fortifies the understanding of ResNet training dynamics through Conditional Optimal Transport and the Polyak-Łojasiewicz condition, establishing conditions that either guarantee or enhance convergence. The insights into metric structuring and the implications of kernel conditioning pave a clearer path for future research and practical implementations in deep learning architectures.