Layer Importance for Mathematical Reasoning is Forged in Pre-Training and Invariant after Post-Training
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2025
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866915599479734272 |
|---|---|
| author | Nepal, Aadim Shrestha, Safal Shrestha, Anubhav Kim, Minwu Naghiyev, Jalal Shwartz-Ziv, Ravid Ross, Keith |
| author_facet | Nepal, Aadim Shrestha, Safal Shrestha, Anubhav Kim, Minwu Naghiyev, Jalal Shwartz-Ziv, Ravid Ross, Keith |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2506_22638 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Layer Importance for Mathematical Reasoning is Forged in Pre-Training and Invariant after Post-Training Nepal, Aadim Shrestha, Safal Shrestha, Anubhav Kim, Minwu Naghiyev, Jalal Shwartz-Ziv, Ravid Ross, Keith Machine Learning Artificial Intelligence 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. |
| title | Layer Importance for Mathematical Reasoning is Forged in Pre-Training and Invariant after Post-Training |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2506.22638 |