Jet-Nemotron: Efficient Language Model with Post Neural Architecture Search

Q1: Throughput Comparison with Equal Batch Size and Sequence Length

First, we would like to clarify that the throughput comparison using chunk-prefilling—where the batch size is maximized under a fixed GPU memory constraint—constitutes a fair, apples-to-apples evaluation. This setting reflects the maximum achievable throughput for each model.

To provide a more comprehensive comparison, we also evaluate the throughput of JetLM-2B and Qwen2.5-1.5B on an H100 GPU using identical batch sizes and sequence lengths. The results are summarized below:

Under this setting, JetLM-2B delivers up to a 5.55× speedup compared to Qwen2.5-1.5B.

Q2: Comparison with MobileLLM

Thank you for bringing this relevant work to our attention. We will include comparisons with MobileLLM in the revised manuscript. The results are provided below:

JetLM-2B achieves higher throughput than MobileLLM-1.5B while also delivering superior accuracy on downstream tasks, including MMLU, Math, Commonsense, and Retrieval.

Q3: Throughput Results on Lower-End Hardware

We measure the throughput of JetLM-2B and Qwen2.5-1.5B on the NVIDIA Jetson Orin (32GB) and NVIDIA RTX 3090 GPUs with a context length of 64K:

JetLM achieves 8.84x and 6.50x speedups over Qwen2.5 on the Jetson Orin and RTX 3090 GPUs, respectively.

Q4: Training Details

The total training time and number of tokens used to train JetLM-2B are summarized below:

Q5: Related Work

Thank you for pointing out these references. We will include and discuss them in the revised manuscript.

Q1: On the Study Comparing Different Attention Blocks

In this study, we ensure that all models use the same cache size to guarantee fair comparisons. As discussed in Section 2.5, cache size is the primary factor affecting generation throughput in long-context settings.

Q2: Finetuning Baseline Models on JetLM Training Dataset

Following your advice, we conducted experiments to provide more comprehensive comparisons with the baseline models. The results are summarized below:

JetLM still provides a substantial accuracy improvement compared to the fine-tuned baseline models.

Q3: Experiment Running Time

The search and training times for JetLM are summarized below:

Q4: PostNAS as a Proxy Testbed for Architecture Exploration (L92)

In the neural architecture search literature [1], leveraging results from a proxy task for architecture selection is a widely adopted strategy. This supports our conjecture that PostNAS can serve as an effective proxy task for architecture selection. We will revise the manuscript accordingly.

Q5: SWA Layer Placement and Final Attention Placement Configuration

We determine the attention placement configuration based on our super-network search results, as detailed in Sections 2.2 and 3.1. Specifically, we use the Retrieval task to guide the placement of full attention layers and the MMLU task to guide the placement of sliding window attention (SWA) layers. As a result, we assign full attention to the top two critical layers for Retrieval (layers 15 and 20), and SWA to the top two critical layers for MMLU (layers 21 and 22).

Q6: Kernel Implementation

Kernel implementations can impact throughput results. To ensure a fair comparison, we adopt the unified Triton implementations from flash-linear-attention [2] as a controlled experimental setting.

Q7: Apply to a 7B Model

While our method significantly reduces the cost of model architecture exploration, experiments with 7B-scale models still require substantial computational resources. We plan to open-source additional accelerated models with 7B+ parameters using our approach. However, due to hardware constraints, it is not feasible to complete these experiments within the rebuttal period. We will include the new results in the next version of our work.

[1] Zoph, Barret, et al. “Learning transferable architectures for scalable image recognition.” Proceedings of the IEEE conference on computer vision and pattern recognition. 2018.

[2] Yang, Songlin and Zhang, Yu. “FLA: A Triton-Based Library for Hardware-Efficient Implementations of Linear Attention Mechanism.” GitHub. 2024.

Q1: Regarding Comparison Against SOTA Models

We would like to clarify that JetLM is a general-purpose language model and has not undergone any domain-specific fine-tuning on downstream tasks such as MMLU, Math, or Commonsense. To ensure fair, apples-to-apples comparisons, we evaluate JetLM only against other state-of-the-art general-purpose language models. Comparisons with specialized or non-general-purpose models are beyond the scope of this work.

Q2: Controlled Study to Exclude the Influence of Training Data

We have conducted comprehensive controlled ablation studies (see Figure 3, Figure 5b, Table 1, and Table 2) to verify the effectiveness of each design component in JetLM. All models in these experiments were trained on the same dataset to ensure a fair comparison.

Additionally, we fine-tuned the baseline models (Qwen2.5, RWKV-7, and Mamba-2) on JetLM’s training dataset to provide a more comprehensive evaluation.

JetLM-2B outperforms all these finetuned baseline models by a significant margin.

Q3: Comparison Between DynaConvGDN and Gated DeltaNet

We have included results for the JetLM variant that incorporates the Gated DeltaNet block in the manuscript. In Table 1, "Gated DeltaNet" refers to this JetLM variant, not the original standalone Gated DeltaNet model. We will revise the text in the next version to clarify this point.

Beyond the MMLU, Math, Retrieval, and Commonsense tasks presented in Table 1, we further compare the DynaConvGDN block with the Gated DeltaNet block on four additional, more challenging benchmarks: MMLU-Pro, BBH, MBPP, and HumanEval.

The DynaConvGDN block outperforms the Gated DeltaNet block across all evaluated tasks, achieving an accuracy improvement of 1.9 to 3.0. These results clearly demonstrate the effectiveness of the DynaConvGDN design.

Moreover, our contributions extend well beyond the DynaConvGDN block. As shown in Figure 3, the searched attention layer placement, linear attention selection, and hardware-aware architecture search each contribute significantly to overall accuracy gains. Additionally, our approach substantially reduces both the cost and risk of LLM architecture exploration, offering practical benefits to the broader language model community.

Q4: Reasons for the Attention Layer Placement

We determine the attention placement configuration based on our super-network search results, as detailed in Sections 2.2 and 3.1. Specifically, we use the Retrieval task to guide the placement of full attention layers and the MMLU task to guide the placement of sliding window attention (SWA) layers. As a result, we assign full attention to the top two critical layers for Retrieval (layers 15 and 20), and SWA to the top two critical layers for MMLU (layers 21 and 22).

Dear Reviewer Ndha,

We are truly grateful for your valuable feedback and the time you’ve taken to review our work. We would be honored to hear if our responses have satisfactorily addressed your concerns, and we warmly welcome any further suggestions you may have.

With sincere thanks and best regards,

The Authors

Dear Reviewer Ndha,

Thank you once again for your thoughtful suggestions and comments. As the discussion period deadline approaches, we would greatly appreciate it if you could let us know whether our responses have adequately addressed your concerns.

Thank you!

Sincerely,

The Authors

Q1: Novelty and Technical Contributions

We respectfully disagree. The technical contributions presented in our work go well beyond the boundaries of the NAS community and address broader challenges in efficient language model design.

• First, we introduce PostNAS, a novel model architecture exploration paradigm for language models. A key distinction between PostNAS and prior NAS methods for image classification or object detection is its ability to flexibly explore attention block designs while reusing pre-trained LLMs. In contrast, previous NAS approaches require training from scratch to achieve similar objectives, which renders them ineffective for large language models.
• Second, our work offers novel insights into language model architecture design, such as the task-specific importance of attention layers and the finding that KV cache size is a more critical factor than parameter count for generation throughput.
• Third, our work introduces a novel linear attention block, DynaConvGDN, which integrates linear attention with dynamic convolution and hardware-aware architecture search. It consistently delivers significant accuracy improvements over previous linear attention blocks while maintaining comparable generation throughput.
• Fourth, our work introduces JetLM, a novel hybrid-architecture language model family that achieves superior accuracy and significantly higher generation throughput than prior state-of-the-art full-attention models (e.g., Qwen2.5, Gemma3, and Llama3.2). With its strong accuracy and exceptional inference efficiency, JetLM offers practical benefits for a wide range of applications requiring efficient language models.
• Fifth, our work significantly reduces the cost and risk associated with LLM architecture exploration, enabling faster and more efficient innovation in language model design. This contribution can benefit the broader language model community.

Q2: Final Model Architecture

The final JetLM models are composed of a stack of blocks, each containing a Multi-Layer Perceptron (MLP) layer and an attention layer. The attention layer is selected from one of three types: full attention, sliding window attention, or DynaConvGDN. The detailed architecture configurations are presented below:

The full attention and sliding window attention layers use grouped-query attention and are configured as follows:

For sliding window attention layers, the window size is set to 1,152 in JetLM-2B and 2,048 in JetLM-4B.

The DynaConvGDN layers are configured as follows:

Q3: Search and Training Cost

The search and training costs for JetLM are summarized below:

Q4: Search Algorithm in Hardware-Aware Architecture Search

We employed grid search for the hardware-aware architecture search.

Q5: Search Space

In Section 2.2, the search space consists of all combinations of selecting 2 layers out of 28, resulting in a total size of $C_{28}^2$.

In Section 2.5, the search space is detailed in Table 2, with $d_K \in [64, 96, 128, 192, 256]$ and $d_V \in [144, 192, 256, 288, 384, 576]$.