[論文レビュー] Efficient Softmax Approximation for Deep Neural Networks with Attention Mechanism

Ihor Vasyltsov, Wooseok Chang|arXiv (Cornell University)|Nov 21, 2021

Advanced Neural Network Applications参考文献 33被引用数 6

ひとこと要約

本稿では、注意機構に基づく深層ニューラルネットワークにおける効率的なソフトマックス近似のための2つのLUTベースの手法を提案する。この手法により、高コストな除算器を必要とせず、ハードウェアアクセラレーションが可能となる。多様なNLPおよびビジョンタスクにおいて、8ビット量子化下でも1%未満の精度損失を達成し、強力な一般化性能と低メモリオーバーヘッド（700バイト）を示している。

ABSTRACT

There has been a rapid advance of custom hardware (HW) for accelerating the inference speed of deep neural networks (DNNs). Previously, the softmax layer was not a main concern of DNN accelerating HW, because its portion is relatively small in multi-layer perceptron or convolutional neural networks. However, as the attention mechanisms are widely used in various modern DNNs, a cost-efficient implementation of softmax layer is becoming very important. In this paper, we propose two methods to approximate softmax computation, which are based on the usage of LookUp Tables (LUTs). The required size of LUT is quite small (about 700 Bytes) because ranges of numerators and denominators of softmax are stable if normalization is applied to the input. We have validated the proposed technique over different AI tasks (object detection, machine translation, sentiment analysis, and semantic equivalence) and DNN models (DETR, Transformer, BERT) by a variety of benchmarks (COCO17, WMT14, WMT17, GLUE). We showed that 8-bit approximation allows to obtain acceptable accuracy loss below $1.0\%$.

研究の動機と目的

• Transformer や BERT などの注意機構ベースのモデルにおけるソフトマックス層の計算コストの増大に対処すること。特にリソース制限のあるエッジ環境において顕著である。
• エッジアクセラレータにおいて、ソフトマックス計算をホストCPUにオフロードする非効率性を克服すること。これは、遅延に敏感なNLPワークロードの性能を制限する要因となる。
• 複雑な除算を回避し、デバイス内推論における消費電力と遅延を低減するハードウェアフレンドリーなソフトマックス近似手法を開発すること。
• さまざまなAIタスクおよびモデルにおいて手法を検証し、シーケンス長やモデルアーキテクチャの変動に対しても一般化性能と耐障害性を確認すること。

提案手法

• 事前に計算されたルックアップテーブル（LUT）を用いてソフトマックス関数を区分的定数近似することで、直接計算を置き換える。
• 入力の正規化を活用して、分子および分母の範囲を安定化させ、多様なモデルにわたってコンパクトなLUTサイズ（約700バイト）を実現する。
• 2つのバリエーションを実装する：1つは全ソフトマックス計算に1つのLUTを用いるもの、もう1つは分子と分母に別々のLUTを用いることで精度を向上させるもの。
• ハードウェアアクセラレータ（例：NPUs）にLUTベースのソフトマックスを統合することで、高面積を要する除算器ユニットの必要性を排除し、消費電力と遅延を低減する。
• 再トレーニングを必要とせず、Post-Training Quantization（PTQ-D）の段階で手法を適用することで、既存の量子化ワークフローとの互換性を確保する。
• 精度とメモリフットプリントのバランスを図るため、LUTサイズと解像度（例：256〜512バイト）を最適化する。特に、DETR+DC5のような高マグニチュードの活性化を示すモデルにおいて有効である。

実験結果

リサーチクエスチョン

• RQ1 コンパクトなLUTベースのソフトマックス近似は、エッジAIアクセラレータにおいて、ハードウェア複雑性を最小限に抑えつつ高い精度を達成できるか？
• RQ2 提案手法は、多様な注意機構ベースのモデル（例：BERT、DETR、Transformer）およびタスク（NLP、オブジェクト検出）において、どのように性能を発揮するか？
• RQ3 LUTサイズと量子化精度（例：8ビット）が、高マグニチュードのソフトマックス入力を持つモデルにおける精度損失に与える影響は何か？
• RQ4 標準的なPTQ-D量子化と比較して、LUTベースの近似は精度劣化と一般化性能の面でどのように差をつけるか？
• RQ5 活性化分布が非対称なモデル、例えば拡張（dilation）を備えたDETR（DETR with DC5）のようなモデルに対しても、この手法はどの程度一般化可能か？

主な発見

• 評価された全タスク（オブジェクト検出、機械翻訳、センチメント分析、意味同等性）において、8ビット量子化下で1.0%未満の精度損失を達成した。
• 機械翻訳（WMT14/WMT17）において、BLEUスコアのベースラインからのずれは0.5%未満であり、性能劣化が最小限に抑えられた。
• センチメント分析および意味同等性（GLUEベンチマーク）において、LUT手法による精度低下は、PTQ-D量子化単体によるものよりも常に小さかった。
• 拡張（dilation）を備えたDETRモデル（DETR+DC5）では、LUTサイズを256バイトから512バイトに増加させることで、8ビット精度下での精度損失を9%から3%に削減した。これはスケーラビリティを示している。
• 本手法はモデル間で良好に一般化される：DETR、BERT、Transformerともに、最小限の精度損失で一貫した性能を示しており、耐障害性が確認された。
• LUTサイズは小さく（約700バイト）維持され、メモリ制限のあるエッジデバイスへのデプロイに実用的であり、顕著な面積的・消費電力的オーバーヘッドを伴わない。