In-Context Learning Distillation for Efficient Few-Shot Fine-Tuning

TL;DR

Context Distillation addresses memory and token-context limitations of in-context learning by transferring context information into a compact 125M student distilled from a 1.3B teacher. The approach blends fine-tuning and ICL with parameter-efficient training (LoRA/BitFit) and a KL–CE loss balance, achieving notably better out-of-domain generalization and substantial memory savings on MNLI. Key findings include approximately a 50% gain in out-of-domain accuracy and up to a 60% reduction in peak memory compared with pattern-based fine-tuning, while enabling on-device, few-shot adaptation without consuming input token budgets. These results suggest context distillation as a practical pathway to scalable, memory-efficient deployment of LLMs for long-context tasks and on-device applications, with promising variants like CD+LoRA for faster training and robust performance.

Abstract

We applied few-shot in-context learning on the OPT-1.3B model for the natural language inference task and employed knowledge distillation to internalize the context information, reducing model parameter from 1.3B to 125M and achieving a size reduction from 2.5GB to 0.25GB. Compared to using in-context learning alone on similarly sized models, this context distillation approach achieved a nearly 50% improvement in out-of-domain accuracy, demonstrating superior knowledge transfer capabilities over prompt-based methods. Furthermore, this approach reduced memory consumption by up to 60% while delivering a 20% improvement in out-of-domain accuracy compared to conventional pattern-based fine-tuning.

Figures (6)

• Figure 1: In-domain and out-of-domain accuracy for different task adaptation techniques applied to OPT-125M model vs. numbers of support examples.
• Figure 2: Peak GPU memory allocated and training time for different task adaptation techniques applied to OPT-125M model vs. numbers of support examples. Note that the baseline and ICL do not involve fine-tuning; their memory usage reflects only the inference process, and the training time is plotted as 0 for reference
• Figure 3: In-domain and out-of-domain accuracy for PBFT OPT-125M model using LoRA and BitFit.
• Figure 4: Peak GPU memory allocated and training time for PBFT OPT-125M model using LoRA and BitFit.
• Figure 5: In-domain and out-of-domain accuracy for OPT-125M student model distilled from OPT-1.3B teacher model with and without LoRA or BitFit.