Not All Layers Need Tuning: Selective Layer Restoration Recovers Diversity, Bowen Zhang★, Meiyi Wang★, and Harold Soh★, Neural Information Processing Systems (NeurIPS)
Links:
Post-training makes large language models more helpful and better at following instructions, but it can also substantially reduce the diversity of their outputs, a phenomenon known as mode collapse. This is undesirable whenever many distinct responses are valid, such as in creative writing, open-ended question answering, and reasoning. This raises a simple question: can we recover the diversity of the pre-trained model without retraining?
Motivated by evidence that LLM layers play distinct functional roles, we hypothesize that mode collapse is localized to specific layers. We propose Selective Layer Restoration (SLR), a training-free method that restores a carefully selected range of layers in a post-trained model back to their pre-trained weights. The resulting hybrid model keeps the same architecture and parameter count, and adds no inference-time cost. To decide which layers to restore, we design a cheap proxy task, Constrained Random Character (CRC), which has an explicit validity set and a natural diversity objective.
Across Llama, Qwen, and Gemma, and across creative writing, open-ended question answering, and multi-step reasoning, SLR consistently recovers substantial output diversity while preserving strong output quality, and achieves higher pass@k across sampling budgets on reasoning. Interestingly, much of the effect of post-training on diversity can be reversed by changing only a small part of the network, suggesting that useful behaviors from pre-training are not erased by post-training, but suppressed in ways that can be selectively recovered.
Resources
You can find our paper here.
Citation
Please consider citing our paper if you build upon our results and ideas.
Bowen Zhang★, Meiyi Wang★, and Harold Soh★, “Not All Layers Need Tuning: Selective Layer Restoration Recovers Diversity”, Neural Information Processing Systems (NeurIPS)
@inproceedings{zhang2026notall, title={Not All Layers Need Tuning: Selective Layer Restoration Recovers Diversity}, author={Zhang, Bowen and Wang, Meiyi and Soh, Harold}, booktitle={Advances in Neural Information Processing Systems (NeurIPS)}, year={2026} }
Contact
If you have questions or comments, please contact Bowen or Harold.
