arXiv cs.AI

ニューラルODEの混合精度トレーニング

Mixed Precision Training of Neural ODEs

http://arxiv.org/abs/2510.23498v1


本論文では、ニューラル常微分方程式(Neural ODEs)のための混合精度トレーニングフレームワークを提案しています。深層学習においては、計算コストの増加に対応するために低精度計算が一般的に利用されていますが、すべてを低精度で実行することは丸め誤差や不安定さを引き起こす可能性があります。したがって、混合精度トレーニングでは通常、高精度で重みを保存し、かつ特定の操作のみで低精度計算を使用します。本研究では、ニューラルODEにおけるトレーニングの2つの主要課題に対処し、低精度計算を用いて速度を評価し、途中の状態を保存しながら、安定性を高精度で解の累積を行うことで確保しています。この手法は、画像分類や生成モデルなどの実際のアプリケーションで約50%のメモリ削減と最大2倍のスピードアップを達成し、単一精度トレーニングと同等の精度を維持しています。さらに、可拡張なオープンソースのPyTorchパッケージも公開しています。