
書記が数学やるだけ#897 Reformer,Longformer(長文に対するメモリ削減)
Transformerが長い文章に対してメモリ消費量が急激に増加する問題を解決するモデルの例を示す。
問題

説明
従来のTransformerは,Scaled Dot-Product Attentionにおいて必要なメモリ数が文章の長さの2乗に比例(O(L^2))する。ReformerはattentionをLocality-Sensitive Hashing(LSH)という手法で効率化することで,メモリの必要量をO(nlogn)に抑える。
Reformer: The Efficient Transformer https://arxiv.org/pdf/2001.04451

通常の ResNet(Residual Network)では,順伝播時の各層の中間アクティベーションを保存する必要があるため,メモリ使用量は O(nld) となる。一方でRevNet(Reversible Network)では、各層の入力を出力から再計算できるため、順伝播の中間値を保存せずに逆伝播が可能である。層数 l が大きい場合、従来の ResNet に比べて メモリ使用量を約 l 倍削減できる。ReformerではRevNetを導入することでメモリ使用量をさらに削減している。
The Reversible Residual Network: Backpropagation Without Storing Activations https://arxiv.org/pdf/1707.04585

Longformer では,Sliding Window Attention により各トークンが自分の近傍トークンにのみ attention を向け,さらに Global Attention により特定の重要なトークンについては文書全体のすべてのトークンに attention を向ける。これにより,計算量を大幅に削減しつつ,長距離依存関係も考慮できる。
Longformer: The Long-Document Transformer https://arxiv.org/pdf/2004.05150

似たような手法としてSparse Transformerが挙げられる。
Generating Long Sequences with Sparse Transformers https://arxiv.org/pdf/1904.10509

解答
Locality-Sensitive Hashing (LSH) では,ベクトル xを固定されたランダム射影行列 R に投影し,正負の符号を見てハッシュ値 h(x) を決定する。類似したベクトルは同じハッシュ値になる確率が高いため,attention の計算対象を同じハッシュバケット内に制限できる。


16 個のトークンに対してハッシュを計算すると、各トークンがどのバケットに属するかが分かる。チャンク長 m=4 に分割した後、各チャンクは 自分のチャンク内のトークン と 1 チャンク前のトークンにだけ attention を計算する。これにより計算量を抑えることができる。

Longformer では、長い文書の attention を効率的に計算するために以下を組み合わせる:
・Sliding Window Attention
各トークンは左右 w個のトークンに対してのみ attention を計算する。
例:トークン 3 は [[2,4] に attention。
・Global Attention
特定の重要なトークン(例:トークン 1 と 8)は文書全体に対して attention を計算。
例:
トークン 1(グローバル) → 全トークン [1,2,3,4,5,6,7,8]
トークン 2 → ウィンドウ内 [1,2,3]

本記事のもくじはこちら: