Nyström形式：ニュストローム法による線形時間とメモリでのセルフアテンションの近似

Nyström形式：線形時間とメモリでのセルフアテンションの近似方法、ニュストローム法

はじめに

トランスフォーマーは、さまざまな自然言語処理やコンピュータビジョンのタスクで優れた性能を発揮しています。その成功は、自己注意メカニズムによるものであり、入力のすべてのトークン間のペアワイズな相互作用を捉えることができます。しかし、標準の自己注意メカニズムは、入力シーケンスの長さ n n n （ここで n n n は入力シーケンスの長さ）に対して O ( n 2 ) O(n^2) O ( n 2 ) の時間とメモリの複雑さを持ち、長い入力シーケンスでのトレーニングには高コストです。

Nyströmformer は、標準の自己注意を O ( n ) O(n) O ( n ) の複雑さで近似する効率的なトランスフォーマーモデルの一つです。Nyströmformer は、標準の自己注意の効率を改善しながら、さまざまな後続の自然言語処理やコンピュータビジョンのタスクで競争力のある性能を発揮します。このブログ投稿の目的は、読者に Nyström 法の概要と、どのように自己注意を近似するために適応できるかを説明することです。

行列近似のための Nyström 法

Nyströmformer の核心にあるのは、行列近似のための Nyström 法です。この方法により、行列の一部の行と列をサンプリングすることで行列を近似することができます。まず、計算が非常にコストのかかる n × n n \times n n × n の行列 P P P を考えてみましょう。その代わりに、Nyström 法を使用して近似することができます。まず、P P P から m m m の行と列をサンプリングします。次に、サンプリングされた行と列を次のように配置します：

ブロック行列として P を表現する

• プロキシマルポリシーオプティマイゼーション（PPO）
• Sentence Transformersモデルのトレーニングと微調整
• TF Servingを使用してKubernetes上に🤗 ViTをデプロイする

これで、4つの部分行列 A P , B P , F P , C P A_P, B_P, F_P, C_P A P ​ , B P ​ , F P ​ , C P ​ が得られました。それぞれのサイズは m × m , m × ( n − m ) , ( n − m ) × m m \times m, m \times (n – m), (n – m) \times m m × m , m × ( n − m ) , ( n − m ) × m および ( n − m ) × ( n − m ) (n – m) \times (n – m) ( n − m ) × ( n − m ) です。m m m 個のサンプリングされた列は A P A_P A P ​ と F P F_P F P ​ に含まれており、m m m 個のサンプリングされた行は A P A_P A P ​ と B P B_P B P ​ に含まれています。したがって、A P , B P , F_P, A P ​ , B P ​ , F P ​ のエントリはわかっていますが、C P C_P C P ​ を推定する必要があります。Nyström 法によると、C P C_P C P ​ は次のようになります：

C P = F P A P + B P C_P = F_P A_P^+ B_P C P ​ = F P ​ A P + ​ B P ​

ここで、+ + + は Moore-Penrose 逆行列（または擬似逆行列）を意味します。したがって、行列 P P P の Nyström 近似は次のように書くことができます：

P の Nyström 近似

2行目に示されているように、P ^ \hat{P} P ^ は3つの行列の積として表現することができます。その理由は後ほど明らかになります。

Nyström 法で自己注意を近似できるか？

私たちの目標は、標準の自己注意における softmax 行列を近似することです：S = softmax Q K T d \frac{QK^T}{\sqrt{d}} d ​ Q K T ​

ここで、Q Q Q と K K K はクエリとキーを示します。上記で説明した手順に従って、S S S から m m m の行と列をサンプリングし、4つの部分行列を形成し、S ^ \hat{S} S ^ を取得します：

S の Nyström 近似

しかし、S S S から列をサンプリングするとはどういう意味でしょうか？それは、各行から1つの要素を選択することを意味します。S の計算方法を思い出してみてください：最終的な操作は行ごとの softmax です。1つの行のエントリを見つけるには、他のすべてのエントリ（softmax の分母のため）にアクセスする必要があります。したがって、1つの列をサンプリングするには、行列の他のすべての列を知る必要があります。したがって、softmax 行列を直接的に Nyström 法で近似することはできません。

ニューストローム法を使って自己注意を近似する方法はどのように適応できますか？

著者らは、S S S からサンプリングする代わりに、クエリとキーからランドマーク（またはニューストロームポイント）をサンプリングすることを提案しています。クエリのランドマークとキーのランドマークをそれぞれ Q ~ \tilde{Q} Q ~ と K ~ \tilde{K} K ~ と表記します。Q ~ \tilde{Q} Q ~ と K ~ \tilde{K} K ~ は、S S S のニューストローム近似に対応する3つの行列を構築するために使用できます。以下の行列を定義します：

F ~ = s o f t m a x ( Q K ~ T d ) A ~ = s o f t m a x ( Q ~ K ~ T d ) + B ~ = s o f t m a x ( Q ~ K T d ) \tilde{F} = softmax(\frac{Q\tilde{K}^T}{\sqrt{d}}) \hspace{40pt} \tilde{A} = softmax(\frac{\tilde{Q}\tilde{K}^T}{\sqrt{d}})^+ \hspace{40pt} \tilde{B} = softmax(\frac{\tilde{Q}K^T}{\sqrt{d}}) F ~ = s o f t m a x ( d ​ Q K ~ T ​ ) A ~ = s o f t m a x ( d ​ Q ~ ​ K ~ T ​ ) + B ~ = s o f t m a x ( d ​ Q ~ ​ K T ​ )

F ~ \tilde{F} F ~ , A ~ \tilde{A} A ~ , および B ~ \tilde{B} B ~ のサイズは、それぞれ ( n × m , m × m , and m × n ) です。ニューストローム近似の3つの行列を、新しく定義した行列に置き換えて、代替ニューストローム近似を得ます：

S ^ = F ~ A ~ B ~ = s o f t m a x ( Q K ~ T d ) s o f t m a x ( Q ~ K ~ T d ) + s o f t m a x ( Q ~ K T d ) \begin{aligned}\hat{S} &= \tilde{F} \tilde{A} \tilde{B} \\ &= softmax(\frac{Q\tilde{K}^T}{\sqrt{d}}) softmax(\frac{\tilde{Q}\tilde{K}^T}{\sqrt{d}})^+ softmax(\frac{\tilde{Q}K^T}{\sqrt{d}}) \end{aligned} S ^ ​ = F ~ A ~ B ~ = s o f t m a x ( d ​ Q K ~ T ​ ) s o f t m a x ( d ​ Q ~ ​ K ~ T ​ ) + s o f t m a x ( d ​ Q ~ ​ K T ​ )

これは、自己注意メカニズムのソフトマックス行列のニューストローム近似です。この行列に値（ V V V ）を掛けることで、自己注意の線形近似を得ます。Q K T QK^T Q K T の積を計算せずに、O ( n 2 ) O(n^2) O ( n 2 ) の計算量を回避しています。

ランドマークはどのように選択されますか？

著者らは、Q Q Q と K K K から m m m 行をサンプリングする代わりに、セグメント平均を使用して Q ~ \tilde{Q} Q ~ と K ~ \tilde{K} K ~ を構築することを提案しています。この手順では、n n n トークンが m m m セグメントにグループ化され、各セグメントの平均が計算されます。理想的には、m m m は n n n よりもはるかに小さくなります。論文の実験によると、標準の自己注意や他の効率的な注意機構と比較して、たった 32 32 3 2 または 64 64 6 4 のランドマークを選択するだけでも、長いシーケンスの長さ（ n = 4096 n=4096 n = 4 0 9 6 または 8192 8192 8 1 9 2 ）でも競争力のあるパフォーマンスが得られます。

アルゴリズム全体は、以下の図によって論文で要約されています：

ニュストレム法による効率的なセルフアテンション

上記の3つのオレンジ色の行列は、キーとクエリのランドマークを使用して構築した3つの行列に対応しています。また、DConvボックスがあることにも注意してください。これは、1D深さ畳み込みを使用して値に追加されるスキップ接続に対応しています。

ニュストレムフォーマーはどのように実装されていますか？

ニュストレムフォーマーの元の実装はこちらで見つけることができ、HuggingFaceの実装はこちらで見つけることができます。HuggingFaceの実装からいくつかのコードの行を見てみましょう（いくつかのコメントを加えました）。正規化、アテンションマスキング、および深さ畳み込みなどの詳細は、単純化のために省略されていることに注意してください。

key_layer = self.transpose_for_scores(self.key(hidden_states)) # K
value_layer = self.transpose_for_scores(self.value(hidden_states)) # V
query_layer = self.transpose_for_scores(mixed_query_layer) # Q

q_landmarks = query_layer.reshape(
    -1,
    self.num_attention_heads,
    self.num_landmarks,
    self.seq_len // self.num_landmarks,
    self.attention_head_size,
).mean(dim=-2) # \tilde{Q}

k_landmarks = key_layer.reshape(
    -1,
    self.num_attention_heads,
    self.num_landmarks,
    self.seq_len // self.num_landmarks,
    self.attention_head_size,
).mean(dim=-2) # \tilde{K}

kernel_1 = torch.nn.functional.softmax(torch.matmul(query_layer, k_landmarks.transpose(-1, -2)), dim=-1) # \tilde{F}
kernel_2 = torch.nn.functional.softmax(torch.matmul(q_landmarks, k_landmarks.transpose(-1, -2)), dim=-1) # \tilde{A} before pseudo-inverse

attention_scores = torch.matmul(q_landmarks, key_layer.transpose(-1, -2)) # \tilde{B} before softmax

kernel_3 = nn.functional.softmax(attention_scores, dim=-1) # \tilde{B}
attention_probs = torch.matmul(kernel_1, self.iterative_inv(kernel_2)) # \tilde{F} * \tilde{A}
new_value_layer = torch.matmul(kernel_3, value_layer) # \tilde{B} * V
context_layer = torch.matmul(attention_probs, new_value_layer) # \tilde{F} * \tilde{A} * \tilde{B} * V

HuggingFaceでのニュストレムフォーマーの使用方法

マスクされた言語モデリング（MLM）用のニュストレムフォーマーは、HuggingFaceで利用できます。現在、さまざまなシーケンス長に対応する4つのチェックポイントがあります：nystromformer-512、nystromformer-1024、nystromformer-2048、およびnystromformer-4096。ランドマークの数mは、NystromformerConfigのnum_landmarksパラメータを使用して制御することができます。ニュストレムフォーマーをMLMに使用する最小の例を見てみましょう：

from transformers import AutoTokenizer, NystromformerForMaskedLM
import torch

tokenizer = AutoTokenizer.from_pretrained("uw-madison/nystromformer-512")
model = NystromformerForMaskedLM.from_pretrained("uw-madison/nystromformer-512")

inputs = tokenizer("Paris is the [MASK] of France.", return_tensors="pt")

with torch.no_grad():
    logits = model(**inputs).logits

# retrieve index of [MASK]
mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]

predicted_token_id = logits[0, mask_token_index].argmax(axis=-1)
tokenizer.decode(predicted_token_id)

または、パイプラインAPIを使用することもできます（すべての複雑さを処理してくれます）：

from transformers import pipeline
unmasker = pipeline('fill-mask', model='uw-madison/nystromformer-512')
unmasker("Paris is the [MASK] of France.")

結論

ニュストレムフォーマーは、他の線形セルフアテンション手法を上回りながら、標準のセルフアテンションメカニズムへの効率的な近似を提供します。このブログ記事では、ニュストレム法の概要とセルフアテンションへの活用方法について概説しました。ニュストレムフォーマーを展開またはファインチューニングして下流タスクに使用したい読者は、こちらのHuggingFaceのドキュメントを参照してください。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!