この記事では DeepSeek-R1 および DeepSeek-V3 のアーキテクチャについて解説します。DeepSeek-R1 は DeepSeek-V3 の事前学習済みモデルに追加学習を施したものなので、アーキテクチャはどちらも同じです。
主に DeepSeek-V3 Technical Report の §2 の内容に沿ってまとめていきます。理解に必要な基礎知識 (特に Attention について) は補足するので、ニューラルネットワークの基本的な知識があれば読めると思います。また、なるべく集合論的な曖昧さがないように書いているので、数学に慣れている人は読みやすいと思います。
元々は V3 の Technical Report および R1 の論文をまとめるつもりでしたが、あまりにも長くなってしまうので分割して、この記事ではアーキテクチャのみを解説することにしました。以下のタイトルの記事も執筆予定です。
- DeepSeek-R1 の凄いところまとめ
- DeepSeek-R1 (V3) の学習方法の解説
目次
語句説明
この記事を読むにあたって重要となる語句をまとめます。
- DeepSeek-V3
- DeepSeek から 2024年12月に発表された LLM モデル
- DeepSeek-R1
- DeepSeek から 2025年1月に発表されて話題になった LLM モデル
- DeepSeek-V3 の事前学習済モデルをベースに学習方法を工夫したもの
- MLA (Multi-head Latent Attention)
- Multi-head Attention に Latent space を加えたもの
- 推論時に必要なキャッシュの量を削減させる
- MoE (Mixture-of-Experts)
- 入力に応じて適切なネットワークに処理を振り分けるもの
- 学習を効率化させる役割を持つ
アーキテクチャ全体像
DeepSeek-V3 のアーキテクチャは基本的に以下の図の通りで、Transformer Block というものがいくつか連なったものになっています。正確には入力側に文章をベクトル化するための Tokenizer と Embedding があり、出力側にはその逆操作をする層が加わります。
ポイントは上図の右にある MLA と DeepSeekMoE です。MLA では、MultiHead Attention の性能を落とさずに
- 推論時の計算量を抑えるための KV-Cache という仕組みに必要なメモリ量の大幅な削減
に成功したようです。DeepSeekMoE では、入力に応じて処理をさせるネットワーク (専門家と呼ばれる) を選択する MoE (専門家の混在) と呼ばれる仕組みの
- 各専門家が共通の知識を獲得しようとする懸念の解消
- 学習時に特定の専門家にのみ振り分けられる問題の解消方法 (Load Balancing) の改善
に成功したようです。
この記事では、
- Tokenizer
- Embedding
- RMS Normalization
- Multi-Head Latent Attention (MLA)
- DeepSeekMoE (Mixture of Experts)
- 最後の出力の処理
についてまとめます。1, 2, 3, 6 は割と一般的なもので、4, 5 は DeepSeek オリジナルの要素が含まれています。token 間の関係についての処理は Attention のみで行い、それ以外は token ごとに処理をすることを念頭におくと理解しやすいと思います。
また、上の図のようにブロックの入力の前で分岐して出力後に合流する矢印がありますが、これは skip connection といって学習の際の勾配消失問題を解消する役割を持ち、これによって深い層のネットワークの学習が可能になります。skip connection についてはこの記事では解説しません。
Tokenizer について
まずは入力の最初の処理である Tokenizer について説明します。
Tokenizer とは
Tokenizer とは文章を token と呼ばれる単位に分割する役割を持つものの事です。例えば単語を token の単位として文章を分割すると
"吾輩は猫である。名前はまだ無い。"
↓
['吾輩', 'は', '猫', 'で', 'ある', '。', '名前', 'は', 'まだ', '無い', '。']のように文章が分割されます。他にも文字単位で分割する方法も考えられます。文章を token の列に分割することを token 化 (tokenize) と言います。ニューラルネットワーク上ではこのように文章を token 化した後、各 token に識別番号をふり、文章を番号列として扱います。
一般に、ネットワークが扱うことのできる token 数が多くなると計算コストや必要なメモリが増加します。一方、token 数が少なくなると
- 文章が細かく分割されすぎて意味のあるテキスト表現を学習しにくい
- 分割後の token 列が長くなる
などの問題が起きます。
その間をとって、粗すぎず細かすぎず、いい感じの分割を行う方法がいくつか知られています (例えば npaka さんの記事を見てください)。
また、文字列を単語の列ではなく byte の列とみなす手法もあり、そのようにすることで原理的にどんな言語でも受け付けられるようになります。
DeepSeek の tokenizer ではありませんが、Llama の tokenizer のデモが公開されています。難しめの漢字を入力すると byte 単位で分割されることが確認できます。
https://belladoreai.github.io/llama3-tokenizer-js/example-demo/build
DeepSeek-V3 の Tokenizer
Tokenizer についての論文中の記述はそれほど詳しくないので、わかる部分だけまとめます。この節は読み飛ばしてもアーキテクチャの理解に影響はありません。
まず、DeepSeek-V3 の Tokenizer は BBPE というアルゴリズムを採用しています。分割の最小単位は文字ではなく byte なので原理的にはどんな言語でも受け付けます。
V3 の Technical Report には V2 の Tokenizer からの変更点のみ書かれているので、まずは V2 の Tokenizer ([DS-24b §2.1] に書かれているもの) についてまとめます。
- 改行、句読点、CJK記号の併合を防ぐために、GPT-2 と同様に事前トークン化を行なった。
- 数字は1桁ずつに分割した。
- 約 24 GB の多言語コーパスでトレーニングした。
- これまでの経験から、トークンの数を 10万 (102,400) とした。
- さらに15個の特殊なトークンを追加した。
V3 の Tokenizer は上記から以下のような変更を加えたようです。
- トレーニングコーパスについて
- 数学サンプルとプログラミングサンプルの比率向上
- 英語、中国語以外の言語の比率向上
- プレトークナイザーとトレーニングデータを多言語圧縮効率を最適化するように変更
- トークンの数を 128K (コードをみる限り 129280) とした。
- V2 とは異なり、改行、句読点が含まれる token を禁止しなかった
- これによりトークン境界バイアス (詳しくはこちら) が発生する可能性があるため、改行、句読点などが含まれる token を学習時にランダムで分割した
Tokenizer の数学的な記号の準備
この後の説明のために、Tokenizer を数学的に、というよりも集合論的に、集合と写像を用いて表します。まず文章 $doc$ とは byte の有限列
$$doc = (b_1, \cdots, b_n), \quad (b_i \textrm{ は byte})$$
のことであるとし、文章全体の集合を $Doc$ とします。ただし byte とは $0$ 以上 $255$ 以下の整数であるとします。
有限部分集合 $T \subset Doc$ で、長さ $1$ の byte 列を全て含むものを token の集合といい、$t \in T$ を token といいます。$t$ は byte の有限列です。$|T|$ を語彙数と言います。
token の集合 $T$ を一つ固定します。$|T| < \infty$ なので、token $t$ に対して番号 $1 \leq j \leq |T|$ を 1 対 1 に与えることができ、それによってtoken と番号を同一視できます。この同一視により番号を token と見なす場合は $[j]$ のように $[]$ をつけることとします。
token の集合 $T$ に対して $T^{\oplus}$ を token の有限列全体の集合とおきます (この節だけの記号です)。Tokenizer とは、token の集合 $T$ と写像 $\mathrm{Tk}: Doc \to T^{\oplus}$ の組 $(T, \mathrm{Tk})$ のことであると定義します。
\[\xymatrix@R=2.8pt@M=8pt{ Doc \ar[r]^{\mathrm{Tk}} & T^{\oplus} \\ (b_1, \cdots, b_n) \ar@{|->}[r] \ar@{}[u]|{\style{display: inline-block; transform: rotate(-90deg)}{\in}} & ([j_1], \cdots, [j_m]) \ar@{}[u]|{\style{display: inline-block; transform: rotate(-90deg)}{\in}} }\]
token の集合 $T$ には、文章の開始を表す token ([SOS] (start of sentence)) や文章の終了を表す token ([EOS] (end of sentence)) など、byte 列とは対応しない特殊なトークンを追加する場合がありますが、$T$ を番号の集合とみなせば、それ用の番号をふるのみで追加できます。
Embedding (単語埋め込み)
Embedding 層では token を低次元のベクトル空間 $\mathbb{R}^{d_{emb}}$ に埋め込みます。$d_{emb}$ はベクトル空間の次元です。埋め込みの方法は次のとおりです。
まず token $[j]$ を、$j$ 番目のみ $1$、それ以外は $0$ である $|T|$ 次元のベクトルに変換します。
\[\xymatrix@R=3.4pt@M=8pt{ T \ar[r] & \mathbb{R}^{|T|} \\ [j] \ar@{|->}[r] \ar@{}[u]|{\style{display: inline-block; transform: rotate(-90deg)}{\in}} & (0, \cdots, 0, 1, 0, \cdots 0) \ar@{}[u]|{\style{display: inline-block; transform: rotate(-90deg)}{\in}} }\]
このようなベクトルを one-hot ベクトルといい、
$$\widehat{j} = {}^t (0, 0, \cdots, \overset{j番目}{1}, \cdots 0) $$
で表すこととします。${}^t$ は転置を意味し、$\widehat{j}$ は縦ベクトルであるとします。Attention の初期の論文ではベクトルを横ベクトルで表していましたが、DeepSeek の論文は縦ベクトルなのでこの記事では縦ベクトルで表します。 (縦か横かは行列を右から掛けるか左から掛けるかにのみ影響し、それ以外に本質的な差異はありません。)
Embedding は線形写像 $W^{emb}: \mathbb{R}^{|T|} \to \mathbb{R}^{d_{emb}}$、つまり $|T| \times d_{emb}$ 行列 $W^{emb}$ を用いて
\[\xymatrix@R=3.4pt@M=8pt{ T \ar[r] \ar@/^15pt/[rr]^-{\textrm{Emb}} & \mathbb{R}^{|T|} \ar[r]|-{\mathrm{w}} & \mathbb{R}^{d_{emb}} \\ [j] \ar@{|->}[r] \ar@{}[u]|{\style{display: inline-block; transform: rotate(-90deg)}{\in}} & \widehat{j} \ar@{|->}[r] \ar@{}[u]|{\style{display: inline-block; transform: rotate(-90deg)}{\in}} & W^{emb} \widehat{j} \ar@{}[u]|{\style{display: inline-block; transform: rotate(-90deg)}{\in}} }\]
で与えられます。ここで、$W^{emb}$ はバイアスなしの線形層で与えられます。矢印に重なるように $w$ がついているのは学習時に値が更新されることを意味ます。以後、線型写像を $W$ を用いて表すときは、バイアスなしの線形層で与えられ、学習時に値が更新されるものとします。
$\mathrm{Emb}$ により token $t \in T$ は $d_{emb}$ 次元のベクトルで表されます。tonek の列に対しては、token 毎に埋め込みます。
\[\xymatrix@R=3.4pt@M=8pt{ T^{L} \ar[r] \ar[r]^-{\textrm{Emb}}|-{\mathrm{w}} & (\mathbb{R}^{d_{emb}})^{L} \\ ([j_1], [j_2], \cdots, [j_{L}]) \ar@{|->}[r] \ar@{}[u]|{\style{display: inline-block; transform: rotate(-90deg)}{\in}} & (W^{emb} \widehat{j_1},\ W^{emb}\widehat{j_2},\ \cdots, \ W^{emb} \widehat{j_{L}}) \ar@{}[u]|{\style{display: inline-block; transform: rotate(-90deg)}{\in}} }\]
単語埋め込みが線形であることの補足
ニューラルネットワークでは非線形性を持たせるために活性化関数と呼ばれるものを用いますが、言語の意味に関する処理を行う場合は活性化関数を用いずに、線形な変換のみを行う場合が多いです。それはおそらく、単語間の意味がある種の線形性を持つと考えられているからです。
それは2013年に発表された Word2Vec と呼ばれるモデルにおいて、おそらく最初に確認されました。有名なのは
「王」-「男」+「女」=「女王」
という式で、Word2Vec により埋め込まれた単語ベクトル空間の中で左辺を計算すると、「女王」の単語ベクトルに近い値になります。(ちなみに近さはベクトルの成す角度のみで測るのでベクトルの大きさは関係なく、正確には線形性とはいえません。)
活性化関数を用いると線形性が崩れてしまい、単語間の意味が失われる可能性があるため、この後説明する Attention においても線形な変換が主に用いられるのだと思います。
RMS Normalization について
ここからは Transformer Block の中身の解説になります。Transformer Block の入力と出力の形式は、長さ $L$ の token の列を埋め込んだ
$$(W^{emb}\widehat{j_1} \ , W^{emb}\widehat{j_2} \ , \cdots, W^{emb} \widehat{j_{L}} ) \in (\mathbb{R}^{d_{emb}})^{L}$$
と同じ形式ですが、Attention 以外は token の位置毎に処理されます。以後、Transformer Block の入力データについては位置の添字になるべく $1 \leq \ell \leq L$ を用いるようにしますが、それ以外の場合はその限りではありません (この $L$ は論文中の記号での Transformer Block の数と被っていますが、別物です)。
Normalization はデータのスケールを調整して学習を安定させる役割を持ちます。まずは、RMS Normalization のベースとなった Layer Normalization について説明します。
Layer Normalization とは
一般にデータ
$$(x_1, x_2, \cdots, x_n) \in \mathbb{R}^n$$
に対して、平均 $\mu$ および分散 $\sigma$
\begin{align}\mu &= \frac{1}{n}\sum_{i=1}^{n}x_i \\ \sigma &= \frac{1}{n}\sum_{i=1}^n (x_i -\mu)^2 \end{align}
を計算し、$x_i$ を
$$\overline{x}_i = \frac{x_i -\mu}{\sqrt{\sigma}}$$
に置き換えると $(\overline{x}_1, \cdots, \overline{x}_n)$ の平均は $0$、分散は $1$ になります。Layer Nomalization ではさらに、$n$ 個の学習可能なパラメータ $\gamma_i$ と $1$ 個の学習可能なパラメータ $\beta$ と、小さい定数 $\varepsilon$ を用いて
$$\widetilde{x}_i = \frac{x_i -\mu}{\sqrt{\sigma + \varepsilon}} \gamma_i + \beta$$
と変換します。$\varepsilon$ は分母が $0$ にならないように便宜的につけています。
RMS Normalization とは
RMS Normalization は Layer Normalization の性能を落とさずに処理を簡略化したものです。RMS は Root Mean Square の略で、平均を取ることによるデータの中心化を省略し、分散の代わりに 2乗の平均の平方根を取ったものになります。つまり、入力データ
$$(x_1, x_2, \cdots, x_n) \in \mathbb{R}^{n}$$
と $n$ 個の学習可能なパラメータ $\gamma_i$ と小さい定数 $\varepsilon$ を用いて
$$\overline{x}_i = \frac{x_i}{\sqrt{\varepsilon + {\displaystyle \frac{1}{n} \sum_{i=1}^n x_i^2}}} \gamma_i$$
と変換します。
DeepSeek のアーキテクチャにおける RMS Normalization への入力データの形は、入力された token の数を $L$ として
$$(x_1, \cdots, x_{L}) \in (\mathbb{R}^{d_{emb}})^{L}$$
となります。例えば Embedding 直後では、各 token の埋め込みベクトル
$$(W^{emb} \widehat{j_1} \ , W^{emb} \widehat{j_2} \ , \cdots, W^{emb} \widehat{j_{L}}) \in (\mathbb{R}^{d_{emb}})^{L}$$
が入力データとなります。これに対して RMS Normalization を、各 token の位置毎に適用します。つまり
$$x_\ell = \begin{pmatrix} x_{\ell,1} \\ x_{\ell, 2} \\ \vdots \\ x_{\ell, d_{emb}} \end{pmatrix} \in \mathbb{R}^{d_{emb}}$$
に対して
$$\mathrm{RMS}(x_j)=\begin{pmatrix} {\displaystyle \frac{x_{\ell, 1}} {\sqrt{\varepsilon + {\displaystyle \frac{1}{d_{emb}} \sum_{i=1}^{d_{emb}} x_{\ell, i}^2} }} } \gamma_1 \\ \vdots \\ {\displaystyle \frac{x_{\ell, d_{emb}}} {\sqrt{\varepsilon + {\displaystyle \frac{1}{d_{emb}} \sum_{i=1}^{d_{emb}} x_{\ell, i}^2} }} } \gamma_{d_{emb}} \end{pmatrix}\in \mathbb{R}^{d_{emb}}$$
と定めます。ただし $\gamma_1, \cdots, \gamma_{d_{emb}}$ は学習可能なパラメータで、$(x_1, \cdots, x_{L}) \in (\mathbb{R}^{d_{emb}})^{L}$ に対して共通です。
このようにして写像
\[\xymatrix@R=3.4pt@M=8pt{ (\mathbb{R}^{d_{emb}})^{L} \ar[r]^-{\textrm{RMS}}|-{\mathrm{w}} & (\mathbb{R}^{d_{emb}})^{L} }\]
が定まります。
MultiHead Latent Attention (MLA)
次は、MultiHead Latent Attention (MLA) について解説します。
Attention や MultiHead Attention は自然言語モデルでよく用いられるものなので、最初にそれらを説明します。Attention は構造上、入力の順番を無視してしまうので、token に位置を表す情報を付加する必要があります。それを Positional Encoding と言います。MLA では RoPE (Rotary Positional Embedding) という手法を使っています。それについても解説します。
Attention とは
Attention とは、どの情報に注目すべきかを判断するための仕組みです。例えば「これはペンです。」という文章を英語に訳するタスクを行い、”This is a” まで訳できたとします。この次に来るであろう単語 “pen” を推測するのに、「これはペンです。」のうちどの部分に注目すべきか、ということを判断します。
このようなタスクをこなす場合、ネットワークには下の図の下部のような token の列を入力し、図の上部のような token の列を出力させます。
\[\xymatrix@C=2.2pt@M=8pt{ {\color{lightgray}\textrm{これ}} & {\color{lightgray}\textrm{は}} & {\color{lightgray}\textrm{ペン}} & {\color{lightgray}\textrm{です}} & {\color{lightgray}\textrm{。}} & \color{lightgray}{[\textrm{delim}]} & \textrm{this} & \textrm{is} & \textrm{a} & {\color{red}\textrm{?}} \\ [\textrm{sos}] \ar[u] & \textrm{これ} \ar[u] & \textrm{は} \ar[u] & \textrm{ペン} \ar[u] & \textrm{です} \ar[u] & \textrm{。} \ar[u] & [\textrm{delim}] \ar[u] & \textrm{this} \ar[u] & \textrm{is} \ar[u] & \textrm{a} \ar[u] }\]
基本的には左からひとつずつ token 入力し、入力した token の 1 つ次の token を予測させ、予測した token を右に繋げて再びその次の token を予測させ、、、ということを繰り返すことで文章を生成します。
ここで、[sos] は文章の始まりを表す特殊な token、[delim] は文章の区切りを表す特殊な token です。グレーの部分は出力させない場合もあります。
Attention の役割は、”a” の次の単語を予測するのに、図の下部の入力のどの部分に注目すれば良いのかを与えることです。
簡単のため日本語部分と英語部分を分け、日本語部分を token 化 (+ 埋め込み) をしたベクトル列を
$$(x_1, \cdots, x_n) \in (\mathbb{R}^{d_{emb}})^{n}$$
英語部分を token 化したものを
$$(y_1, \cdots, y_m) \in (\mathbb{R}^{d_{emb}})^{m}$$
とおき、”a” の埋め込みベクトルは $y_m$ であるとします。
単語埋め込み空間では内積により単語の近さを測ることができると考えられているので、$y_m$ と $x_i$ の内積をとって大きいものに注目する、という方法が考えられますが、必要なのは “a” そのものではなく “a” の次に来そうな単語です。
そこで、$y_m$ を別のベクトル空間 $\mathbb{R}^{d_q}$ に埋め込み、$x_1, \cdots, x_n$ も同じ次元のベクトル空間に埋め込み、そこで内積を取ることで近さを測ります。$y_m$ の $\mathbb{R}^{d_q}$ への埋め込みは次の単語の情報を持つことが期待されますが、$x_i$ の埋め込みは次の単語と関係ないので、$x$ と $y$ で別の埋め込み方が与えられるべきです。ただし内積を取るので同じ次元である必要があります。
よって 2 つの線型写像 $W^Q, W^K: \mathbb{R}^{d_{emb}} \to \mathbb{R}^{d_q}$ を与え、$W^Q y_m$ と $W^K x_i$ との内積を取ります。内積の値を $a_i = W^Q y_m \cdot W^K x_i$ とおくと
$$(a_1, \cdots, a_n) \in \mathbb{R}^n$$
が得られ、$a_i$ が大きいほど注目度が高いと考えます。$(a_1, \cdots, a_n)$ の $\textrm{softmax}$ を取とって $0 \sim 1$ の値に正規化したものを $(\overline{a}_1, \cdots, \overline{a}_n)$ とおきます。$\sum_{i=1}^{n} \overline{a}_i x_i$ は入力 token のうち $\overline{a}_i$ が大きい token が強調されたベクトルになります。
それに線形写像 $W^V: \mathrm{R}^{d_{emb}} \to \mathrm{R}^{d_v}$ を適用して
$$W^V \left(\sum_{n=1}^{n} \overline{a}_i x_i \right) = \sum_{n=1}^{n} \overline{a}_i W^V x_i$$
を出力とします (次の token を予測するには、この出力から token を予測する層が必要ですが、省略します)。$W^V$ による変換は注目度を計算する前に行っても後に行っても同じです。
Attention の式
一般的な用語に合わせて Attention を説明し、式を整理します。
$(y_1, \cdots, y_m) \in (\mathbb{R}^{d_{Input}})^m$ を Input (さっきの例だと英語の文章), $(x_1, \cdots, x_n) \in (\mathbb{R}^{d_{Mem}})^n$ を Momery (さっきの例だと日本語の文章) と呼びます。$y_i$ は $d_{Input}$ 次元、$x_i$ は $d_{Mem}$ 次元ベクトルとします。
線形写像 $W^Q: \mathbb{R}^{d_{Input}} \to \mathbb{R}^{d_q}$, $W^K: \mathbb{R}^{d_{Mem}} \to \mathbb{R}^{d_q}$, $W^V: \mathbb{R}^{d_{Mem}} \to \mathbb{R}^{d_v}$ に対して
\begin{align} q_i &= W^Q y_i \quad (1 \leq i \leq m) \\ k_i &= W^K x_i \quad (1 \leq i \leq n) \\ v_i &= W^V x_i \quad (1 \leq i \leq n) \end{align}
とおいて、$q_i$ を query、$k_i$ を key、$v_i$ を value といいます。$W^Q, W^K, W^V$ はバイアスなしの線形層で与えられます。このとき
\begin{align} Q &= (q_1, q_2, \cdots, q_m) \\[.3em] K &= (k_1, k_2, \cdots, k_n) \\[.3em] V &= (v_1, v_2, \cdots, v_n) \end{align}
(縦ベクトルを横に並べた行列) とおくと、Attention の出力は
$$\mathrm{Attention}(Q, K, V) = V \ \mathrm{softmax} \left(\frac{{}^t K Q}{\sqrt{d_q}}\right) $$
で与えられます。${}^t K$ は $K$ の転置行列です。(埋め込みベクトルを縦ベクトルとした影響で式の形が標準のものと少し違っています。)
ひとつひとつ見ていくと ${}^t K Q$ は
$${}^t K Q = \begin{pmatrix} k_1 \cdot q_1 & k_1 \cdot q_2 & \cdots & k_1 \cdot q_m\\ k_2 \cdot q_1 & k_2 \cdot q_2 & \cdots & k_2 \cdot q_m\\ \vdots & & \ddots & \vdots \\ k_n \cdot q_1 & k_n \cdot q_2 & \cdots & k_n \cdot q_m \end{pmatrix}$$
であり、最後の列 (縦の並び) に注目すれば、さっきの例でいう $y_m$ に対応する注目度に相当します。
$\sqrt{d_q}$ で割っていることに関しては、$\mathrm{softmax}$ を取る際に中身が大きくなると学習効率が落ちる (勾配が小さくなる) ことを防ぐためのようです。$1$ が $d_q$ 個並んだベクトル $(1, 1, \cdots, 1)$ のベクトルの長さが $\sqrt{d_q}$ であり、次元が大きいほど (適当に取った) ベクトルの長さが長くなる傾向があります。
$\mathrm{softmax}$ は列毎に取ります。
最後に
$$\mathrm{softmax} \left(\frac{{}^t K Q}{\sqrt{d_q}}\right) = \begin{pmatrix} \overline{a}_{11} & \overline{a}_{12} & \cdots & \overline{a}_{1m}\\ \overline{a}_{21} & \overline{a}_{22} & \cdots & \overline{a}_{2m}\\ \vdots & & \ddots & \vdots \\ \overline{a}_{n1} & \overline{a}_{n2} & \cdots & \overline{a}_{nm} \end{pmatrix}$$
とおけば
$$V\ \mathrm{softmax} \left(\frac{{}^t K Q}{\sqrt{q}}\right) = \left(\sum_{i=1}^n \overline{a}_{i1}v_i , \sum_{i=1}^n \overline{a}_{i2}v_i, \cdots, \sum_{i=1}^n \overline{a}_{im}v_i \right) $$
となります。$q_j$ に対応するベクトルが $\sum_{i=1}^n \overline{a}_{i j}v_i$ になります。
出力は $d_v$ 次元ベクトルが $m$ 個並んだものになるので、入力と型を合わせるなら $d_v = d_{Input}$ とすれば良いです。skip connection を付加することが多いため、$d_{Input} = d_{emb}$ (token の埋め込みの次元) であることが多いようです。
ちなみに $v_i = W^V x_i$ であり、$X$ を $x_i$ を横に並べた行列とすると、行列の積の結合性から
$$(W^V X) \ \mathrm{softmax} \left(\frac{{}^t K Q}{\sqrt{q}}\right) = W^V \left(X \ \mathrm{softmax} \left(\frac{{}^t K Q}{\sqrt{q}}\right) \right) $$
が成り立つので、$W^V$ の役割は単に最後に線形変換をかけているだけであり、本質は key と query により注目度を計算するところにあります。
self-Attention
Input と Memory が同じとき、self-Attention と言います。
MultiHead Attention とは
MultiHead Attention とは、Attention を複数並列に用意して、その出力を連結して、線形変換を施すものです。具体的には以下のようにします。
先ほどと同様に Input を $y \in (\mathbb{R}^{d_{Input}})^m$ $(y = (y_1, \cdots, y_m))$、Memory を $x \in (\mathbb{R}^{d_{Mem}})^n$ $(x = (x_1, \cdots, x_n))$ とします。そして
\begin{align}W^Q_h&: \mathbb{R}^{d_{Input}} \to \mathbb{R}^{d_{q_h}} \quad (1 \leq h \leq n_h) \\ W^K_h&: \mathbb{R}^{d_{Mem}} \to \mathbb{R}^{d_{q_h}} \quad (1 \leq h \leq n_h) \\ W^V_h &: \mathbb{R}^{d_{Mem}} \to \mathbb{R}^{d_{v_h}} \quad (1 \leq h \leq n_h) \end{align}
に対して
$$o_h = \mathrm{Attention}(W^Q_h y , W^K_h x, W^V_h x) \quad (1 \leq h \leq n_h)$$
を考えます。$o_h$ は $d_{v_h}$ 次元の縦ベクトルを横に $n$ 個並べたものです。$o_h$ を縦に連結した
$$A = \mathrm{vconcat}(o_1, \cdots, o_{n_h})$$
は $\sum_{h=1}^{n_h} d_{v_h}$ 次元の縦ベクトルを横に $n$ 個並べたものになります。これの各列成分に線形変換 $W^O: \mathbb{R}^{\sum_{h=1}^{n_h} d_{v_h}} \to \mathbb{R}^{d_o}$ を施したものが MultiHead Attention の出力となります。$W_O$ はバイアスなしの線形層として与えられます。
普通は $d_o = d_{emb}$, $d_{v_h} = d_{emb} / h$ とします。
RoPE (Rotary Positional Embeddings) とは
attention は前述したように、各 token を並列に処理するので token の順序を認識できません。そこで Positional Embedding という方法で位置情報を付加します。まずは初期に考えられた手法を説明し、その後 RoPE について説明します。
Absolute Positional Embedding
初期に考えられた Absolute Positional Embedding は以下のとおりです。
$(x_1, \cdots, x_n) \in (\mathbb{R}^{d_{emb}})^n$ を token の埋め込みの列とし、$j$ 番目のベクトル $x_j$ を
$$x_j = \begin{pmatrix}x_{j, 1} \\ \vdots \\ x_{j, d_{emb}} \end{pmatrix}$$
と表します。また、$p_{j, i} \in \mathbb{R}$ を
\begin{align} p_{j, 2i} &= \sin \left(\frac{j}{10000^{2i/d_{emb}}}\right) \\ p_{j, 2i+1} &= \cos \left(\frac{j}{10000^{2i/d_{emb}}}\right) \end{align}
とおいて位置を表すベクトル $p_j \in \mathbb{R}^{d_{emb}}$ を
$$p_j = \begin{pmatrix}p_{j, 1} \\ p_{j, 2}\\ \vdots \\ p_{j, d_{emb}} \end{pmatrix}$$
と定めます。そして
$$(x_1, \cdots, x_n) + (p_1, \cdots, p_n) = (x_1 + p_1, \cdots, x_n + p_n)$$
と位置を表すベクトルを足すことで token の埋め込みの列に位置情報を付加します (右辺の $+$ は $\mathbb{R}^{d_{emb}}$ 上のベクトルとしての和です)。
RoPE
Absolute Positional Embedding は Attention の入力前に行いますが、実際に Attention に重要なのは query と key の位置関係だけで、value には必要ありません。また、文章が長くなるほど絶対的な位置よりも相対的な位置関係が明確になる方が好ましいと考えられます。RoPE はこの 2 つの性質を持ちます。
RoPE は Attention の内部で行います。まず、以下の行列
$$\begin{pmatrix} \cos \theta & -\sin \theta \\ \sin \theta & \cos \theta \end{pmatrix}$$
は、縦ベクトルに左から作用させることで、平面を反時計回りに $\theta$ (ラジアン) 回転させることに注意します。また、この行列の転置は逆回転の行列になります。
query と key の次元 $d_q$ が偶数であると仮定して
$$\theta_i = \frac{1}{10000^{2(i-1) / d_q}}$$
とおき、
$$R_{j} = \begin{pmatrix} \cos j \theta_1 & -\sin j \theta_1 & 0 & 0 & \cdots & 0 & 0 \\ \sin j \theta_1 & \cos j \theta_1 & 0 & 0 & \cdots & 0 & 0 \\ 0 & 0 & \cos j \theta_2 & -\sin j\theta_2 & \cdots & 0 & 0 \\ 0 & 0 & \sin j \theta_2 & \cos j \theta_2 & \cdots & 0 & 0 \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & 0 & 0 & \cdots & \cos j \theta_{\frac{d_q}{2}} & -\sin j \theta_{\frac{d_q}{2}} \\ 0 & 0 & 0 & 0 & \cdots & \sin j \theta_{\frac{d_q}{2}} & \cos j \theta_{\frac{d_q}{2}} \end{pmatrix}$$
という行列を考えます。これは $d_q$ 次元のベクトルを $2$ 次元ずつに分けて、それぞれ時計回りに $j \theta_i$ 回転させる行列を表します。
$(x_1, \cdots, x_n) \in (\mathbb{R}^{d_{emb}})^n$ を self-Attention の Input (かつ Memory) とし、$W^Q: \mathbb{R}^{d_{emb}} \to \mathbb{R}^{d_q}$ を query への線形写像、$W^K: \mathbb{R}^{d_{emb}} \to \mathbb{R}^{d_q}$ を key への線形写像とします。$j$ 番目の query、$\ell$ 番目の key をそれぞれ
\begin{align} \overline{q}_j &= R_j q_j = R_j W^Q x_j \\ \overline{k}_{\ell} &= R_{\ell}k_{\ell} = R_{\ell} W^K x_{\ell} \end{align}
と $R_j$, $R_{\ell}$ で回転させます。すると $\overline{k}_{\ell}$ と $\overline{q}_j$ の内積は
\begin{align} \overline{k}_{\ell} \cdot \overline{q}_j &= {}^t(R_{\ell} k_{\ell}) R_j q_j\\ &= {}^t k_{\ell} \, {}^tR_{\ell} R_{j} q_j \\ &= {}^t k_{\ell} R_{j-\ell} \, q_{j} \\ \end{align}
となり、$k_{\ell}$ と $q_j$ の間に行列 $R_{j -\ell}$ が挟まります。ただし転置は逆回転なので ${}^tR_\ell = R_{-\ell}$ です。
$R_{j -\ell}$ は $j$ と $\ell$ の相対的な位置で決まるという特徴があり、また $|j -\ell|$ が大きくなるほど $\overline{k}_{\ell} \cdot \overline{q}_j $ の値が小さくなることが知られています (long-term decay)。これにより、比較的遠い位置にある単語同士の attention は小さくなります。long-term decay は [RoPE] で式変形により証明されていますが、完全な証明にはなっていません。
MultiHead Latent Attention (MLA) とは
ここからは本題の MultiHead Latent Attention について説明します。上の図を見ると query の前と key-value の前に Latent という層が挟まっていて、RoRE を適用したものとしていないものを結合 (concat) するという形になっています。また、斜線が引かれている部分は推論時にキャッシュされます。
MLA の特徴
MLA は DeepSeek-V2 [DS-24c] で導入された方法です。一般に文章生成においては、途中まで生成された文章 (token 列) を入力として次の token を予測する、というタスクを繰り返します。例えば以下の図だと、”a” まで入力して次の単語 “pen” を予測し、”pen” の次の単語を予測するときには下の入力の末尾に “pen” を加えてその次の単語を予測します。
\[\xymatrix@C=2.2pt@M=8pt{ {\color{lightgray}\textrm{これ}} & {\color{lightgray}\textrm{は}} & {\color{lightgray}\textrm{ペン}} & {\color{lightgray}\textrm{です}} & {\color{lightgray}\textrm{。}} & \color{lightgray}{[\textrm{delim}]} & \textrm{this} & \textrm{is} & \textrm{a} & {\color{red}\textrm{?}} \\ [\textrm{sos}] \ar[u] & \textrm{これ} \ar[u] & \textrm{は} \ar[u] & \textrm{ペン} \ar[u] & \textrm{です} \ar[u] & \textrm{。} \ar[u] & [\textrm{delim}] \ar[u] & \textrm{this} \ar[u] & \textrm{is} \ar[u] & \textrm{a} \ar[u] }\]
Attention の計算において上の図の ? を予測するとき、”a” の位置に対応する出力の計算のためには、”a” の query, key, value と “a” より前の token の key, value のみわかっていれば十分ですが、”a” より前の token の key, value は前の token の予測時に計算されているので、キャッシュすることで計算効率が向上します。これを KVcache といいます。しかし文章が長くなるとキャッシュに必要なメモリが多くなります。
KVcache の削減の試みは以前にもあったようですが、パフォーマンスが低下してしまうという問題がありました。MLA では cache するべき情報を低い次元のベクトル空間 (Latent Space) に埋め込むことで KVcache を削減し、かつパフォーマンスを維持することに成功しました。
MLA の仕組み
MLA の入力を $(x_1, \cdots, x_L) \in (\mathbb{R}^{d_{emb}})^L$ とします。Head の数は $n_h$ とします。MLA については論文の図を見ながら式を見ることをお勧めします。
まず key と value については、線型写像
$$W^{DKV}: \mathbb{R}^{d_{emb}} \to \mathbb{R}^{d^{KV}_c}$$
により $x_\ell$ を $\mathbb{R}^{d^{KV}_c}$ (Latent Space) に埋め込んで、Normalize したものを
$$c^{KV}_\ell = \mathrm{RMS}(W^{DKV} x_\ell)$$
とおき (論文には書いていませんが、潜在空間への埋め込み時は $\mathrm{RMS}$ を行うようです)、線型写像
\begin{align}& W^{UK}_h: \mathbb{R}^{d^{KV}_c} \to \mathbb{R}^{d_q} \quad (1 \leq h \leq n_h) \\ & W^{UV}_h: \mathbb{R}^{d^{KV}_c} \to \mathbb{R}^{d_v} \quad (1 \leq h \leq n_h)\end{align}
により
\begin{align} k^c_{\ell,h} &= W^{UK}_h c^{KV}_\ell \quad (1 \leq h \leq n_h) \\ v^c_{\ell,h} &= W^{UV}_h c^{KV}_\ell \quad (1 \leq h \leq n_h) \end{align}
と key, value を定めます。$c^{KV}_i$ は全ヘッド共通で、 推論時にキャッシュされます。$W_h^{DKV}$ の $D$ は down-projection の $D$, $W_h^{UK}$ の $U$ は up-projection の $U$ だと思われます。
query も線型写像
$$W^{DQ}: \mathbb{R}^{d_{emb}} \to \mathbb{R}^{d^{Q}_c}$$
により $x_\ell$ を Latent Space $\mathbb{R}^{d^{Q}_c}$ に埋め込んで、Normalization したものを
$$c^{Q}_\ell = \mathrm{RMS}(W^{DQ} x_\ell)$$
とおき、線型写像
$$W^{UQ}_h: \mathbb{R}^{d^{Q}_c} \to \mathbb{R}^{d_q} \quad (1 \leq h \leq n_h)$$
により
\begin{align} q^c_{\ell,h} = W^{UQ}_h c^Q_\ell \quad (1 \leq h \leq n_h)\end{align}
と query を定めます。
次に RoPE ですが、$k^c_{\ell,h}$ に行列 $R_\ell$ を掛ける方式にすると、$k^c_{\ell,h}$ をキャッシュしていないので推論時に毎回 $R_\ell$ を掛ける必要があり効率が悪く、$c^{KV}_\ell$ に $R_\ell$ を掛けると、query $q^c_{j ,h}$ との内積をとるときに
$${}^t (W^{UK} R_{\ell} c^{KV}_\ell) R_j q^c_{j ,h} = {}^t c^{KV}_\ell R_{-\ell} {}^t W^{UK} R_j q^c_{j ,h}$$
と、$R_{-\ell}$ と $R_{j}$ 間に行列 $W^{UK}_h$ が入ってしまい、相対位置がうまく表現されません。
そこで、RoPE を適用した query と key を別に用意し、さっき定義した $k^c_{i,h}$ や $q^c_{i,h}$ の後ろにくっつけます。RoPE を適用した key は推論時にキャッシュし、(おそらく) キャッシュ量の削減のため各ヘッドで共有します。具体的には以下の通りです。線型写像
\begin{align} W^{KR} & : \mathbb{R}^{d^{Q}_c} \to \mathbb{R}^{d_R} \\ W^{QR}_h & : \mathbb{R}^{d^{Q}_c} \to \mathbb{R}^{d_R} \quad (1 \leq h \leq n_h) \end{align}
を用いて
\begin{align} k^R_\ell &= R_\ell W^{KR}_h x_\ell \\ q^R_{\ell, h} &= R_\ell W^{QR}_h c^Q_\ell \\ k_{\ell, h} &= \mathrm{vconcat}(k^c_{\ell,h}, k^R_\ell) \\ q_{\ell, h} &= \mathrm{vconcat}(q^c_{\ell,h}, q^R_{\ell, h}) \\ \end{align}
とします。ただし $\mathrm{vconcat}$ は token 列の方向ではなく、埋め込みベクトルの方向に結合します。そして
\begin{align} Q_h &= (q_{1, h}, \cdots, q_{n, h}) \\ K_h &= (k_{1, h}, \cdots, k_{n, h}) \\ V_h &= (v^c_{1, h}, \cdots, v^c_{n, h}) \end{align}
として $h$ 番目の Head の Attention
$$o_{h} = \mathrm{Attention}(Q_h, K_h, V_h)$$
を計算して、$o_h$ を縦方向に結合して、線型写像 $W^O: \mathbb{R}^{d_v n_h} \to \mathbb{R}^{d_{emb}}$ を縦ベクトル毎に作用させて MLA の出力を得ます。
\begin{gather} u_i = \mathrm{vconcat}(o_1, \cdots, o_{n_h}) \\ \mathrm{MLA}(x_1, \cdots, x_n) = (W^O u_1, \cdots, W^O u_n) \end{gather}
DeepSeek-V3 の一番大きいモデル (671B) では
\begin{align} d_{emb} &= 7168 && (\textrm{ 埋め込み次元 }) \\ n_h &= 128 && (\textrm{ Headの数})\\ d^{KV}_c &= 512 && (\textrm{ KVの潜在空間の次元}) \\ d^{Q}_c &= 1536 && (\textrm{ query の潜在空間の次元}) \\ d_q &= 128 && (\textrm{ RoPE を作用させないベクトルの次元}) \\ d_R &= 64 && (\textrm{ RoPE を作用させるベクトルの次元}) \\ \end{align}
とされています。また、一番小さいモデルでは query の Latent Space は使われていません。
DeepSeekMoE (Mixture of Experts)
次は MoE (Mixture of Experts) について説明します。通常の Transformer では 1 つの 2 層の Feed-Forward Network
$$\mathrm{FNN}(x) = W_2 \varphi(W_1 x)$$
を全ての位置で使い回します。MoE では $\mathrm{FNN}$ を複数用意し、その中から入力に応じてうまく処理できる $\mathrm{FNN}$ をいくつか選び、それらに処理をさせます。$\mathrm{FNN}$ 達は専門家 (expert) に例えられ、専門家の集まりの中から適した専門家を選び、処理させる、と解釈できます。
MoE のメリットを ChatGPT に聞いたところ
- 計算コスト削減
- 一部のエキスパートのみを使用し、効率的な推論が可能
- スケーラビリティ向上
- 巨大モデルを運用しやすく、高性能を維持できる
- 専門的な学習が可能
- 各エキスパートが異なる知識を学習し、マルチタスクに強い
- 推論が高速化
- 必要なエキスパートのみ使用するため、応答速度が向上
- 分散処理が容易
- エキスパートを異なるGPUやサーバーに配置し、効率的に計算可能
とのことらしいです。
DeepSeekMoE の特徴
DeepSeekMoE の特徴のひとつ目は、普遍的な情報の処理を行う expert を備えていることです。つまり、どのような入力に対しても処理を行う expert (shared experts) がいくつか存在し、さらに入力に応じて他の expert を選択する、という仕組みになっています。
従来の MoE では shared experts が存在しないため、複数の expert が共通の知識を獲得しようとしてしまう懸念がありましたが、shared experts によって他の expert がより専門化され、その冗長性が解消されるという考えのようです。
一般に MoE の学習の際、特定の expert のみが選ばれてしまう問題が発生するようです。それを避けるために従来の方法では Auxiliary-Loss というものを導入し、学習の損失関数に加えていました。しかし、Auxiliary-Loss が大きいとモデルの学習が妨げられパフォーマンスが低下し、小さいと特定の expert のみが選ばれる、という問題がありました。
DeepSeekMoE の特徴のふたつ目は、選択される expert の偏りの解消 (load balancing) を Auxiliary-Loss なしに行うことです。これにより上述の問題が解消されたようです ([DS-MoE])。
正確には Auxiliary-Loss-Free な方法をメインとしつつ、1 つの入力列に対して expert が極端に偏ることを防ぐために、Auxiliary-Loss に非常に小さな係数をかけて損失関数に加えたようです。
DeepSeekMoE の入力と出力
まずは DeepSeekMoE の入力から出力までを説明します。
$x \in \mathbb{R}^{d_{emb}}$ を MoE の入力とし (token の位置毎に計算します)、2 層の Feed-Forward Network
$$\mathrm{FNN}_i(x) = W^2_i \varphi(W^1_i x)$$
を $N_s + N_r$ 個用意します。ただし $\varphi$ は活性化関数で、ベクトルの要素毎に適用されます。DeepSeek-V3 では
$$\varphi(x) = \mathrm{silu}(x) = x \cdot \frac{1}{1 +e^{-x}}$$
が使われています。$W^1_i: \mathbb{R}^{d_{emb}} \to \mathbb{R}^{d_1}$, $W^2_i: \mathbb{R}^{d_1} \to \mathbb{R}^{d_{emb}}$ は線型写像です。(DeepSeek-V3 のコードでは $W^1_i$ と同じを形の行列 $W^3_i$ も定義し、$W^1_i x$ の代わりに要素毎の積 $W^1_i x \cdot W^3_i x$ を使っていますが理由は分かりません。)
$1 \leq i \leq N_s$ のとき
$$\mathrm{FNN}^{(s)}_i = \mathrm{FNN}_i$$
を shared expert、$N_s + 1 \leq i \leq N_r$ のとき
$$\mathrm{FNN}^{(r)}_{i-N_s} = \mathrm{FNN}_i$$
を routed expert と呼びます。つまり最初の $N_s$ 個を shared expert、残りの $N_r$ を routed expert とし、routed expert の添字を $1$ から振り直します。
次に $N_r$ 個の routed experts から $K_r$ 個の expert を選択する方法を説明します。まず線型写像 $W^G: \mathbb{R}^{d_{emb}} \to \mathbb{R}^{N_r}$ をとり、$W^G x$ の各成分の sigmoid を取ります。
$$\begin{pmatrix}s_1 \\ \vdots \\ s_{N_r} \end{pmatrix} = \mathrm{sigmoid}(W^G x)$$
$s_i$ を $i$ 番目の expert のスコアと呼びます。DeepSeek-V2 では softmaxを取っていましたが、V3 では sigmoid にしたようです。ちなみに論文では 「$i$ 番目の expert の centroid ベクトル $e_i$ との内積をとる」と書いていますが、$W^G$ の $i$ 行目が $e_i$ です。
そして $s_1, \cdots s_{N_r}$ のうち、値が大きいもの上位 $K_r$ 個の添字の集合を
$$\mathrm{Topk}((s_1, \cdots, s_{N_r}), K_r)$$
とおき、
\begin{align} g_i^{\prime} &= \begin{cases} s_i & (i \in \mathrm{Topk}((s_1, \cdots, s_{N_r}), K_r)) \\ 0 & (それ以外) \end{cases} \\ g_i &= \frac{g_i^{\prime}}{\sum_{j=1}^{N_r} g_j^{\prime}} \end{align}
とします。つまり、$\mathrm{Topk}$ に含まれないものは $0$ とし、和が $1$ になるように正規化します。入力 $x$ から $(g_1, \cdots, g_{N_r})$ を得る処理を
$$\mathrm{Gate}(x) = (g_1, \cdots, g_{N_r})$$
とおきます。$W^G$ の G は Gate の G です。
そして MoE の出力を
$$\mathrm{MoE}(x) = \sum_{i=1}^{N_s} \mathrm{FFN}^{(s)}_i (x) + \sum_{i=1}^{N_r} g_i \mathrm{FFN}^{(r)}_i (x)$$
とします。
DeepSeek V3 の一番大きいモデル (671B) では
\begin{align} N_s &= 1 && (\textrm{shared expert の数})\\ N_r &= 256 && (\textrm{routed expert の数})\\ K_r &= 8 && (\textrm{選ばれる expert の数})\\ d_1 &= 2048 && (\textrm{FNN の中間層の次元})\\ \end{align}
とされています。
コード上では、routed expert をいくつかのグループに分け、グループを選んでからその中の expert を選ぶ処理になっています。また、61 個あるうちの最初の 3 つの Transformer Block では MoE の代わりに普通の $\mathrm{FFN}$ が使われています。
Load Balancing
次は expert の偏りの解消について説明します。まずは従来の方法で用いられていた Auxiliary-Loss について説明します。
Auxiliary-Loss とは
入力値を $(x_1, \cdots, x_L) \in (\mathbb{R}^{d_{emb}})^n$ とします。入力の $\ell$ 番目の値 $x_{\ell}$ に対してスコアを
$$\begin{pmatrix}s_{\ell, 1} \\ \vdots \\ s_{\ell, N_r} \end{pmatrix} = \mathrm{sigmoid}(W^G x_{\ell})$$
とおきます。そして $i$ 番目の expert が選ばれた入力の位置の集合を
$$C_i = \{\ell \mid 1 \leq \ell \leq L, \ i \in \mathrm{Topk}((s_{\ell, 1}, \cdots, s_{\ell, N_r}), K_r)\}$$
とおき
\begin{align} f_i &= \frac{N_r}{K_r} \frac{|C_i|}{L} \\ s^{\prime}_{\ell, i} &= \frac{s_{\ell, i}}{\sum_{j=1}^{N_r} s_{\ell, j}} \\ P_i &= \frac{1}{T} \sum_{\ell = 1}^L s^{\prime}_{\ell, i} \end{align}
とします。$f_i$ は $i$ 番目の expert が選ばれた割合に $\frac{N_r}{K_r}$ を掛けたもので、$P_i$ は $i$ 番目の expert のスコア (を正規化したもの) の平均です。softmax でスコアを計算する場合は正規化が必要ありませんが、V3 では sigmoid を使っているので正規化しているものと思われます。$\frac{N_r}{K_r}$ を掛けている理由は分かりません。
そして、Auxiliary-Loss $\mathcal{L}_{Balance}$ を
$$\mathcal{L}_{Balance} = \alpha \sum_{t=1}^T f_i P_i$$
と定めます。$\alpha$ はハイパーパラメータで、DeepSeek-V3 では非常に小さな値に設定されたようです。$\mathcal{L}_{Balance}$ は expert の振り分け方が偏るほど大きくなるようですが、詳細は省略します。
Auxiliary-Loss-Free Load Balancing
DeepSeek の方法は単純で、学習中に Expert が選択される回数を監視し、選択される回数が少ない場合はスコアに大きめバイアスを加え、多い場合はバイアスを小さくするという方法です。ただしバイアスは $\mathrm{Topk}$ の計算にのみ使用され、最終的な出力には使用されません。
バイアスは学習のバッチ毎に更新され、以下のように更新されるようです。
- $b_i$ $(1 \leq i \leq N_r)$ の初期値を $0$ とする.
- バッチ内で $i$ 番目の Expert が選ばれた回数 $c_i$ を計算する.
- $\bar{c}$ を $N_r$ / (バッチ内の token 数) とし、$e_i = c_i -\bar{c}$ を計算する.
- $b_i \leftarrow b_i + u \cdot \mathrm{sign}(e_i)$ と更新する. ただし $u > 0$ はハイパーパラメータ.
つまり、平均より多く選ばれたものはバイアスを $u$ 減らし、平均より少なく選ばれたものはバイアスを $u$ 増やします。ちょうど平均だった場合はそのままです。
出力部分
以上で Transfomer Block 内の説明が終わったので、最後に出力の部分の処理について説明します。最後に出力の部分では、Transfomer Block の出力
$$(x_1, \cdots, x_L) \in (\mathbb{R}^{d_{emb}})^L$$
の各 $x_\ell$ $(1 \leq \ell \leq L)$ を、token の集合 $T$ 上の確率分布に変換します。
そのために、線型写像 $W: \mathbb{R}^{d_{emb}} \to \mathbb{R}^{|T|}$ ($|T|$ は全ての token の数) を用意して
$$W (\mathrm{RMS}(x_\ell))$$
と、$x_\ell$ を $\mathbb{R}^{|T|}$ 次元のベクトルに変換されます ($\mathrm{RMS}$ は RMS Normalization です)。それに $\mathrm{softmax}$ を適用させて、$T$ 上の確率分布 $p$ が得られます。
$$p([j]) = \mathrm{softmax}(W (\mathrm{RMS}(x_\ell)))_j$$
ただし $[j]$ は $j$ 番目の token を表します。
学習時はこの確率分布を使って損失関数を計算し、推論時は確率の最も高い token を選ぶか、サンプリングを行います。
$\mathrm{softmax}$ は温度と呼ばれるパラメータ $\tau$ を用いて
$$\mathrm{softmax}_{\tau} (x)_j = \frac{\exp(x_j /\tau)}{\sum_{i=1}^{|T|} \exp(x_j /\tau)}$$
変形される場合があります。$\tau = 1$ の場合は $\mathrm{softmax}$ と同じです。$\tau < 1$ の場合は確率の高いものがより強調されランダム性が下がり、$\tau > 1$ の場合は確率が全体的になだらかになりランダム性が上がります。
以上で、DeepSeek-R1 (V3) のアーキテクチャを一通り説明できました。
ご支援のお願い
記事を読んで、「支援してもいいよ」と思っていただけましたら、ご支援いただけると幸いです。サーバー維持費などに充てさせていただきます。登録不要で、100円から寄付でき、金額の90%がクリエイターに届きます。
参考文献
[DS-V2] DeepSeek-AI. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
[DS-V3] DeepSeek-AI. DeepSeek-V3 Technical Report.
[DS-24a] DeepSeek-AI. DeepSeek-Coder-V2: Breaking the Barrier of Closed-Source Models in Code Intelligence.
[DS-24b] DeepSeek-AI. DeepSeek LLM Scaling Open-Source Language Models with Longtermism.
[DS-24c] DeepSeek-AI. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model.
[DS-git] https://github.com/deepseek-ai/DeepSeek-V3
[DS-MoE] Daiほか. DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models.
[RoPE] Jianlin Suほか. ROFORMER: ENHANCED TRANSFORMER WITH ROTARY POSITION EMBEDDING.