Mistral/Mixtralのアーキテクチャ — Sliding Window AttentionとMoEの融合

7Bのパラメータで13Bクラスのモデルを上回る性能を出せるとしたら、どうでしょうか。推論に必要なGPUメモリは半分以下になり、レスポンスは速くなり、エッジデバイスへのデプロイすら視野に入ります。2023年9月、フランスのAIスタートアップMistral AIが公開したMistral 7Bは、まさにそれを実現しました。LLaMA 2 13Bを多くのベンチマークで上回り、一部のタスクではLLaMA 1 34Bにすら匹敵する性能を示したのです。

さらに2023年12月には、Mistral 7Bの設計を発展させたMixtral 8x7Bが登場しました。8つのエキスパートネットワークを持ちながら、推論時には2つだけを使うSparse Mixture of Experts（MoE）により、46.7Bのパラメータを持ちつつ推論コストは12.9B相当に抑えるという離れ業を成し遂げています。

「なぜ7Bのモデルが13Bを超えられるのか？」「MoEでパラメータを増やしても推論が遅くならないのはなぜか？」 — 本記事ではこれらの疑問に答えながら、Mistral/Mixtralのアーキテクチャを数式とコードで完全に理解します。

Mistral/Mixtralのアーキテクチャを理解することは、以下のような場面で直接役立ちます。

• オープンソースLLMの選定と運用: Mistral/Mixtralはllama.cppやvLLMで動作するため、アーキテクチャの特性を理解すればメモリ見積もりやバッチサイズの最適化が正確に行えます
• エッジ推論の設計: Sliding Window Attentionによるメモリ上限の保証は、GPUメモリが限られるエッジ環境で特に重要です
• MoEモデルの理解: Mixtralの成功を受けて、Qwen-MoE、DBRX、Grok-1など多くのモデルがMoEを採用しています。Mixtralの設計を理解すれば、これらのモデルの差分を追うだけで済みます
• 後続モデルの基盤知識: Mistral Large、Codestral、Pixtralなど、Mistral AIの後続モデルはすべてこのアーキテクチャの延長線上にあります

本記事の内容

• Mistral 7BはLLaMAから何を変えたか — 全体像の把握
• Sliding Window Attention（SWA）の直感と数学的定式化
• Grouped-Query Attention（GQA）によるメモリ効率化
• Rolling Buffer KV Cacheの仕組み
• Mixtral 8x7BのSparse MoEアーキテクチャ
• MoEルーターの数学 — Top-2ゲーティングとロードバランス損失
• PyTorchによるSWAマスクとMoEルーターの実装
• Mistral/Mixtralの性能比較と位置づけ

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

それでは、まずMistral 7Bの全体像から見ていきましょう。LLaMAのアーキテクチャをベースに、どこをどう変えたのかを明確にします。

Mistral 7Bの全体像 — LLaMAからの変更点

LLaMAを土台とした設計

Mistral 7BはLLaMA 2と同じDecoder-only Transformerをベースとしています。RMSNorm、SwiGLU活性化関数、RoPE（Rotary Position Embedding）といったLLaMAの設計要素はそのまま引き継がれています。つまり、Mistral 7Bの革新性は「全く新しいアーキテクチャを発明した」ことにあるのではなく、「既知の技術を最適に組み合わせ、さらに独自の工夫を加えた」ことにあります。

LLaMAからの主な変更点を表にまとめます。

ここで特に注目すべきは3つの変更です。Sliding Window Attention（SWA） は注意の範囲を局所的に限定してメモリと計算量を削減し、GQA はKey-Valueヘッドを共有してメモリ使用量を減らし、Rolling Buffer KV Cache はSWAと連携してKVキャッシュのサイズに上限を設けます。

パラメータ数の比較

Mistral 7BとLLaMA 2のパラメータ数を比較してみましょう。

Mistral 7Bは7.3Bパラメータで、パラメータ数が約2倍のLLaMA 2 13Bをすべてのベンチマークで上回っています。この結果は、アーキテクチャの改良と学習データの質が、単純なパラメータ数の増加よりも効果的であることを示しています。

それでは、Mistral 7Bの最大の特徴であるSliding Window Attentionの仕組みを詳しく見ていきましょう。

Sliding Window Attention（SWA）の直感

全トークンに注目する必要はあるか？

標準的なSelf-Attentionでは、各トークンが系列内の全てのトークンに注目します。系列長を $n$ とすると、注意行列は $n \times n$ のサイズになり、計算量もメモリも $O(n^2)$ で増加します。これは系列長が長くなるほど深刻な問題になります。

しかし、自然言語を考えてみてください。文章を読むとき、私たちは全ての単語を等しく参照しているでしょうか？実際には、直近の数語〜数十語が最も重要で、遠くの単語の影響は薄れていきます。「昨日の会議で田中さんが言った」という文を処理するとき、「言った」の意味を理解するために最も重要なのは近くにある「田中さんが」であり、10段落前の単語ではありません。

Sliding Window Attention（SWA）はこの直感に基づいています。各トークンが注目する範囲を、直近の $W$ トークンに制限するのです。窓（window）をスライドさせながら、各位置で局所的な注意のみを計算します。

窓サイズと情報の到達距離

「でも、窓の外の情報は完全に失われるのでは？」という疑問が自然に生まれます。ここがSWAの巧妙な点です。

1層のSWAでは、各トークンは直近 $W$ 個のトークンしか見えません。しかし、Transformerは複数の層を重ねる構造です。2層目のトークンは、1層目で窓内のトークンが集約した情報を受け取ります。つまり、2層目では間接的に $2W$ 個のトークンの情報にアクセスできるのです。

$L$ 層のTransformerでは、情報の理論的な到達距離は次のようになります。

$$ \text{理論的到達距離} = L \times W $$

Mistral 7Bでは $L = 32$ 層、$W = 4{,}096$ ですから、理論的な到達距離は次のように計算できます。

$$ 32 \times 4{,}096 = 131{,}072 \text{ トークン} $$

つまり、窓サイズは4,096に限定されていますが、層を重ねることで約13万トークンの文脈情報を間接的に利用できるのです。これは、直接的な全注意を行わなくても、深いネットワークが長距離の依存関係を捉えられることを意味しています。

この仕組みは「伝言ゲーム」に似ています。一人一人は隣の数人としか話せませんが、複数のラウンドを経ることで、最初の人のメッセージは全員に伝わります。ただし、実際には情報は層を重ねるごとに薄まるため、理論上の到達距離がそのまま実効的な文脈窓になるわけではありません。それでも、標準的な全注意に比べて大幅に計算量を削減しながら、実用上十分な文脈を捉えられることが実験的に確認されています。

それでは、SWAの数学的な定式化に進みましょう。注意マスクの構造を明確にし、計算量の削減効果を定量的に理解します。

SWAの数学的定式化

注意行列のマスク構造

標準的なDecoder-only TransformerのSelf-Attentionでは、因果マスク（causal mask）を用いて未来のトークンへの注意を遮断します。位置 $i$ のトークンが位置 $j$ のトークンに注意できる条件は $j \leq i$ です。

SWAでは、これに加えて窓サイズ $W$ による制約が加わります。位置 $i$ のトークンが位置 $j$ のトークンに注意できる条件は次のようになります。

$$ i – W < j \leq i $$

すなわち、位置 $i$ から見て $W$ ステップ以内かつ過去のトークンのみが注意対象です。これをマスク行列 $\bm{M} \in \{0, -\infty\}^{n \times n}$ として表現すると、各成分は次のように定義されます。

$$ M_{ij} = \begin{cases} 0 & \text{if } i – W < j \leq i \\ -\infty & \text{otherwise} \end{cases} $$

SWA付きのアテンション出力は、Query $\bm{Q}$、Key $\bm{K}$、Value $\bm{V}$ を用いて次のように書けます。

$$ \text{Attention}(\bm{Q}, \bm{K}, \bm{V}) = \text{softmax}\!\left(\frac{\bm{Q}\bm{K}^\top}{\sqrt{d_k}} + \bm{M}\right)\bm{V} $$

ここで $d_k$ はKeyの次元数です。マスク $\bm{M}$ の $-\infty$ の成分は、softmaxを通すとゼロになるため、対応する位置のValueは出力に寄与しません。

標準注意との計算量比較

マスク行列の構造を視覚的に理解しましょう。系列長 $n = 8$、窓サイズ $W = 3$ の場合を考えます。

標準因果マスクでは下三角行列（注意できるのは自分以前の全トークン）になりますが、SWAでは下三角行列のうち対角線から $W$ 行以内の帯（バンド）だけが残ります。

$$ \bm{M}_{\text{causal}} = \begin{pmatrix} 0 & -\infty & -\infty & -\infty & \cdots \\ 0 & 0 & -\infty & -\infty & \cdots \\ 0 & 0 & 0 & -\infty & \cdots \\ 0 & 0 & 0 & 0 & \cdots \\ \vdots & & & & \ddots \end{pmatrix} $$

$$ \bm{M}_{\text{SWA}} = \begin{pmatrix} 0 & -\infty & -\infty & -\infty & \cdots \\ 0 & 0 & -\infty & -\infty & \cdots \\ 0 & 0 & 0 & -\infty & \cdots \\ -\infty & 0 & 0 & 0 & \cdots \\ \vdots & & & & \ddots \end{pmatrix} $$

標準因果注意の計算量は $O(n^2 \cdot d_k)$ です。一方、SWAでは各トークンが注目するのは最大 $W$ 個のトークンですから、計算量は次のようになります。

$$ O(n \times W \times d_k) $$

$W$ が $n$ に比べて十分小さい場合（$W \ll n$）、計算量は系列長に対して線形 $O(n)$ になります。Mistral 7Bでは $W = 4{,}096$ ですから、系列長が4,096以下の場合は標準注意と同じですが、系列長がそれを超えると大きな計算量の削減が得られます。

具体的な数値例

系列長 $n = 32{,}768$（32Kトークン）、窓サイズ $W = 4{,}096$ の場合を考えてみましょう。

標準因果注意のソフトマックス計算回数（非ゼロ要素数）は下三角行列の要素数に相当します。

$$ \frac{n(n+1)}{2} = \frac{32{,}768 \times 32{,}769}{2} \approx 5.37 \times 10^8 $$

一方、SWAのソフトマックス計算回数は近似的に次のようになります。

$$ n \times W = 32{,}768 \times 4{,}096 \approx 1.34 \times 10^8 $$

ここで比率を取ると次のようになります。

$$ \frac{n \times W}{\frac{n(n+1)}{2}} \approx \frac{2W}{n} = \frac{2 \times 4{,}096}{32{,}768} = 0.25 $$

つまり、32Kトークンの系列に対して、SWAは標準注意の約25%の計算量で済むことがわかります。系列長が長くなるほどこの削減効果は大きくなります。

SWAによって計算量を大幅に削減できることがわかりました。次に、もう一つの重要な改良であるGrouped-Query Attention（GQA）について見ていきましょう。GQAはアテンション機構のメモリ効率を改善する技術であり、SWAとは異なるアプローチでLLMの効率化に貢献します。

Grouped-Query Attention（GQA）

MHA・MQA・GQAの3つのアプローチ

大規模言語モデルの推論では、KVキャッシュがメモリの大きなボトルネックになります。Auto-Regressive（自己回帰）生成では、過去のトークンのKey（K）とValue（V）を保持して再利用する必要があるためです。このKVキャッシュのメモリ消費を削減するアプローチとして、アテンションヘッドの構成に3つの選択肢があります。

Multi-Head Attention（MHA） は標準的な構成です。全てのヘッドがそれぞれ独立したQuery（Q）、Key（K）、Value（V）を持ちます。32ヘッドあれば、32組のKVペアがキャッシュに保存されます。表現力は最も高いですが、KVキャッシュのメモリ消費も最大です。

Multi-Query Attention（MQA） は対極にある設計です。全てのヘッドがQueryだけを独立に持ち、Key-Valueは全ヘッドで1組だけを共有します。KVキャッシュのメモリは$1/h$（$h$はヘッド数）に減りますが、表現力が低下して性能が落ちることがあります。

Grouped-Query Attention（GQA） はこの2つの中間に位置する設計です。$h$個のQueryヘッドを$g$個のグループに分け、各グループ内でKey-Valueを共有します。$g = h$ ならMHA、$g = 1$ ならMQAに一致するため、GQAはMHAとMQAを一般化した枠組みと言えます。

GQAの数式

$h$個のアテンションヘッドを$g$個のグループに分けます。各グループ内のヘッド数は $h / g$ です。ヘッド$i$が属するグループを $\lfloor ig / h \rfloor$ とすると、ヘッド$i$のアテンション出力は次のように計算されます。

$$ \text{head}_i = \text{Attention}(\bm{Q}_i, \bm{K}_{\lfloor ig/h \rfloor}, \bm{V}_{\lfloor ig/h \rfloor}) $$

ここで $\bm{Q}_i \in \mathbb{R}^{n \times d_k}$ はヘッド $i$ 固有のQuery射影、$\bm{K}_j, \bm{V}_j \in \mathbb{R}^{n \times d_k}$ はグループ $j$ で共有されるKey・Value射影です。

全ヘッドの出力を結合して最終出力を得ます。

$$ \text{GQA}(\bm{Q}, \bm{K}, \bm{V}) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) \bm{W}^O $$

KVキャッシュのメモリ削減効果

Mistral 7Bでは $h = 32$（Queryヘッド数）、$g = 8$（KVグループ数）です。KVキャッシュに必要なメモリを比較しましょう。

1トークンあたりのKVキャッシュサイズは次のように計算できます。ヘッド次元を $d_k = d_{\text{model}} / h = 4{,}096 / 32 = 128$、層数を $L = 32$、数値型をbfloat16（2バイト）とします。

MHA（$g = 32$）の場合、KVそれぞれ $g$ 組が必要です。

$$ \text{KV}_{\text{MHA}} = 2 \times g \times d_k \times L \times 2 = 2 \times 32 \times 128 \times 32 \times 2 = 524{,}288 \text{ bytes} \approx 0.5 \text{ MB/token} $$

GQA（$g = 8$）の場合、KV組数が $g = 8$ に減ります。

$$ \text{KV}_{\text{GQA}} = 2 \times g \times d_k \times L \times 2 = 2 \times 8 \times 128 \times 32 \times 2 = 131{,}072 \text{ bytes} \approx 0.125 \text{ MB/token} $$

したがってメモリ削減率は次のようになります。

$$ \frac{\text{KV}_{\text{GQA}}}{\text{KV}_{\text{MHA}}} = \frac{8}{32} = 0.25 $$

GQAにより、KVキャッシュのメモリはMHAの25%に削減されます。8,192トークンの系列全体では、MHAが約4GBのKVキャッシュを必要とするのに対し、GQAでは約1GBで済みます。この差は、バッチサイズを大きくしたいサービング環境では極めて大きな意味を持ちます。

GQAがKVキャッシュのメモリを削減することがわかりました。しかし、系列が長くなるとKVキャッシュ自体が際限なく成長する問題は依然として残ります。この問題に対するMistral独自の解決策が、次に解説するRolling Buffer KV Cacheです。

Rolling Buffer KV Cache

KVキャッシュの成長問題

標準的なKVキャッシュでは、新しいトークンが生成されるたびにそのKとVがキャッシュに追加されます。系列長が $n$ に達すると、キャッシュサイズは $n$ に比例して増大します。これは長い文書を処理する際に深刻なメモリ問題を引き起こします。

ここで思い出してください。SWAでは、各トークンが注目するのは直近 $W$ 個のトークンだけです。つまり、$W$ ステップより前のトークンのKVはもう二度と参照されないのです。参照されないデータを保持し続けるのは無駄です。

Rolling Bufferの仕組み

Rolling Buffer KV Cacheは、この観察に基づいてKVキャッシュを固定サイズの循環バッファとして実装します。バッファのサイズは窓サイズ $W$ に等しく、位置 $i$ のトークンのKVはバッファの位置 $i \bmod W$ に格納されます。

イメージとしては、レコードプレーヤーの溝のように、バッファの末尾に達したら先頭に戻って上書きします。こうすることで、常に直近 $W$ 個のトークンのKVだけがバッファに保持されます。

位置 $i$ のトークンを処理するとき、バッファへの書き込み位置は次のように決まります。

$$ \text{pos}_{\text{buffer}} = i \bmod W $$

例えば $W = 4$ の場合を考えましょう。

位置4のトークンが入ると、バッファ位置0のt0が上書きされてt4に置き換わります。位置4のSWA窓は $[1, 2, 3, 4]$ ですから、t0はもう不要です。

メモリ使用量の上限保証

Rolling Bufferの最大の利点は、KVキャッシュのメモリ使用量に上限が存在することです。GQAと組み合わせた場合のメモリ上限を計算しましょう。

$$ \text{KV}_{\text{max}} = 2 \times g \times d_k \times L \times W \times 2 $$

Mistral 7Bの具体的な数値を代入します（$g = 8$、$d_k = 128$、$L = 32$、$W = 4{,}096$）。

$$ \text{KV}_{\text{max}} = 2 \times 8 \times 128 \times 32 \times 4{,}096 \times 2 = 536{,}870{,}912 \text{ bytes} = 512 \text{ MB} $$

つまり、Mistral 7BのKVキャッシュはどれだけ長い系列を処理しても512MBを超えないのです。これは標準的なKVキャッシュ（系列長に比例して成長）と比べて、予測可能なメモリ管理を可能にする大きな利点です。

Pre-fill時の注意点

Rolling Bufferは推論時（トークンを1つずつ生成するフェーズ）では自然に機能しますが、Pre-fill（プロンプト全体を一度に処理するフェーズ）では注意が必要です。Pre-fillではSWAのマスクを明示的に適用しつつ、バッファの各位置を正しく埋める必要があります。

具体的には、Pre-fill時にはプロンプト全体に対してSWAマスク付きのアテンションを一括計算し、最後の $W$ トークン分のKVをRolling Bufferに格納します。それ以降の自己回帰生成では、新しいトークンが $i \bmod W$ の位置にKVを書き込みながら進みます。

ここまでで、Mistral 7Bのコアとなる3つの技術 — SWA、GQA、Rolling Buffer KV Cache — を理解しました。これらはすべて「効率化」の技術です。次に、Mistral AIが効率化と性能向上の両方を同時に追求したMixtral 8x7Bのアーキテクチャに進みましょう。ここでは新たにSparse Mixture of Experts（MoE）が登場します。

Mixtral 8x7B — Sparse MoEの導入

MoEの基本的なアイデア

Mixtral 8x7Bの名前にある「8x7B」は、8つのエキスパートネットワークがあり、各エキスパートが7Bクラスの構造を持つことを示しています。ただし、推論時に8つ全てを使うわけではありません。各トークンに対して2つだけを選んで使います。

これは「大きな病院にたとえると理解しやすいでしょう。病院には内科、外科、小児科、皮膚科など多くの診療科（エキスパート）がありますが、一人の患者が全ての科を受診するわけではありません。症状に応じて適切な科に振り分けられます。Mixtralも同様に、入力トークンの特性に応じて最も適切なエキスパートに「診てもらう」のです。

Mixtralのアーキテクチャ

Mixtral 8x7BはMistral 7Bのアーキテクチャをベースとしていますが、FFN（Feed-Forward Network）レイヤーをMoEレイヤーに置き換えています。Self-Attentionの部分は全く同じで、SWAとGQAがそのまま使われます。

各Transformerブロックの構成は次のようになります。

1. RMSNorm → SWA（GQA付き） → 残差接続
2. RMSNorm → MoEレイヤー → 残差接続

MoEレイヤーの内部には8つの独立したFFN（エキスパート）とルーター（ゲートネットワーク）が存在します。ルーターが入力に基づいて2つのエキスパートを選択し、それぞれの出力を重み付き平均して最終出力とします。

パラメータ数と推論コスト

Mixtral 8x7Bの総パラメータ数は約46.7Bですが、推論時に各トークンで活性化されるパラメータ数は約12.9Bです。この差はどこから来るのでしょうか。

Self-Attention層とEmbedding層のパラメータは全トークンで共有されるため、常に活性化されます。一方、MoEレイヤーの8つのエキスパートのうち2つだけが活性化されるため、FFN部分のパラメータ使用量は $2/8 = 25\%$ に抑えられます。

具体的に概算してみましょう。1つのFFNエキスパートのパラメータ数を $P_{\text{FFN}}$ とすると、以下の関係になります。

Mistral 7Bの全パラメータ数（FFN含む）が約7.3Bですから、Attention部分のパラメータを $P_{\text{attn}}$、FFN部分を $P_{\text{FFN}}$ とすると次の関係があります。

$$ P_{\text{attn}} + P_{\text{FFN}} \approx 7.3\text{B} $$

Mixtralでは8つのFFNエキスパートを持つため、総パラメータ数は次のようになります。

$$ P_{\text{total}} = P_{\text{attn}} + 8 \times P_{\text{FFN}} \approx 46.7\text{B} $$

推論時に活性化されるパラメータ数は、2つのエキスパートのみが選ばれるため次のようになります。

$$ P_{\text{active}} = P_{\text{attn}} + 2 \times P_{\text{FFN}} \approx 12.9\text{B} $$

これらの連立方程式から、$P_{\text{FFN}} \approx 5.63\text{B}$、$P_{\text{attn}} \approx 1.67\text{B}$ と推定できます。FFNのパラメータが全体の大部分を占めていることがわかります。

この「パラメータは多いが計算量は少ない」という特性がMoEの最大の魅力です。46.7Bのモデル容量（知識量）を持ちながら、推論速度はMistral 7Bの約2倍弱に抑えられます。GPUメモリさえ足りれば、13Bクラスの推論速度で70Bクラスの性能を得られるのです。

それでは次に、MoEレイヤーの心臓部であるルーターの数学的な仕組みを詳しく見ていきましょう。

MoEの数学 — ルーターとゲーティング機構

Top-2ゲーティングの定式化

MoEレイヤーには $N = 8$ 個のエキスパート $\{E_1, E_2, \dots, E_8\}$ と、1つのルーターネットワーク $G$ が含まれます。入力トークンの隠れ状態を $\bm{x} \in \mathbb{R}^{d_{\text{model}}}$ とすると、ルーターはまず各エキスパートに対するスコア（ロジット）を計算します。

$$ \bm{s} = \bm{W}_g \bm{x} $$

ここで $\bm{W}_g \in \mathbb{R}^{N \times d_{\text{model}}}$ はルーターの重み行列であり、$\bm{s} \in \mathbb{R}^N$ は各エキスパートに対するスコアベクトルです。ルーターのパラメータは $\bm{W}_g$ だけで、バイアス項はありません。

次に、Top-2の選択を行います。スコアが最も高い2つのエキスパートのインデックスを $i_1, i_2$（$s_{i_1} \geq s_{i_2}$）とします。

$$ \{i_1, i_2\} = \text{Top-2}(\bm{s}) $$

選ばれた2つのエキスパートのスコアに対してソフトマックスを適用し、ゲーティング重みを計算します。

$$ g_{i_k} = \frac{\exp(s_{i_k})}{\exp(s_{i_1}) + \exp(s_{i_2})}, \quad k \in \{1, 2\} $$

この操作では、選ばれなかった6つのエキスパートのスコアはソフトマックスの計算から除外されている点に注意してください。全8エキスパートに対してソフトマックスを取ってからTop-2を選ぶのではなく、Top-2を選んでからソフトマックスを取ることで、選ばれた2つの重みの和が常に1になることが保証されます。

最終的なMoEレイヤーの出力は、2つのエキスパートの出力を重み付き平均したものです。

$$ \text{MoE}(\bm{x}) = g_{i_1} \cdot E_{i_1}(\bm{x}) + g_{i_2} \cdot E_{i_2}(\bm{x}) $$

この式は直感的には「2人の専門家の意見を、信頼度に応じて重み付けして統合する」ことに相当します。ゲーティング重み $g_{i_1}, g_{i_2}$ はルーターが判断した「この入力にどれだけ適しているか」の度合いを反映しています。

各エキスパートの構造

各エキスパート $E_i$ はMistral 7BのFFNと同じSwiGLU構造を持っています。入力 $\bm{x}$ に対して次のように計算します。

$$ E_i(\bm{x}) = (\text{SiLU}(\bm{W}_1^{(i)} \bm{x}) \odot \bm{W}_3^{(i)} \bm{x}) \bm{W}_2^{(i)} $$

ここで $\bm{W}_1^{(i)}, \bm{W}_3^{(i)} \in \mathbb{R}^{d_{\text{ff}} \times d_{\text{model}}}$、$\bm{W}_2^{(i)} \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ff}}}$ はエキスパート $i$ 固有の重み行列、$\odot$ は要素ごとの積、SiLUは $\text{SiLU}(x) = x \cdot \sigma(x)$ です。$d_{\text{ff}} = 14{,}336$ はFFNの中間次元で、Mistral 7Bと同じ値です。

8つのエキスパートはそれぞれ独立した重みを持つため、FFNのパラメータは8倍になります。しかし各トークンでは2つだけが計算されるため、推論時のFLOPsは2倍に留まります。

ロードバランス損失

MoEの学習において最も深刻な問題はルーターの偏りです。ルーターが特定のエキスパートばかりを選ぶようになると、そのエキスパートだけが学習の機会を得て強くなり、さらにルーターがそのエキスパートを選ぶようになる — という正のフィードバックループが形成されます。極端な場合、8つ中1〜2個のエキスパートしか使われなくなり、MoEの意味が失われます。これをエキスパートの崩壊（expert collapse）と呼びます。

この問題を防ぐために、ロードバランス損失（load balancing loss）を学習の損失関数に加えます。バッチ内の $T$ 個のトークンに対して、エキスパート $i$ が選ばれた頻度 $f_i$ と、ルーターがエキスパート $i$ に割り当てたゲーティング確率の平均 $p_i$ を次のように定義します。

$$ f_i = \frac{1}{T} \sum_{t=1}^{T} \mathbb{1}[i \in \text{Top-2}(\bm{s}_t)] $$

$$ p_i = \frac{1}{T} \sum_{t=1}^{T} \text{softmax}(\bm{s}_t)_i $$

ここで $\mathbb{1}[\cdot]$ は指示関数です。$f_i$ は「エキスパート $i$ が実際に選ばれた割合」、$p_i$ は「ルーターがエキスパート $i$ を選びたいと思っている度合いの平均」と解釈できます。

ロードバランス損失は次のように定義されます。

$$ \mathcal{L}_{\text{balance}} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot p_i $$

ここで $\alpha$ は損失の強さを制御するハイパーパラメータです。Mixtralでは $\alpha = 0.01$ 程度が使われます。係数 $N$ は正規化のためのもので、全エキスパートが均等に使われた場合に $\sum f_i \cdot p_i = 1/N$ となることから、$N$ を掛けて $\mathcal{L}_{\text{balance}} = \alpha$ となるように調整しています。

この損失関数の直感を理解しましょう。もし全エキスパートが均等に使われていれば、$f_i = 2/N$（Top-2なので各トークンで2つ選ばれる）、$p_i = 1/N$ となり、損失は最小値に近くなります。一方、特定のエキスパートに偏ると、偏ったエキスパートの $f_i$ と $p_i$ が大きくなり、損失が増大します。この損失を最小化する方向に学習を進めることで、エキスパートの均等な利用が促進されます。

全体の学習損失は、通常の言語モデル損失（交差エントロピー）とロードバランス損失の和になります。

$$ \mathcal{L} = \mathcal{L}_{\text{LM}} + \mathcal{L}_{\text{balance}} $$

ここまでで、SWA、GQA、Rolling Buffer KV Cache、MoEルーター、ロードバランス損失という、Mistral/Mixtralの全てのコアコンポーネントの数学を理解しました。次は、これらをPyTorchで実装して理解を確認しましょう。

PyTorch実装

SWAマスクの実装

まず、Sliding Window Attentionのマスクを生成し、標準因果マスクとの違いを可視化します。

import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np

def create_causal_mask(seq_len):
    """標準因果マスク（下三角行列）を生成"""
    mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1)
    mask = mask.masked_fill(mask == 1, float('-inf'))
    return mask

def create_swa_mask(seq_len, window_size):
    """Sliding Window Attentionマスクを生成"""
    # まず因果マスクを作成
    mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1)
    # 窓の外（window_sizeより前）もマスク
    window_mask = torch.tril(torch.ones(seq_len, seq_len), diagonal=-(window_size))
    mask = mask + window_mask
    mask = mask.masked_fill(mask >= 1, float('-inf'))
    return mask

# マスクの生成と可視化
seq_len = 16
window_size = 4

causal_mask = create_causal_mask(seq_len)
swa_mask = create_swa_mask(seq_len, window_size)

fig, axes = plt.subplots(1, 2, figsize=(12, 5))

# 標準因果マスク
im0 = axes[0].imshow(
    (causal_mask == 0).float().numpy(),
    cmap='Blues', aspect='equal'
)
axes[0].set_title('Standard Causal Mask', fontsize=13)
axes[0].set_xlabel('Key position (j)')
axes[0].set_ylabel('Query position (i)')

# SWAマスク
im1 = axes[1].imshow(
    (swa_mask == 0).float().numpy(),
    cmap='Blues', aspect='equal'
)
axes[1].set_title(f'SWA Mask (W={window_size})', fontsize=13)
axes[1].set_xlabel('Key position (j)')
axes[1].set_ylabel('Query position (i)')

plt.tight_layout()
plt.savefig('swa_mask_comparison.png', dpi=150, bbox_inches='tight')
plt.show()

# 非ゼロ要素数の比較
causal_nonzero = (causal_mask == 0).sum().item()
swa_nonzero = (swa_mask == 0).sum().item()
print(f"標準因果マスクの注意対象数: {causal_nonzero}")
print(f"SWAマスクの注意対象数:       {swa_nonzero}")
print(f"計算量削減率: {1 - swa_nonzero / causal_nonzero:.1%}")

上のコードを実行すると、2つのマスクの構造の違いが一目でわかります。標準因果マスクは下三角行列全体が注意対象（青色）になりますが、SWAマスクでは対角線から窓サイズ $W = 4$ の範囲だけがバンド状に残ります。系列長16、窓サイズ4の場合、標準因果マスクの注意対象数は136（$= 16 \times 17 / 2$）、SWAの注意対象数は58（$\approx 16 \times 4 – \text{端の補正}$）となり、約57%の計算量が削減されることが確認できます。系列長が大きくなるほどこの削減率は増大します。

SWA付きアテンションの実装

次に、SWAマスクを実際のアテンション計算に組み込みます。

import torch
import torch.nn as nn
import torch.nn.functional as F

class SlidingWindowAttention(nn.Module):
    """Sliding Window Attention（簡略版）"""
    def __init__(self, d_model, n_heads, n_kv_heads, window_size):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim = d_model // n_heads
        self.window_size = window_size
        # GQA: n_headsのQueryに対してn_kv_headsのKVグループ
        self.n_rep = n_heads // n_kv_heads

        self.wq = nn.Linear(d_model, n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(n_heads * self.head_dim, d_model, bias=False)

    def _repeat_kv(self, x, n_rep):
        """KVヘッドをQueryヘッド数に合わせて繰り返す"""
        if n_rep == 1:
            return x
        bs, n_kv, seq_len, head_dim = x.shape
        x = x.unsqueeze(2).expand(bs, n_kv, n_rep, seq_len, head_dim)
        return x.reshape(bs, n_kv * n_rep, seq_len, head_dim)

    def forward(self, x):
        bs, seq_len, _ = x.shape

        # Q, K, Vの射影
        q = self.wq(x).view(bs, seq_len, self.n_heads, self.head_dim)
        k = self.wk(x).view(bs, seq_len, self.n_kv_heads, self.head_dim)
        v = self.wv(x).view(bs, seq_len, self.n_kv_heads, self.head_dim)

        # (bs, heads, seq_len, head_dim)に転置
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # GQA: KVを繰り返してQueryヘッド数に合わせる
        k = self._repeat_kv(k, self.n_rep)
        v = self._repeat_kv(v, self.n_rep)

        # スケーリング付き内積
        scale = self.head_dim ** 0.5
        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / scale

        # SWAマスクの適用
        swa_mask = create_swa_mask(seq_len, self.window_size).to(x.device)
        attn_weights = attn_weights + swa_mask.unsqueeze(0).unsqueeze(0)

        attn_weights = F.softmax(attn_weights, dim=-1)
        output = torch.matmul(attn_weights, v)

        # ヘッドの結合
        output = output.transpose(1, 2).contiguous().view(bs, seq_len, -1)
        return self.wo(output)

この実装では、GQAの仕組みも組み込まれています。_repeat_kv メソッドが、$g$ 個のKVグループを $h$ 個のQueryヘッドに合わせて複製する処理を行います。Mistral 7Bの場合、8個のKVヘッドをそれぞれ4回複製して32個のQueryヘッドに対応させます。この複製はメモリ上では効率的なビューで実現でき、実際にデータをコピーするわけではありません。

SWAの注意パターンの可視化

SWAが実際にどのような注意パターンを生み出すかを可視化して確認しましょう。

import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np

def visualize_swa_attention_pattern(seq_len=32, window_size=8, d_model=64, n_heads=4):
    """SWA付きアテンションの重みパターンを可視化"""
    torch.manual_seed(42)

    # ランダムなQ, K を生成
    head_dim = d_model // n_heads
    q = torch.randn(1, n_heads, seq_len, head_dim)
    k = torch.randn(1, n_heads, seq_len, head_dim)

    # スケーリング付き内積
    scores = torch.matmul(q, k.transpose(-2, -1)) / (head_dim ** 0.5)

    # SWAマスクを適用
    swa_mask = create_swa_mask(seq_len, window_size)
    scores_masked = scores + swa_mask.unsqueeze(0).unsqueeze(0)

    # ソフトマックスで重みに変換
    attn_weights = F.softmax(scores_masked, dim=-1)

    # ヘッド0の注意重みを可視化
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    # マスクなし（標準因果注意）
    causal_mask = create_causal_mask(seq_len)
    scores_causal = scores + causal_mask.unsqueeze(0).unsqueeze(0)
    attn_causal = F.softmax(scores_causal, dim=-1)

    im0 = axes[0].imshow(
        attn_causal[0, 0].detach().numpy(),
        cmap='viridis', aspect='equal', vmin=0
    )
    axes[0].set_title('Standard Causal Attention (Head 0)', fontsize=12)
    axes[0].set_xlabel('Key position')
    axes[0].set_ylabel('Query position')
    plt.colorbar(im0, ax=axes[0], shrink=0.8)

    # SWAマスクあり
    im1 = axes[1].imshow(
        attn_weights[0, 0].detach().numpy(),
        cmap='viridis', aspect='equal', vmin=0
    )
    axes[1].set_title(f'SWA Attention (Head 0, W={window_size})', fontsize=12)
    axes[1].set_xlabel('Key position')
    axes[1].set_ylabel('Query position')
    plt.colorbar(im1, ax=axes[1], shrink=0.8)

    plt.tight_layout()
    plt.savefig('swa_attention_pattern.png', dpi=150, bbox_inches='tight')
    plt.show()

visualize_swa_attention_pattern()

左のグラフ（標準因果注意）では、各トークンの注意重みが過去の全トークンに分散しています。特に後方のトークンほど注意対象が多いため、個々のKey位置に割り当てられる重みが薄くなる傾向が見られます。右のグラフ（SWA）では、注意が窓サイズ内のバンド状に集中し、窓外の領域は完全にゼロです。窓内に注意が集約されるため、各Keyに対する重みが相対的に大きくなり、局所的な情報をより強く捉えられていることがわかります。これは「重要な局所情報に集中する」というSWAの設計意図と一致しています。

簡易MoEルーターの実装

次に、MixtralのMoEレイヤーの核心であるルーターを実装します。

import torch
import torch.nn as nn
import torch.nn.functional as F

class MoERouter(nn.Module):
    """Mixtral style Top-2 MoE Router"""
    def __init__(self, d_model, n_experts, top_k=2):
        super().__init__()
        self.n_experts = n_experts
        self.top_k = top_k
        # ルーターの重み行列（バイアスなし）
        self.gate = nn.Linear(d_model, n_experts, bias=False)

    def forward(self, x):
        """
        Args:
            x: (batch_size, seq_len, d_model)
        Returns:
            weights: (batch_size, seq_len, top_k) 正規化された重み
            indices: (batch_size, seq_len, top_k) 選択されたエキスパートのインデックス
        """
        # ルーターのスコア計算
        logits = self.gate(x)  # (bs, seq_len, n_experts)

        # Top-k選択
        top_k_logits, top_k_indices = torch.topk(logits, self.top_k, dim=-1)

        # 選ばれたエキスパートのスコアにソフトマックスを適用
        top_k_weights = F.softmax(top_k_logits, dim=-1)

        return top_k_weights, top_k_indices

    def load_balance_loss(self, logits):
        """ロードバランス損失の計算"""
        # logits: (bs, seq_len, n_experts)
        bs, seq_len, n_experts = logits.shape
        T = bs * seq_len

        # 各エキスパートが選ばれた頻度 f_i
        _, top_k_indices = torch.topk(logits, self.top_k, dim=-1)
        # one-hotに変換して頻度を計算
        one_hot = F.one_hot(top_k_indices, n_experts).float()
        # top_k分を合計（各トークンでtop_k個選ばれる）
        f = one_hot.sum(dim=2).sum(dim=(0, 1)) / T  # (n_experts,)

        # ルーターの確率分布 p_i
        probs = F.softmax(logits, dim=-1)
        p = probs.mean(dim=(0, 1))  # (n_experts,)

        # ロードバランス損失
        loss = n_experts * (f * p).sum()
        return loss

この実装のポイントは3つあります。第一に、ルーターは単純な線形層（nn.Linear）で、バイアスを持ちません。第二に、Top-2選択の後にソフトマックスを適用することで、選ばれた2つのエキスパートの重みの和が常に1になります。第三に、ロードバランス損失は頻度 $f_i$ と確率 $p_i$ の内積として計算され、偏りが大きいほど損失が増大する仕組みです。

MoEレイヤー全体の実装

ルーターとエキスパートを組み合わせて、MoEレイヤー全体を実装します。

class SwiGLUExpert(nn.Module):
    """SwiGLU FFN（1つのエキスパート）"""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
        self.w3 = nn.Linear(d_model, d_ff, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class MixtralMoELayer(nn.Module):
    """Mixtral style MoE Layer"""
    def __init__(self, d_model, d_ff, n_experts=8, top_k=2):
        super().__init__()
        self.n_experts = n_experts
        self.top_k = top_k
        self.router = MoERouter(d_model, n_experts, top_k)
        self.experts = nn.ModuleList([
            SwiGLUExpert(d_model, d_ff) for _ in range(n_experts)
        ])

    def forward(self, x):
        bs, seq_len, d_model = x.shape

        # ルーターで重みとインデックスを取得
        weights, indices = self.router(x)

        # 出力の初期化
        output = torch.zeros_like(x)

        # 各エキスパートの出力を重み付き加算
        for k in range(self.top_k):
            expert_indices = indices[:, :, k]   # (bs, seq_len)
            expert_weights = weights[:, :, k]   # (bs, seq_len)

            for i in range(self.n_experts):
                # エキスパートiが選ばれたトークンのマスク
                mask = (expert_indices == i)
                if mask.any():
                    # 該当トークンを抽出してエキスパートに通す
                    expert_input = x[mask]  # (n_selected, d_model)
                    expert_output = self.experts[i](expert_input)
                    # 重みをかけて加算
                    output[mask] += expert_weights[mask].unsqueeze(-1) * expert_output

        return output

上の実装はわかりやすさを優先したナイーブな実装です。実際のMixtralの推論実装では、トークンをエキスパートごとにバッチ化して効率的に処理する手法（Expert Parallelism）が使われます。しかし、数学的な処理は同じです。

ルーターの動作確認

実装したMoEルーターが正しく動作するかを確認しましょう。

import torch
import matplotlib.pyplot as plt
import numpy as np

torch.manual_seed(42)

d_model = 256
n_experts = 8
seq_len = 64
batch_size = 4

# MoEレイヤーの初期化
moe_layer = MixtralMoELayer(d_model, d_ff=512, n_experts=n_experts, top_k=2)

# ランダムな入力
x = torch.randn(batch_size, seq_len, d_model)

# ルーターのスコアを取得
with torch.no_grad():
    logits = moe_layer.router.gate(x)  # (bs, seq_len, n_experts)
    weights, indices = moe_layer.router(x)

# エキスパート選択頻度の集計
all_indices = indices.reshape(-1, 2)  # (bs*seq_len, 2)
expert_counts = torch.zeros(n_experts)
for i in range(n_experts):
    expert_counts[i] = (all_indices == i).sum().item()

# ロードバランス損失
lb_loss = moe_layer.router.load_balance_loss(logits)

# 可視化
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# エキスパート選択頻度
colors = plt.cm.Set2(np.linspace(0, 1, n_experts))
axes[0].bar(range(n_experts), expert_counts.numpy(), color=colors)
axes[0].axhline(y=expert_counts.sum().item() / n_experts, color='red',
                linestyle='--', label='Ideal uniform')
axes[0].set_xlabel('Expert ID')
axes[0].set_ylabel('Selection count')
axes[0].set_title('Expert Selection Frequency (before training)')
axes[0].legend()

# Top-2ゲーティング重みの分布
w1 = weights[:, :, 0].flatten().detach().numpy()
w2 = weights[:, :, 1].flatten().detach().numpy()
axes[1].hist(w1, bins=30, alpha=0.7, label='Top-1 weight', color='steelblue')
axes[1].hist(w2, bins=30, alpha=0.7, label='Top-2 weight', color='coral')
axes[1].set_xlabel('Gating weight')
axes[1].set_ylabel('Count')
axes[1].set_title('Distribution of Gating Weights')
axes[1].legend()

plt.suptitle(f'Load Balance Loss: {lb_loss.item():.4f}', fontsize=13, y=1.02)
plt.tight_layout()
plt.savefig('moe_router_analysis.png', dpi=150, bbox_inches='tight')
plt.show()

print(f"ロードバランス損失: {lb_loss.item():.4f}")
print(f"エキスパート選択頻度: {expert_counts.tolist()}")
print(f"理想的な均等頻度: {expert_counts.sum().item() / n_experts:.1f}")

左のグラフはエキスパートの選択頻度を示しています。初期化直後のランダムな重みでは、各エキスパートの選択頻度にばらつきがありますが、赤い点線（理想的な均等頻度）からの乖離は比較的小さいことがわかります。これはルーターの重みがランダム初期化されているためです。学習が進むと、ロードバランス損失が各エキスパートの利用を均等化する方向に働きます。右のグラフはゲーティング重みの分布です。Top-1の重みは0.5以上に集中し、Top-2の重みは0.5以下に集中していることがわかります。これは「メインの専門家が大きな影響力を持ち、サブの専門家が補助的な役割を果たす」というTop-2ゲーティングの設計意図と合致しています。

Rolling Buffer KV Cacheの実装

最後に、Rolling Buffer KV Cacheの動作を実装して確認します。

import torch
import matplotlib.pyplot as plt
import numpy as np

class RollingBufferKVCache:
    """Rolling Buffer KV Cache for Sliding Window Attention"""
    def __init__(self, max_batch_size, n_kv_heads, head_dim, window_size, n_layers):
        self.window_size = window_size
        self.n_layers = n_layers
        # 固定サイズの循環バッファ
        self.k_cache = torch.zeros(
            n_layers, max_batch_size, n_kv_heads, window_size, head_dim
        )
        self.v_cache = torch.zeros(
            n_layers, max_batch_size, n_kv_heads, window_size, head_dim
        )

    def update(self, layer_idx, k, v, position):
        """
        新しいKVを循環バッファに書き込む
        Args:
            layer_idx: 層のインデックス
            k, v: (batch_size, n_kv_heads, 1, head_dim)
            position: 現在のトークン位置
        """
        buf_pos = position % self.window_size
        self.k_cache[layer_idx, :, :, buf_pos, :] = k.squeeze(2)
        self.v_cache[layer_idx, :, :, buf_pos, :] = v.squeeze(2)

    def get(self, layer_idx, current_position):
        """
        現在のSWA窓内のKVを取得
        """
        start = max(0, current_position - self.window_size + 1)
        length = min(self.window_size, current_position + 1)

        # バッファ内の有効な位置を計算
        positions = torch.arange(start, current_position + 1)
        buf_positions = positions % self.window_size

        k = self.k_cache[layer_idx, :, :, buf_positions, :]
        v = self.v_cache[layer_idx, :, :, buf_positions, :]
        return k, v

# Rolling Bufferの動作を可視化
window_size = 4
n_steps = 12
buffer_states = []

# 各ステップでバッファの中身を記録
buffer = ['_'] * window_size
for pos in range(n_steps):
    buf_pos = pos % window_size
    buffer[buf_pos] = f't{pos}'
    buffer_states.append(list(buffer))

# 可視化
fig, ax = plt.subplots(figsize=(10, 6))
cell_width = 1.5
cell_height = 0.5

for step, state in enumerate(buffer_states):
    y = (n_steps - 1 - step) * cell_height
    for slot, val in enumerate(state):
        x = slot * cell_width
        is_new = (step % window_size == slot) and val != '_'
        color = '#4CAF50' if is_new else ('#E3F2FD' if val != '_' else '#F5F5F5')
        rect = plt.Rectangle((x, y), cell_width - 0.05, cell_height - 0.05,
                              facecolor=color, edgecolor='gray', linewidth=0.5)
        ax.add_patch(rect)
        ax.text(x + cell_width/2, y + cell_height/2, val,
                ha='center', va='center', fontsize=9,
                fontweight='bold' if is_new else 'normal')
    ax.text(-0.8, y + cell_height/2, f'pos={step}',
            ha='right', va='center', fontsize=9)

for slot in range(window_size):
    ax.text(slot * cell_width + cell_width/2, n_steps * cell_height + 0.1,
            f'buf[{slot}]', ha='center', va='bottom', fontsize=10, fontweight='bold')

ax.set_xlim(-1.2, window_size * cell_width + 0.2)
ax.set_ylim(-0.2, n_steps * cell_height + 0.5)
ax.set_aspect('equal')
ax.axis('off')
ax.set_title(f'Rolling Buffer KV Cache (W={window_size})', fontsize=14, pad=10)
plt.tight_layout()
plt.savefig('rolling_buffer.png', dpi=150, bbox_inches='tight')
plt.show()

この可視化では、窓サイズ4のRolling Bufferが12ステップにわたってどのように更新されるかを示しています。緑色のセルは各ステップで新しく書き込まれた位置を表します。位置0〜3ではバッファが順に埋まり、位置4で先頭に戻ってt0がt4に上書きされる様子が確認できます。以降も循環的に上書きが繰り返され、バッファには常に直近4トークンのKVだけが保持されます。バッファのサイズは窓サイズで固定されるため、どれだけ長い系列を処理してもメモリ使用量は一定です。

ここまでで、Mistral/Mixtralの主要コンポーネントをPyTorchで実装しました。最後に、これらのモデルの性能を既存モデルと比較して、アーキテクチャの改良がどれだけの効果をもたらしたかを確認しましょう。

性能比較

Mistral 7B vs LLaMA 2

Mistral 7Bの公開時に報告されたベンチマーク結果を見てみましょう。

この結果から、2つの重要なことが読み取れます。

第一に、Mistral 7Bは全てのベンチマークでLLaMA 2 13Bを上回っています。パラメータ数が約半分であるにもかかわらずです。これはSWA、GQA、Rolling Buffer KV Cacheによる効率化が、モデルの性能を犠牲にしていないことを示しています。

第二に、GSM8K（数学的推論）においてMistral 7Bの改善は特に顕著です。LLaMA 2 7Bの14.6から35.4へと大幅に向上しています。これはアーキテクチャの改良だけでなく、学習データの品質と量が大きく貢献していると考えられます。

Mixtral 8x7B vs 他のモデル

Mixtral 8x7Bは、さらに上のクラスのモデルと比較されます。

Mixtral 8x7Bは多くのベンチマークでLLaMA 2 70Bと同等以上の性能を示しています。LLaMA 2 70Bのパラメータ数は70Bであるのに対し、Mixtralの活性パラメータ数は12.9Bです。推論時に実際に計算が走るパラメータ数はLLaMA 2 70Bの約5分の1であり、推論速度が大幅に速いことを意味します。

特筆すべきはGSM8K（数学）でのスコアです。Mixtral 8x7Bは74.4と、LLaMA 2 70B（56.8）やGPT-3.5（57.1）を大きく引き離しています。MoEによってモデルの「知識容量」が増大し、特に複雑な推論タスクで効果を発揮していることが示唆されます。

推論効率の観点

性能だけでなく、推論効率も重要な比較軸です。

Mixtral 8x7Bは総パラメータが46.7Bであるため、モデル全体をGPU VRAMに載せるには約90GB（BF16）が必要です。しかし、推論のFLOPsはMistral 7Bの約2倍程度に抑えられています。つまり、メモリに余裕があるサーバー環境では、70Bクラスの性能を13Bクラスの速度で得られるという、極めて魅力的なトレードオフを実現しています。

まとめ

本記事では、Mistral 7BとMixtral 8x7Bのアーキテクチャを解説しました。

• Sliding Window Attention（SWA） は注意の範囲を直近 $W$ トークンに限定することで、計算量を $O(n^2)$ から $O(n \times W)$ に削減します。層を重ねることで間接的に $L \times W$ の情報到達距離を確保し、実用上十分な長距離依存を捉えます
• Grouped-Query Attention（GQA） はKVヘッドをグループ化して共有することで、KVキャッシュのメモリを $g/h$ に削減します。Mistral 7Bでは8KVヘッド/32Queryヘッドで、MHAの25%のKVキャッシュで済みます
• Rolling Buffer KV Cache はSWAと連携して、KVキャッシュを窓サイズ $W$ の循環バッファとして実装します。これにより、系列長によらずKVキャッシュのメモリに上限が設定されます
• Mixtral 8x7B はMistral 7BのFFNを8つのSwiGLUエキスパートに置き換えたSparse MoEアーキテクチャです。Top-2ゲーティングにより、46.7Bのパラメータを持ちながら推論コストは12.9B相当に抑えられ、ロードバランス損失がエキスパートの均等な利用を保証します

Mistral/Mixtralの設計は、「利用可能な計算資源の中で最大の性能を引き出す」という実用的な目標に徹底的に最適化されています。SWAとRolling Bufferによる予測可能なメモリ管理、GQAによる推論効率化、MoEによる知識容量の拡大 — これらの技術は個別には既存の研究に基づくものですが、それらを適切に組み合わせた設計の完成度の高さがMistral/Mixtralの真価です。

次のステップとして、以下の記事も参考にしてください。