Attention・FFN・MoE

このページの目次

logits の使い方を学んだ後は、その前段の表現計算を見ます。ある位置のベクトルを追い、注意力による情報の集約、FFN による変換、MoE の専門家選択を説明します。重みの配置と KV の容量は次の記事で扱います。

1つのDecoder層が表現をどのように更新するか

本稿では、自己回帰型Transformer Decoderの一般的な構造を例に挙げます。トークンIDはまずベクトルにマッピングされ、モデルは位置情報も表現する必要があります。位置情報は、位置ベクトルや注意機構造内のベクトルに対する位置依存変換などの方法で導入され、すべてのモデルが単純に2つの埋め込み(embedding)を加算するわけではありません。

1つの層には通常、注意機構造、前馈ネットワーク(FFN)、残差接続、正規化が含まれます。以下に示すのは理解を容易にするためのPre-norm形式です。元のTransformerやその後のモデルでは、正規化の位置、活性化関数、位置メカニズムに差異があります。Attention Is All You Need

入力表現 x各位置に1つのベクトルu = x + Attention(Norm(x))因果的注意機構造によりアクセス可能な位置を混合y = u + FFN(Norm(u))前馈ネットワークによる位置ごとの変換;MoEでここに置き換え可能簡略化されたPre-norm Decoder層複数の層を繰り返した後、出力投影を経てlogitsが得られます。正規化の位置や具体的なFFNの形式はアーキテクチャによって異なります。

注意機構造:現在の表現が他の位置をどのように活用するか

入力行列Xに対して、線形射影によりQ=XWQ、K=XWK、V=XWVが得られます。ある位置のクエリ(Query)が、アクセスが許可された位置のキー(Key)と内積を行い、スケーリング、マスク、ソフトマックスを経て重みが得られ、その後バリュー(Value)に対して重み付き和が計算されます。単一ヘッドの場合、以下のように表されます:

Attention(Q,K,V) = softmax(QKᵀ/√dₖ + M)V

因果的マスクMにより、未来の位置は現在の位置の注意計算に参加できず、通常現在の位置およびそれ以前の位置へのアクセスが許可されます。位置iの表現は未来のトークンを予測するために使用されるため、現在の位置が含まれていても未知の未来を「のぞき見」ているわけではありません。マルチヘッド注意機構造では、この計算を複数のグループに分割し、結果を結合します。ある1つの注意マップだけを見て、特定のヘッドが固定された意味を担っている、あるいはモデルの判断を完全に説明できると断定することはできません。

3つの位置を持つ教育的な例でマスクを見てみましょう。もし2番目の位置が3つの位置に対する元のスコアがすべて0の場合、マスク後、3番目の位置はアクセス不可となり、1番目と2番目の位置はそれぞれ1/2の重みを持ちます。もし対応するバリューが [2,0]、[0,4]、[99,99] の場合、結果は [1,2] となり、未来の位置の大きな値は混入しません。3番目の位置だけが、それまでの3つの位置を結合する資格を持ちます。

FFN、残差接続、層の深さ

FFNは通常、各位置のベクトルに対して線形変換、非線形活性化を行い、モデルの次元に射影します。現代の実装ではゲート制御がよく使用されます。ここで言う「位置ごと」とは、このサブレイヤーが位置間を直接混合しないことを指しますが、入力にはすでに以前の注意計算で集約された情報が含まれています。注意機構造とFFNは複数の層にわたって表現計算を構成し、「注意は見るだけで、FFNだけが考える」と単純に分割することはできません。

残差接続はサブレイヤーの増分を元の表現に加算し、情報と勾配に直接のパスを提供します。正規化は数値のスケールを制御するのに役立ちます。層数を増やすことは順序変換の深さを変更しますが、エクスパート数(MoEの場合)を増やすことは主に選択可能なパラメータの集合を変更するもので、これら2つの「拡大」は同等ではありません。活性化パラメータ数は、モデル内の推論ステップ数を測る指標でもありません。

マスクは加重結果をどう変えるか

内積と縮小は相対スコア、マスクは参加資格、softmax は配分を決めます。q=0 で因果境界を確認してから q を変え、寄与を見ます。重みの強調はモデル全体の説明ではありません。

図を準備しています
マスクは加重結果をどう変えるか

アクセスできない位置の重みは 0 で、利用可能な位置間で確率を配分して value を加重します。

稠密FFNからスパースエクスパートへ

稠密なFFN層では、各位置が同じ前馈重みを使用します。一方、スパースなMoE層では、複数のエクスパートネットワークが用意され、ルーターが現在の位置の表現に基づいてその一部を選択します。注意機構造などの他のコンポーネントは独自の計算を行うため、モデル全体が単に互いに排他的なサブモデルに分割されるわけではありません。

ルーターが3つのエクスパートに対して [0.6,0.3,0.1] の重みを出力し、この層でTop-2が選択され、選択された集合内で正規化されると、[2/3,1/3] が得られます。もし2つのエクスパートの出力がそれぞれ [3,0] と [0,6] であれば、結合結果は [2,2] になります。3番目のエクスパートはこの位置のエクスパート計算には参加しませんが、次の位置や次の層で選択される可能性があります。

現在の位置表現 hルーター重み 0.6 / 0.3 / 0.1エクスパート 1E₁(h) = [3, 0]エクスパート 2E₂(h) = [0, 6]エクスパート 3この位置では選択されず重み付き結果 [2, 2]重みの再正規化:2/3、1/3例:ルーターによるエクスパート選択と出力の結合これはTop-2かつ再正規化された教育的な例です。MoEのゲート制御や共有エクスパートの設計は多様です。

これは1つのゲート制御形式に過ぎません。Top-kの選択、重みの正規化方法、容量制御、MoEが適用される層などは、具体的な設計によって異なります。Switch TransformerはTop-1のスパースルーティングを示しました。Mixtralの公開されているアーキテクチャでは、各層で各トークンに対して2つの前馈エクスパートが選択されます。Switch Transformers、Mixtral of Experts

エクスパートが自動的に人間が定義した分野になるわけではない

「エクスパート」はサブネットワークの名前であり、システムが事前に数学エクスパート、法律エクスパート、コードエクスパートなどを設定していることを意味しません。トレーニング中に何らかの専門化が生じる可能性はありますが、実際のルーティングは表現、位置、最適化目標の影響を受けます。ある1回の選択を、信頼性の高いトピック分類結果と見なすことはできません。

トレーニング中の負荷の偏りにより、少数のエクスパートが大量のトークンを受け取り、他のエクスパートのトレーニングが不十分になることがあります。サービス時にはホットスポットや待機が発生することもあります。設計では補助損失やルーティング制約などの方法で改善できますが、「バッチごとに絶対的に均等である必要がある」とは限りません。一部のアーキテクチャでは、常に参加する共有エクスパートを導入し、共通知識を吸収して冗長性を減らそうとしています。これは特定のアーキテクチャの選択であり、MoEに必須の構成要素ではありません。DeepSeekMoE

1つのパラメータ計算で容量と計算の分離を見る

モデルのエクスパート以外の部分が1Bパラメータを持ち、すべてのMoE層で合計1パスあたりのエクスパートが2Bパラメータ、合計8パスあり、各位置で2パスを選択すると仮定します。総パラメータ数は 1+8×2=17B となり、この簡略化した数え方では、トークンあたりの参加パラメータ数は約 1+2×2=5B です。パス数を16に増やしてもなお2パスを選択する場合、総量は33Bになりますが、トークンあたりの参加パラメータ数は約5Bのままです。

この例は、候補容量を増やしても、エクスパートの計算量を同じ比率で増やす必要がないことを示していますが、計算量が完全に不変である証明にはなりません。ルーター、通信、負荷分散は変化します。実際のモデルのパラメータ統計には、埋め込み(embedding)、出力ヘッド、共有エクスパートなどを集計に含めるかという問題もあるため、モデルレポートを読み、単に「エクスパート数×ある稠密モデルのサイズ」で総量を推測してはいけません。

専門家の選択と重み付け

ルーターの値と最終混合係数は異なり、top-k の選択後に正規化します。E1 を増やすと混合が変わり、k を増やすと実行数も変わります。専門家出力を固定してルーティングだけを見ます。

図を準備しています
専門家の選択と重み付け

選ばれた専門家の間で正規化して出力を合成します。未選択でも重みは保存されています。

アーキテクチャ情報とタスク能力を分離して検証

ツール呼び出しの失敗については、まずモデルの出力、チャットテンプレート、ツールプロトコルの解析、コンテキストの欠如、ツール実行の例外を区別する必要があります。テキスト形式の関数名が表示されたとしても、有効な呼び出し構造になっていなければ、MoEの活性化パラメータが少なすぎることを証明するものではありません。モデルが対応するプロトコルでトレーニングされていないか、推論エンジンが正しいテンプレートを使用していない可能性があります。

「7Bなら安定して1ステップ、20Bなら安定して複数ステップ」という、活性化パラメータに共通する閾値はありません。トレーニングデータ、ファインチューニング手法、コンテキスト、デコーディング設定、タスクの難易度、ツールのフィードバックなどが結果に影響します。パラメータ規模をフィルタリングの指標として使うことはできますが、成功を保証するものとして扱うには証拠が不足しています。

次に読む:推論段階・KV cache・メモリ。