Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Distilling LLMs' Decomposition Abilities into Compact Language Models

Denis Tarasov, Kumar Shridhar

TL;DR

This work tackles the scalability challenge of complex reasoning by distilling LLMs' sub-question decomposition abilities into compact language models via offline RL. It introduces the GSM8K-AI-SubQ dataset, generated from GSM8K with AI-driven sub-questions and feedback, to train and evaluate smaller models against a ChatGPT baseline. Across multiple backbones (e.g., GPT-2 small/medium, DistilGPT, Mistral, LLaMA), the study finds a substantial gap to ChatGPT and mixed results for offline RL methods, with Filtered BC often outperforming vanilla BC but offline RL not consistently beating supervised baselines. The dataset and baselines provide a foundation for future offline NLP reasoning research, highlighting the need for more effective offline RL techniques and larger, diverse benchmarks for sub-questioning and decomposition tasks.

Abstract

Large Language Models (LLMs) have demonstrated proficiency in their reasoning abilities, yet their large size presents scalability challenges and limits any further customization. In contrast, compact models offer customized training but often fall short in solving complex reasoning tasks. This study focuses on distilling the LLMs' decomposition skills into compact models using offline reinforcement learning. We leverage the advancements in the LLM`s capabilities to provide feedback and generate a specialized task-specific dataset for training compact models. The development of an AI-generated dataset and the establishment of baselines constitute the primary contributions of our work, underscoring the potential of compact models in replicating complex problem-solving skills.

Distilling LLMs' Decomposition Abilities into Compact Language Models

TL;DR

This work tackles the scalability challenge of complex reasoning by distilling LLMs' sub-question decomposition abilities into compact language models via offline RL. It introduces the GSM8K-AI-SubQ dataset, generated from GSM8K with AI-driven sub-questions and feedback, to train and evaluate smaller models against a ChatGPT baseline. Across multiple backbones (e.g., GPT-2 small/medium, DistilGPT, Mistral, LLaMA), the study finds a substantial gap to ChatGPT and mixed results for offline RL methods, with Filtered BC often outperforming vanilla BC but offline RL not consistently beating supervised baselines. The dataset and baselines provide a foundation for future offline NLP reasoning research, highlighting the need for more effective offline RL techniques and larger, diverse benchmarks for sub-questioning and decomposition tasks.

Abstract

Large Language Models (LLMs) have demonstrated proficiency in their reasoning abilities, yet their large size presents scalability challenges and limits any further customization. In contrast, compact models offer customized training but often fall short in solving complex reasoning tasks. This study focuses on distilling the LLMs' decomposition skills into compact models using offline reinforcement learning. We leverage the advancements in the LLM`s capabilities to provide feedback and generate a specialized task-specific dataset for training compact models. The development of an AI-generated dataset and the establishment of baselines constitute the primary contributions of our work, underscoring the potential of compact models in replicating complex problem-solving skills.
Paper Structure (22 sections, 2 figures, 12 tables)

This paper contains 22 sections, 2 figures, 12 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. LM Distillation
  4. LLMs Reasoning
  5. Sub-Questioning
  6. Preliminaries
  7. Language Modeling
  8. Offline RL task formulation
  9. GSM8K-AI-SubQ Dataset
  10. Dataset Collection Process
  11. Dataset Analysis
  12. Baselines
  13. Evaluation Protocol
  14. Applied Approaches
  15. Experimental Results
  16. ...and 7 more sections

Figures (2)

  • Figure 1: Number of questions distributions in train set. (a) Distribution in the entire train set with mean 4.0 $\pm$ 1.5 and median 4, (b, c, d) Comparisons of distributions in number of questions between problems with 3 vs. 0, 1, 2 sets of sub-questions that lead to the correct final answer.
  • Figure 2: (a) distribution of the usefulness feedback over individual sub-questions, (b) distribution of the usefulness feedback averaged over sets of sub-questions, (c) confusion matrix between the presence of negative feedback in the set of sub-questions and correctness of the final answer based on corresponding sub-questions.