[Paper Review] Implicit Bias of Gradient Descent for Two-layer ReLU and Leaky ReLU Networks on Nearly-orthogonal Data

Yiwen Kou, Zixiang Chen+1 more|arXiv (Cornell University)|Oct 29, 2023

Machine Learning and ELM4 citations

TL;DR

This paper establishes the implicit bias of gradient descent in two-layer ReLU and Leaky ReLU networks on nearly-orthogonal data, proving that gradient descent converges to solutions with stable rank 1 for Leaky ReLU and bounded stable rank for ReLU, while all training examples achieve the same normalized margin asymptotically. The analysis reveals logarithmic weight norm growth and $\Theta(t^{-1})$ loss convergence, extending prior gradient flow results to gradient descent in non-smooth networks.

ABSTRACT

The implicit bias towards solutions with favorable properties is believed to be a key reason why neural networks trained by gradient-based optimization can generalize well. While the implicit bias of gradient flow has been widely studied for homogeneous neural networks (including ReLU and leaky ReLU networks), the implicit bias of gradient descent is currently only understood for smooth neural networks. Therefore, implicit bias in non-smooth neural networks trained by gradient descent remains an open question. In this paper, we aim to answer this question by studying the implicit bias of gradient descent for training two-layer fully connected (leaky) ReLU neural networks. We showed that when the training data are nearly-orthogonal, for leaky ReLU activation function, gradient descent will find a network with a stable rank that converges to $1$, whereas for ReLU activation function, gradient descent will find a neural network with a stable rank that is upper bounded by a constant. Additionally, we show that gradient descent will find a neural network such that all the training data points have the same normalized margin asymptotically. Experiments on both synthetic and real data backup our theoretical findings.

Motivation & Objective

• To understand the implicit bias of gradient descent in non-smooth neural networks, particularly ReLU and Leaky ReLU networks, where such bias remains poorly understood.
• To analyze how gradient descent optimizes two-layer fully connected networks when the training data are nearly-orthogonal, a condition that enables theoretical tractability.
• To establish convergence properties of the weight matrix, including stable rank and normalized margin, under gradient descent for both ReLU and Leaky ReLU activations.
• To extend prior results on gradient flow to gradient descent, particularly in non-smooth settings, and provide rigorous theoretical guarantees.

Proposed method

• Theoretical analysis of gradient descent dynamics on two-layer ReLU and Leaky ReLU networks with nearly-orthogonal data, defined by $\|\mathbf{x}_i\|_2^2 \geq Cn \max_{k \neq i} |\langle \mathbf{x}_i, \mathbf{x}_k \rangle|$.
• Proof of stable rank convergence to 1 for Leaky ReLU networks and bounded stable rank for ReLU networks via analysis of neuron activation patterns and weight matrix evolution.
• Derivation of weight norm growth at rate $\Theta(\log t)$ and loss decay at rate $\Theta(t^{-1})$ using differential inequality techniques and logarithmic integration.
• Proof that all training data points asymptotically achieve the same normalized margin, leveraging properties of the loss function and gradient dynamics.
• Use of sub-Gaussian matrix concentration (Lemma H.3) and logarithmic inequality lemmas (H.4, H.5) to bound key terms in the convergence analysis.
• Application of differential inequality techniques (Lemmas H.1 and H.2) to control the growth of logarithmic weight norms and derive convergence rates.

Experimental results

Research questions

• RQ1 What is the implicit bias of gradient descent in two-layer ReLU and Leaky ReLU networks trained on nearly-orthogonal data?
• RQ2 How does the stable rank of the weight matrix evolve during gradient descent for ReLU and Leaky ReLU networks?
• RQ3 Does gradient descent drive all training samples to have the same normalized margin asymptotically in these networks?
• RQ4 What are the convergence rates of weight norm and training loss under gradient descent for these non-smooth networks?
• RQ5 How do the theoretical results for gradient descent compare to prior gradient flow analyses in the same setting?

Key findings

• For two-layer Leaky ReLU networks, gradient descent drives the stable rank of the weight matrix to converge to 1, matching the gradient flow result of Frei et al. (2022b).
• For two-layer ReLU networks, the stable rank of the weight matrix is upper bounded by a constant, with an illustrative example showing convergence to approximately 2 under fully orthogonal data.
• Gradient descent increases the weight norm at a rate of $\Theta(\log t)$, where $t$ is the number of iterations.
• The training loss converges to zero at a rate of $\Theta(t^{-1})$, improving upon the $O(t^{-1/2})$ rate in prior work for smoothed networks.
• All training data points asymptotically achieve the same normalized margin, indicating convergence to a max-margin solution.
• The theoretical results are empirically validated on both synthetic and real-world datasets, confirming the predicted convergence behaviors.