[論文レビュー] Online normalizer calculation for softmax

Maxim Milakov, Natalia Gimelshein|arXiv (Cornell University)|May 8, 2018

Parallel Computing and Optimization Techniques参考文献 8被引用数 20

ひとこと要約

この論文では、Softmax関数のオンライン正規化子計算を提案し、各要素あたりのメモリアクセス回数を4回から3回に削減することで、より高速な推論を実現する。Tesla V100でのベンチマークでは、Softmax単体で最大1.3倍の高速化、Softmax+TopKの統合処理では最大5倍の高速化が達成され、主にメモリ帯域幅の圧迫が軽減されたことによるものである。

ABSTRACT

The Softmax function is ubiquitous in machine learning, multiple previous works suggested faster alternatives for it. In this paper we propose a way to compute classical Softmax with fewer memory accesses and hypothesize that this reduction in memory accesses should improve Softmax performance on actual hardware. The benchmarks confirm this hypothesis: Softmax accelerates by up to 1.3x and Softmax+TopK combined and fused by up to 5x.

研究の動機と目的

• ディープラーニング推論におけるパフォーマンスボトルneckであるSoftmax計算におけるメモリアクセスのオーバーヘッドを低減すること。
• 多数の代替手法が存在する中で、古典的なSoftmax関数に対する特化した最適化が不足しているという問題に取り組むこと。
• 正規化と選択の論理を同時に処理することで、SoftmaxとTopKの操作を効率的に統合すること。
• 1パスアルゴリズムにより、数値的安定性を維持しながらメモリ帯域幅の使用量を削減すること。
• 特にGPU上で現代のハードウェアにおいて測定可能なパフォーマンス向上を示すこと。

提案手法

• 各要素あたりのメモリアクセス回数を4回から3回に削減する1パスアルゴリズムを導入し、最大値とSoftmax正規化項を1回の反復で同時に計算する。
• オーバーフロー／アンダーフローを防ぐために、指数関数計算の前にすべてのログチットから最大値を差し引く数値的に安定した定式化を採用する。
• 正規化項の段階的更新ルールを用いる：$ d_j = d_{j-1} \cdot e^{m_{j-1} - m_j} + e^{x_j - m_j} $、ここで$ m_j $ は累積的最大値を表す。
• 処理される各要素に対して、累積的最大値 $ m_j $ と正規化済み和 $ d_j $ を段階的に更新する。
• 同じパスでTop-k値を追跡することで、SoftmaxとTopKの操作を統合し、重複するメモリアクセスを排除する。
• メモリ帯域幅の使用量を最小限に抑え、カーネル統合を可能にするようにGPUパフォーマンスを最適化する。

実験結果

リサーチクエスチョン

• RQ1 Softmax計算におけるメモリアクセス回数の削減は、現代のハードウェア上で測定可能なパフォーマンス向上をもたらすか？
• RQ2 数値的安定性を損なわずに、最大値とSoftmax正規化項を1パスで計算することは可能か？
• RQ3 提案されたオンライン正規化子は、標準の2パスまたは3パス実装と比較して、パフォーマンスと正確性の面でどのように異なるか？
• RQ4 SoftmaxとTopKの統合処理は、エンドツーエンドの推論速度をどの程度向上できるか？
• RQ5 GPUアーキテクチャにおいて、ベクトルサイズとバッチサイズが増加するに従い、パフォーマンス向上がスケーリングするか？

主な発見

• 提案されたオンライン正規化子により、各要素あたりのメモリアクセス回数が4回から3回に削減され、Tesla V100で大規模なベクトルサイズの場合、Softmax単体で最大1.3倍の高速化が達成された。
• TopKと統合された場合、統合処理による2.5倍の高速化とオンライン正規化子による2倍の高速化を組み合わせ、Softmax+TopKの合計で最大5倍の高速化が達成された。
• パフォーマンス向上は、特にメモリ帯域幅が制限要因となる大バッチ環境で顕著に現れた。
• 小バッチ設定でも、レイテンシーとメモリアクセスのオーバーヘッドが低減されたことから、オンラインSoftmaxは1.5倍～2.5倍の高速化を達成した。
• 数値的安定性を維持しており、既存のディープラーニングフレームワークと互換性があり、即座に適用可能な最適化である。
• 他のSoftmax最適化技術（例：Hierarchical Softmax, SVD-Softmax, Importance Sampling）とは直交的に機能し、組み合わせることでさらなる高速化が可能である。