Zenn
💨

self-attentionを代替する各種手法について

2023/09/22に公開
#
transformer
#
attention
tech

概要

Transformerはvision, textなど各種ドメインで成功を収めているが、計算量とメモリ消費が系列長の2乗に比例するため、長い系列を扱うのに課題がある。本稿ではこのような計算量の課題に対応することを目的として提案されたTransFormerの代替アーキテクチャについてまとめる。

以下に挙げる手法は自分が過去に論文を読んだことのあるものをまとめたものであり、網羅的ではないことに注意。

代替パターン

概ね以下のいずれかに分類される。

  1. attentionの間引き
  2. 近似計算
  3. 計算過程の最適化

表1に画像、言語の各ドメインについてこれらの手法をまとめる。

Attentionの間引き

Attentionは長さNの系列について、全てのペアについてattentionを計算しているが、関連性の強いペアのみに限定することで計算量を減らすというのがこの種の手法の方針である。

画像ドメインにいては、「画像の近傍のピクセルほど関連性が強い」という仮定に基づき、局所的なattentionが用いられることが多い[1-4]。一方で、attentionの適用範囲を近傍のピクセルに限ることで受容野が狭まる問題が生じる。これに対して、多くのアーキテクチャでは層を重ねるごとに受容野が広くなるようにすることでこの問題を解決している(Convolutional networkと似ている)。

言語ドメインでは、attentionを計算するペアをデータサンプルごとにdynamicに決定する手法も提案されている[9]。

近似計算

dot-product attentionに変わる代替、あるいは近似アルゴリズムについても提案されている。これには、Mlp-Mixerをはじめとする、Multi-Head Self Attention(MHSA)ブロックを任意のブロックで置き換える一連の手法がある[5-7]。

計算過程の最適化

FlashAttention[10]は、間引きなしの通常のAttention(full-attentionと呼ばれることが多い)を計算上の工夫によって効率化する試みである。近年では巨大言語モデルでも採用されている[11]。

表1. self-attentionの代替手法

タスクドメイン 分類 手法 代表的なアーキテクチャ
画像 attentionの間引き spatial-reduction attention PVT[1]
画像 attentionの間引き Block/Grid attention MaxViT[2]
画像 attentionの間引き Shifted Window Swin Transformer[3]
画像 attentionの間引き Deformable Attention DAT[4]
画像 近似計算 MLP Mlp-mixer[5]
画像 近似計算 mean pooling PoolFormer[6]
画像 近似計算 Factorized Attention CoaT[7]
言語 attentionの間引き Sparse Attention BigBird[8]
言語 近似計算, attentionの間引き Locality-Sensitive Hash Reformer[9]
言語 計算過程の最適化 FlashAttention[10] Falcon[11]

参考文献

GitHubで編集を提案

Discussion