Training Task Experts through Retrieval Based Distillation

TL;DR

This paper introduces Retrieval Based Distillation (ReBase), a framework that overcomes data scarcity for domain-specific model fine-tuning by retrieving diverse labeled data from a large pool of online datasets and transforming it with an LLM to produce task-specific (query, answer) data, including Chain-of-Thought reasoning for complex tasks. The method builds a large, column-wise datastore from Hugging Face Datasets, retrieves highly relevant items via instruction and example-guided cosine similarity scoring, and then transforms the retrieved data into usable training data for a task-expert model. Across MNLI, SQuAD, MCoNaLa, and BBH, ReBase achieves notable improvements over baselines such as data synthesis and zero-shot generation, demonstrating enhanced data diversity, coverage, and reasoning ability in the downstream model. The work highlights the practical impact of combining multisource retrieval with structured data transformation to enable cost-effective training of domain-specific models and provides insights into data diversity, CoT distillation, and the tradeoffs of different generation architectures.

Abstract

One of the most reliable ways to create deployable models for specialized tasks is to obtain an adequate amount of high-quality task-specific data. However, for specialized tasks, often such datasets do not exist. Existing methods address this by creating such data from large language models (LLMs) and then distilling such knowledge into smaller models. However, these methods are limited by the quality of the LLMs output, and tend to generate repetitive or incorrect data. In this work, we present Retrieval Based Distillation (ReBase), a method that first retrieves data from rich online sources and then transforms them into domain-specific data. This method greatly enhances data diversity. Moreover, ReBase generates Chain-of-Thought reasoning and distills the reasoning capacity of LLMs. We test our method on 4 benchmarks and results show that our method significantly improves performance by up to 7.8% on SQuAD, 1.37% on MNLI, and 1.94% on BigBench-Hard.

Figures (9)

• Figure 1: Motivation of ReBase. Previous methods either uses manually annotated data or use LLMs to generate synthetic data. This is either too costly or lacks diversity/quality. ReBase retrieves data from existing examples then uses an LLM to create new domain-specific data based on the retrieved content.
• Figure 2: Pipeline of ReBase. First, ReBase iterates over a large number of datasets available on Hugging Face Datasets and encodes each item in this datasets to build a large datastore. Then, ReBase uses the instruction and few-shot examples provided by the new task to retrieve the relevant items from the datastore. Finally, ReBase uses an LLM to generate new data for the target task from the retrieved data.
• Figure 3: Examples of ReBase transformations on BBH. In the data transformation stage, ReBase takes in the original full row of the retrieved data and use the content to generate a new data with the field query and answer. The LLM need to identify the necessary fields in the row. For the BBH task, the transformation contains chain-of-thought reasoning.
• Figure 4: Qualitative Examples on ReBase (Transformed) compared to directly synthesized data (Generated). ReBase outputs data that are more diverse while directly synthesized data tend to be simpler and replicate. In MCoNaLa, we found that ReBase generates samples that contains dynamic programming, counting, mathematical calculations whereas directly synthesized dataset is limited to simpler commands such as printing or simple list operation. In SquAD, we found that ReBase generates samples that contain diverse and harder questions whereas directly synthesized data asks simpler questions.
• Figure 5: Embedding Visualization Result of MNLI and SQuAD. The data generated by ReBase are more widely scattered across the embedding space compared to the synthesized data.