Instruction-Following Pruning for Large Language Models

We thank Reviewer bo9C for the support and the valuable suggestions. Below we address each point raised.

Q1: Clarification on actual model speedup.

Since inference speed is concerned, we would like to first clarify that our method is motivated and designed for on-device models (e.g. on smartphone / laptop / desktop), where the inference typically samples a few responses given the same user query (or the same task). In this case, the same activated parameters are selected and cached as a dense model, therefore achieving the same speedup as the static pruning and dense baseline. We discussed the limitations of our work and possible extensions to batch inference in Section 5. We will better clarify our limitations in the next version. Thank you!

We evaluate the speedups by pruning the open-sourced LLaMA-3.1-8B-Instruct model to 3B. Although the tests are done on GPUs, we used batch size 1 and 4 generations per query, reflecting on-device usage. We report the time-to-first-token (TTFT) and the decoding time, both measured in seconds.

For dense models (8B and 3B), the TTFT consists of pre-filling only. For our method, we break down TTFT into its components:

• Sub-network selection (via the sparsity predictor)
• Parameter loading (load the selected parameters and cache the sub-network as a dense model)
• Pre-filling using the 3B sub-network

We highlight the following observation:

• Practical inference efficiency Gains: TTFT decreased by up to 57%, decoding time decreased by up to 41%.
• Minimal overhead from dynamic pruning & parameter caching: Overhead from dynamic pruning & parameter caching is negligible (~0.05s, ~2% of the total generation time)
• Despite dynamic masking, runtime of IFPruning is on par with static pruning, while offering input-specific adaptivity and superior accuracy.

We will include this analysis in the final version of our paper.

Q2: Sub-network overlap rate calculation

Thank you for asking. We clarify the process in three steps:

• 1: Data: We sample 128 inputs per dataset (MMLU, GSM8K, CodeAlpaca-20K, GPTTeacher). Inputs to the sparsity predictor are formatted with in-context examples (MMLU: 5-shot, GSM8K: 8-shot) or raw prompts (CodeAlpaca, GPTTeacher).
• 2: Sub-network Generation: Each input is passed through the sparsity predictor, which selects a fixed number of FFN units (a binary mask) at each layer, producing 128 sub-networks per dataset.
• 3. Overlap Calculation: We compare every pair of sub-networks within the 128 examples. For each pair, we compute the fraction of selected FFN units they share. The final overlap rate is the average of these pairwise overlaps.

We will include these details in the final version of our paper.

Q3: More details on SFT datasets

We will expand the final version with a detailed breakdown of the SFT datasets, including data sources, instruction formats, and size per domain.

Q4: Sparsity predictor architecture

The predictor is a lightweight model with 302M parameters, built on a pre-trained LM backbone. It consists of:

• A feature extractor (last hidden state of the final input token)
• Two-layer MLP:
  • Linear(hidden_dim → 128)
  • Linear(128 → num_layers × ffn_dim)

The output is a tensor of FFN importance scores per layer. The SoftTopK operator then converts the scores into structured binary masks.

Q5: Related work on SFT for efficient LLMs

Thank you for the pointers. We will add these references in the final paper to strengthen the related work section.

Q6: Will the code be published upon acceptance?

We will make our best effort to release the full implementation on top of an open-source model upon acceptance.

We thank Reviewer Kh2x for the support and the valuable suggestions. Please see our response below.

Q1: Why does continued pre-training help and what if we remove it?

Thank you for the question. We first explain our motivation followed by the ablation study on continued pre-training.

• Intuition of continued pre-training. In continued pretraining, we split long text into chunks; the model predicts the next chunk after selecting a sparse sub-network from the current one. Consider the extreme case where we split a pretraining text into just two chunks. the first acts as a prompt to select parameters, and the second as the target for prediction—closely mirroring the instruction–response format of SFT. By training the model in this way across millions of natural examples, the sparsity predictor learns to select the right sub-networks given different input contexts.

• Empirical Results: We ran an ablation study with and without continued pre-training (6B → 3B, train for 400B tokens)

We observe consistent and notable gains across all benchmarks. We will add these results and clarify this design choice in the final version.

Q2: Configuration details for the mask predictor network and the LLMs

The predictor consists of:

• An LLM with 302M parameter as the feature extractor (last hidden state of the final input token)
• Two-layer MLP:
  • Linear(hidden_dim → 128)
  • Linear(128 → num_layers × ffn_dim)

Regarding the LLMs used in this work: they follow standard LLM design, such as grouped-query attention and RMSNorm, with no custom components—similar to LLaMA. We will include these details in the final version.

Q3: Why no comparison with ShearedLLaMA or Minitron? What is the model distillation objective?

Thank you for your helpful question. We clarify both aspects:

• On comparison: We did not include ShearedLLaMA or Minitron due to differences in training data and model setup. Instead, we implemented a fairer pruning + distillation baseline using similar techniques in ShearedLLaMA and Minitron: structured pruning using learned masks, and logit distillation in continued pre-training. As shown below, our baseline outperforms the results in ShearedLLaMA.

Minitron results are only partially available, but our pruning baseline already outperforms its reported MMLU score. We will add these comparisons as a reference..

• Distillation objective: We apply KL divergence between the output distributions of the student and the teacher model for each output token. The teacher distribution only keeps the highest-scoring tokens, similar to Minitron. We minimize a combined loss of the standard next-token prediction and KL divergence loss.

Q4: Comparison with MoE.

Thank you for the question. We did not include MoE models due to a fundamental difference in inference scenarios. We will improve the clarity of our paper:

• Our method is designed for edge devices (e.g., smartphones) where memory and compute resources are limited, and the inference batch size is small.
• While MoE models are great for server-side large batch size inference, they are not efficient for on-device inference when generating responses given a single query.
• In this case, decoding is bottlenecked by weight loading. Since MoE requires reading many expert weights (e.g., 1-2 for each token), the cost of MoE is multiple times higher than a dense model and our method.

To illustrate, we compare our method with the open-source model, Qwen1.5-MoE-A2.7B. It activates 2.7B parameters per token. For our method, we prune LLaMA-3-8B to 3B parameters.

We report time-to-first-token (TTFT) and decoding time with input length = 4k, generation length = 100, and sample 4 responses for each query. For our method, we also report the latency for sub-network selection and loading parameters for selected sub-network.

We can see the dense baseline and our method have significantly better latency and throughput than MoE.

Finally, we agree with the reviewer's point that MoE and our method share the same spirit by dynamically activating parameters. In this regard, our model is a sparse model designed for on-device scenarios.

We thank Reviewer epS3 for the support and the valuable feedback. Please see our response below.

Q1: What is the reasoning behind designing a method that requires task-specific pruning for general-purpose LLMs?

We totally agree with the reviewer that having “general-purpose LLMs and adaptability across tasks” is essential. Indeed, our method aims to improve the general applicaiblity of smaller LMs. Please allow us to clarify a few things:

1. Given a user instruction (can be a general-purpose question such as a coding problem, a travel suggestion, and a knowledge-seeking question etc.), our method dynamically prunes the model and uses the most suited parameters for inference.
2. As the pruning mask is generated by “reading” the input instruction, our model can handle questions/tasks that are not seen during training. See for example our evaluation results on AlpacaEval and Arena-Hard.

In other words, our approach offers the following advantages:

1. Improve inference efficiency over large dense model (see Q3);
2. Improve the model quality over static pruning;
3. Pruning based on natural language description enables zero-shot generalization to unseen tasks and instructions, making the model still a general-purpose LLM.

In contrast, naive task-specific fine-tuning can have problems: no model available for unseen tasks during inference; memory/storage overhead for deploying many task-specific models; additional cost to collect training data for each task, etc.

Q2: Included baseline is not strong enough. Need to compare with baselines including LLM-Pruner or SliceGPT.

We would like to clarify that our pruning+distill baseline reflects the SOTA practice in large-scale model development. It first prunes the model then continued pre-trains on trillions of tokens. This method combined with logit distillation is consistent with recent model developments like LLaMA 3.2, Nvidia Minitron, and Gemma 3.

To show the effectiveness of our method and our baseline, the table below summarizes the (relative) performance drop compared to the source model with different pruning methods:

Our baseline and IFPruning achieves much higher sparsity and smaller accuracy degradation, validating the effectiveness of our method.

Q3: Overhead of generating masks for each new input

We evaluate the latency with LLaMA-3.1 8B as the source model and prune it to 3B. Due to space limit, please also see our response to Q1 by reviewer Lv3b. In brief, mask generation adds only ~0.1s overhead. Specifically, with input length 4000, output length 100, and response sample size 4, we show:

• Time-to-first token (TTFT) decreased by up to 57%
• Decoding time decreased by up to 41%
• Comparable latency to static 3B models

The overhead from sub-network selection (mask generation) & parameter loading is negligible (~0.1s, ~2% of the total generation time).

Q4: Concerns about JAX/AXLearn and framework portability

Our core method is framework-agnostic and can be easily implemented in PyTorch, as shown in our latency experiments using LLaMA-3.1-8B-Instruct and PyTorch. We will make our best effort to release the full implementation on top of an open-source model.

Q5: Sparsity predictor architecture

The predictor consists of:

• An LLM with 302M parameter as the feature extractor (last hidden state of the final input token)
• Two-layer MLP:
  • Linear(hidden_dim → 128)
  • Linear(128 → num_layers × ffn_dim)

Training is end-to-end with the masked LLM using standard language modeling loss.

Q6: Missing citations, discussions with other dynamic pruning papers, and figure issues

Thank you for your suggestion. We promise that we will address your comments in the final revision of our paper.

We thank Reviewer Lv3b for the support and the valuable suggestions for our paper. We will address the writing feedback, such as adding statistics of the dataset and discussing additional related work, in the next version. Please see our response to the questions and/or major comments below.

Q1: Comparison with static pruning, and latency numbers.

We agree with the reviewer that static pruning enjoys better batch parallelism and lower memory usage compared to dynamic pruning. We discussed the downside of our work in Section 5 and will improve the clarification in our next version. Thank you!

As far as inference speed is concerned, we would like to clarify that our method is designed for on-device models (e.g. on smartphone / laptop / desktop), where the inference typically samples a few responses given the same user query (or the same task). In this case, the same activated parameters are selected and cached as a dense model, therefore achieving the same speedup as the static pruning and dense baseline.

To illustrate the real inference speedup, we test the inference latency by pruning the open-sourced LLaMA-3.1-8B-Instruct model to 3B. Although the tests are done on GPUs, we used batch size 1 and 4 generations per query, reflecting on-device usage. We report the time-to-first-token (TTFT) and the decoding time, both measured in seconds.

For dense models (8B and 3B), the TTFT consists of pre-filling only. For our method, we break down TTFT into its components:

• Sub-network selection (via the sparsity predictor)
• Parameter loading (load the selected parameters and cache the sub-network as a dense model)
• Pre-filling using the 3B sub-network

Key takeaways:

• TTFT decreased by up to 57%, decoding time decreased by up to 41%. In total, we achieve 1.8x speedup compared to Llama-8b.
• Overhead from dynamic pruning & parameter caching is negligible (~0.05s, ~2% of the total generation time)
• Despite dynamic masking, runtime of IFPruning is on par with static pruning, while offering input-specific adaptivity and superior accuracy.

We will include this analysis in the final version of our paper.

Q2: Choice of SoftTopK over HardConcrete.

A2: We appreciate the suggestion. Both SoftTopK and HardConcrete are mask generation operators introduced in previous work, and not a contribution of our work. We chose SoftTopK for simplicity for two reasons:

1. Both methods achieve similar task performance.
2. SoftTopK does not require another auxiliary loss, whereas HardConcrete needs some tuning (e.g. on the auxiliary loss weight and its learning rate).
3. SoftTopk seems more stable in our preliminary study.

Q3: More details on SFT datasets

We will expand the final version with a detailed breakdown of the SFT datasets, including data sources, instruction formats, and size per domain.

Q4: Missing related work.

Thank you for highlighting these valuable references—we will include them in the final revision to strengthen the related work section.