DeepSeek V3をJAX約200行で再実装
原題: Show HN: DeepSeek V3 in ~200 lines of pure Jax
なぜ重要か
大規模モデルのアーキテクチャを最小コードで理解・再現できる教材として、研究者や学習者にとって参照価値が高い。
個人開発者のVaibhaav Tiwari氏が、DeepSeek V3のアーキテクチャ（MLA、DeepSeekMoE、aux-loss-free balancing、MTP）をわずか190行の純粋なJAXコードで実装し、GitHubで公開した。外部ライブラリ（Flax、Equinox、PyTorch）に依存せず、MacBook Air上でスクラッチから学習を完了している。リポジトリは公開後に142のスターを獲得した。
Vaibhaav Tiwari氏がGitHubで公開したリポジトリ「deepseek-v3-jax」は、DeepSeek-V3の技術レポートに基づくアーキテクチャを190行の純粋なJAXだけで実装したものだ。FlaxもEquinoxもPyTorchも使わない。モデルのパラメータはプレーンな配列の辞書として扱われ、すべての処理が1つのファイル「deepseek_v3.py」に収まる関数として記述されている。
実装に含まれる主な要素は4つ。まず「Multi-head Latent Attention（MLA）」は、トークンごとにフルサイズのキー・バリューをキャッシュする代わりに小さな潜在ベクトルを保持し、1トークン1レイヤーあたりのキャッシュを480から144数値に削減する。次に「Decoupled RoPE」は全ヘッドで共有する回転位置エンコーディングキーを圧縮潜在の外に保持する手法。「DeepSeekMoE」は常に動作する1つの共有エキスパートと16のルーティングエキスパートを組み合わせる混合エキスパート構造。最後に補助損失なしのバランシングと「Multi-Token Prediction（MTP）」が含まれる。
学習データはパブリックドメインの探偵小説を使用し、MacBook Air単体でスクラッチからの事前学習を実施した。Tiwari氏はあわせて詳細な解説記事（vaibhaav.in）を公開しており、図解やコードの逐次説明、学習結果も掲載している。リポジトリのライセンスはMITで、公開直後から14フォーク・142スターを集めている。