TL;DR
Standard model distillation treats a teacher network as a black box, forcing a smaller student to mimic only final outputs. Circuit distillation instead opens the box, directly transferring the teacher's internal computational mechanisms ("circuits") by aligning functionally corresponding attention heads via Centered Kernel Alignment (CKA). By updating only a small, targeted subset of student parameters (11–15% of heads), this method outperforms full-model behavioral distillation and converges significantly faster.
What Problem Does This Paper Address?
Knowledge distillation (KD) usually relies on behavioral mimicry: training a smaller student model to reproduce the teacher’s output logits. However, this creates a representational bottleneck. The student sees what the teacher predicted, but receives no direct guidance on how to compute it, forcing it to reverse-engineer algorithms from scratch.
Meanwhile, mechanistic interpretability has revealed that neural networks implement specialized algorithms via sparse subgraphs of attention heads and MLP layers ("circuits"). The authors ask: Can we distill internal algorithms directly into a student, rather than just matching outputs?
Methods
To transfer internal mechanisms between models of different sizes (e.g., LLaMA-3-8B down to 1B/3B), the authors tackle two key challenges:
-
Functional Component Mapping: Because dimensions and layers do not line up 1:1 across architectures, the authors match teacher and student heads based on ablation impact similarity. They measure how much task accuracy drops when a specific head is mean-ablated: $$\Delta P_s(h_s) = P_{s,\text{base}} - P_{s,\text{abl}}(h_s)$$ $$\Delta P_t(h_t) = P_{s,\text{base}} - P_{t,\text{abl}}(h_t)$$ Student heads are mapped to the teacher heads that produce the most similar relative drop in baseline capability.
-
Representational Alignment Loss: Because raw activations cannot be compared directly due to arbitrary basis rotations, differing dimensions, and isotropic scaling, the authors use linear Centered Kernel Alignment (CKA). For each paired head $c$, the loss minimizes: $$\mathcal{L}_{\text{CKA}}(K_s^{(c)}, K_t^{(c)}) = 1 - \text{CKA}(K_s^{(c)}, K_t^{(c)})$$
- Targeted Parameter Updates: The total training objective combines standard task cross-entropy with the sum of CKA alignment losses over the paired heads. Crucially, only the parameters of the identified student circuit heads are updated, leaving the rest of the model frozen.
Datasets
The authors evaluate on two tasks with previously documented internal circuits:
- Entity Tracking: Based on Kim & Schuster (2023) and Prakash et al. (2024), where models track object-container bindings (e.g., "The apple is in box F... Box F contains the [apple]"). Phrases are structured to avoid shallow n-gram heuristics. The 8B teacher is fine-tuned on the GOAT arithmetic dataset (Liu & Low, 2023), which strengthens the underlying tracking circuit.
- Theory of Mind (ToM): Evaluated on the CausalToM benchmark (Prakash et al., 2025). Models read stories involving two agents manipulating objects with partial visibility and must infer an agent's true belief rather than ground-truth reality. The 8B teacher is fine-tuned on Stanford Alpaca (Taori et al., 2023).
Key Insights and Ideas
- Principled alignment is mandatory: Aligning randomly paired heads (
CE + Rand CKA) drastically degrades performance (dropping below the untrained baseline), proving that the CKA regularizer is not just a generic smoother—it strictly depends on correct functional correspondence. - Updating 11–15% of parameters beats updating 100%: Mechanistically guiding only the relevant attention heads yields higher task accuracy than traditional full-parameter distillation.
- Denser supervisory signal: While standard cross-entropy provides a sparse training signal on the final output token, intermediate circuit alignment provides dense gradients directly into computational subgraphs, cutting convergence time roughly in half.
Key Results
Experiments used LLaMA-3-8B as the teacher, distilling into LLaMA-3-1B and 3B students:
- Entity Tracking (3B Student):
- Base-3B student: 0.72 accuracy
- Standard Full Distillation (CE only, all parameters): 0.78
- Circuit Distillation (CE only, ~11% heads): 0.77
- Circuit Distillation (CE + Align CKA): 0.82 (teacher achieves 0.85).
- Random Head Control: Drops to 0.63.
- Theory of Mind (3B Student):
- Base-3B student: 0.58 accuracy
- Standard Full Distillation (CE only): 0.63
- Circuit Distillation (CE + Align CKA): 0.65 (teacher achieves 0.76).
- Random Head Control: Drops to 0.49.
- Training Efficiency: For the 3B student on entity tracking at step 1,500, the cross-entropy loss under circuit distillation was 0.368, compared to 0.689 for standard CE training alone.
Limitations and Future Work
- Prerequisite Circuit Discovery: The framework requires knowing the teacher's task circuit in advance. Automated circuit discovery is still an open and difficult research area.
- Heuristic Head Matching: Mapping heads via 1D ablation impact similarity is a coarse heuristic that may miss complex multi-head interactions or polysemantic features.
- Scale and Task Scope: Experiments were restricted to two cognitive tasks and relatively small open models (up to 8B) due to compute budgets; scaling to 70B+ teachers and open-domain generative tasks remains untested.