[論文レビュー] Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers

Krzysztof Choromański, Valerii Likhosherstov+9 more|arXiv (Cornell University)|Jun 5, 2020

Machine Learning in Bioinformatics被引用数 40

ひとこと要約

Performerを導入する、Fast Attention Via Orthogonal Random features（FAVOR）を用いたTransformerの派生モデルで線形スケールのアテンションを実現し、理論的保証と後方互換性を備えた長文脈のタンパク質配列モデリングを可能にする。

ABSTRACT

Transformer models have achieved state-of-the-art results across a diverse range of domains. However, concern over the cost of training the attention mechanism to learn complex dependencies between distant inputs continues to grow. In response, solutions that exploit the structure and sparsity of the learned attention matrix have blossomed. However, real-world applications that involve long sequences, such as biological sequence analysis, may fall short of meeting these assumptions, precluding exploration of these models. To address this challenge, we present a new Transformer architecture, Performer, based on Fast Attention Via Orthogonal Random features (FAVOR). Our mechanism scales linearly rather than quadratically in the number of tokens in the sequence, is characterized by sub-quadratic space complexity and does not incorporate any sparsity pattern priors. Furthermore, it provides strong theoretical guarantees: unbiased estimation of the attention matrix and uniform convergence. It is also backwards-compatible with pre-trained regular Transformers. We demonstrate its effectiveness on the challenging task of protein sequence modeling and provide detailed theoretical analysis.

研究の動機と目的

• スパース性の事前知識に依存せず、長い生体シーケンス（例：タンパク質）に対してTransformerモデルをスケールさせる動機付け。
• 理論的保証と事前学習済みTransformersとの実用的な互換性を持つ線形時間のアテンション機構を開発する。
• タンパク質配列モデリングとImageNet64での有効性を示し、複雑さと収束を分析する。

提案手法

• Attentionがカーネルベースになり得る Generalized Attention (GA) を提示する。
• 正規のアテンションを Fast Attention Via Orthogonal Random features (FAVOR) に置換し、アテンション行列の偏りのない低ランク近似を得る。
• クエリとキー間のカーネル類似度を推定するためにランダム特徴写像を用い、全ての A 行列を作成せずに計算形 Q'K'ᵀ を得る。
• 分散を減らし近似品質を向上させるために Orthogonal Random Features (ORFs) を組み込む。
• O(Ld log d) の空間、O(Ld² log d) の時間という時間・空間計算量分析を示し、通常のアテンションの O(L²d) と比較する。
• 他のTransformerコンポーネントをそのままにして、アテンション機構だけを置換することで後方互換性を示す。

実験結果

リサーチクエスチョン

• RQ1 FAVOR は Generalized kernel-based attention に対して標準のアテンション行列を偏りのない推定で近似できるのか？
• RQ2 FAVOR ベースのアテンションは長いタンパク質配列に対して精度を維持しつつ、シーケンス長 L に対して線形にスケールするのか？
• RQ3 通常の Transformers およびスパースアテンション手法と比較して、長大なタンパク質配列モデリングとImageNet64における Performer の性能はどうか？
• RQ4 GAフレームワークにおけるFAVORの理論的保証（偏りなし性と一様収束）は何か？

主な発見

• FAVOR はトークン数に対して線形時間を達成し、空間は二次以上でなく準二次以下を実現し、二次のアテンションを低ランク近似に置換する。
• この手法はアテンション行列の偏りのない推定と一様収束保証を提供する。
• FAVORを用いて訓練されたPerformers はファインチューニング後にタンパク質配列タスクでTransformerの性能を回復またはほぼ一致させ、タンパク質におけるReformerのような一部のスパースアテンションベースのベンチマークより優れていることがある。
• Orthogonal random features は非構造的特徴と比較して近似誤差を低減し、下流の性能を向上させる。
• この手法は事前学習済みの通常のTransformersとAPI互換であり、通常のアテンションのスケーラブルな置換として使用できる。