Not All Layers Need Tuning: Selective Layer Restoration Recovers Diversity

2026-02-06

AI Summary

This work addresses the mode collapse commonly observed in large language models after post-training, which severely degrades output diversity in open-ended generation. The authors propose Selective Layer Restoration (SLR), a method that reveals— for the first time—that mode collapse is localized to specific network layers. By reverting only these layers to their pre-trained weights, SLR effectively balances diversity and generation quality without incurring additional inference costs. To guide the selection of optimal layers for restoration, the study introduces a proxy task based on Constrained Random Character (CRC) generation to evaluate the trade-off between diversity and effectiveness. Experiments across mainstream models—including Llama, Qwen, and Gemma—demonstrate that SLR consistently enhances output diversity in creative writing, open-ended question answering, and multi-step reasoning tasks while preserving high-quality generation.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work

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.

Problem

• mode collapse
• generation diversity
• post-training
• large language models
• repetitive outputs

Innovation

• Selective Layer Restoration
• mode collapse
• diversity recovery
• post-training
• constrained random character