The success of deep learning comes at a tremendous computational and energy cost, and the scalability of training massively overparametrized neural networks is becoming a real barrier to the progress of artificial intelligence (AI). Despite the popularity and low cost-per-iteration of traditional backpropagation via gradient decent, stochastic gradient descent (SGD) has prohibitive convergence rate in non-convex settings, both in theory and practice. To mitigate this cost, recent works have proposed to employ alternative (Newton-type) training methods with much faster convergence rate, albeit with higher cost-per-iteration. For a typical neural network with $m=\mathrm{poly}(n)$ parameters and input batch of $n$ datapoints in $\mathbb{R}^d$, the previous work of [Brand, Peng, Song, and Weinstein, ITCS'2021] requires $\sim mnd + n^3$ time per iteration. In this paper, we present a novel training method that requires only $m^{1-\alpha} n d + n^3$ amortized time in the same overparametrized regime, where $\alpha \in (0.01,1)$ is some fixed constant. This method relies on a new and alternative view of neural networks, as a set of binary search trees, where each iteration corresponds to modifying a small subset of the nodes in the tree. We believe this view would have further applications in the design and analysis of deep neural networks (DNNs).

翻译：深度学习的成功伴随着巨大的计算和能源成本，训练大规模超参数化神经网络的可扩展性正成为人工智能发展的现实障碍。尽管传统的基于梯度下降的反向传播方法广受欢迎且单次迭代成本低廉，但在非凸场景下，随机梯度下降（SGD）的理论与实践收敛速度均不理想。为降低成本，近期研究提出采用牛顿类替代训练方法，其收敛速度更快，但单次迭代成本更高。针对具有 $m=\mathrm{poly}(n)$ 个参数且输入批量包含 $\mathbb{R}^d$ 中 $n$ 个数据点的典型神经网络，[Brand, Peng, Song, and Weinstein, ITCS'2021] 的先前工作需要每次迭代 $\sim mnd + n^3$ 时间。本文提出一种新型训练方法，在相同超参数化场景下仅需 $m^{1-\alpha} n d + n^3$ 均摊时间，其中 $\alpha \in (0.01,1)$ 为固定常数。该方法基于对神经网络的全新替代视角——将其视为一组二叉搜索树，每次迭代对应修改树中少量节点。我们相信这一视角将在深度神经网络（DNN）的设计与分析中产生更多应用。