Technical Summary: Selective Layer Restoration Recovers Diversity

Problem Statement

Post-training techniques, such as instruction tuning and Reinforcement Learning from Human Feedback (RLHF), significantly enhance the instruction-following capabilities and helpfulness of Large Language Models (LLMs). However, these methods often induce mode collapse, a phenomenon where the model assigns disproportionately high probability mass to a narrow range of responses, suppressing diverse but valid alternatives. This leads to repetitive outputs in open-ended settings (e.g., creative writing, open-ended QA, and reasoning), limiting the model's utility in applications that value distinct responses.

Existing mitigation strategies fall into three categories, each with limitations:

1. Decoding-based methods (e.g., temperature scaling, min-p sampling) increase stochasticity but cannot recover diversity lost due to parameter changes during post-training.
2. Prompt-based methods (e.g., requesting multiple answers) increase inference costs and are sensitive to prompt wording.
3. Training-based methods (e.g., objective redesign) require expensive retraining and data collection.

The authors seek a method that effectively recovers diversity, is computationally efficient during adaptation, and introduces no additional inference-time cost.

Methodology

The proposed approach, Selective Layer Restoration (SLR), is a training-free weight-space intervention based on two empirical observations:

1. LLM layers exhibit distinct functional roles (e.g., factual associations are often localized to mid-depth layers).
2. Post-training primarily steers which pre-trained behaviors are expressed, often suppressing others (catastrophic forgetting), rather than acquiring entirely new knowledge.

The authors hypothesize that mode collapse is localized to specific layers and that restoring a carefully chosen range of layers in a post-trained model ($M_{post}$) to their pre-trained weights ($M_{pre}$) can recover diversity while maintaining quality.

The SLR Framework

Given a model with $N$ layers, SLR constructs a hybrid model $M_{i:j}$ by replacing a contiguous interval of layers $[i, j]$ in $M_{post}$ with the corresponding layers from $M_{pre}$. The goal is to find the interval $(i, j)$ that maximizes diversity ($D$) subject to a minimum quality threshold ($Q \ge q_{min}$).

The CRC Proxy Task

Directly optimizing the restoration interval on complex downstream tasks is computationally expensive and difficult to evaluate due to the subjectivity of "quality." To address this, the authors introduce Constrained Random Character (CRC) as a proxy task.

• Mechanism: CRC involves generating random characters (digits or letters) within a strictly defined valid set (e.g., "Generate a random integer in [0, 5]").
• Metrics:
  • Quality ($Q_{CRC}$): The probability mass assigned to valid outputs (measuring instruction adherence).
  • Diversity ($D_{CRC}$): The entropy of the distribution over valid outputs (measuring variety).
• Selection Process: The authors perform a grid search over all possible contiguous layer intervals using CRC. They select the interval that maximizes $D_{CRC}$ while satisfying $Q_{CRC} \ge 0.9 \times Q_{CRC}(M_{post})$. This selected interval is then applied to downstream tasks without further tuning.

Key Contributions

1. Selective Layer Restoration (SLR): A training-free method that restores a specific contiguous interval of layers in a post-trained LLM to their pre-trained parameters. It preserves the original architecture and parameter count, incurring no additional inference cost. It is shown to be complementary to decoding and prompt-based diversity methods.
2. Constrained Random Character (CRC): A proxy task with an explicit validity set and natural diversity objective. CRC enables efficient, low-cost exploration of restoration intervals and provides empirical evidence for the layer-localization hypothesis.
3. Empirical Validation: Comprehensive evaluation across three model families (Llama, Qwen, Gemma) and three downstream tasks (creative writing, open-ended QA, multi-step reasoning).

Experimental Results

The authors evaluated SLR against the post-trained baseline and a "Model Soup" baseline (weighted parameter averaging of pre-trained and post-trained models).

• Creative Writing: SLR consistently improved diversity (average +42.5% increase in embedding-based dissimilarity) with minimal quality loss (average -2.1% in LLM-judge scores). In contrast, Model Soup showed smaller diversity gains (+27.3%) and more significant quality degradation (-6.0%).
• Open-Ended QA: SLR substantially increased answer diversity, showing a +112.7% increase in entropy and a +169.5% increase in coverage (fraction of unique correct answers), while maintaining high precision (average -1.8% drop).
• Reasoning (GSM8K): On problems where post-trained models failed under greedy decoding, SLR improved Pass@k across all sampling budgets ($k \in \{1, \dots, 64\}$), indicating a shift in probability mass toward diverse but correct reasoning paths.
• Composability: SLR remained effective when combined with higher temperature decoding ($T=1.5$) and verbalized sampling (prompt-based methods), suggesting that SLR modifies the underlying model distribution while other methods manipulate the sampling process or input.
• Ablation Studies: Restoring layers guided by the CRC proxy yielded superior results compared to restoring arbitrary intervals (e.g., the earliest or latest layers). Restoring arbitrary layers often failed to improve diversity or degraded quality, confirming that the specific interval selection is critical.

Significance and Claims

The paper claims that SLR offers a simple yet effective path to mitigating mode collapse without the computational overhead of retraining or the inference costs of complex prompting. By demonstrating that diversity can be recovered through targeted weight-space interventions, the work supports a modular view of post-trained LLMs:

1. Functionalities and behaviors (including mode collapse) are localized to specific layers rather than being distributed uniformly.
2. Post-training primarily steers the expression of pre-trained behaviors rather than fundamentally altering the model's core knowledge base.

The authors conclude that fine-grained weight-space interventions like SLR provide a promising avenue for controlling the quality-diversity trade-off in LLMs, allowing for the recovery of desirable pre-trained properties (like diversity) while retaining the benefits of post-training (like instruction adherence). The paper remains modest regarding scalability, noting that results are currently limited to ~7–9B parameter models and that future work is needed to determine if these findings hold at larger scales or with more granular interventions (e.g., sub-layer restoration).