Few-Shot Knowledge Distillation of LLMs With Counterfactual Explanations

We thank the reviewer for the review of the paper!

---

On practical significance few-shot knowledge distillation and CFE-infusion

Distilling LLMs into smaller resource-constrained environments is of utmost importance. The unprecedented surge of AI comes with ever-increasing model size and complexity [1,4], leading to an unprecedented demand for energy. As models get larger, deploying them on resource-constrained environments such as mobile phones, Internet of Things (IoT), and edge devices gets particularly challenging due to their limited processing power, memory, and battery life [3,4]. To this end, this work addresses the broader question: can we develop efficient reduced-order models that can achieve high performance with minimal training data?

In fact, a recent survey [1] on LLM compression highlights the significance and timely nature of our work:

"Crucially, the success of Knowledge Distillation (KD) in LLMs hinges on Dataset Distillation (DD) techniques, which enable the creation of compact, information-rich synthetic datasets that encapsulate the diverse and complex knowledge of the teacher LLMs." [1]

Our CFE-infused approach embodies this principle by injecting precisely those high-value samples to better approximate the teacher’s decision boundary.

---

Data-efficient machine learning is also rapidly growing field of research [7]. Two primary motivations for data-efficient ML is: (i) data acquisition or labeling cost [5]; and (ii) training or fine-tuning costs. In this context, our work shows how it is possible to distill using so few carefully selected examples – essentially, how few examples hold the key to knowledge transfer. Few-shot distillation can have applications for on‑device personalization (keyboard, voice dictation), low‑resource or dialectal languages in under-served regions, clinical or legal settings, finance applications with privacy rules restricting data sharing data, deployment in hospitals/firms that have very few local examples.

---

Classification with LLMs still have various applications today beyond sentiment analysis, e.g., security/safety, toxicity detection, hate-speech/abuse detection, spam email detection, fraud detection, topic categorization, and tabular dataset prediction in high-stakes applications such as finance, healthcare, education, etc [2,3, 6, 8].

---

On the extension to generative LLMs: One possible way to extend this to generative LLMs is to define counterfactual explanation as a minimal change to the input (prompt) that flips a chosen property of the generated text while keeping everything else as similar as possible. Formally, given a generative model $f$, a prompt $x$, and a binary attribute function $g(\text{output})$ (e.g., sentiment, toxicity, factuality, topic relevance), a counterfactual explanation prompt $x^{\star}$ would be a small semantic perturbation of $x$ such that the generated output $f(x^{\star})$ flips the value of $g(f(x^{\star}))$.

Another possible extension is to rethink counterfactual generation in terms of model sensitivity: identifying minimal semantic perturbations to the input $x$ that cause large shifts in the output distribution or likelihood. For generative models, this corresponds to large changes in the sequence-level probability: $p_\theta(y \mid x) = \prod_{t=1}^{T} p_\theta(y_t \mid x, y_{<t})$, where $y = (y_1, y_2, \ldots, y_T)$ is the generated sequence for an input prompt $x$. In this framing, CFEs could reveal parts of the input space where the model is uncertain, making them especially informative for distillation. While outside the scope of the current paper, this extension would be an interesting future direction, potentially inspiring novel and efficient pathways for data selection for a broad range of tasks, including distillation, supervised finetuning (SFT), post-training, etc.

---

We thank the reviewer for raising these important points; we will update the paper to include these clarifications and believe that the paper will be significantly strengthened as a result of this discussion.

---

References

[1] Knowledge Distillation and Dataset Distillation of Large Language Models: Emerging Trends, Challenges, and Future Directions, Luyang Fang, Xiaowei Yu, Jiazhang Cai, Yongkai Chen, et al.

[2] Smart Expert System: Large Language Models as Text Classifiers Zhiqiang Wang, Yiran Pang, Yanbin Lin.

[3] Enabling On-Device Large Language Model Personalization with Self-Supervised Data Selection and Synthesis, Ruiyang Qin, Jun Xia, Zhenge Jia, Meng Jiang et al.

[4] Optimizing LLMs for Resource-Constrained Environments: A Survey of Model Compression Techniques, Sanjay Surendranath Girija, Shashank Kapoor, Lakshit Arora, Dipen Pradhan, Aman Raj, Ankit Shetgaonkar.

[5] Investigating Cost-Efficiency of LLM-Generated Training Data for Conversational Semantic Frame Analysis, Shiho Matta, Yin Jou Huang, Fei Cheng, Hirokazu Kiyomaru, Yugo Murawaki.

[6] Less is More: Task-aware Layer-wise Distillation for Language Model Compression Chen Liang, Simiao Zuo, Qingru Zhang, Pengcheng He, Weizhu Chen, Tuo Zhao.

[7] Foundations of Data-efficient Machine Learning, Siddharth Joshi · Baharan Mirzasoleiman, ICML Tutorial 2025

[8] TabLLM: Few-shot Classification of Tabular Data with Large Language Models Stefan Hegselmann, Alejandro Buendia, Hunter Lang, Monica Agrawal, Xiaoyi Jiang, David Sontag.

Thank you for your suggestion. We will update the title and manuscript to better clarify that our focus is on classification tasks, while preserving the broader motivation and possible extensions to generative LLMs.

Please let us know if our rebuttal has addressed your other concerns. We are happy to clarify further or respond to any additional questions you may have.

We thank the reviewer for taking out the time to review this paper, and appreciate the positive opinion about our work!

---

On template designing and prompt choices

Thank you for the thoughtful question. We conducted an experiment varying four prompt templates for generating CFEs and observed that our method (CoD) is robust to prompt choices, showing low standard deviation across variants (e.g., ≤0.018) and consistently outperforming the KD baseline low shot cases (see table below (SST2 dataset)). This suggests CoD is not overly sensitive to prompt used. Automatic prompt generation (e.g., AutoPrompt, RL-based search) are typically more compute-intensive. Given our already strong and stable performance using simple manually designed prompts, such complex techniques may not be necessary. We will include this prompt robustness study and the varying prompts used in the updated version.

---

On soft-label computational and memory requirements

As the sample $k$ increases, we observe a consistent improvement in model performance (accuracy) for both KD+CoD and LWD+CoD (see table below (SST2 dataset)). However, this comes at the cost of increased computational demands. We use codecarbon package to track energy usage across compute components, excluding the counterfactual generation step (which is API-based and not locally measured). Runtime and energy consumption (CPU, GPU, RAM) grow with larger $k$, indicating more processing and memory usage. LWD+CoD is more computationally intensive than KD+CoD at every $k$, due to its additional computational steps of aligning the teacher and students intermediate representations. We will include this ablation in the revised manuscript.

KD + CoD

---

LWD + CoD

---

---

On the quality of soft label calibration on downstream performance

Thank you for the suggestion. We have added an ablation that isolates the effect of the soft labels. As shown in the table below, removing the soft label term ($\alpha$ = 0) leads to a substantial drop in performance across all shot levels. Although CoD still improves over KD in this setting, the gains are significantly reduced. This highlights that the soft label calibration from the teacher is a key contributor to the effectiveness of counterfactual explanation data. Additionally, when replacing soft labels with random values, performance degrades sharply, likely due to inconsistency with the hard labels, introducing conflicting supervision signals in the training objective. We will include this analysis in the revised paper.

---

On the extension to generative LLMs and structured tasks

One possible way to extend this to generative LLMs is to define counterfactual explanation as a minimal change to the input (prompt) that flips a chosen property of the generated text while keeping everything else as similar as possible. Formally, given a generative model $f$, a prompt $x$, and a binary attribute function $g(\text{output})$ (e.g., sentiment, toxicity, factuality, topic relevance), a counterfactual explanation prompt $x^{\star}$ would be a small semantic perturbation of $x$ such that the generated output $f(x^{\star})$ flips the value of $g(f(x^{\star}))$.

Another possible extension is to rethink counterfactual generation in terms of model sensitivity: identifying minimal semantic perturbations to the input $x$ that cause large shifts in the output distribution or likelihood. For generative models, this corresponds to large changes in the sequence-level probability: $p_\theta(y \mid x) = \prod_{t=1}^{T} p_\theta(y_t \mid x, y_{<t})$, where $y = (y_1, y_2, \ldots, y_T)$ is the generated sequence for an input prompt $x$. In this framing, CFEs could reveal parts of the input space where the model is uncertain, making them especially informative for distillation. While outside the scope of the current paper, this extension would be an interesting future direction, potentially inspiring novel and efficient pathways for data selection for a broad range of tasks, including distillation, supervised finetuning (SFT), post-training, etc.

---

On RAG comparison

We thank the reviewer for highlighting the conceptual connection to RAG. However, our setup and that of RAG systems differ fundamentally. Our method is focused on classification using supervision from soft labels and counterfactual explanations, whereas RAG augments data from a memory bank for generation using instruction-tuned models. A direct comparison would require switching to instruction-tuned LLMs capable of ingesting augmented data (e.g., "Given this input and context, classify "). However, the models we use are classification-specific (added classification head, similar to related works [1]), making such a comparison incompatible. Moreover, instruction-tuned LLMs used in RAG systems are typically very large (billions of parameters), which would make the comparison unbalanced and outside the scope of our focus on few-shot distillation under resource constraints. We believe bridging these paradigms is a promising direction for future work.

---

On limitations

We appreciate the reviewer’s for their suggestion. We will expand the limitations section to explicitly discuss the engineering overhead involved in our system.

---

[1] Less is More: Task-aware Layer-wise Distillation for Language Model Compression Chen Liang, Simiao Zuo, Qingru Zhang, Pengcheng He, Weizhu Chen, Tuo Zhao

We thank the reviewer for the thoughtful review and are delighted to hear that the reviewer has found our work “well-defined and formulated” and the method to be “clear and elegant”!

Below, we provide detailed responses to the questions asked:

---

On the effect of more samples ($k$) in the distillation process

Firstly, note that for a fixed distillation dataset budget $k$, our method effectively uses half the number of real samples for distillation. E.g., for $k=8$, KD uses 8 randomly selected labeled examples, while our strategy only uses 4 labeled examples complemented with their counterfactuals (CFEs) and still significantly outperforms regular KD. Thus, it achieves more with less in the few-shot regime.

Increasing the number of samples ($k$) does not necessarily worsen our approach, but it could also make regular KD/LWD without CFE more competitive (reducing the performance gap). We hypothesize that this could be because as $k$ increases, real datapoints start to densely cover the decision space and hence we see diminishing returns from CFEs. Indeed, if the entire training data would be available during distillation ($k \rightarrow \infty$), the student’s performance is expected to keep getting better using KD/LWD.

---

On new experiments for multi-class classification

Yes, the strategy could be extended to multi-class classification setting. One way to extend this to multiclass could be as follows: Consider a multi-class setting with classes ${1,2,...C}$. For a data point $x$ belonging to a specific class $c$, we can generate several CFEs to flip the label to each of the other classes, excluding $c$.

As an initial demonstration, we implemented this variant of our proposed strategy on SST5 Dataset. Stanford Sentiment Treebank with 5 labels: very positive, positive, neutral, negative, very negative. We filtered it to $3$ classes to have a good separation among labels. When regular KD/LWD uses $k$ real samples, our proposed strategy effectively uses $k/3$ real samples + $2k/3$ CFEs (so basically $1/3$ of what regular KD/LWD uses).

---

$k=9$

KD: $0.614$ KD+CoD (Ours): $\bf0.622$

LWD: $0.655$ LWD+ CoD (Ours): $\bf0.696$

---

Notably, there could be several other variants of our strategy for the multi-class setting. For instance, one could also consider only one CFE per datapoint corresponding to the smallest perturbation leading to any changed class, or even considering the counterfactual of the counterfactual, etc. It will be an exciting direction of future research to study the performance tradeoffs among different variants of our strategy for the multi-class setting.

---

On extension to generative LLMs

One possible way to extend our approach to generative LLMs is to define counterfactual explanation as a minimal change to the input (prompt) that flips a chosen property of the generated text while keeping everything else as similar as possible. Formally, given a generative model $f$, a prompt $x$, and a binary attribute function $g(\text{output})$ (e.g., sentiment, toxicity, factuality, topic relevance), a counterfactual explanation prompt $x^{\star}$ would be a small semantic perturbation of $x$ such that the generated output $f(x^{\star})$ flips the value of $g$.

Another possible extension is to rethink counterfactual generation in terms of model sensitivity: identifying minimal semantic perturbations to the input $x$ that cause large shifts in the output distribution or likelihood. For generative models, this corresponds to large changes in the sequence-level probability: $p_\theta(y \mid x) = \prod_{t=1}^{T} p_\theta(y_t \mid x, y_{<t})$, where $y = (y_1, y_2, \ldots, y_T)$ is the generated sequence for an input prompt $x$. In this framing, CFEs could reveal parts of the input space where the model is uncertain, making them especially informative for distillation. While outside the scope of the current paper, this extension would be an interesting future direction, potentially inspiring novel and efficient pathways for data selection for a broad range of tasks, including distillation, supervised finetuning (SFT), post-training, etc.

We will include these discussions in the revised paper.