## Memo ## Memo with LLM ### 論文情報 - **論文タイトル**: Jet-Nemotron: Efficient Language Model with Post Neural Architecture Search - **著者と所属**: Yuxian Gu, Qinghao Hu, Shang Yang, Haocheng Xi, Junyu Chen, Song Han, Han Cai（NVIDIA） - **カンファレンス/ジャーナル名**: arXiv preprint - **発表年**: 2025年 ### 論文概要 本論文では、既存のフル注意機構モデルと同等以上の精度を保ちながら生成スループットを大幅に向上させるハイブリッドアーキテクチャ言語モデル「Jet-Nemotron」を提案している。事前学習済みのフル注意モデルから出発してMLPの重みを固定し、効率的に注意ブロックの設計を探索するPost Neural Architecture Search（PostNAS）と呼ばれる新しいアーキテクチャ探索パイプラインを開発した。 ### 詳細解説 #### 問題設定 大規模言語モデル（LLM）は優れた性能を示すが、特に長文脈での生成や推論において、自己注意機構による計算複雑度O(n²)と大きなKey-Value（KV）キャッシュにより効率性が課題となっている。従来の線形注意機構を用いたモデルは効率的だが、精度が最先端のフル注意モデルに劣る問題があった。 #### 提案手法 **PostNAS（Post Neural Architecture Search）**: 1. **フル注意層の配置と除去**: 事前学習されたフル注意モデルから開始し、MLPの重みを固定。once-for-all超ネットワークを構築し、ビーム探索によって最適なフル注意層の配置を学習する。 2. **線形注意ブロック選択**: RWKV7、RetNet、Mamba2、GLA、DeltaNet、Gated DeltaNetなど6つの最先端線形注意ブロックを評価し、最適なものを選択。 3. **新しい注意ブロック設計**: 動的畳み込みを統合した新しい線形注意ブロック「JetBlock」を提案。入力に基づいて動的に畳み込みカーネルを生成するカーネル生成器モジュールを導入。 4. **ハードウェア対応アーキテクチャ探索**: パラメータ数ではなく生成スループットを直接的な指標として、キー/バリュー次元と注意ヘッド数を最適化。 **JetBlock**: Gated DeltaNetをベースに、動的畳み込みカーネルを生成するモジュールを追加。静的畳み込みと異なり、入力特徴に基づいて畳み込みカーネルを適応的に生成し、より表現力の高い特徴抽出を可能にする。 #### 新規性 - **PostNAS**: 従来のゼロから学習する手法と異なり、事前学習済みモデルから開始してMLPを固定し、注意メカニズムのみを探索することで、コストとリスクを大幅に削減 - **JetBlock**: 既存の線形注意ブロックが静的畳み込みカーネルを使用するのに対し、入力に応じて動的にカーネルを生成する新しいアプローチ - **ハードウェア対応設計**: パラメータ数ではなく、実際のハードウェアでの生成スループットを最適化指標として使用 #### 実験設定 **データセット**: Nemotron-CCとRedstone-QAを50Bトークンで第1段階、数学とコーディングドメインの高品質データを追加して350Bトークンで第2段階の学習を実施。 **評価指標**: MMLU(-Pro)、数学推論（GSM8K、MATH、MathQA）、常識推論（ARC-c、ARC-e、PIQA、Winograd）、検索タスク（FDA、SWDE、Squad）、コーディング（EvalPlus、CRUXEval）、長文脈タスク（LongBench）で総合評価。 **ベースライン**: Qwen2.5、Qwen3、Gemma3、Llama3.2などの最先端フル注意モデルと、Mamba2、RWKV7などの線形注意モデル、ハイブリッドモデルと比較。 #### 実験結果 **Jet-Nemotron-2B**: Qwen3-1.7B-Baseと比較して47倍高速な生成スループットを実現し、MMLU、MMLU-Pro、数学推論など全ベンチマークで同等以上の精度を達成。特にMMLD-Proでは15B総パラメータ、2.2B活性化パラメータを持つDeepSeek-V3-SmallやMoonlightを上回る性能を示した。 **効率性**: H100 GPU上で256Kコンテキスト長において、Qwen3-1.7B-Baseに対して事前充填で6.14倍、デコーディングで53.6倍の高速化を達成。KVキャッシュサイズの削減により、より大きなバッチサイズでの処理が可能となり、スループットが向上した。 論文の提案手法は、既存の事前学習済みモデルを活用して効率的にアーキテクチャを探索する新しいパラダイムを提示し、実用的なLLMの展開において大きな意義を持つ研究である。 ## Abstract 我々は、主要なフル注意モデルの精度に匹敵するかそれを上回りながら、生成スループットを大幅に改善するハイブリッドアーキテクチャ言語モデルの新しいファミリーであるJet-Nemotronを提案する。Jet-NemotronはPost Neural Architecture Search（PostNAS）を使用して開発されており、これは効率的なモデル設計を可能にする新しい神経アーキテクチャ探索パイプラインである。従来のアプローチとは異なり、PostNASは事前学習されたフル注意モデルから開始し、そのMLPの重みを固定することで、注意ブロック設計の効率的な探索を可能にする。パイプラインには4つの重要なコンポーネントが含まれている：（1）最適なフル注意層の配置と除去の学習、（2）線形注意ブロックの選択、（3）新しい注意ブロックの設計、（4）ハードウェア対応ハイパーパラメータ探索。我々のJet-Nemotron-2Bモデルは、包括的なベンチマークスイート全体でQwen3、Qwen2.5、Gemma3、およびLlama3.2と同等またはそれ以上の精度を達成しながら、最大53.6倍の生成スループット高速化と6.1倍の事前充填高速化を実現している。