# Optimal Depth of Neural Networks

arXiv: https://arxiv.org/abs/2506.16862 · HTML: https://mlsanity.com/paper/2506.16862/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 9 of its 13 references that are in the map.
- First posted on arXiv 20 Jun 2025.

## Builds on (most cited first, 9 of 9)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Visualizing the Loss Landscape of Neural Nets](https://mlsanity.com/paper/1712.09913/) (2017)
- [Deep Networks with Stochastic Depth](https://mlsanity.com/paper/1603.09382/) (2016)
- [MINE: Mutual Information Neural Estimation](https://mlsanity.com/paper/1801.04062/) (2018)
- [Deep Learning via Dynamical Systems: An Approximation Perspective](https://mlsanity.com/paper/1912.10382/) (2019)
- [Universal Approximation Theorem of Deep Q-Networks](https://mlsanity.com/paper/2505.02288/) (2025)

## Abstract (arXiv)

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, Lrm 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.