Layer Importance for Mathematical Reasoning is Forged in Pre-Training and Invariant after Post-Training
This work addresses the problem of understanding how large language models adapt to specialized tasks like math, revealing that core capabilities are forged early and remain invariant, which is incremental but clarifies model interpretability for researchers.
The study investigated whether post-training methods like instruction tuning alter the transformer layer structure for mathematical reasoning, finding that a few critical layers formed during pre-training remain important and stable, with removal reducing math accuracy by up to 80%.
Large language models improve at math after instruction tuning, reinforcement learning, or knowledge distillation. We ask whether these gains come from major changes in the transformer layers or from smaller adjustments that keep the original structure. Using layer-wise ablation on base and trained variants, we find that math reasoning depends on a few critical layers, which stay important across all post-training methods. Removing these layers reduces math accuracy by as much as 80%, whereas factual recall tasks only show relatively smaller drops. This suggests that specialized layers for mathematical tasks form during pre-training and remain stable afterward. As measured by Normalized Mutual Information (NMI), we find that near these critical layers, tokens drift from their original syntactic clusters toward representations aligned with tokens less syntactically related but potentially more useful for downstream task.