Layer Importance for Mathematical Reasoning is Forged in Pre-Training and Invariant after Post-Training

2025-06-27

AI Summary

It remains unclear whether enhancing large language models’ mathematical reasoning capabilities requires substantive architectural reconfiguration of Transformer layers. Method: We conduct systematic layer-wise ablation experiments and information-theoretic analyses, comparing pre-trained base models against post-training variants—including instruction tuning, knowledge distillation, and reinforcement learning—on mathematical reasoning benchmarks. Contribution/Results: We discover a task-specific, highly stable layer importance structure for mathematical reasoning, which emerges during pre-training and persists unchanged across all post-training methods. Removing critical layers causes up to an 80% drop in reasoning accuracy, whereas non-reasoning tasks (e.g., factual recall) exhibit no such sensitivity. This is the first empirical evidence demonstrating innate, layer-wise functional specialization for mathematical reasoning—challenging the prevailing assumption that post-training fundamentally reshapes model architecture. Our findings establish a new paradigm for interpretable modeling and efficient inference optimization.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Knowledge Representation and Reasoning: Computational Complexity of Reasoning
• Natural Language Processing: (Large) Language Models

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Web Mining and Content Analysis: Large pretrained models with web data
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Large language models can exhibit improved mathematical reasoning capabilities following post-training with instruction tuning, reinforcement learning, or knowledge distillation. However, it remains unclear whether these improvements are driven by major changes in transformer layers or from minor adjustments that leave the relative layer importance structures of the base model largely unchanged. We investigate this question through systematic layer-wise ablation experiments, examining base, instruction-tuned, knowledge-distilled, and reinforcement learning variants on mathematical reasoning benchmarks. Our findings show that mathematical reasoning gives rise to a specific layer importance structure, and this structure persists across all post-training paradigms. Removal of such layers causes accuracy drops of up to 80%. In contrast, non-mathematical tasks like factual recall exhibit no critical layers. This distinction suggests that mathematical reasoning requires specialized layers that emerge during pre-training, while other non-reasoning tasks do not. From an information-theoretic perspective, we also observe that these critical layers are the same layers where major representational transformation occurs.

Problem

• Examines if post-training changes transformer layer importance for math reasoning
• Identifies critical layers for math tasks that persist after post-training
• Shows non-math tasks lack specialized layers unlike math reasoning

Innovation

• Layer-wise ablation experiments analyze model variants
• Mathematical reasoning requires specialized pre-trained layers
• Critical layers show major representational transformations