Distilling LLM Reasoning into Graph of Concept Predictors

Ziyang Yu; Liang Zhao

cs.AIFeb 3, 2026

Distilling LLM Reasoning into Graph of Concept Predictors

Authors: Ziyang Yu, Liang Zhao

Organizations: Department of Computer Science Emory University Atlanta, USA

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 small 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 in which the teacher's reasoning is elicited as a directed acyclic graph of intermediate concepts and mirrored in the student. GCP enhances sample efficiency through a graph-aware acquisition strategy that weights per-concept uncertainty, gradient diversity, and coverage by node centrality. 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.

Figures & tables

Panel of Figure 1 — Figure 1: Mirroring a teacher’s reasoning DAG into a Graph of Concept Predictors, shown on a Yelp review. Left: the LLM teacher judges the food ( $R_{1}$ ) and the service ( $R_{2}$ ), combines them into an overall impression ( $R_{3}$ ), and labels the review negative. Right: GCP gives each conclusion its own concept predictor $f_{j}$ and wires the predictors with the same edges; every $f_{j}$ also reads $x$ . The teacher’s intermediate answers serve as labels for the matching concepts $c_{j}$ .

Figure 2: Overview of the GCP loop: the student (lower left), graph-aware acquisition (top), and sub-module retraining by counterfactual reruns (lower right).

Panel of Figure 3 — Figure 3: A correct but under-confident prediction, on the review from Figure 1 . The student’s service predictor $f_{2}$ misses the complaint about waiting (red), and the error passes to $c_{3}$ and $\hat{y}$ (dashed). The student still predicts negative , but with lower confidence and higher loss.

Classic uncertainty- and diversity-based baselines
Method	MNLI	GoEmotions	SemEval	MIMIC-III	AG News	Amazon	IMDB	Yelp
Random	51.21	23.06	30.33	75.21	82.87	86.60	79.21	89.19
Least Confidence	57.42	25.01	35.18	80.32	83.05	88.94	82.37	91.84
Entropy Sampling	59.30	25.98	37.46	81.51	83.27	89.63	83.41	92.56
Disagreement (Vote Entropy)	58.91	25.44	36.87	81.03	83.11	89.21	83.02	92.19
Representation- and gradient-based strong baselines

Table 1: Baseline comparison under a fixed 20% token budget. All methods are evaluated in a pool-based active learning setting and spend the same number of teacher tokens, so the label-only baselines annotate more texts than GCP. Best results are in bold .

Figure 4: Performance curves of different sample selection methods for active learning. The y-axis denotes classification accuracy, and the x-axis denotes the teacher-token budget as a percentage of the tokens needed to label the full training set with labels only. At the same budget, GCP annotates fewer texts than the baselines because each of its queries also returns concept values.

Architecture	MNLI	GoEmotions	SemEval	MIMIC-III	AG News	Amazon	IMDB	Yelp
GCP	68.41	31.92	47.83	89.27	87.96	93.42	91.18	95.63
Flat CBM	65.73	30.84	44.26	86.12	87.21	92.67	89.94	94.71
End-to-End MLP	63.18	29.97	41.83	84.03	86.74	92.11	88.92	94.02

Table 2: Architecture ablation. Comparison of GCP with a flat CBM (no dependencies) and an end-to-end MLP at 50% annotation budget.

Setting	MNLI	GoEmotions	SemEval	MIMIC-III	AG News	Amazon	IMDB	Yelp
GCP (Full)	68.41	31.92	47.83	89.27	87.96	93.42	91.18	95.63
w/o SWU	66.82	30.74	45.61	87.43	87.21	92.81	90.32	94.91
w/o Grad	66.41	30.28	45.02	87.11	87.08	92.64	89.97	94.74
w/o Coverage	67.02	30.91	46.13	87.86	87.35	92.93	90.51	95.02
w/o Intersection (Union + tie-break)	65.94	30.12	44.58	86.73	86.89	92.41	89.63	94.38
w/o Sub-module Retraining	64.87	29.96	43.21	85.02	86.71	92.18	89.11	94.09

Table 3: Ablation study on major components of GCP . Accuracy (%) at 50% annotation budget across eight datasets. Each row disables one component; the full model includes all components.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	# Instances	# Classes	Avg. Length	Task Type
AG News	127,600	4	38	Topic
Amazon Review Polarity	4,000,000	2	90	Sentiment
IMDB	50,000	2	230	Sentiment
Yelp Review Polarity	598,000	2	150	Sentiment
MNLI	433,000	3	33	NLI
GoEmotions	58,000	27	30	Emotion

Appendix

Table 4: Summary of dataset statistics. # Instances is the size of the full dataset, with all splits combined.

Model	$\widehat{\mu}_{\min}$	median $\widehat{\mu}_{v}$	$v^{\star}$
MNLI ( $\|V\|=16$ )
MLP	0.63	1.9	negation scope
flat CBM	1.10	3.2	negation scope
GCP	1.80	4.6	negation scope
MIMIC-III ( $\|V\|=18$ )
MLP	0.51	1.5	lab abnormality

Appendix

Table 5: Estimated node-level PL constants (all values

\times 10^{-2}

\widehat{\mu}_{\min}=\min_{v}\widehat{\mu}_{v}

is the quantity that controls the rate in Theorem 2 under uniform

\lambda_{v}

;

v^{\star}

is the node attaining it.

Figure 5: Ablation study of targeted sub-module retraining. The y-axis denotes classification accuracy and the x-axis denotes teacher-token usage (millions; thousands for SemEval). Budgets are set per dataset and span roughly 8–50% of the tokens needed to label its full training set with labels only (cf. the 5–50% range of Figure 4 ).

Figure 6: Compute cost versus annotation scale on a logarithmic axis. Comparison between the LLaMA-3-70B teacher and GCP.

Method	MNLI	SemEval	MIMIC-III	Yelp
Entropy	38	21	64	73
CoreSet	62	34	102	121
BADGE	94	51	158	187
CAL	112	67	174	213
GCP	137	82	211	248

Appendix

Table 6: Per-round runtime (seconds) at 20% budget.

Topic / sentiment
Dataset	$\|C\|$	$\|V\|$	$\|E_{\mathrm{cand}}\|$	$\|E\|$
AG News	32	6	11	7
Amazon	35	8	15	9
IMDB	31	5	9	6
Yelp	33	7	13	9
Multi-factor reasoning

Appendix

Table 7: Concept graph statistics per dataset.

|C|

: candidate concepts elicited in Stage 1;

|V|

: nodes retained after pruning;

|E_{\mathrm{cand}}|

: teacher-proposed dependencies whose endpoints both survive pruning, i.e. those actually submitted to the CMI test;

|E|

: edges retained after that test. Node count and edge density both track reasoning depth rather than input length.

Dataset	Example concepts
AG News	geopolitical entity, market terminology, scientific jargon
Amazon	product defect mention, comparative phrasing, sentiment polarity
IMDB	cast praise, plot critique, recommendation cue
Yelp	service quality, food descriptor, price sentiment
MNLI	lexical overlap, negation scope, temporal mismatch
GoEmotions	intensifier, second-person address, gratitude cue

Appendix

Table 8: Representative concepts induced for each dataset.

Explore similar work

Distribution Corrected Offline Data Distillation for Large Language Models
May 13, 2026Yumeng Zhang, Zhengbang Yang, Yevin Nikhel Goonatilake +1Probe-Logit Distillation Mathematical Reasoning Benchmarks
George Mason University, Fairfax, VA, USA
Structural Rationale Distillation via Reasoning Space Compression
May 8, 2026Jialin Yang, Jiankun Wang, Jiajun Wu +3Reasoning Paths Dataset Distillation
University of Calgary · Calgary, Canada · University of Michigan +1
Distilled Reinforcement Learning for LLM Post-training
Jul 19, 2026Chen Wang, Zhaochun Li, Jionghao Bai +4Post-Training Offline Reinforcement Learning
College of Elite Engineers, Nankai University · Zhongguancun Academy · Beijing Institute of Technology +4

May 13, 2026, cs.CLJ/K move · Enter open · S save

Distribution Corrected Offline Data Distillation for Large Language Models

Yumeng Zhang, Zhengbang Yang, Yevin Nikhel Goonatilake, Zhuangdi Zhu

George Mason University, Fairfax, VA, USA

Probe-Logit Distillation Mathematical Reasoning Benchmarks Reasoning Traces Traces

Open paper AskPDF arXiv

Distilling reasoning traces from strong large language models into smaller ones is a promising route to improve intelligence in resource-constrained settings. Existing approaches face a fundamental trade-off: offline distillation from teacher-generated traces provides high-quality, sample-efficient supervision but suffers from distributional drift: during training, the student model conditions on teacher-generated prefixes, whereas during inference the student autoregresses on self-generated prefixes, leading to compounding errors over long reasoning trajectories. Meanwhile, on-policy or self-distillation methods better match the student's inference-time distribution, but require costly online sampling and often produce low-quality traces in early training. We propose a principled offline reasoning distillation framework that preserves the efficiency and supervision quality of offline teacher-generated data while correcting teacher-student distribution drift. It adaptively emphasizes teacher supervision that is better aligned with the student's on-policy distribution. Evaluations on mathematical reasoning benchmarks of GSM8K, MATH, MATH500, and harder held-out competition-style tasks, including AMC, AIME, and OlympiadBench, show that our method improves reasoning accuracy over prior offline distillation algorithms and yields more stable reasoning traces while preserving instruction-following capabilities. Our work shows that lightweight, distribution-correction-aware training can substantially strengthen offline reasoning distillation without online rollouts.

Abstract

Explore similar work

Distribution Corrected Offline Data Distillation for Large Language Models

Structural Rationale Distillation via Reasoning Space Compression

Distilled Reinforcement Learning for LLM Post-training