HUOXIU

大規模モデルアーキテクチャをより深く理解したいですか?まずは計算パラメータの数を確認することから始めましょう。

編集者注:新しい機械学習アーキテクチャ(そして他の新しいテクノロジー)を理解する最も効率的な方法は、実際にゼロから実装してみることです。しかし、さらに簡単な方法があります。それは、パラメータの数を計算することです。

パラメータの数を計算することで、読者はモデルアーキテクチャをより深く理解し、ソリューションに未発見のエラーがないか確認することができます。

この記事では、Transformer モデルのパラメータの数を計算するための数式と、それほど正確ではない簡略化されたバージョンを示します。これにより、読者は Transformer ベースのモデル内のパラメータの数を迅速に推定できます。

翻訳はこちらです。お楽しみください!

著者|ドミトロ・ニコライエフ(ディミッド)

編纂者:岳陽

新しい機械学習アーキテクチャ(そして他のあらゆる新しいテクノロジー)を理解する最も効率的な方法は、それをゼロから実装することです。これは非常に複雑で時間がかかり、時にはほぼ不可能になることもありますが、あらゆる技術的な詳細を理解するのに最適な方法です。例えば、類似のコンピューティングリソースやデータがなければ、私たちのソリューションに未発見のエラーがないと確信することはできません。

しかし、さらに簡単な方法があります。それは、パラメータの数を計算することです。これは論文を読むよりもはるかに難しいですが、より深く掘り下げて、新しいアーキテクチャのコンポーネント(この記事では、Transformerのエンコーダとデコーダのコンポーネント)を完全に理解しているかどうかを確認することができます。

この問題は、新しい ML アーキテクチャを理解する 3 つの方法を示す次の図を使用して考えることができます。円の大きさは、アーキテクチャの理解レベルを示しています。

本論文では主にTransformerアーキテクチャを研究し、 PyTorchのTransformerEncoderLayer [1]およびTransformerDecoderLayer [2]クラスにおけるパラメータ数の計算方法について考察します。そのため、このアーキテクチャの構成がもはや謎に包まれていないことを確認する必要があります。

TLDR(要約)

(この記事はかなり長いです。詳しく読みたくない場合や時間がない場合は、概要セクションまで飛ばしてください。)

パラメータを計算するためのすべての式がまとめられている「結論」セクションをお読みください。

この記事では、パラメータの数を計算するための数式だけでなく、これらの数の精度は低いものの近似バージョンも提供し、Transformer ベースのモデル内のパラメータの数を迅速に見積もることができます。

01 トランスフォーマーアーキテクチャ

Transformerアーキテクチャは、2017年に論文「Attention Is All You Need」[3]で提案され、長距離の依存関係を効果的に捉える能力があるため、自然言語処理やコンピュータービジョンのタスクにおける標準アーキテクチャとなっています。

2023年初頭には、テキスト画像生成モデル[5]の人気を受けて、拡散モデル[4]が急速に普及しました。おそらく近い将来、拡散モデルはTransformer、LSTM、CNNと同様に、様々なタスクにおける最先端技術となるでしょう。まずはTransformerについて見ていきましょう。

この記事は、既に多くの優れた記事で述べられているように、Transformerアーキテクチャを詳しく説明するものではありません。この記事では、Transformerを様々な視点から考察したり、詳細について議論したりすることを目指しています。このアーキテクチャに関する学習リソースをさらに探している方は、いくつかお勧めできるものがあるので、もしそうでない場合は、このまま読み進めてください。

1.1 Transformerリソースの詳細

Transformer アーキテクチャのより詳細な概要を知りたい場合は、次の資料をお読みください (インターネット上には多くの技術的なコンテンツがあることに注意してください。私は個人的にこれらを好みます)。

  • まず、公式論文[3]を読んでみてください。Transformerを初めて使う場合は、論文を読むのが最善の方法ではないかもしれませんが、見た目ほど複雑ではありません。Explainpaper(マウスでマークしたテキストをAIが解釈するツール)を使って、この論文[6]や他の論文を読むのも良いでしょう。

  • ジェイ・アラマーの「グレート・イラストレイテッド・トランスフォーマー」[7]。記事を読むのが苦手な方は、同じ著者によるYouTube動画[8]を視聴することもできます。

  • Lukasz KaiserによるGoogle Brainでの講義「Awesome Tensor2Tensor」 [9]。

  • 実践的な経験を積み、さまざまなTransformerモデルを使用してアプリケーションを構築したい場合は、Hugging Faceコース[10]を参照してください。

1.2 オリジナルトランスフォーマー

まず、Transformer の基本を確認しましょう。

Transformerアーキテクチャは、エンコーダ(左側)とデコーダ(右側)という2つのコンポーネントで構成されています。エンコーダは入力トークンシーケンスを受け取り、隠れ状態シーケンスを生成し、デコーダはこの隠れ状態シーケンスを受け取り、出力トークンシーケンスを生成します。

Transformer のアーキテクチャ図 (https://arxiv.org/pdf/1706.03762.pdf より)

エンコーダとデコーダはどちらも同一の層で構成されています。エンコーダの層は、マルチヘッドアテンション(1 - ここと下の数字は下の図の番号付き部分に対応しています)と、層正規化(3)とスキップ接続を備えたフィードフォワードニューラルネットワーク(2)で構成されています。

デコーダーはエンコーダーに似ていますが、最初のマルチヘッドアテンション(4)(機械翻訳タスクではマスクされるため、デコーダーは将来のトークンを見て不正行為を行うことはできません)とプレフィックスネットワーク(5)に加えて、 2番目のマルチヘッドアテンション機構(6)も備えています。これにより、デコーダーは出力を生成する際に、エンコーダーから提供されたコンテキストを使用することができます。エンコーダーと同様に、デコーダーにもレイヤー正規化(7)とスキップ接続コンポーネントが備わっています。

番号付きコンポーネントを含むトランスフォーマーアーキテクチャ図

https://arxiv.org/pdf/1706.03762.pdf より

入力埋め込み層(位置エンコーディングを含む)と最終出力層(線形 + ソフトマックス)はTransformerのコンポーネントとはみなさず、エンコーダーとデコーダーのブロックのみに焦点を当てます。これは、これらのコンポーネントが特定のタスクや埋め込み手法に適しているのに対し、エンコーダーとデコーダーのスタックは他のアーキテクチャの基盤となるためです。

このアーキテクチャの例には、エンコーダ用の BERT ベースのモデル ( BERT、RoBERTa、ALBERT、DeBERTa など)、デコーダ用の GPT ベースのモデル ( GPT、GPT-2、GPT-3、ChatGPT )、完全なエンコーダ/デコーダ フレームワーク上に構築されたモデル ( T5、BART など) などがあります。

このアーキテクチャでは 7 つのコンポーネントにラベルを付けていますが、そのうち本当に一意なのは 3 つだけであることがわかります。

  • マルチヘッドアテンション;

  • フィードフォワードネットワーク。

  • レイヤーの正規化。

Transformer コンポーネントは、論文 https://arxiv.org/pdf/1706.03762.pdf からのものです。

02 変圧器コンポーネントブロック

各モジュールの内部構造と必要なパラメータ数について考えてみましょう。このセクションでは、PyTorch[11]を用いて計算結果を検証していきます。

モデルブロックのパラメータ数を確認するには、次の関数[12]を使用します。

輸入トーチ
# https://discuss.pytorch.org/t/how-do-i-check-the-number-of-parameters-of-a-model/4325/9
def count_parameters(モデル: torch.nn.Module) -> int:
""" PyTorch モデルの学習可能なパラメータの数を返します """
p.requires_grad の場合、model.parameters() 内の p について sum(p.numel() を返す)

始める前に、すべてのビルディングブロックは標準化されており、スキップ接続を使用していることに注意してください。つまり、すべての入力と出力の形状(より正確には、バッチサイズとトークン数は異なる可能性があるため、最後の番号は同じである必要があります元の論文では、この番号(d_model)は512です。

2.1 多頭注意

アテンションメカニズムはTransformerアーキテクチャの鍵となります。しかし、設計の動機や技術的な詳細にかかわらず、アテンションメカニズムは数回の行列乗算のみを必要とします。

Transformer マルチヘッドアテンションアーキテクチャ図

論文より https://arxiv.org/pdf/1706.03762.pdf

各ヘッドのアテンションを計算した後、すべてのヘッドを連結し、線形層(W_O行列)に渡します。逆に、各ヘッドはスケーリングされたドット積アテンション層であり、クエリ、キー、値(それぞれW_Q、W_K、W_V行列)という3つの独立した行列が乗算されます。これらの3つの行列はヘッドごとに異なるため、添え字iが付いています。

最後の線形層(W_O)はd_modelからd_modelまでの形状を持ちます。他の3つの行列(W_Q、W_K、W_V)もd_modelからd_qkvまでの形状を持ちます。

上の画像では、d_qkv は原論文では d_k または d_v と表記されていることに注意してください。これらの行列は形状は異なるものの、ほぼ同じであるため、この名前の方が直感的だと思います。

さらに、d_qkv = d_model / num_heads(本文中のh)であることに注意してください。これが、後続の連結が正しく行われるようにするために、d_modelがnum_headsで割り切れる必要がある理由です。

上の画像ですべての中間段階の形状を確認することで、これを自分でテストできます (正しい形状は右下隅にマークされています)。

したがって、各ヘッドごとに3つの小さな行列と1つの大きな最終行列が必要になります。では、バイアスを無視しない場合、いくつのパラメータが必要でしょうか?

Transformerのアテンションモジュールにおけるパラメータ数を計算する式。画像は著者提供。

この式があまり複雑にならないことを願っています。導出をできるだけ明確にしようと努力しました。ご安心ください!今後の式はさらに短くなります。

パラメータのおおよその数は、4 * d_model と 4 * d_model^2 を比較して無視できるほどです。PyTorch でテストしてみましょう。

トーチインポートnnから
d_model = 512
n_heads = 8 # `d_model` の約数でなければなりません
マルチヘッドアテンション = nn.MultiheadAttention(埋め込み次元 = d_model、数値ヘッド = n_heads)
print(count_parameters(multi_head_attention)) # 1050624
印刷(4 * (d_model * d_model + d_model)) # 1050624

数字が一致しました。素晴らしい仕事ができたということです!

2.2 フィードフォワードネットワーク

Transformerのフィードフォワードネットワークは、ReLU活性化関数を挟んだ2つの完全結合層で構成されています。このネットワークの内部部分は、入力と出力(入力と出力は同じである必要があります)よりも表現力に優れています。

一般的には、MLP(d_model, d_ff) -> ReLU -> MLP(d_ff, d_model) となり、元の論文では d_ff = 2048 となっています

フィードフォワードニューラルネットワークの説明図は、論文 https://arxiv.org/pdf/1706.03762.pdf からの引用です。

少し視覚化しても害はありません。

Transformerのフィードフォワードネットワーク。画像は著者提供。

パラメータの計算は非常に簡単です。重要なのは、それらを混同しないことです。

Transformerフィードフォワードネットワークのパラメータ数を計算する式。画像は著者提供。

このような単純なネットワークは、次のコードを使用して記述し、そのパラメータの数を確認できます(公式の PyTorch 実装では、ドロップアウトも使用されており、これについては後でエンコーダ/デコーダ コードで説明します。ただし、ご存知のとおり、ドロップアウト レイヤーにはトレーニング可能なパラメータがないため、ここでは簡単にするために省略します)。

トーチインポートnnから
クラスTransformerFeedForward(nn.Module):
__init__(self, d_model, d_ff): を定義します。
super(TransformerFeedForward、self).__init__()
self.d_model = d_model
自己.d_ff = d_ff
self.linear1 = nn.Linear(self.d_model, self.d_ff)
self.relu = nn.ReLU()
self.linear2 = nn.Linear(self.d_ff, self.d_model)
定義forward(self, x):
x = 自己.線形1(x)
x = 自己.relu(x)
x = 自己.線形2(x)
xを返す
d_model = 512
d_ff = 2048
feed_forward = TransformerFeedForward(d_model, d_ff)
print(count_parameters(feed_forward)) # 2099712
プリント(2 * d_model * d_ff + d_model + d_ff) # 2099712

画像内の数字をもう一度見てみると、まだ説明されていないコンポーネントが 1 つだけ残っています。

2.3 レイヤーの正規化

Transformerアーキテクチャの最後の構成要素は、レイヤー正規化です。簡単に言えば、これはスケーリング機能を備えたインテリジェントな(学習可能な)正規化手法であり、トレーニングプロセスの安定性を向上させます。

Transformer レイヤーの正規化。画像は著者提供。

ここでのトレーニング可能なパラメーターは、それぞれ d_model の次元を持つ gamma と beta の 2 つのベクトルです。

Transformerレイヤーの正規化モジュールにおけるパラメータ数を計算する式。画像は著者提供。

コードを使用して仮説をテストしてみましょう。

トーチインポートnnから
d_model = 512
レイヤー正規化 = nn.LayerNorm(d_model)
print(count_parameters(layer_normalization)) # 1024
プリント(d_model * 2) # 1024

素晴らしい!近似計算では、レイヤーに正規化されたパラメータの数は、フィードフォワードネットワークやマルチヘッドアテンションブロックの数よりもはるかに少ないため、この数は無視できます(ただし、このモジュールは数回登場します)。

03 完全な式を導出する

これですべての準備が整いましたので、エンコーダー/デコーダー モジュール全体のパラメーターを計算できます。

3.1 PyTorchで実装されたエンコーダとデコーダ

エンコーダーは、アテンション ブロック、フィードフォワード ネットワーク、および 2 つの正規化レイヤーで構成されていることに注意してください。

トランスフォーマーエンコーダ。出典:論文 https://arxiv.org/pdf/1706.03762.pdf

PyTorchコードの詳細を調べて、すべてのコンポーネントが適切に配置されていることを確認できます。マルチヘッドアテンション機構は赤(左)、フィードフォワードネットワークは青レイヤー正規化は緑(PyCharmのPythonコンソールのスクリーンショット)で示されています。

PyTorch TransformerEncoderLayer。画像は著者提供。

3.2 最終式

確認ができたら、引数の数を計算する次の関数を記述できます。これはたった3行のコードで、1行にまとめることもできます。関数の残りの部分は説明用のドキュメント文字列です。

 def transformer_count_params(d_model=512, d_ff=2048, エンコーダー=True, approx=False):
「」
Transformer Encoder/Decoder のパラメータの数を計算します。
数式は次のとおりです。
マルチヘッドアテンション: 4*(d_model^2 + d_model)
approx=Falseの場合、4*d_model^2、それ以外の場合
フィードフォワード: 2*d_model*d_ff + d_model + d_ff
approx=Falseの場合、それ以外の場合は2*d_model*d_ff
レイヤーの正規化: approx=False の場合は 2*d_model、それ以外の場合は 0
エンコーダー ブロックは次の要素で構成されます。
1つのマルチヘッドアテンションブロック、
1つのフィードフォワードネット、および
2 層の正規化。
デコーダー ブロックは次の要素で構成されます。
2つのマルチヘッドアテンションブロック、
1つのフィードフォワードネット、および
3 層の正規化。
:param d_model: (int) モデルの次元
:param d_ff: (int) フィードフォワードニューラルネットワークの内部次元
:param encoder: (bool) Trueの場合、エンコーダのパラメータの数を返します。
それ以外の場合はデコーダー
:param approx: (bool) Trueの場合、結果は近似値です(数式を参照)
:return: (int) Transformer Encoder/Decoderの学習可能なパラメータの数
「」
attention = 4 * (d_model ** 2 + d_model) 近似でない場合は 4 * d_model ** 2
feed_forward = 2 * d_model * d_ff + d_model + d_ff (近似でない場合は 2 * d_model * d_ff)
layer_norm = 2 * d_model 近似でない場合は 0
注意 + フィードフォワード + 2 * レイヤーノルムを返す \
エンコーダの場合、そうでない場合、2 * attention + feed_forward + 3 * layer_norm

今こそそれをテストする時です。

トーチインポートnnから
エンコーダーレイヤー = nn.TransformerEncoderLayer(d_model=512, nhead=8)
print(count_parameters(encoder_layer)) #3152384
print(transformer_count_params(d_model=512, d_ff=2048, encoding=True, approx=False)) # 3152384
print(transformer_count_params(d_model=512, d_ff=2048, encoding=True, approx=True)) # 3145728
# 約0.21%の差
デコーダーレイヤー = nn.TransformerDecoderLayer(d_model=512, nhead=8)
print(count_parameters(デコーダーレイヤー)) # 4204032
print(transformer_count_params(d_model=512, d_ff=2048, encoding=False, approx=False)) # 4204032
print(transformer_count_params(d_model=512, d_ff=2048, encoding=False, approx=True)) # 4194304
# 約0.23%の差

正確な式は正しいです。つまり、すべての構成要素を正しく特定し、個々の要素に分解できたということです。興味深いことに、近似式では比較的小さな値(数百万ではなく数千)を無視したため、結果に対する誤差は約0.2%です!しかし、これらの式をより単純化する別の方法があります。

注意ブロックのパラメータ数は、おおよそ4 * d_model^2です。d_modelが重要なハイパーパラメータであることを考慮すると、計算的には単純に思えます。しかし、フィードフォワードネットワークの場合、d_ffを知る必要があります。式は2 * d_model * d_ffです。

d_ffは、式の中で覚えておく必要がある別のハイパーパラメータなので、それを取り除く方法を考えてみましょう。上で見たように、d_model = 512のとき、d_ff = 2048なので、d_ff = 4 * d_modelとなります。

多くのTransformerモデルでは、この仮定は理にかなっており、パラメータのおおよその数を提供しながら、式を大幅に簡素化します。結局のところ、正確な数を知りたい人は誰もいません。数十万なのか数千万なのかを知りたいだけなのです。

エンコーダとデコーダのおおよその式。画像は著者提供。

扱うべき桁数を把握するために、乗数を丸めることもできます。これにより、各エンコーダ/デコーダ層には10 * d_model^2のパラメータが与えられます。

04 結論

今日導き出したすべての式をまとめてみましょう。

著者提供の画像付きの公式概要。

この記事ではTransformerのエンコーダ/デコーダブロックのパラメータ数を計算していますが、もちろん、すべての新しいモデルに対してパラメータを計算することは推奨しません。このアプローチを選択した理由は、Transformerについて調べ始めた際に、このトピックに関する記事が全く見つからなかったことに驚いたからです。

パラメータの数はモデルの複雑さや学習に必要なデータ量を示す指標ですが、モデルのアーキテクチャをより深く理解するための方法の一つに過ぎません。ぜひ、様々なハイパーパラメータを使ってコードを見て、実装して、実行してみるなど、探求と実験をしてみてください。さあ、学び続けて、人工知能の楽しさを味わってください!

終わり

参考文献

1.https://pytorch.org/docs/stable/generated/torch.nn.TransformerEncoderLayer.html

2.https://pytorch.org/docs/stable/generated/torch.nn.TransformerDecoderLayer.html

3.https://arxiv.org/abs/1706.03762

4. https://techcrunch.com/2022/12/22/a-brief-history-of-diffusion-the-tech-at-the-heart-of-modern-image-generating-ai/

5.https://www.washingtonpost.com/technology/interactive/2022/ai-image-generator/

6.https://www.explainpaper.com/papers/attention

7.https://jalammar.github.io/illustrated-transformer/

8.https://youtu.be/-QH8fRhqFHM

9.https://www.youtube.com/watch?v=rBCqOTEfxv

10.https://huggingface.co/course/chapter1/1

11. https://pytorch.org/

12.https://discuss.pytorch.org/t/how-do-i-check-the-number-of-parameters-of-a-model/4325/9

この記事は、原著者の許可を得てBaihai IDPによって翻訳されました。翻訳の転載をご希望の場合は、お問い合わせください。

オリジナルリンク:

https://towardsdatascience.com/トランスフォーマーモデルのパラメータ数を推定する方法-ca0f57d8dff0