PyImageSearchサイトが、深層学習ライブラリであるPyTorchを使用してDeepSeek-V3のアーキテクチャを構築するチュートリアルを公開した。6回シリーズの第5回となるこのチュートリアルは、Puneet Mangla氏によるもので、これまでのレッスンで紹介された構成要素を組み合わせて、モデルが処理するテキストの単位であるトークンを、可能な各トークンに割り当てられるスコアであるロジットに変換する、完全な言語モデルを形成する。
これまでのレッスンでは、モデルの設定、トークンの位置をエンコードするRotary Position Embeddings (RoPE)、注意機構であるMultihead Latent Attention (MLA)、計算を複数の専門サブネットワークに分散させるMixture of Experts (MoE)（エキスパートの混合）を扱った。最後に、Multi-Token Prediction (MTP)は、一度に複数のトークンを予測することを可能にする。チュートリアルによれば、これは学習効率を向上させ、投機的デコーディングにより推論を高速化する可能性があるという。
トランスフォーマーの各ブロックは、MLAとMoEを「pre-norm」と呼ばれる方式で組み合わせる。これは、各サブレイヤーの前に正規化が行われる方式である。残差接続により信号が層から層へ直接伝播し、著者によれば、これは深いネットワークにおける勾配消失問題の解決に役立つ。完全なモデルはその後、トークン埋め込みと位置埋め込み、ドロップアウト層、ブロック群、最終正規化、語彙予測ヘッド、そしてオプションでMTPヘッドを連結する。
重み共有と複数トークン予測
チュートリアルではまた、トークン埋め込み行列と予測ヘッド間の重み共有についても説明している。語彙サイズ50,259トークン、埋め込み次元256の場合、この技術により約1,290万パラメータを節約できる。著者によれば、入力空間と出力空間が整合性を保つため、これはしばしば性能向上につながるという。
学習中、MTPヘッドの出力は複数の将来のタイムステップに対して計算され、それらの損失は主要な損失と、エキスパート間の負荷分散を目的としたMoEの補助損失と組み合わされる。提示されたコードでは、MTPの平均損失は主要な損失に対して0.3の重みで加重される。MTPヘッドは、推論時、つまりテキスト生成時には使用されない。
モデルサイズとメモリフットプリント
著者によれば、このモデルは約3,000万から3,500万のパラメータを6層に分散しており、層あたり200万から300万パラメータとなる。これに加えて、MTPヘッドごとに200万から300万パラメータが追加される。32ビットのFP32精度では、パラメータだけで120から140メガバイトを占める。勾配とAdamWオプティマイザの状態を含む学習には、約500から600MBが必要であり、これに活性化が加わる。1,024トークンの4サンプルからなるバッチの場合、活性化は200から300MBを占め、合計で約800MB近くになる。チュートリアルでは、この量は現在のグラフィックスプロセッシングユニット（GPU）で扱えるとしている。
推論時、FP16精度を使用するとメモリ使用量が約半分に削減される。この実装はキー・バリューキャッシュを使用せず、生成の各ステップでコンテキストを再計算する。シリーズの次のレッスンでは、実際のデータを用いたこのモデルの学習について扱う予定である。
