Lite Transformer with Long-Short Range Attention

shimopino commented 4 years ago

論文へのリンク

[arXiv:2004.11886] Lite Transformer with Long-Short Range Attention

著者・所属機関

Zhanghao Wu, Zhijian Liu, Ji Lin, Yujun Lin, Song Han

MIT
Shanghai Jiao Tong University

投稿日時（YYYY-MM-DD）

2020-04-24

1. どんなもの？

NLPでのデファクトスタンダードとなったTransformerは計算リソースを多く必要としておりモバイル端末などの制限された環境で動作させることが難しい。

本研究ではMulti-Head Attention部分を、局所的な特徴量の重みを畳み込みを使用して計算するモジュールと、長期的な依存関係をとらえるAttentionを使用して計算するモジュールに分けることで、Transformerを上回る精度を発揮しながらも、モバイル端末でも稼働できる軽量なモデルを提案した。

2. 先行研究と比べてどこがすごいの？

重量級であるTransformerモデルを軽量化する単純な方法としては、単語IDから単語ベクトルに変換する際に次元数を落とすことであるが、このやり方はモデルが単語の局所的・長期的な依存関係を補足する能力を落としてしまう。

Attention計算は自然言語処理や画像認識など多岐にわたって使用されているが、計算コストが非常に高く、入力するベクトルの次元数がNであり特徴量の次元数がdの場合計算コストはdN^2で増加していく。

この問題に対して、Attention層に入力する前段階でチャンネル数を小さくしておくことで計算コストを下げる手法がよく採用されており、元のTransformerの設計ではAttention層に入力する前にチャンネル数を4分の1にまで落としている。しかしこの手法ではAttentionが捉えることのできる依存関係を補足する能力までも落としてしまう。

計算コストを考えた際にAttention層はO(4Nd^2+N^2d)であり、FFNではO(2x4Nd^2)になる。機械翻訳タスクでは一度に入力される文章のトークン数は20から30であることを考えると計算コストを担っているのはFFNである可能性が高い。

実際に計算してみるとFFNの計算コストのほうが高かった。

本研究はTransformerのFFNが計算コストの大部分を占めており、ボトルネック構造が非効率であることを発見した。また局所的な依存関係を畳み込み層でとらえ、長期的な依存関係をAttention層でとらえるLSRAモジュールを導入したLite Transformerを提案した。

本手法は一般的な圧縮技術である枝刈りや量子化と組み合わせることで、約18倍モデルを圧縮できることを示した。

3. 技術や手法の"キモ"はどこにある？

BERTで学習したAttention層の重みの可視化は以下であり、それぞれの重みはSourceの文中の単語がTargetの文中のどの単語に着目しているのかを表している。

得られたAttentionの重みには特徴的なパターンが表れており、疎であり直交性を有していることがわかる。疎であることはつまり長期的な依存関係を捉えていることを示し、直交性は隣り合った単語同士の関係性を表している。

本研究では前者を"global"な関係、後者を"local"な関係と捉えている。

今回提案しているLSRAにより局所的な関係性を捉えるモジュールと長期的な関係性を捉えるモジュールを明確に分けることで、Attentionのみを使ったモデルと比較してより計算コストを落としモバイル端末でも実行可能なモデルを構築することを可能となる。

LSRAの計算の流れは以下のように2つに分岐する。またそれぞれのbranchに入力する際に入力ベクトルをチャンネル次元に沿って半分に分け計算し、続くFFNで2つの出力を混ぜることで計算コストを削減することができる。

さらに通常の畳み込み層ではなく、線形結合層とdepthwise畳み込み層を使用したモジュールを採用することでさらに計算コストを削減している。

これらの計算によって得られたLSRAのAttention層の重みは以下のように長期的な依存関係を補足するものになる。

4. どうやって有効だと検証した？

4.1 動作環境

トークン数が30個の文の対して500M Mult-Adds（あるいは1G FLOPs）を目安とする
MobileNetは7M程度の容量なので、今回は容量の上限を10Mに設定

4.2 結果

ドイツ語から英語への機械翻訳タスクであるIWSLTでの比較結果は以下になる。それぞれ同程度のパラメータ数や計算コストに対して、Lite Transformerはより高いBLUEを達成していることがわかる。

また英語からドイツ語への機械翻訳、英語からフランス語への機械翻訳タスクであるWMTでの比較結果は以下になる。先ほどと同じく同程度のパラメータ数と計算コストではLite Transformerがより高いBLUEを達成している。

同タスクでの計算コストとBLUEスコアのトレードオフの関係は以下のようになっており、元のTransformerからおよそ2.5倍の計算コストの削減に達成している。

Lite Transformerは枝刈りや量子化といった圧縮技術に対しても頑強な性能を発揮しており、無視できる程度の性能悪化でありながらモデルサイズを約18倍程度にまで圧縮できている。

AutoMLの探索によって得られたモデルであるETと比較すると、同程度のパラメータ数と計算コストでより高い精度をより安く達成していることがわかる。

また長い入力文を取り扱う要約タスクであるCNN-DailyMailでは、同程度の性能を約2分の1程度の計算コストとパラメータ数で達成している。

言語モデリングタスクであるWIKITEXT-103ではLite Transformerはより低いPPLを達成しており、Transformerに対して同程度の精度を約2分の1の計算コストで達成している。

5. 議論はあるか？

shimopino commented 4 years ago

https://github.com/mit-han-lab/lite-transformer

shimopino commented 4 years ago

BERTのAttention層の可視化と理解

Olga Kovaleva, Alexey Romanov, Anna Rogers, and Anna Rumshisky. Revealing the Dark Secrets of BERT. In EMNLP, 2019. 4
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D Manning. What Does BERT Look At? An Analysis of BERT’s Attention. In BlackboxNLP, 2020. 4

計算効率のいいConv

Bichen Wu, Xiaoliang Dai, Peizhao Zhang, Yanghan Wang, Fei Sun, Yiming Wu, Yuandong Tian, Peter Vajda, Yangqing Jia, and Kurt Keutzer. FBNet: Hardware-Aware Efficient ConvNet Design via Differentiable Neural Architecture Search. In CVPR, 2019a. 3

shimopino / papers-challenge