Not All Layers Need Tuning: Selective Layer Restoration Recovers Diversity

TL;DR

This paper introduces Selective Layer Restoration (SLR), a post-training method that selectively restores certain layers of large language models to their pre-trained weights, improving output diversity without sacrificing quality.

Contribution

The paper proposes a novel, training-free approach called SLR that identifies and restores specific layers to recover diversity in LLM outputs, addressing mode collapse.

Findings

SLR increases diversity across multiple tasks and models.

Restoring certain layers balances diversity and quality effectively.

SLR incurs no additional inference cost.

Abstract

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.