GQA、大規模モデルで一般的に使用されるアテンションメカニズム、および Pytorch コードの実装の詳細な説明-AI-php.cn

GQA、大規模モデルで一般的に使用されるアテンションメカニズム、および Pytorch コードの実装の詳細な説明

WBOY

リリース： 2024-04-03 17:40:09

転載

997 人が閲覧しました

グループ化されたクエリアテンションは、大規模な言語モデルにおけるマルチクエリアテンション手法であり、その目標は、MQA の速度を維持しながら MHA の品質を達成することです。グループ化されたクエリアテンションは、各グループ内のクエリが同じアテンションの重みを共有するようにクエリをグループ化します。これにより、計算の複雑さが軽減され、推論速度が向上します。

この記事では、GQA の考え方とそれをコードに変換する方法について説明します。

GQA は、「GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints」という論文で提案されています。これは非常にシンプルでクリーンなアイデアであり、マルチヘッドに基づいて構築されています。注意、強度以上です。

GQA、大規模モデルで一般的に使用されるアテンションメカニズム、および Pytorch コードの実装の詳細な説明

GQA

標準のマルチヘッドアテンションレイヤー (MHA) は、H クエリヘッドで構成されます。キーヘッドと値ヘッダーの構成。各ヘッドの寸法はDです。 Pytorch コードは次のとおりです。

from torch.nn.functional import scaled_dot_product_attention  # shapes: (batch_size, seq_len, num_heads, head_dim) query = torch.randn(1, 256, 8, 64) key = torch.randn(1, 256, 8, 64) value = torch.randn(1, 256, 8, 64)  output = scaled_dot_product_attention(query, key, value) print(output.shape) # torch.Size([1, 256, 8, 64])

ログイン後にコピー

クエリヘッダーごとに、対応するキーがあります。このプロセスを次の図に示します。

GQA、大規模モデルで一般的に使用されるアテンションメカニズム、および Pytorch コードの実装の詳細な説明

そして、GQA はクエリヘッダーを G 個のグループに分割し、各グループがキーと値を共有します。

GQA、大規模モデルで一般的に使用されるアテンションメカニズム、および Pytorch コードの実装の詳細な説明

視覚的な表現を使用すると、上で述べたように、GQA の動作原理を明確に理解できます。 GQA は非常にシンプルでクリーンなアイデアです。

Pytorch コードの実装

クエリヘッダーを G 個のグループに分割し、各グループが key と value を共有するコードを作成しましょう。 einops ライブラリを使用すると、テンソルに対して複雑な演算を効率的に実行できます。

まず、クエリ、キー、値を定義します。次に、アテンションヘッドの数を設定します。この数は任意ですが、num_heads_for_query % num_heads_for_key = 0、つまり割り切れる必要があることを確認する必要があります。私たちの定義は次のとおりです:

import torch  # shapes: (batch_size, seq_len, num_heads, head_dim) query = torch.randn(1, 256, 8, 64) key = torch.randn(1, 256, 2, 64) value = torch.randn(1, 256, 2, 64)  num_head_groups = query.shape[2] // key.shape[2] print(num_head_groups) # each group is of size 4 since there are 2 kv_heads

ログイン後にコピー

効率を向上させるために、seq_len 次元と num_heads 次元を交換するために、einops は次のように簡単に実行できます。

from einops import rearrange  query = rearrange(query, "b n h d -> b h n d") key = rearrange(key, "b s h d -> b h s d") value = rearrange(value, "b s h d -> b h s d")

ログイン後にコピー

次に、「グループ化」の概念をクエリマトリックスに導入する必要があります。

from einops import rearrange query = rearrange(query, "b (h g) n d -> b g h n d", g=num_head_groups) print(query.shape) # torch.Size([1, 4, 2, 256, 64])

ログイン後にコピー

上記のコードを使用して、2D を 2D に再形成します。定義したテンソルの場合、元の次元 8 (クエリ内のヘッドの数) は現在、 (キーと値のヘッドの数を一致させるため) 2 つのグループに分割し、各グループのサイズは 4 です。

最後で最も難しい部分は、注意スコアの計算です。しかし実際には、insum 演算を通じて 1 行で実行できます。

from einops import einsum, rearrange # g stands for the number of groups # h stands for the hidden dim # n and s are equal and stands for sequence length scores = einsum(query, key, "b g h n d, b h s d -> b h n s") print(scores.shape) # torch.Size([1, 2, 256, 256])

ログイン後にコピー

スコアテンソルは、上記の値テンソルと同じ形状をしています。仕組みを見てみましょう

einsum は次の 2 つのことを行います:

1. クエリとキーの行列乗算。私たちの場合、これらのテンソルの形状は (1,4,2,256,64) と (1,2,256,64) であるため、最後の 2 次元に沿った行列の乗算は (1,4,2,256,256) になります。

2. 2 番目の次元 (次元 g) の要素を合計します - 指定された出力形状で次元が省略されている場合、einsum はこの作業を自動的に完了します。ヘッダー内のキーと値の数を一致させます。

最後に、分数と値の標準的な乗算に注目してください:

import torch.nn.functional as F  scale = query.size(-1) ** 0.5 attention = F.softmax(similarity / scale, dim=-1)  # here we do just a standard matrix multiplication out = einsum(attention, value, "b h n s, b h s d -> b h n d")  # finally, just reshape back to the (batch_size, seq_len, num_kv_heads, hidden_dim) out = rearrange(out, "b h n d -> b n h d") print(out.shape) # torch.Size([1, 256, 2, 64])

ログイン後にコピー

最も単純な GQA 実装が完了しました。必要な Python コードは 16 行未満です。

GQA、大規模モデルで一般的に使用されるアテンションメカニズム、および Pytorch コードの実装の詳細な説明