Source-linked AI summary
FOCUS & RePAIR: Mitigating Text Degeneration via Token-Level Guidance for Pruned Large Language Models
Junyoung Lee, Sehyeon Park, Shinhyoung Jang, Seonha Ryu, Hojeong Kim, Hyunsei Lee, Il Hong Suh, Yeseong Kim
TL;DR
Pruning can amplify repetition-loop degeneration, exposing a gap in evaluations centered on perplexity and task performance. The paper models degeneration through loop entry and persistence, then introduces FOCUS and RePAIR for post-pruning fine-tuning; both consistently reduce repetition and improve generation quality. The main scope boundary is that perceptual quality is assessed with an additional LLM-as-a-judge evaluation because automatic metrics may not fully capture it.
Problem
Pruning can increase repetition-loop degeneration even when perplexity and task performance remain largely intact, motivating analysis beyond knowledge-preservation metrics.
Method
The paper analyzes decoding as dynamics over recurrent contexts and proposes FOCUS distillation reweighting plus RePAIR onset-centered pairwise guidance for post-pruning fine-tuning.
Results
Across open-ended continuation and instruction-based generation, FOCUS and RePAIR consistently reduce repetition and improve generation quality with only a small perplexity increase.
Takeaways & Limitations
Token-level control of leakage and plausible alternatives provides a practical way to mitigate repetition after pruning while preserving broad generation quality improvements.
Takeaways & Limitations
Automatic metrics may not fully capture perceptual quality, so the evaluation adds an LLM-as-a-judge assessment.
Abstract
from arXiv · showhide
Pruning is a practical approach to compress large language models (LLMs), but it can amplify text degeneration, especially repetition loops, even when perplexity and task accuracy remain largely unchanged. In this work, we present a token-level analysis of this failure mode by viewing decoding as a dynamical process that enters and persists in a small set of recurrent contexts. Our analysis decomposes degeneration into loop entry risk and loop persistence, and shows that persistence is controlled by the escape mass assigned to plausible alternatives within the token sampling set. Motivated by these findings, we propose two token-level guidance objectives for post-pruning fine-tuning. FOCUS reweights distillation toward high-confidence teacher regions to suppress leakage, while RePAIR uses onset-centered positive/negative continuation pairs with a margin loss to promote plausible alternatives and prevent early commitment to repetition loops. Experiments on open-ended continuation and instruction-based generation show that both methods consistently reduce repetition and improve generation quality.
1. Introduction
The paper argues that pruning can worsen repetition-loop degeneration despite conventional metric recovery, and develops token-level analysis and guidance methods to mitigate it.
- Motivation: Pruning can intensify repetitive generation even when fine-tuning partially restores knowledge.The paper frames repetition as a side effect that standard perplexity and accuracy evaluations may miss.
- Token-level analysis: The analysis separates degeneration into loop entry risk and loop persistence, with persistence compounding across consecutive in-loop transitions.A small number of high-probability token changes at loop onset can steer generation away from repetition.
- Proposed methods: FOCUS emphasizes high-confidence teacher regions during distillation to suppress probability leakage into low-support areas.Its objective addresses token-probability drift under the capacity constraints of pruned models.
- Results: Across open-ended continuation and instruction-based generation, both methods reduce degeneration and improve generation quality with only a small perplexity increase.Reported quality measures include MAUVE, CREP, EAD1, and BERTScore.
- Proposed methods: RePAIR uses onset-centered positive and negative continuation pairs to promote plausible alternatives at contexts vulnerable to repetition loops.The method is designed to prevent early commitment to repetitive trajectories.
2. Related Work
Prior work addresses pruning mainly through knowledge preservation and addresses degeneration through decoding-time or training-time repetition penalties, but these approaches involve limitations in diversity, guidance, or perplexity.
- Pruning and knowledge preservation: Pruning research commonly uses knowledge distillation during post-pruning fine-tuning to transfer teacher knowledge to the pruned student.Earlier evaluations emphasize perplexity, zero-shot accuracy, and downstream task performance.
- Decoding-time strategies: Greedy and beam search can limit diversity, while top-k and top-p sampling improve diversity by restricting generation to high-probability tokens.Top-p adapts the candidate set to the sharpness of the token distribution.
- Training-time strategies: Training-time methods such as unlikelihood training and ScaleGrad discourage repetition, but unlikelihood-based approaches often degrade perplexity.These methods modify token probabilities or gradients to promote novel tokens.
3. Token-level Guidance: Repetition-loop Dynamics and Supervision Signals
The paper models repetition degeneration as a token-level dynamical process involving loop entry and persistence, then derives supervision targets that control leakage and preserve plausible escape alternatives. Experiments show that small onset interventions can sharply reduce repetition, motivating guidance objectives based on nucleus-level probability allocation.
- Coverage-based Repetition Metrics and Onset Sensitivity: Coverage measures how much of a generated sequence is explained by occurrences of its dominant recurring N-gram.Dataset-level degeneration is summarized using the Coverage-based REPetition rate (CREP).
- Coverage-based Repetition Metrics and Onset Sensitivity: Replacing only two onset tokens reduces CREP from 5.9% to 0.7% under sampling and from 26.6% to 10.8% under greedy decoding.The replacements use high-probability alternatives, preserving plausible local continuation while shifting generation away from repetition.
- Decoding Dynamics: Entry Risk and Persistence: The decoding-dynamics view represents generation as transitions among contexts and defines degeneration as entering and remaining within a recurrent repetition-prone region L.This framework separates degeneration into loop entry risk and loop persistence.
- Nucleus Alternatives and Escape Mass: Persistence is controlled by nucleus-level escape mass: increasing plausible alternatives within the nucleus reduces persistence, whereas full-vocabulary entropy alone is insufficient.Under nucleus sampling, only tokens in the nucleus are sampleable, so persistence depends on the allocation between loop-continuing and loop-exiting tokens.
- How Pruning and Naive Distillation Reshape Risks: Forward KL distillation can mode-average a multimodal teacher, redistributing student probability into intermediate or low-density regions and increasing leakage into teacher-suppressed tokens.Such leaked tokens may enter the nucleus and increase loop-entry risk under top-p decoding.
- How Pruning and Naive Distillation Reshape Risks: Pruning can also collapse teacher-supported near-tie alternatives, reducing nucleus escape mass and increasing persistence; the proposed targets therefore suppress leakage and preserve plausible escape alternatives.These targets instantiate complementary guidance objectives for controlling entry and persistence after pruning.
4. Method
The method combines confidence-weighted distillation with repetition-aware pairwise alignment to guide token probabilities away from degeneration while preserving plausible alternatives.
- FOCUS: FOCUS uses teacher and student predictive distributions over the vocabulary as the basis for its token-wise weighting and distillation objective.The teacher logits define the teacher distribution, while the student and teacher predictive distributions are used in the objective.
- FOCUS: FOCUS reweights distillation toward tokens where the teacher is highly confident, mitigating probability leakage into low-support regions.It preserves the standard knowledge-distillation form by effectively replacing the teacher distribution with a reweighted version.
- RePAIR: RePAIR supplies token-level corrective signals at repetition onset, where runtime guidance cannot detect loops without access to future tokens.Its training-time design targets the point at which degeneration begins rather than attempting to intervene during decoding.
- RePAIR: RePAIR constructs negative samples from repetitive pruned-model outputs and positive samples from unpruned-model top-p continuations after the first repetition.Sequences exceeding Coverage threshold 0.3 are used to identify repetition, and regeneration begins from the prefix preceding the first repetition.
- RePAIR: A token-level pairwise margin loss trains the model to assign higher confidence to positive continuations than to negative ones.The final objective combines continuation likelihood and margin-ranking terms, with α1 = 0.05 and α2 = 1 unless otherwise specified.
5. Experiments
Experiments evaluate FOCUS and RePAIR on open-ended continuation and instruction-based generation using pruned Llama models, alongside performance and degeneration metrics. Both methods reduce repetition and improve generation quality with only small perplexity increases.
- Experimental Setup: 25% width pruning is applied to Llama models, followed by two epochs of LoRA fine-tuning on Alpaca with self-distillation.Additional pruning rates of 35% and 45% are evaluated in Appendix I.
- Experimental Setup: The evaluation covers 1,000 WikiText-103 continuations and 1,000 Self-Instruct prompts under top-p sampling with p = 0.9.Both task performance and degeneration metrics are reported to assess repetition mitigation without excessive task degradation.
- Open-ended Generation: RePAIR alone achieves the highest MAUVE score of 0.68 and the lowest CREP score of 2.23 in open-ended generation.Its unique n-gram distribution is also comparable to the original WikiText-103 dataset.
- Instruction-based Generation: RePAIR reaches a CREP score of 0.63 and an EAD1 score of 0.32 in instruction-based generation while maintaining comparable perplexity to UL.Combining FOCUS with the methods improves CREP, EAD1, and BERTScore while keeping zero-shot accuracy stable.
- LLM-as-a-Judge Evaluation: LLM-as-a-judge evaluation reports that the proposed method consistently outperforms baselines in generation quality.This complements the automatic metrics, which may not fully capture perceptual quality.
6. Analysis
The analysis examines FOCUS distribution alignment, likelihood instability, loop persistence, supervision granularity, comparisons with ToDi, and RePAIR data efficiency. These studies connect token-level guidance to reduced repetition and preserved generation quality.
- FOCUS Distribution Study: FOCUS increases head agreement and tail correlation with the teacher while preserving tail agreement, indicating structural alignment of low-confidence probabilities.The teacher head is defined by top-p with p = 0.9.
- Robustness to Likelihood Instability: Despite increasing perplexity on WikiText, FOCUS improves TruthfulQA performance, while RePAIR achieves the largest overall gain.The analysis examines whether likelihood instability affects factual behavior.
- Supervision Granularity: RePAIR achieves larger repetition-mitigation gains than DPO under WikiText-103 continuation, consistent with its onset-level rather than sequence-level supervision.Both methods are compared under the same continuation setting to isolate supervision granularity.
- Loop Persistence: FOCUS delays the increase and saturation of loop-continuing probability mass, and combining it with RePAIR further suppresses the persistence metric ρ(c).The persistence analysis uses synthetic contexts formed by repeating WikiText-103 sentence pairs ten times.
- Distribution-shaping Distillation: FOCUS reduces repetition more effectively than ToDi, and FOCUS combined with RePAIR achieves the lowest repetition scores with MAUVE comparable to ToDi.The comparison uses the same Llama 3.1-8B setting as Table 3.
- Pairwise Data Efficiency: About 4k of 12k RePAIR pairwise samples suffice to achieve repetition rates comparable to those from the full dataset.This makes RePAIR more data-efficient than DITTO and preserves more data for standard training.
7. Conclusion
The paper concludes that pruning increases repetition in language models and that token-level guidance mitigates this degeneration. FOCUS and RePAIR reduce repetition while improving generation quality.
- Findings: Pruning increases repetition in language models, motivating token-level guidance as a mitigation strategy.The conclusion frames repetition as a side effect of pruning that requires targeted intervention.
- Methods: FOCUS uses token-probability-weighted distillation, whereas RePAIR uses a pairwise margin objective to guide models toward better alternative tokens.The methods provide complementary token-level guidance mechanisms.
- Conclusion: Both methods consistently reduce repetition and improve generation quality.This is the paper’s overall reported outcome.
Impact Statement
The work is presented as advancing the field of Machine Learning. The authors identify no societal consequences that require specific highlighting.
- Stated Impact: The paper’s stated goal is to advance the field of Machine Learning.
- Stated Impact: The impact statement acknowledges that the work may have potential societal consequences.
- Stated Impact: No societal consequences are identified as requiring specific emphasis.
B. Gradient Analysis of FOCUS
This section frames FOCUS as a reweighted distillation objective and analyzes its gradient relative to the student and teacher token distributions.
- B. Gradient Analysis of FOCUS: FOCUS is analyzed as modifying the knowledge-distillation objective rather than using the standard teacher distribution directly.The analysis introduces a reweighted teacher distribution for the FOCUS loss.
- B. Gradient Analysis of FOCUS: The gradient analysis differentiates the loss through the softmax logits using the softmax Jacobian and the chain rule.The derivation separates the i = k and i ≠ k cases before expressing the result in vector form.
- B. Gradient Analysis of FOCUS: The resulting gradient has the form Z(p − ˜q), where ˜q is formed by weighting and renormalizing teacher probabilities.The displayed formulation defines ˜q from w(q) ⊙ q and a normalization term.
- B. Gradient Analysis of FOCUS: The section also describes constructing RePAIR pairs from a pruned-model degeneration onset and an unpruned-model positive continuation.The prefix is identified using coverage, then paired with negative and positive continuations for margin-loss training.
E. Implementation Details
The implementation uses standard Hugging Face-based baselines, token-level repetition objectives, coverage-based detection, and parameter studies balancing diversity against perplexity.
- E. Implementation Details: All baseline methods are implemented with the Hugging Face framework using official implementations where available.The implementation section states this as the common baseline setup.
- E. Implementation Details: Knowledge distillation transfers teacher knowledge to the pruned student through soft probability supervision with temperature T = 2.The stated temperature is intended to balance training stability and knowledge transfer.
- E. Implementation Details: Unlikelihood training uses token-level negative candidates to suppress repeated tokens while combining the objective with maximum-likelihood training.The experiments set α = 0.5 as a trade-off between repetition suppression and fluency or perplexity preservation.
- E. Implementation Details: FOCUS parameter studies show that larger β and γ improve unique n-gram scores but slightly increase perplexity, while α1 = 0.05 balances perplexity and repetition reduction.Perplexity deteriorates sharply as α1 rises from 0.05 to 0.1, whereas the n-gram rate saturates from 0.1 onward.
- E. Implementation Details: The CREP procedure sweeps n-gram lengths from 4 to 16 and marks a sentence degenerate when maximum repeated-segment coverage exceeds a threshold.It identifies recurring n-grams, reconstructs repeated segments, and computes their global coverage.
H. Experiment on Other Pruning Methods
Experiments across pruning methods, model families, and aggressive sparsity settings indicate that FOCUS and RePAIR consistently reduce repetition beyond the original LLaMA and pruning configurations.
- H. Experiment on Other Pruning Methods: FOCUS and RePAIR consistently reduce repetition across LLMPruner, other width- and depth-pruning methods, and the Qwen model family.The reported robustness extends beyond the LLaMA family and the LLMPruner pruning method.
- H. Experiment on Other Pruning Methods: At 35% and 45% pruning ratios, both methods improve repetition-related metrics, including n-gram rate, across all tested pruning levels.These settings evaluate the methods under more aggressive sparsity where degeneration is more prominent.
- H. Experiment on Other Pruning Methods: Increasing pruning ratios lowers MAUVE, reflecting reduced model capacity and distributional fidelity under aggressive compression.The text distinguishes this decline from the repetition-mitigation behavior of the proposed methods.
J. DPO Training Details
The DPO comparison isolates supervision granularity by training sequence-level preference learning on the same positive and negative continuation pairs used for RePAIR.
- J. DPO Training Details: DPO is implemented as a baseline using the positive non-repetitive and negative repetitive pairs constructed for RePAIR.This keeps the data and supervision source comparable between the two methods.
- J. DPO Training Details: DPO and RePAIR are evaluated separately under the same WikiText-103 continuation protocol without combining either method with KD or FOCUS.The setup isolates sequence-level versus token-level supervision.
- J. DPO Training Details: The DPO temperature parameter is set to βdpo = 0.1, consistent with prior work.
- J. DPO Training Details: Human-style judging compares anonymized candidate pairs as A, B, or TIE, emphasizing fluency, coherence, and repetition while randomizing order.Preference reversals after swapping positions are conservatively counted as ties.