Optimal Depth of Neural Networks

TL;DR

This paper presents a theoretical framework modeling the optimal depth of neural networks as an optimal stopping problem, providing insights into early exiting strategies that improve efficiency without sacrificing accuracy.

Contribution

It introduces a formal optimal stopping framework for neural network depth, proves finiteness of optimal depth under certain conditions, and proposes a regularizer to encourage early exiting.

Findings

The regularizer effectively induces early exiting behavior.

Empirical results show significant computational savings on ImageNet.

The framework extends to Transformers and continuous-depth models.

Abstract

Determining the optimal depth of a neural network is a fundamental yet challenging problem, typically resolved through resource-intensive experimentation. This paper introduces a formal theoretical framework to address this question by recasting the forward pass of a deep network, specifically a Residual Network (ResNet), as an optimal stopping problem. We model the layer-by-layer evolution of hidden representations as a sequential decision process where, at each layer, a choice is made between halting computation to make a prediction or continuing to a deeper layer for a potentially more refined representation. This formulation captures the intrinsic trade-off between accuracy and computational cost. Our primary theoretical contribution is a proof that, under a plausible condition of diminishing returns on the residual functions, the expected optimal stopping depth is provably finite, even in an infinite-horizon setting. We leverage this insight to propose a novel and practical regularization term, $\mathcal{L}_{\rm depth}$, that encourages the network to learn representations amenable to efficient, early exiting. We demonstrate the generality of our framework by extending it to the Transformer architecture and exploring its connection to continuous-depth models via free-boundary problems. Empirical validation on ImageNet confirms that our regularizer successfully induces the theoretically predicted behavior, leading to significant gains in computational efficiency without compromising, and in some cases improving, final model accuracy.