arXiv:2510.23498cs.LGcs.AI2025-10被引 1

为神经微分方程设计混合精度训练方案,显著节省内存与计算时间。

Mixed Precision Training of Neural ODEs

  • 用低精度计算速度和中间状态,高精度保存解与梯度。
  • 实现约50%内存减少,最高2倍加速,精度接近单精度训练。
  • 适合大规模连续时间模型训练,如图像分类与生成模型。

利用低精度计算已成为缓解大模型与大数据集带来的计算成本增长的标准策略。然而,盲目在所有计算中使用低精度可能导致舍入误差和不稳定性。因此,混合精度训练通常将权重存储于高精度,仅对白名单操作使用低精度计算。尽管如此,这些原则对连续时间架构(如神经微分方程,Neural ODEs)尚不可靠。本文提出一种面向 Neural ODEs 的混合精度训练框架,包含显式微分方程求解器与定制反向传播方案,并在多个学习任务中验证其有效性。该方案对由神经网络参数化的速度采用低精度计算,对中间状态也使用低精度,而通过自定义动态伴随缩放和在更高精度下累积解与梯度来保障数值稳定性。该方法解决了 Neural ODEs 训练的两大挑战:重复网络评估带来的计算开销,以及随时间步数或层数增加的内存增长问题。我们还发布了可扩展的开源 PyTorch 工具包 rampde,语法类似主流库,可直接替换现有代码。在具有挑战性的测试案例及图像分类、生成模型等 Neural ODE 应用中,验证了方案的可靠性与高效性,实现了约50%内存减少和最高2倍加速,同时保持与单精度训练相当的精度。

原文摘要 · Abstract (English)

Exploiting low-precision computations has become a standard strategy in deep learning to address the growing computational costs imposed by ever larger models and datasets. However, naively performing all computations in low precision can lead to roundoff errors and instabilities. Therefore, mixed precision training schemes usually store the weights in high precision and use low-precision computations only for whitelisted operations. Despite their success, these principles are currently not reliable for training continuous-time architectures such as neural ordinary differential equations (Neural ODEs). This paper presents a mixed precision training framework for neural ODEs consisting of explicit ODE solvers and a custom backpropagation scheme and shows their effectiveness in a range of learning tasks. Our scheme uses low-precision computations for evaluating the velocity, parameterized by the neural network, and for storing intermediate states, while numerical reliability is provided by custom dynamic adjoint scaling and by accumulating the solution and gradients in higher precision. These contributions address two key challenges in training neural ODEs: the computational cost of repeated network evaluations and the growth of memory requirements with the number of time steps or layers. Along with the paper we publish our extendable, open-source PyTorch package \texttt{rampde}, whose syntax resembles that of leading packages to provide a drop-in replacement in existing codes. We demonstrate the reliability and effectiveness of our scheme using challenging test cases and on neural ODE applications in image classification and generative models, achieving approximately 50\% memory reduction and up to 2x speedup while maintaining accuracy comparable to single-precision training.

神经微分方程混合精度内存优化PyTorch

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。