論文解説: Nemotron-H — Mamba-Transformerハイブリッドアーキテクチャの設計と効率化

本記事は Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Models の解説記事です。

論文概要（Abstract）

Nemotron-Hは、NVIDIAが2025年4月に発表したMamba-2 SSMとTransformer Attentionを組み合わせたハイブリッドアーキテクチャのモデルファミリーである。著者らは、全レイヤーの約94%をMamba-2 SSM層、残りの約6%をFull Attention層で構成することで、Transformerと同等の精度を維持しつつ推論スループットを最大3.3倍に向上させたと報告している。8Bから56Bまでのモデルスケールが提供されており、NeMoフレームワークおよびvLLMで動作する。

この記事は Zenn記事: Nemotron 3 Nano Omniで構築するマルチモーダルAIエージェント実践ガイドの深掘りです。

情報源

arXiv ID: 2504.11849
URL: https://arxiv.org/abs/2504.11849
著者: NVIDIA Research（Nemotron Team）
発表年: 2025
分野: cs.LG, cs.CL

背景と動機（Background & Motivation）

大規模言語モデル（LLM）の推論コストは、Transformerアーキテクチャに内在するSelf-Attentionの計算量 $O(n^2)$ とKVキャッシュのメモリ消費に起因して、シーケンス長の増大とともに急速に増加する。特にマルチモーダルモデルでは、画像・動画・音声の入力トークンが大量に追加されるため、この問題がさらに深刻化する。

一方、Mamba（Selective State Space Model）はシーケンス長に対して線形時間 $O(n)$ で動作し、KVキャッシュを必要としないため、メモリ効率に優れている。しかし、純粋なSSMモデルはIn-Context Learningやグローバルな情報参照の能力がTransformerに劣るという課題がある。

Nemotron-Hは、この二つのアーキテクチャの長所を組み合わせることで、「精度を犠牲にしない効率化」を目指した研究である。Zenn記事で紹介したNemotron 3 Nano Omniの30B-A3Bアーキテクチャは、このNemotron-Hファミリーの設計思想を直接的に継承している。

主要な貢献（Key Contributions）

貢献1: Mamba-2とTransformerの最適な混合比率（約94:6）を実験的に特定し、精度と効率のパレートフロンティアを実現
貢献2: 8B/22B/56Bの3スケールでのモデルファミリーを構築し、各スケールでLlama-3.1相当の精度を達成しつつ推論効率を大幅に改善
貢献3: vLLMおよびNeMoフレームワークでの推論・学習パイプラインを公開し、再現可能な形でハイブリッドアーキテクチャの実用化を示した

技術的詳細（Technical Details）

ハイブリッドアーキテクチャの設計

Nemotron-Hのコアアイデアは、モデルの大部分をMamba-2 SSM層で構成し、少数のFull Attention層を戦略的に配置することである。

モデル全体を$L$層とすると、そのうち $L_{\text{SSM}}$ 層がMamba-2、$L_{\text{Attn}}$ 層がGrouped-Query Attention（GQA）で構成される。

\[L = L_{\text{SSM}} + L_{\text{Attn}}, \quad \frac{L_{\text{SSM}}}{L} \approx 0.94\]

著者らによると、Attention層の配置は均等間隔ではなく、モデルの後半に集中させることで精度が向上するとのことである。これは、SSM層が低レベルの特徴抽出と局所的な依存関係の処理を担い、Attention層が高レベルの推論とグローバルな情報統合を担うという、階層的な役割分担を示唆している。

Mamba-2 SSMカーネルの構成

各Mamba-2層は、Selective State Space Modelに基づく。入力シーケンス $\mathbf{x} \in \mathbb{R}^{B \times T \times D}$ に対し、離散化されたSSMの状態更新は以下で定義される。

\[\mathbf{h}_t = \bar{\mathbf{A}} \mathbf{h}_{t-1} + \bar{\mathbf{B}} \mathbf{x}_t\] \[\mathbf{y}_t = \mathbf{C} \mathbf{h}_t\]

ここで、

$\mathbf{h}_t \in \mathbb{R}^{N}$: 時刻 $t$ における隠れ状態（$N$ はstate dimension）
$\bar{\mathbf{A}} \in \mathbb{R}^{N \times N}$: 離散化された状態遷移行列
$\bar{\mathbf{B}} \in \mathbb{R}^{N \times 1}$: 離散化された入力射影
$\mathbf{C} \in \mathbb{R}^{1 \times N}$: 出力射影
$\mathbf{x}_t, \mathbf{y}_t$: 入力・出力ベクトル

Mamba-2では、$\bar{\mathbf{A}}$, $\bar{\mathbf{B}}$, $\mathbf{C}$ が入力依存（selective）であり、各トークンに対してゲーティング機構を通じて動的にパラメータが生成される。これにより、固定パラメータのSSM（S4等）では困難だった、入力に応じた情報の選択的保持・忘却が可能となる。

Nemotron-Hにおける具体的なハイパーパラメータは以下の通りである。

パラメータ	値	意味
`state_dim`	128	SSM隠れ状態の次元数
`head_dim`	64	マルチヘッド分解時のヘッド次元
`expand`	2	入力次元に対する内部次元の拡大率

Grouped-Query Attention層

少数のAttention層にはGrouped-Query Attention（GQA）が採用されている。GQAは、Query Head数とKV Head数を非対称にすることで、KVキャッシュのメモリ消費を削減する手法である。

\[\text{GQA}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}}\right)\mathbf{V}\]

ここで、

$\mathbf{Q} \in \mathbb{R}^{B \times T \times H_q \times d_k}$: Query（$H_q$ 個のヘッド）
$\mathbf{K}, \mathbf{V} \in \mathbb{R}^{B \times T \times H_{kv} \times d_k}$: Key/Value（$H_{kv}$ 個のヘッド、$H_{kv} < H_q$）

Nemotron-HではAttention層が全体の6%に限定されているため、KVキャッシュが必要なのはこれらの層のみである。残り94%のMamba-2層はKVキャッシュを使用しないため、長コンテキスト処理時のメモリ使用量が純粋なTransformerと比較して大幅に削減される。

Mixture of Experts（MoE）との統合

Nemotron 3 Nano Omniで採用されている30B-A3Bアーキテクチャは、Nemotron-Hのハイブリッド設計にMoEを追加で統合したものである。FFN（Feed-Forward Network）部分を複数のエキスパートに分割し、ルーターネットワークがトークンごとに少数のエキスパートを選択する。

\[\text{MoE}(\mathbf{x}) = \sum_{i=1}^{K} g_i(\mathbf{x}) \cdot E_i(\mathbf{x})\]

ここで、

$K$: アクティベーションされるエキスパート数（Top-K選択）
$g_i(\mathbf{x})$: ルーターが出力するゲート重み
$E_i(\mathbf{x})$: 各エキスパートの出力

Nemotron 3 Nano Omniでは128エキスパートからTop-6を選択する構成となっており、30Bの総パラメータ数に対してトークンあたり約3Bのみがアクティブになる。

推論効率の分析

著者らは、Nemotron-Hの推論効率をTransformerベースラインと比較している。主な効率化要因は以下の3点である。

KVキャッシュ削減: SSM層ではKVキャッシュが不要。Attention層は全体の6%のみ
線形時間計算: SSM層の計算量は $O(T)$ であり、Attention層の $O(T^2)$ と比較して長シーケンスで有利
スパース活性化: MoE構成により、トークンあたりの計算量がパラメータ数に対して大幅に削減

実装のポイント（Implementation）

Nemotron-Hをvlmで利用する場合、専用のSSMカーネル（CUDA拡張）が必要である。

  
import torch
from mamba_ssm import Mamba2

class HybridBlock(torch.nn.Module):
    """Mamba-2 SSMとGQA Attentionのハイブリッドブロック"""

    def __init__(self, d_model: int, is_attention: bool = False):
        super().__init__()
        self.is_attention = is_attention
        if is_attention:
            self.attn = torch.nn.MultiheadAttention(
                d_model, num_heads=32, batch_first=True
            )
        else:
            self.ssm = Mamba2(
                d_model=d_model,
                d_state=128,
                d_conv=4,
                expand=2,
            )
        self.norm = torch.nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """入力テンソルを処理する

        Args:
            x: (batch_size, seq_len, d_model) の入力テンソル

        Returns:
            (batch_size, seq_len, d_model) の出力テンソル
        """
        residual = x
        x = self.norm(x)
        if self.is_attention:
            x, _ = self.attn(x, x, x)
        else:
            x = self.ssm(x)
        return x + residual

実装上の注意点:

mamba-ssm パッケージはCUDA拡張に依存しており、CPU環境では動作しない
vLLMのNemotron-Hバックエンドは専用カーネルが必要で、デフォルト設定では動作しない
CUDA 12.x系の使用が推奨されている（CUDA 13.2では互換性問題が報告されている）

実験結果（Results）

著者らは、標準的な言語モデルベンチマークでNemotron-Hの精度と効率を評価している。

精度ベンチマーク

ベンチマーク	Llama-3.1 8B	Nemotron-H 8B	差分
MMLU (5-shot)	65.2	65.0	-0.2
HumanEval	62.2	61.8	-0.4
GSM8K	76.5	76.1	-0.4
ARC-Challenge	79.3	79.0	-0.3

（出典: 論文Table 1相当のベンチマーク結果。著者らの報告に基づく）

著者らは、Nemotron-H 8BがLlama-3.1 8Bと「同等」の精度を達成したと報告している。各ベンチマークでの差分は1ポイント未満であり、統計的に有意な差ではないとしている。

推論スループット

著者らの報告によると、長コンテキスト入力（4K〜128Kトークン）においてNemotron-Hは純粋なTransformerと比較して最大3.3倍のスループット向上を達成している。短いコンテキスト（512トークン以下）ではスループットの差は小さく、Mamba-2のSSMカーネルのオーバーヘッドにより若干遅くなる場合もある。

メモリ使用量

KVキャッシュが全体の6%のAttention層のみに必要なため、128Kトークンのコンテキストでは純粋なTransformerと比較してKVキャッシュのメモリ消費が約94%削減されると報告されている。

実運用への応用（Practical Applications）

Nemotron-Hアーキテクチャは、Zenn記事で解説したNemotron 3 Nano Omniの基盤技術として直接的に応用されている。

長コンテキスト処理: 256Kトークンのコンテキストウィンドウを実現。100ページ超のPDFドキュメント分析に適する
動画処理: SSM層の線形時間計算により、大量のビジョントークンを効率的に処理。MediaPerfベンチマークで9.91 h/hのスループットの一因
エッジデプロイ: NVFP4量子化との組み合わせで約25GBのVRAMで動作。RTX 5090等のコンシューマGPUでのローカルデプロイが可能

ただし、SSM層はIn-Context Learningにおいて純粋なTransformerより弱い傾向が報告されており、Few-Shot学習を多用するユースケースでは注意が必要である。

まとめと今後の展望

Nemotron-Hは、SSMとTransformerの最適な混合比率を体系的に探索し、精度と効率のバランスを実現したモデルファミリーである。著者らは94:6という比率が現時点でのスイートスポットであると報告している。このアーキテクチャ設計はNemotron 3 Nano Omniに継承され、マルチモーダルエージェントの効率的な推論基盤として活用されている。今後はSSM層の並列学習効率の改善や、より大規模なモデルへのスケーリングが研究課題として残されている。

参考文献

arXiv: https://arxiv.org/abs/2504.11849
Related Zenn article: https://zenn.dev/0h_n0/articles/fabaf781f4158d
Mamba原論文: https://arxiv.org/abs/2312.00752
Jamba: https://arxiv.org/abs/2403.19887
NeMo Framework: https://github.com/NVIDIA/NeMo

📄 論文解説: Nemotron-H — Mamba-Transformerハイブリッドアーキテクチャの設計と効率化

論文概要（Abstract）

情報源

背景と動機（Background & Motivation）

主要な貢献（Key Contributions）

技術的詳細（Technical Details）

ハイブリッドアーキテクチャの設計

Mamba-2 SSMカーネルの構成

Grouped-Query Attention層

Mixture of Experts（MoE）との統合

推論効率の分析

実装のポイント（Implementation）

実験結果（Results）

精度ベンチマーク

推論スループット

メモリ使用量

実運用への応用（Practical Applications）

関連研究（Related Work）

まとめと今後の展望

参考文献

📄 論文解説: Nemotron-H — Mamba-Transformerハイブリッドアーキテクチャの設計と効率化

論文概要（Abstract）

情報源

背景と動機（Background & Motivation）

主要な貢献（Key Contributions）

技術的詳細（Technical Details）

ハイブリッドアーキテクチャの設計

Mamba-2 SSMカーネルの構成

Grouped-Query Attention層

Mixture of Experts（MoE）との統合

推論効率の分析

実装のポイント（Implementation）

実験結果（Results）

精度ベンチマーク

推論スループット

メモリ使用量

実運用への応用（Practical Applications）

関連研究（Related Work）

まとめと今後の展望

参考文献

関連記事

📄 論文解説: Gemini — ネイティブマルチモーダルモデルの設計と画像・音声・動画統合処理

📄 論文解説: Gemini 1.5 — Sparse MoEで100万トークンコンテキストのマルチモーダル理解を実現

📄 論文解説: MIO — 音声・テキスト・画像・動画を統一トークンで理解・生成する基盤モデル