[論文レビュー] Sharing Attention Weights for Fast Transformer

Tong Xiao, Yinqiao Li+3 more|arXiv (Cornell University)|Jun 26, 2019

Topic Modeling参考文献 20被引用数 7

ひとこと要約

本稿では、自己注意およびエンコーダデコーダ注意メカニズムの両方において隣接する層間で注意重みを共有することで、Transformerの推論を高速化する手法である共有注意ネットワーク（San）を提案する。隠れ状態の縦方向再利用と共有ポリシーの共同学習を可能にすることで、キャッシュベースラインに対して1.3倍の高速化を達成し、キャッシュなしでは最大16倍の高速化を実現した。BLEUスコアの低下は最小限に抑えられた。

ABSTRACT

Recently, the Transformer machine translation system has shown strong results by stacking attention layers on both the source and target-language sides. But the inference of this model is slow due to the heavy use of dot-product attention in auto-regressive decoding. In this paper we speed up Transformer via a fast and lightweight attention model. More specifically, we share attention weights in adjacent layers and enable the efficient re-use of hidden states in a vertical manner. Moreover, the sharing policy can be jointly learned with the MT model. We test our approach on ten WMT and NIST OpenMT tasks. Experimental results show that it yields an average of 1.3X speed-up (with almost no decrease in BLEU) on top of a state-of-the-art implementation that has already adopted a cache for fast inference. Also, our approach obtains a 1.8X speed-up when it works with the extsc{Aan} model. This is even 16 times faster than the baseline with no use of the attention cache.

研究の動機と目的

• 自己回帰的デコードと繰り返し発生するドット積注意計算によるTransformerモデルの高い推論遅延を解消すること。
• 隣接層間の注意重み分布の構造的類似性を活用して、スタックされた注意層における冗長な計算を低減すること。
• ヒューリスティックな共有ルールではなく、共有ポリシーとMTモデルを同時に最適化する共同学習フレームワークを構築すること。
• 重み共有による隠れ状態の再利用によって、メモリフットプリントを最小限に抑えること。
• 特に既存の高速化技術（例：注意キャッシュ）と組み合わせた場合に顕著な精度損失なしに高速推論を達成すること。

提案手法

• 自己注意およびエンコーダデコーダ注意サブレイヤーの両方において、隣接する層間で垂直方向に注意重みを共有する共有注意ネットワーク（San）を提案する。
• 同じ注意重み行列を複数のスタックされた層にわたって使用するように注意メカニズムを定式化することで、冗長な計算を削減する。
• 共有ポリシーをMTモデルと同時にEnd-to-endで学習する戦略を導入し、ネットワークが最適な共有パターンを自動で発見できるようにする。
• 共有メモリ領域に隠れ状態を格納することで、効率的な隠れ状態の再利用を可能にし、メモリフットプリントを削減するとともに推論を高速化する。
• 注意キャッシュやAanモデルといった既存の高速化技術と統合し、相乗的かつ補完的な改善効果を示す。
• スケールドドット積注意メカニズムをベースとし、層間で重みを共有することで、モデル構造を維持しつつ効率性を向上させる。

実験結果

リサーチクエスチョン

• RQ1 隣接層間で注意重みを共有することで、翻訳品質の低下を伴わずTransformerの推論時間を短縮できるか？
• RQ2 MTモデルと共有ポリシーを同時に学習することで、ヒューリスティックな共有戦略に比べて性能が向上し、収束が速くなるか？
• RQ3 注意キャッシュやAanモデルといった既存の高速化手法と組み合わせた場合、Sanの有効性はどの程度か？
• RQ4 訓練済みのTransformerにおいて、隣接層がどの程度類似した注意重み分布を示すか。これは重み共有の妥当性を裏付けるか？
• RQ5 序列変換モデルにおいて、Sanは推論速度を維持または向上させながらメモリ使用量を削減できるか？

主な発見

• キャッシュを活用する最先端のTransformer実装よりも、Sanは平均で1.3倍の高速化を達成し、BLEUスコアの低下はほとんどない。
• Aanモデルと組み合わせた場合、Sanは1.8倍の高速化を達成し、異なる高速化技術間の相乗効果が明確に示された。
• SanとAanを組み合わせたシステムは、キャッシュなしのベースラインに比べて最大16倍速くなった。
• 訓練過程において、隣接層間の注意重み分布のJSダイバージェンスが顕著に低下しており、層が自然に類似した注意パターンに収束していることが示された。
• Sanの訓練尤度はベースラインと比べてわずかに高い水準を示しており、共有制約下でもモデルが有効かつ安定していることが示された。
• 本手法は実装が簡単であり、既存技術と直交的で、さまざまなTransformerバリアントに直接適用可能である。