Jet-Nemotron：後段ニューラルアーキテクチャ探索を用いた効率的な言語モデル

本稿では、ハイブリッドアーキテクチャを採用した新たな言語モデル族「Jet-Nemotron」を紹介する。Jet-Nemotronは、従来のフルアテンションモデルと同等またはそれ以上の精度を達成しつつ、生成スループットを著しく向上させている。Jet-Nemotronは、効率的なモデル設計を可能にする新しいニューラルアーキテクチャ探索パイプライン「PostNAS（後期神経アーキテクチャ探索）」を用いて開発された。従来の手法とは異なり、PostNASは事前学習済みのフルアテンションモデルを出発点とし、そのMLP重みを固定することで、アテンションブロックの設計を効率的に探索できる。このパイプラインは以下の4つの主要な構成要素からなる：（1）最適なフルアテンション層の配置および削除の学習、（2）線形アテンションブロックの選定、（3）新たなアテンションブロックの設計、（4）ハードウェアに配慮したハイパーパラメータ探索。本研究で提案するJet-Nemotron-2Bモデルは、Qwen3、Qwen2.5、Gemma3、Llama3.2と比較して、包括的なベンチマークセットにおいて同等または優れた精度を達成しつつ、生成スループットで最大53.6倍、プリフィル（prefilling）処理で最大6.1倍の高速化を実現した。また、パラメータ数が150億（合計）・22億（アクティブ）に達する最新の高度なMoEフルアテンションモデル（DeepSeek-V3-SmallやMoonlightなど）と比較しても、MMLUおよびMMLU-Proにおいてより高い精度を達成している。