Improving In-Context Learning with Reasoning Distillation

2025-04-14

AI Summary

Language models rely on semantic priors for in-context learning but struggle to infer implicit input-output rules from few-shot examples in inductive reasoning tasks. To address this, we propose Reasoning Distillation (ReDis), the first framework that jointly integrates synthetic data augmentation, logical consistency filtering, supervised fine-tuning, and feature-space alignment—explicitly enhancing models’ understanding and generalization of inductive rules. Built upon foundation models such as LLaMA-3, ReDis achieves up to a 66.6% absolute improvement over GPT-4o on benchmarks including 1D-ARC, ACRE, and MiniSCAN, substantially outperforming its zero-shot and few-shot capabilities. All code, datasets, and trained models are publicly released to foster reproducibility and further research.

Technology Category

• Knowledge Representation and Reasoning: Common-Sense Reasoning
• Reasoning under Uncertainty: Other Foundations of Reasoning under Uncertainty
• Cognitive Modeling & Cognitive Systems: Conceptual Inference and Reasoning

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Language models rely on semantic priors to perform in-context learning, which leads to poor performance on tasks involving inductive reasoning. Instruction-tuning methods based on imitation learning can superficially enhance the in-context learning performance of language models, but they often fail to improve the model's understanding of the underlying rules that connect inputs and outputs in few-shot demonstrations. We propose ReDis, a reasoning distillation technique designed to improve the inductive reasoning capabilities of language models. Through a careful combination of data augmentation, filtering, supervised fine-tuning, and alignment, ReDis achieves significant performance improvements across a diverse range of tasks, including 1D-ARC, List Function, ACRE, and MiniSCAN. Experiments on three language model backbones show that ReDis outperforms equivalent few-shot prompting baselines across all tasks and even surpasses the teacher model, GPT-4o, in some cases. ReDis, based on the LLaMA-3 backbone, achieves relative improvements of 23.2%, 2.8%, and 66.6% over GPT-4o on 1D-ARC, ACRE, and MiniSCAN, respectively, within a similar hypothesis search space. The code, dataset, and model checkpoints will be made available at https://github.com/NafisSadeq/reasoning-distillation.git.

Problem

• Enhances inductive reasoning in language models
• Improves understanding of input-output rules
• Boosts performance on diverse reasoning tasks

Innovation

• Reasoning distillation enhances inductive reasoning
• Combines data augmentation, filtering, fine-tuning, alignment
• Outperforms GPT-4o on multiple reasoning tasks