Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Distilling LLM Reasoning into Graph of Concept Predictors

Ziyang Yu, Liang Zhao

TL;DR

This work tackles the high cost of deploying LLMs for discriminative tasks by distilling their reasoning into a Graph of Concept Predictors (GCP), a DAG-based framework that extracts and transfers intermediate reasoning signals to a compact student. Key innovations include a graph-aware acquisition strategy that targets uncertainty, gradient diversity, and representativeness across reasoning nodes, and a targeted sub-module retraining procedure that attributes loss to specific concept predictors and updates only the most influential modules. Theoretical results claim GCP is strictly more expressive than flat CBMs/MLPs, enjoys faster linear convergence, and allows optimal top-K retraining with manageable time complexity. Empirically, GCP achieves stronger accuracy under limited annotation budgets across eight NLP benchmarks and substantially lowers compute relative to annotating with the full LLM, enabling scalable and interpretable LLM-assisted deployment. The approach offers practical benefits for domains requiring efficient, interpretable reasoning transfer and diagnostics, while acknowledging risks of propagating teacher biases through the concept graph.

Abstract

Deploying Large Language Models (LLMs) for discriminative workloads is often limited by inference latency, compute, and API costs at scale. Active distillation reduces these costs by querying an LLM oracle to train compact discriminative students, but most pipelines distill only final labels, discarding intermediate reasoning signals and offering limited diagnostics of what reasoning is missing and where errors arise. We propose Graph of Concept Predictors (GCP), a reasoning-aware active distillation framework that externalizes the teacher's decision process as a directed acyclic graph and mirrors it with modular concept predictors in the student. GCP enhances sample efficiency through a graph-aware acquisition strategy that targets uncertainty and disagreement at critical reasoning nodes. Additionally, it improves training stability and efficiency by performing targeted sub-module retraining, which attributes downstream loss to specific concept predictors and updates only the most influential modules. Experiments on eight NLP classification benchmarks demonstrate that GCP enhances performance under limited annotation budgets while yielding more interpretable and controllable training dynamics. Code is available at: https://github.com/Ziyang-Yu/GCP.

Distilling LLM Reasoning into Graph of Concept Predictors

TL;DR

This work tackles the high cost of deploying LLMs for discriminative tasks by distilling their reasoning into a Graph of Concept Predictors (GCP), a DAG-based framework that extracts and transfers intermediate reasoning signals to a compact student. Key innovations include a graph-aware acquisition strategy that targets uncertainty, gradient diversity, and representativeness across reasoning nodes, and a targeted sub-module retraining procedure that attributes loss to specific concept predictors and updates only the most influential modules. Theoretical results claim GCP is strictly more expressive than flat CBMs/MLPs, enjoys faster linear convergence, and allows optimal top-K retraining with manageable time complexity. Empirically, GCP achieves stronger accuracy under limited annotation budgets across eight NLP benchmarks and substantially lowers compute relative to annotating with the full LLM, enabling scalable and interpretable LLM-assisted deployment. The approach offers practical benefits for domains requiring efficient, interpretable reasoning transfer and diagnostics, while acknowledging risks of propagating teacher biases through the concept graph.

Abstract

Deploying Large Language Models (LLMs) for discriminative workloads is often limited by inference latency, compute, and API costs at scale. Active distillation reduces these costs by querying an LLM oracle to train compact discriminative students, but most pipelines distill only final labels, discarding intermediate reasoning signals and offering limited diagnostics of what reasoning is missing and where errors arise. We propose Graph of Concept Predictors (GCP), a reasoning-aware active distillation framework that externalizes the teacher's decision process as a directed acyclic graph and mirrors it with modular concept predictors in the student. GCP enhances sample efficiency through a graph-aware acquisition strategy that targets uncertainty and disagreement at critical reasoning nodes. Additionally, it improves training stability and efficiency by performing targeted sub-module retraining, which attributes downstream loss to specific concept predictors and updates only the most influential modules. Experiments on eight NLP classification benchmarks demonstrate that GCP enhances performance under limited annotation budgets while yielding more interpretable and controllable training dynamics. Code is available at: https://github.com/Ziyang-Yu/GCP.
Paper Structure (26 sections, 4 theorems, 25 equations, 6 figures, 4 tables, 1 algorithm)

This paper contains 26 sections, 4 theorems, 25 equations, 6 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Concept-Based Reasoning
  4. Large Language Models for Annotation
  5. Active Learning with LLM as Annotator
  6. Proposed Method
  7. GCP Framework
  8. Graph of Concept Predictors
  9. Graph-Aware Acquisition
  10. Sub-module Retraining
  11. Theoretical Analysis
  12. Experiment
  13. Experiment Settings
  14. Datasets Details
  15. Evaluation Metric
  16. ...and 11 more sections

Key Result

Theorem 3.1

Let $G=(V,E)$ be a directed acyclic graph (DAG) over $V=\{c_1,\ldots,c_K,y\}$. Assume the data-generating distribution admits a DAG factorization $p^\star(y,\mathbf{c}\mid x) = \prod_{j=1}^K p^\star\!(c_j \mid x,\mathbf{c}_{\mathrm{pa}^\star(c_j)}) \cdot p^\star\!(y \mid x,\mathbf{c}_{\mathrm{pa}^\s ,where $\mathcal{P}_{\mathrm{GCP}}$ consist of all models that factorize according to an arbitrary

Figures (6)

  • Figure 1: Mirror mapping from a teacher’s reasoning DAG to a probabilistic Graph of Concept Predictors for classifying whether a paper belongs to Hyperbolic Reasoning in LLMs. Left: logical dependencies among intermediate conclusions. Right: corresponding concept variables and latent features with conditional links.
  • Figure 2: Overview of our proposed GCP framework (loop view). The left below part shows concept-level reasoning, where unlabeled inputs are processed by a learned Graph of Concept Predictors that explicitly models dependencies among concepts. The top part illustrates graph-aware acquisition, in which structure-weighted uncertainty, topology-aware gradient diversity, and graph-aware representativeness are jointly evaluated and combined through consensus selection to identify informative samples for LLM annotation. The right below part depicts sub-module retraining, where counterfactual reruns of the GCP attribute errors to individual concept predictors, enabling targeted retraining of high-impact sub-modules for more efficient learning.
  • Figure 3: Mirror-structured reasoning and concept graphs for classifying whether a paper belongs to RAG for EHR Question Answering with grounded clinical evidence. Both predict the same label, but weaker internal evidence on the right yields lower confidence and higher loss.
  • Figure 4: Performance curves of different sample selection methods for active learning. The y-axis denotes the accuracy for the classification task, and the x-axis represents the percentage of samples annotated by the LLM for small model training. In this case, 100% denotes that all samples from the training set have been annotated.
  • Figure 5: Compute cost versus annotation scale on a logarithmic axis. We compare the computational cost between teacher model (Llama-3-70B) versus GCP.
  • ...and 1 more figures

Theorems & Definitions (4)

  • Theorem 3.1: Strict performance advantage of GCP
  • Theorem 3.2: Faster linear convergence of GCP
  • Theorem 3.3: Optimality of Top-$K$ Selection
  • Theorem 3.4