Circuit Distillation

Somin Wadhwa; Silvio Amir; Byron C. Wallace

cs.CLSep 29, 2025

Circuit Distillation

Authors: Somin Wadhwa, Silvio Amir, Byron C. Wallace

Organizations: Khoury College of Computer Sciences, Northeastern University, Boston, MA 02115, USA

Abstract

Model distillation typically focuses on behavioral mimicry, where a student model is trained to replicate a teacher's output while treating its internal computations as a black box. In this work we propose an alternative approach: Distilling the underlying computational mechanisms implemented by a teacher model. Specifically, we propose circuit distillation, which introduces an objective to align internal representations between analogous circuit components in teacher and student models. We propose a method to match ``functionally correspondent'' circuit components and introduce a loss reflecting similarities between the representations that these induce. We evaluate circuit distillation on entity tracking and theory of mind (ToM) tasks using models from the Llama3 family. Our results demonstrate that circuit distillation outperforms standard distillation, successfully transferring algorithmic capabilities by adjusting only a small, targeted subset of student model parameters. This work establishes the feasibility of transferring mechanisms, which may in turn allow for efficient distillation of targeted teacher capabilities via interpretable and controllable internal student mechanisms.

Figures & tables

Figure 1: We propose circuit distillation which entails training a small student model to mimic not only the outputs of a larger teacher model (left is the standard distillation setting; Hinton et al. 2015 ), but also its internal circuitry, i.e., how it executes a task of interest (right).

Figure 2: Circuit distillation requires mapping student heads to functionally analogous teacher heads. Here we propose do this by comparing performance drops under ablation (other approaches may also be viable). These correspondences are then used to align model components via the composite loss during distillation.

Figure 3: Illustrations of the evaluated tasks: entity tracking (left), which tests recall of item locations, and theory of mind (right), which tests reasoning about beliefs under partial observability.

			Entity Tracking		Theory of Mind
		Loss	Full Model	Circuit	Full Model	Circuit
Baselines	Base–1B	–	0.68	0.65	0.55	0.54
	Base–3B	–	0.72	0.70	0.58	0.56
	8B-Instruct (Meta)	–	–	–	0.79	0.78
	8B-FT †	CE	0.85	0.81	0.76	0.76
Fully Distilled	Base–1B	CE	0.75	0.66	0.60	0.55

Table 1: Entity tracking (left) and Theory of Mind (ToM; right) results for Llama3 models. We report results for full distillation (all parameters) and circuit distillation (only matched circuit head parameters) under different losses: Cross Entropy (CE), logit distillation (KD; Hinton et al. 2015 ), Centered Kernel Alignment (CKA), CKA with Random Head Assignment, CE + CKA between whole-layer attention outputs at the layers containing circuit heads (Layer CKA; no head mapping), and CE + CKA with ablation-based head alignment loss. † 8B-FT is GOAT-8B for entity tracking and 8B-Alpaca for ToM. Chance accuracy on entity tracking is 0.14. Instruction-tuned models (e.g., Llama3-Instruct) outperform their base counterparts on ToM. On both tasks, the improvements of circuit distillation (1B, 3B) using CE + Align CKA are statistically significantly better than standard full distillation under CE and circuit distillation using CE only (

p

<

0.05

) under McNemar’s test (details in Appendix).

Figure 4: Mechanism transfer on entity tracking ( Llama3-1B student, 8B teacher), measured on held-out prompts. (a) CKA between each student circuit head and its mapped teacher head (left) and between the same student heads and non-mapped teacher circuit heads (control, right), before (Base) and after distillation; large markers show means with 95% confidence intervals. (b) Attention from the final token to the correct object token for the three mapped pairs with the strongest teacher Value Fetcher scores; error bars are standard errors.

Student	VF overlap@20	Spearman $\rho$ (S vs. T)	Mean mapped PP score
Base (pre)	1.00	0.18	$-$ 0.021
CE	0.80	0.27	$-$ 0.026
CE + Align CKA	0.90	0.55	$-$ 0.041

Table 2: Path patching (Value Fetcher stage) re-run on distilled Llama3-1B students. VF overlap@20 : fraction of the 20 most causally important heads at the final position that belong to the student circuit. Spearman

\rho

: rank correlation between the patching scores of mapped student heads and those of their teacher counterparts. Mean mapped PP score : mean patching score of mapped student heads (more negative = more causally important).

		Loss	Full Model	Circuit	Random Circuit	Faithfulness
Baselines	Base-1B	–	0.68	0.65	0.03	0.96
	Base-3B	–	0.72	0.70	0.02	0.97
	Base-8B	–	0.73	0.72	0.01	0.97
	GOAT-1B	CE	0.77	0.69	0.03	0.89
	GOAT-3B	CE	0.83	0.75	0.04	0.90
	GOAT-8B	CE	0.85	0.81	0.02	0.89

Table 3: Extended results on the Entity Tracking task.

		Loss	Full Model	Circuit	Random Circuit	Faithfulness
Baselines	Base-1B	–	0.55	0.54	0.02	0.98
	Base-3B	–	0.58	0.56	0.01	0.97
	Base-8B	–	0.71	0.68	0.02	0.95
	Base-1B-Instruct	–	0.65	0.64	0.02	0.98
	Base-3B-Instruct	–	0.69	0.67	0.01	0.97
	Base-8B-Instruct	–	0.79	0.78	0.02	0.99

Table 4: Extended results on the ToM task.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Path patching isolates which model components store and retrieve entity-attribute bindings.

Figure 6: We compare the Cross-Entropy (CE) loss component for the CircuitDistilled-3B model on the entity tracking task when trained with a standard behavioral objective ( CE Only ) versus our mechanistic distillation method ( CE + Align CKA )

Figure 7: Mean CKA between mapped student and teacher circuit heads on training batches, over training of the Llama3-1B entity tracking student.

Language Model Distillation

Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.

See all 512 papers

What Survives Distillation? A Mechanistic Interpretability Study of GPT-2 and DistilGPT-2
Oct 27, 2026Nabeelah Faumi, Deshan SumanathilakaLanguage Model Distillation Knowledge Distillation
Informatics Institute of Technology, 57, Ramakrishna Road, Colombo 06, Sri Lanka
Your Teacher Can't Help You Here: Combating Supervision Fidelity Decay in On-Policy Distillation
May 29, 2026Yanjiang Liu, Jie Lou, Xinyan Guan +7CoT Distillation Language Model Distillation
University of Chinese Academy of Sciences · 2Chinese Information Processing Laboratory Institute of Software, Chinese Academy of Sciences 2University of Chinese Academy of Sciences, Beijing, China · 3Xiaohongshu
Reference-Based Distillation Detection in LLMs
Jun 19, 2026Rajat Rawat, Sizhe Chen, Akshay Anand +3LLM Auditing Teacher-Student Learning
University of California, Berkeley · University of Southern California · OpenAI

Oct 27, 2026J/K move · Enter open · S save

What Survives Distillation? A Mechanistic Interpretability Study of GPT-2 and DistilGPT-2

Nabeelah Faumi, Deshan Sumanathilaka

Informatics Institute of Technology, 57, Ramakrishna Road, Colombo 06, Sri Lanka

Language Model Distillation Knowledge Distillation Mechanistic Interpretability

Open paper AskPDF arXiv

Knowledge distillation compresses large language models into compact students that preserve aggregate metrics, yet the mechanistic consequences of depth-reduction distillation remain largely uncharacterized. This work applies linear CKA, logit lens analysis, attention pattern decomposition, sparse autoencoder feature statistics, and activation patching to compare GPT-2 and its six-layer student DistilGPT-2 across two canonical circuits: the Indirect Object Identification (IOI) circuit and the Greater-Than (GT) circuit. The results reveal a compression-collapse dichotomy: the IOI circuit suffers complete mechanistic collapse (DistilGPT-2: 35.2% accuracy vs. 94.8%, zero name-mover preservation), while the GT circuit survives with 87.7% task performance recovery through active remapping of the teacher's layer-9 comparison head into student layer~4 (mean Pearson r = 0.945). These findings demonstrate that circuit compressibility governs mechanistic survival under distillation, and that aggregate metrics are insufficient diagnostics for circuit integrity.

Abstract

Language Model Distillation

Explore similar work

What Survives Distillation? A Mechanistic Interpretability Study of GPT-2 and DistilGPT-2

Your Teacher Can't Help You Here: Combating Supervision Fidelity Decay in On-Policy Distillation

Reference-Based Distillation Detection in LLMs