[論文レビュー] TensorCoder: Dimension-Wise Attention via Tensor Representation for Natural Language Modeling

Shuai Zhang, Peng Zhang+4 more|arXiv (Cornell University)|Jul 28, 2020

Topic Modeling参考文献 26被引用数 4

ひとこと要約

TensorCoderは、テンソル演算を用いてクエリ-キー相互作用を再定義することにより、自己注意機構の計算複雑度をO(N²d)からO(Nd²)に低減する次元別注意メカニズムを提案する。これにより、PTBのマスク言語モデル化タスクで13%の性能向上を達成し、標準Transformerに比べて3.5倍少ないFLOPsを実現。機械翻訳タスクでも計算量を削減しながら競争力ある結果を維持する。

ABSTRACT

Transformer has been widely-used in many Natural Language Processing (NLP) tasks and the scaled dot-product attention between tokens is a core module of Transformer. This attention is a token-wise design and its complexity is quadratic to the length of sequence, limiting its application potential for long sequence tasks. In this paper, we propose a dimension-wise attention mechanism based on which a novel language modeling approach (namely TensorCoder) can be developed. The dimension-wise attention can reduce the attention complexity from the original $O(N^2d)$ to $O(Nd^2)$, where $N$ is the length of the sequence and $d$ is the dimensionality of head. We verify TensorCoder on two tasks including masked language modeling and neural machine translation. Compared with the original Transformer, TensorCoder not only greatly reduces the calculation of the original model but also obtains improved performance on masked language modeling task (in PTB dataset) and comparable performance on machine translation tasks.

研究の動機と目的

• 標準Transformerの自己注意機構における2次計算複雑度の問題を解決し、長文系列へのスケーラビリティを向上させること。
• 計算複雑度を低減しながらも、トークン別注意の出力の有用性を維持する次元別注意メカニズムを開発すること。
• テンソル演算と畳み込みを用いて、次元別注意をTransformerアーキテクチャに統合すること。
• マスクを施した次元別注意メカニズムを用いて、デコーダーにおける自己回帰的性質を保持すること。
• FLOPsとメモリ使用量を削減することで、リソース制限のある環境における事前学習モデルの効率的デプロイを可能にすること。

提案手法

• 標準的なクエリ-キードット積注意（QKᵀ）を、QᵀKを用いる次元別注意メカニズムに置き換え、計算複雑度をO(N²d)からO(Nd²)に低減する。
• Khatri-Rao（KR）積を用いて、注意を3次元テンソルとして表現することで、暗黙的および明示的な出力表現を可能にする。
• テンソル表現に畳み込み演算を適用し、多次元の相互作用特徴を抽出・統合する。
• 情報伝播を以前の位置に制限するマスクを施した次元別注意メカニズムを構築し、自己回帰的生成を保持する。
• 標準のマルチヘッド注意を置き換える形で、次元別注意をTransformerのエンコーダーおよびデコーダー層に統合する。
• テンソルベースの計算と表現を用いて、エンコーダーとデコーダー構造における注意の設計と実装を統一する。

実験結果

リサーチクエスチョン

• RQ1 次元別注意メカニズムは、自己注意機構の2次計算複雑度を低減しつつ、その表現力も維持できるか？
• RQ2 次元別注意は、既存のTransformerアーキテクチャと互換性のある出力形式を生成できるように設計できるか？
• RQ3 テンソルベースの表現と演算は、系列モデリングにおける次元間相互作用を効果的に捉えることができるか？
• RQ4 提案されたマスクを施した次元別注意は、自己回帰的デコーディングにおける自己回帰的生成特性を維持できるか？
• RQ5 次元別注意の低複雑度は、長文系列およびリソース制限のあるNLPタスクにおける効率性と性能向上にどの程度寄与するか？

主な発見

• PTBのマスク言語モデル化タスクにおいて、TensorCoderは標準Transformerに比べ13%の性能向上を達成し、推論時間も3分の1に削減された。
• WMT16 En-De翻訳タスクにおいて、TensorCoderはBLEUスコア33.1（標準Transformerは33.5）という同等の性能を達成したが、注意層のFLOPsは3.5倍少ない。
• IWSLT16タスクでは、TensorCoderが標準TransformerのFLOPsを3分の1に削減し、顕著な効率性の向上を示した。
• 次元別注意メカニズムにより、計算複雑度がO(N²d)からO(Nd²)に低減され、シーケンス長に対して線形スケーリングが可能になった。
• モデルは、メモリと計算要求を大幅に削減しながらも、さまざまなタスクで競争力ある性能を維持しており、エッジデバイスへの適性が高まった。
• Khatri-Rao積とテンソル畳み込みの使用により、注意機構内での効果的な特徴抽出と統合が可能となり、低コストで高い性能を実現した。