Distilling LLM Reasoning into Graph of Concept Predictors

TL;DR

This paper introduces Graph of Concept Predictors (GCP), a reasoning-aware active distillation framework that improves LLM training efficiency, interpretability, and performance by externalizing decision processes as a graph and focusing on critical reasoning nodes.

Contribution

GCP is a novel framework that externalizes LLM reasoning as a graph, enhancing sample efficiency, interpretability, and training stability through targeted sub-module retraining.

Findings

GCP improves performance on NLP classification benchmarks with limited annotations.

GCP yields more interpretable and controllable training dynamics.

GCP outperforms traditional distillation methods in efficiency and accuracy.

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.