Computation Alignment for Chain-of-Thought Distillation

Chain-of-thought (CoT) distillation transfers reasoning capabilities from large teacher models to smaller student models by training on teacher-generated rationales. However, existing methods primarily align observable outputs or intermediate representations, leaving a fundamental question unaddressed: whether the student's computation actually follows the reasoning process prescribed by the teacher's CoT. We identify this overlooked problem as computation alignment and formalize it as the overlap between computational circuits underlying intermediate reasoning steps and final-answer generation, recovered via multi-step circuit tracing. Directly optimizing circuit overlap during distillation, however, is computationally prohibitive. To address this, we propose Computation-Aligned Distillation (CAD), which exploits the sensitivity discrepancy between reasoning and answer circuits to construct a lightweight proxy. Across diverse settings, CAD consistently improves distillation performance over corresponding baselines without it.