[论文解读] Bypass Exponential Time Preprocessing: Fast Neural Network Training via Weight-Data Correlation Preprocessing

Josh Alman, Jiehao Liang+3 more|arXiv (Cornell University)|Nov 25, 2022

Machine Learning in Materials Science被引用 4

一句话总结

本文提出了一种新颖的预处理方法，利用权重-数据相关性树来加速过参数化 ReLU 神经网络的训练。通过利用神经元激活的稀疏性，该方法实现了每次迭代 o(nmd) 的训练时间，且仅需 O(nmd) 的预处理时间，显著优于以往需要指数级预处理时间的方法，同时在标准复杂度猜想下保持理论保证。

ABSTRACT

Over the last decade, deep neural networks have transformed our society, and they are already widely applied in various machine learning applications. State-of-art deep neural networks are becoming larger in size every year to deliver increasing model accuracy, and as a result, model training consumes substantial computing resources and will only consume more in the future. Using current training methods, in each iteration, to process a data point $x \in \mathbb{R}^d$ in a layer, we need to spend $Θ(md)$ time to evaluate all the $m$ neurons in the layer. This means processing the entire layer takes $Θ(nmd)$ time for $n$ data points. Recent work [Song, Yang and Zhang, NeurIPS 2021] reduces this time per iteration to $o(nmd)$, but requires exponential time to preprocess either the data or the neural network weights, making it unlikely to have practical usage. In this work, we present a new preprocessing method that simply stores the weight-data correlation in a tree data structure in order to quickly, dynamically detect which neurons fire at each iteration. Our method requires only $O(nmd)$ time in preprocessing and still achieves $o(nmd)$ time per iteration. We complement our new algorithm with a lower bound, proving that assuming a popular conjecture from complexity theory, one could not substantially speed up our algorithm for dynamic detection of firing neurons.

研究动机与目标

• 为解决训练大规模深度神经网络日益增长的计算成本。
• 设计一种具有多项式预处理时间与在数据量和宽度上亚二次每次迭代复杂度的训练算法。
• 克服先前最先进方法中使用最近邻数据结构所导致的指数级预处理瓶颈。
• 为训练过程中动态检测激活神经元提供一种确定性、实用的解决方案。
• 在标准复杂度猜想下建立理论极限。

提出的方法

• 为每个数据点构建 n 棵二叉搜索树，以维护每个数据点与全部 m 个权重之间的内积。
• 每棵树在叶节点存储内积值，并在内部节点向上传播最大值，以支持高效的范围查询。
• 利用树结构通过自顶向下的遍历动态检测激活神经元（即内积 ≥ 阈值 b 的神经元）。
• 支持高效更新：当权重被修改时，每个数据点仅需更新 O(log m) 个节点，耗时 O(nd log m)。
• 采用阈值 b = √(0.4 log m) 来控制激活神经元的数量，确保激活集合的稀疏性。
• 通过仅对激活神经元进行计算，将该数据结构集成到梯度下降中，从而降低每次迭代的计算成本。

实验结果

研究问题

• RQ1 我们能否仅通过多项式预处理时间，实现 o(nmd) 的每次迭代训练时间？
• RQ2 是否可能设计一种确定性、实用的数据结构，用于神经网络训练中动态激活神经元的检测？
• RQ3 在标准复杂度假设下，动态激活神经元检测的理论加速极限是什么？
• RQ4 权重-数据相关性如何实现过参数化网络中高效的稀疏激活检测？
• RQ5 我们能否在确保正确性与效率的前提下，保持亚二次的训练时间？

主要发现

• 所提算法的期望每次迭代运行时间为 Õ(m⁴ᐟ⁵n²d)，当 m ≫ n 时为 o(nmd)。
• 预处理时间为 O(nmd)，显著优于以往需要 O(2^d) 或 O(n^d) 时间的方法。
• 该算法支持动态更新和查询，时间复杂度分别为 O(nd log m) 和 O(min{|Q|, m⁴ᐟ⁵n})。
• 该方法是确定性的，不依赖随机性，从而增强了可复现性与可靠性。
• 在正交向量猜想下证明了下界，表明若不破坏标准复杂度假设，o(nmd) 的每次迭代时间无法被显著改进。
• 通过理论分析展示了经验效率，表明平均每个数据点仅有 O(m⁴ᐟ⁵n) 个神经元被激活，从而支持稀疏计算。