# Not All Layers Need Tuning: Selective Layer Restoration Recovers Diversity

arXiv: https://arxiv.org/abs/2602.06665 · HTML: https://mlsanity.com/paper/2602.06665/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 34 references that are in the map.
- First posted on arXiv 6 Feb 2026.

## Builds on (most cited first, 30 of 30)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Gemini 2.5: Pushing the Frontier with Advanced Reasoning, Multimodality, Long Context, and Next Generation Agentic Capabilities](https://mlsanity.com/paper/2507.06261/) (2025)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Qwen2 Technical Report](https://mlsanity.com/paper/2407.10671/) (2024)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [Multitask Prompted Training Enables Zero-Shot Task Generalization](https://mlsanity.com/paper/2110.08207/) (2021)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [Qwen3 Embedding: Advancing Text Embedding and Reranking Through Foundation Models](https://mlsanity.com/paper/2506.05176/) (2025)
- [BERT Rediscovers the Classical NLP Pipeline](https://mlsanity.com/paper/1905.05950/) (2019)
- [Language Models are Super Mario: Absorbing Abilities from Homologous Models as a Free Lunch](https://mlsanity.com/paper/2311.03099/) (2023)
- [Scaling Synthetic Data Creation with 1,000,000,000 Personas](https://mlsanity.com/paper/2406.20094/) (2024)
- [Understanding the Effects of RLHF on LLM Generalisation and Diversity](https://mlsanity.com/paper/2310.06452/) (2023)
- [Chain-of-Thought Reasoning Without Prompting](https://mlsanity.com/paper/2402.10200/) (2024)
- [The Effect of Sampling Temperature on Problem Solving in Large Language Models](https://mlsanity.com/paper/2402.05201/) (2024)
- [Turning Up the Heat: Min-p Sampling for Creative and Coherent LLM Outputs](https://mlsanity.com/paper/2407.01082/) (2024)
- [Understanding Catastrophic Forgetting in Language Models via Implicit Inference](https://mlsanity.com/paper/2309.10105/) (2023)
- [Verbalized Sampling: How to Mitigate Mode Collapse and Unlock LLM Diversity](https://mlsanity.com/paper/2510.01171/) (2025)
- [Revisiting Catastrophic Forgetting in Large Language Model Tuning](https://mlsanity.com/paper/2406.04836/) (2024)
- [On the Algorithmic Bias of Aligning Large Language Models with RLHF: Preference Collapse and Matching Regularization](https://mlsanity.com/paper/2405.16455/) (2024)
- [Extract, Define, Canonicalize: An LLM-based Framework for Knowledge Graph Construction](https://mlsanity.com/paper/2404.03868/) (2024)
- [Preserving Diversity in Supervised Fine-Tuning of Large Language Models](https://mlsanity.com/paper/2408.16673/) (2024)
- [Diversity-Aware Policy Optimization for Large Language Model Reasoning](https://mlsanity.com/paper/2505.23433/) (2025)
- [Scaling Data Diversity for Fine-Tuning Language Models in Human Alignment](https://mlsanity.com/paper/2403.11124/) (2024)
- [Layer Swapping for Zero-Shot Cross-Lingual Transfer in Large Language Models](https://mlsanity.com/paper/2410.01335/) (2024)
- [Demystifying the Roles of LLM Layers in Retrieval, Knowledge, and Reasoning](https://mlsanity.com/paper/2510.02091/) (2025)
- [Navigating the Alignment-Calibration Trade-off: A Pareto-Superior Frontier via Model Merging](https://mlsanity.com/paper/2510.17426/) (2025)
- [Top-nσ: Not All Logits Are You Need](https://mlsanity.com/paper/2411.07641/) (2024)

## Built on by (most cited first, 1 of 1)

- [Beyond the Hivemind: Escaping LLM Homogeneity via Meta-Persona Anchoring and Sequential Temperature Scaling](https://mlsanity.com/paper/2608.02618/) (2026)

## Abstract (arXiv)

Post-training improves instruction-following and helpfulness of large language models (LLMs) but often reduces generation diversity, which leads to repetitive outputs in open-ended settings, a phenomenon known as mode collapse. Motivated by evidence that LLM layers play distinct functional roles, we hypothesize that mode collapse can be localized to specific layers and that restoring a carefully chosen range of layers to their pre-trained weights can recover diversity while maintaining high output quality. To validate this hypothesis and decide which layers to restore, we design a proxy task -- Constrained Random Character(CRC) -- with an explicit validity set and a natural diversity objective. Results on CRC reveal a clear diversity-validity trade-off across restoration ranges and identify configurations that increase diversity with minimal quality loss. Based on these findings, we propose Selective Layer Restoration (SLR), a training-free method that restores selected layers in a post-trained model to their pre-trained weights, yielding a hybrid model with the same architecture and parameter count, incurring no additional inference cost. Across three different tasks (creative writing, open-ended question answering, and multi-step reasoning) and three different model families (Llama, Qwen, and Gemma), we find SLR can consistently and substantially improve output diversity while maintaining high output quality.