self-attentionを代替する各種手法について
概要
Transformerはvision, textなど各種ドメインで成功を収めているが、計算量とメモリ消費が系列長の2乗に比例するため、長い系列を扱うのに課題がある。本稿ではこのような計算量の課題に対応することを目的として提案されたTransFormerの代替アーキテクチャについてまとめる。
以下に挙げる手法は自分が過去に論文を読んだことのあるものをまとめたものであり、網羅的ではないことに注意。
代替パターン
概ね以下のいずれかに分類される。
- attentionの間引き
- 近似計算
- 計算過程の最適化
表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] |
参考文献
- [1] Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions
- [2] MaxViT: Multi-Axis Vision Transformer
- [3] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
- [4] Vision Transformer with Deformable Attention
- [5] MLP-Mixer: An all-MLP Architecture for Vision
- [6] MetaFormer Is Actually What You Need for Vision
- [7] Co-Scale Conv-Attentional Image Transformers
- [8] Big Bird: Transformers for Longer Sequences
- [9] Reformer: The Efficient Transformer
- [10] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
- [11] Falcon - Hugging Face
Discussion