Gradient Descent Optimizes Infinite-Depth ReLU Implicit Networks with Linear Widths
This provides theoretical guarantees for training implicit networks, addressing a bottleneck in their adoption for memory-efficient deep learning.
The paper tackles the problem of understanding when gradient descent converges for nonlinear ReLU implicit networks, which have infinitely many layers and can suffer from ill-posed equilibrium equations. It proves that with a scaled weight matrix and linear width relative to sample size, both gradient flow and gradient descent converge to a global minimum at a linear rate.
Implicit deep learning has recently become popular in the machine learning community since these implicit models can achieve competitive performance with state-of-the-art deep networks while using significantly less memory and computational resources. However, our theoretical understanding of when and how first-order methods such as gradient descent (GD) converge on \textit{nonlinear} implicit networks is limited. Although this type of problem has been studied in standard feed-forward networks, the case of implicit models is still intriguing because implicit networks have \textit{infinitely} many layers. The corresponding equilibrium equation probably admits no or multiple solutions during training. This paper studies the convergence of both gradient flow (GF) and gradient descent for nonlinear ReLU activated implicit networks. To deal with the well-posedness problem, we introduce a fixed scalar to scale the weight matrix of the implicit layer and show that there exists a small enough scaling constant, keeping the equilibrium equation well-posed throughout training. As a result, we prove that both GF and GD converge to a global minimum at a linear rate if the width $m$ of the implicit network is \textit{linear} in the sample size $N$, i.e., $m=Ω(N)$.