[論文レビュー] Faster Transformer Decoding: N-gram Masked Self-Attention

Ciprian Chelba, Mia Xu Chen+2 more|arXiv (Cornell University)|Jan 14, 2020

Topic Modeling参考文献 8被引用数 13

ひとこと要約

本稿では、自己注意機構のコンテキストを直前の N−1 個のターゲットトークンに制限することで、Transformerのデコードを高速化する N-gram マスク付き自己注意機構を提案する。これにより計算量の複雑性は O(T²) から O(N·T) に低下する。WMT EnDe および EnFr における実験では、N=8 でベースラインの BLEU スコアにほぼ同等の性能（0.3–0.4 BLEU 内）を達成し、最大で 2–3 倍の高速化とメモリ帯域幅の削減を実現した。

ABSTRACT

Motivated by the fact that most of the information relevant to the prediction of target tokens is drawn from the source sentence $S=s_1, \ldots, s_S$, we propose truncating the target-side window used for computing self-attention by making an $N$-gram assumption. Experiments on WMT EnDe and EnFr data sets show that the $N$-gram masked self-attention model loses very little in BLEU score for $N$ values in the range $4, \ldots, 8$, depending on the task.

研究の動機と目的

• 自己回帰的 Transformer デコードにおける計算およびメモリのオーバーヘッドを、性能の著しい低下を伴わずに低減すること。
• 自己注意のコンテキストを固定された N-gram ウィンドウに制限しても、ニューラル機械翻訳における翻訳品質が保持されるかどうかを調査すること。
• 系列生成における推論速度、メモリ効率、モデル精度のトレードオフを評価すること。
• ビームサーチ中に直近の N−1 個のトークンのみを格納する固定サイズのバッファを用いることで、メモリ帯域幅を削減できるかを検討すること。

提案手法

• デコーダー内の自己注意計算を直前の N−1 個のターゲットトークンに制限する N-gram マスク付き自己注意機構を導入する。
• 標準的な因果的自己注意機構を、ターゲット系列において N−1 个より前の位置のトークンを除外するマスクを適用することで修正する。
• Lingvo フレームワークに N-gram 注意を構成可能オプションとして実装し、デコーダー層の完全な因果的注意を置き換える。
• デコード中に逐次的に更新されるサイズ N−1 のスライディングウインドウバッファを用いて、コンテキストトークンを格納・更新することで、メモリアクセスのオーバーヘッドを低減する。
• エンコーダー・デコーダー間の注意メカニズムはそのままで、デコーダーの自己注意のみを N-gram に制限する。
• ワードピeceトークン化と固定ハイパーパramータを用いて、WMT EnDe および EnFr データセット上で標準的な Transformer アーキテクチャを用いてモデルを学習・評価する。

実験結果

リサーチクエスチョン

• RQ1 直前の N−1 個のトークンに自己注意のコンテキストを制限した場合、BLEU スコアの観点から翻訳品質にどのような影響を与えるか？
• RQ2 機械翻訳タスクにおいて、速度と性能のバランスを最適化するのに適した N-gram の順序（N）は何か？
• RQ3 N-gram 自己注意機構は、推論時における計算量とメモリ帯域幅をどの程度低減できるか？
• RQ4 コンテキストの切り詰めが行われても、N-gram メカニズムは長距離依存関係を十分に捉える能力を保持しているか？
• RQ5 固定サイズのバッファを用いて、ビームサーチ中のメモリアクセスを効率的に削減できるか？

主な発見

• WMT EnDe において N=8 の場合、テストセットで BLEU スコア 22.5 を達成し、ベースライン（22.4）よりわずか 0.1 BLEU 低いが、最高性能のベースラインと 0.3–0.4 BLEU 内の差で収束している。
• WMT EnFr では N=8 でテストデータで BLEU スコア 32.7 を得ており、ベースライン（32.9）と 0.2 BLEU 内で一致し、より小さい N-gram モデルを上回る性能を示した。
• N-gram モデルにより、計算量の複雑性は O(T²) から O(N·T) に低下し、T≈16–25 のシーケンス長に対して理論的な 2–3 倍の高速化が可能である。
• WMT EnDe テストデータにおいて、N=8 のモデルは 2.66 のログ周辺度を達成し、ベースラインの 2.64 とわずかに高いが、言語モデルの能力にほとんど劣化がないことを示している。
• N-gram メカニズムにより、長さ N−1 の固定サイズのメモリバッファを用いることができ、O(T/N) の割合でメモリ帯域幅が削減される。これは TPU ハードウェアにおいて特に有益である。
• N≥3 の場合、性能は滑らかに低下し、N=4 でも既に WMT EnDe で 22.2 BLEU を達成しており、より小さい N 値でも速度と精度のトレードオフに有効であることが示された。