Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Distilled Circuits: A Mechanistic Study of Internal Restructuring in Knowledge Distillation

Reilly Haskins, Benjamin Adams

TL;DR

This work investigates how knowledge distillation reshapes internal computations in neural networks by applying mechanistic interpretability to a GPT-2 small teacher and its distilled student, with a replication on a BERT pair to test generality. It introduces an alignment metric that combines component influence with cross-model similarity to quantify functional replication beyond surface output similarity. Case studies on numeral sequence completion reveal that distilled students compress and reallocate computation, relying on fewer, more critical components and displaying brittleness under ablations. The findings highlight a trade-off between parameter efficiency and robustness, offering a principled metric to guide safer deployment of distilled models.

Abstract

Knowledge distillation compresses a larger neural model (teacher) into smaller, faster student models by training them to match teacher outputs. However, the internal computational transformations that occur during this process remain poorly understood. We apply techniques from mechanistic interpretability to analyze how internal circuits, representations, and activation patterns differ between teacher and student. Focusing on GPT2-small and its distilled counterpart DistilGPT2, we find that student models reorganize, compress, and discard teacher components, often resulting in stronger reliance on fewer individual components. To quantify functional alignment beyond output similarity, we introduce an alignment metric based on influence-weighted component similarity, validated across multiple tasks. Our findings reveal that while knowledge distillation preserves broad functional behaviors, it also causes significant shifts in internal computation, with important implications for the robustness and generalization capacity of distilled models.

Distilled Circuits: A Mechanistic Study of Internal Restructuring in Knowledge Distillation

TL;DR

This work investigates how knowledge distillation reshapes internal computations in neural networks by applying mechanistic interpretability to a GPT-2 small teacher and its distilled student, with a replication on a BERT pair to test generality. It introduces an alignment metric that combines component influence with cross-model similarity to quantify functional replication beyond surface output similarity. Case studies on numeral sequence completion reveal that distilled students compress and reallocate computation, relying on fewer, more critical components and displaying brittleness under ablations. The findings highlight a trade-off between parameter efficiency and robustness, offering a principled metric to guide safer deployment of distilled models.

Abstract

Knowledge distillation compresses a larger neural model (teacher) into smaller, faster student models by training them to match teacher outputs. However, the internal computational transformations that occur during this process remain poorly understood. We apply techniques from mechanistic interpretability to analyze how internal circuits, representations, and activation patterns differ between teacher and student. Focusing on GPT2-small and its distilled counterpart DistilGPT2, we find that student models reorganize, compress, and discard teacher components, often resulting in stronger reliance on fewer individual components. To quantify functional alignment beyond output similarity, we introduce an alignment metric based on influence-weighted component similarity, validated across multiple tasks. Our findings reveal that while knowledge distillation preserves broad functional behaviors, it also causes significant shifts in internal computation, with important implications for the robustness and generalization capacity of distilled models.
Paper Structure (54 sections, 4 equations, 21 figures, 2 tables)

This paper contains 54 sections, 4 equations, 21 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Background and related work
  3. Knowledge distillation
  4. Mechanistic interpretability
  5. Case study: Numeral sequence completion
  6. Methodology
  7. Circuit discovery
  8. Component comparison
  9. Findings
  10. Attention analysis
  11. Similar member detection (T-L1-H5).
  12. Numeral detection (T-L4-H4 / S-L2-H4).
  13. Numeral mover (T-L7-H11 / S-L3-H11).
  14. Successor computation (T-L9-H1 / S-L4-H1).
  15. Performance drop per attention head.
  16. ...and 39 more sections

Figures (21)

  • Figure 1: Teacher (left) and student (right) QK attention matrices for the successor head functionality
  • Figure 2: Comparison of PCA projections across different teacher-student MLP pairs.
  • Figure 3: Results of our activation noise-injection experiment on our proposed alignment metric for the numeral sequence completion task, with error bars shown in red.
  • Figure 4: Query-key attention matrix for T-L1-H5 (similar member detection)
  • Figure 5: Teacher (left) and student (right) QK attention matrices for the numeral detection functionality
  • ...and 16 more figures