Source-linked AI summary
Better Fine-Tuning by Reducing Representational Collapse
Armen Aghajanyan, Akshat Shrivastava, Anchit Gupta, Naman Goyal, Luke Zettlemoyer, Sonal Gupta
TL;DR
Fine-tuning pre-trained language models is unstable, and standard updates can degrade generalizable representations. The paper introduces lightweight trust-region-inspired regularization using parametric noise instead of adversarial objectives, achieving competitive or better performance with reduced overhead while alleviating representational collapse.
Problem
Fine-tuning pre-trained language models is unstable across hyper-parameter settings, and standard fine-tuning degrades generalizable representations through representational collapse.
Method
R3F and R4F use trust-region-inspired regularization with parametric normal or uniform noise, with R4F additionally constraining the classification head through spectral normalization.
Results
The methods match or exceed prior trust-region methods across understanding and generation tasks while requiring a fraction of their computational and memory overhead.
Takeaways & Limitations
The methods retain more generalizable pre-trained representations and improve end-task performance across monolingual, multilingual, sentence-prediction, and generation tasks.
Takeaways & Limitations
The method is not evaluated on STS-B because its KL divergence is undefined for that regression task.
Abstract
from arXiv · showhide
Although widely adopted, existing approaches for fine-tuning pre-trained language models have been shown to be unstable across hyper-parameter settings, motivating recent work on trust region methods. In this paper, we present a simplified and efficient method rooted in trust region theory that replaces previously used adversarial objectives with parametric noise (sampling from either a normal or uniform distribution), thereby discouraging representation change during fine-tuning when possible without hurting performance. We also introduce a new analysis to motivate the use of trust region methods more generally, by studying representational collapse; the degradation of generalizable representations from pre-trained models as they are fine-tuned for a specific end task. Extensive experiments show that our fine-tuning method matches or exceeds the performance of previous trust region methods on a range of understanding and generation tasks (including DailyMail/CNN, Gigaword, Reddit TIFU, and the GLUE benchmark), while also being much faster. We also show that it is less prone to representation collapse; the pre-trained models maintain more generalizable representations every time they are fine-tuned.
1 INTRODUCTION
Fine-tuning pre-trained language models is unstable and can damage generalizable representations. The paper proposes lightweight trust-region-inspired methods that reduce representational collapse while preserving or improving performance.
- Fine-tuning pre-trained language models can produce failed runs, seed-sensitive results, over-fitting, and other unwanted consequences across hyper-parameter settings.
- Trust-region methods improve fine-tuning stability and accuracy by constraining changes to initial parameters, but require substantially more computation and memory.
- The proposed lightweight strategy matches or improves SMART and FreeLB performance with a fraction of their computational and memory overhead and no additional backward passes.
- Representational collapse is defined and measured as degradation of generalizable pre-trained representations during fine-tuning.Probing experiments connect this degradation to standard fine-tuning and compare trust-region-based methods.
- The method achieves state-of-the-art or improved results across sentence prediction, summarization, monolingual, and cross-lingual tasks.Reported outcomes include state of the art on GLUE and improvements spanning sentence prediction to summarization.
2 LEARNING ROBUST REPRESENTATIONS THROUGH REGULARIZED FINE-TUNING
The method approximates trust-region optimization by regularizing prediction changes under parametric noise rather than using adversarial optimization. Optional spectral normalization constrains the classification head, reducing computational cost while retaining performance.
- The objective minimizes task loss while constraining movement in the density space of pre-trained representations between update steps.
- Direct natural-gradient optimization over representation density is impractical, motivating an approximation over g·f with a 1-Lipschitz classification head.
- SMART approximates the intractable supremum with adversarial gradient-ascent steps, whereas the proposed method removes those ascent computations.
- R3F regularizes symmetric KL divergence between outputs for an input and a perturbed input, with perturbations sampled from centered normal or uniform distributions.
- The proposed methods are significantly more computationally efficient than adversarial fine-tuning while matching or exceeding FreeLB and SMART across many tasks.
3 EXPERIMENTS
The experiments evaluate R3F/R4F against standard and trust-region fine-tuning methods across GLUE, XNLI, and abstractive summarization, emphasizing performance, stability, and computational cost.
- Experimental setup: The evaluation uses fixed-budget hyperparameter searches and, where tractable, reports median and maximum results plus distributions across multiple runs.GLUE experiments use 10 seeds, while XNLI results average five zero-shot runs.
- Sentence prediction: R3F and R4F are evaluated on RoBERTa-Large single-task models across MNLI, QQP, RTE, QNLI, MRPC, CoLA, and SST-2.Performance is reported on the GLUE development set.
- Sentence prediction: On GLUE, R3F and R4F unanimously improve over Standard and Standard++ fine-tuning while matching or exceeding SMART and FreeLB at a fraction of the computational cost.The relative wall-time advantage is reported as consistent across GLUE tasks.
- Cross-lingual prediction: R4F reaches an average XNLI score of 81.4 across five runs and dominates standard pre-training on 14 of 15 languages.The result uses XLM-R Large in the zero-shot setting; the cited current state of the art requires novel pretraining.
- Summarization: The summarization evaluation reports ROUGE-1, ROUGE-2, and ROUGE-L for each dataset, following PEGASUS’s convention for highlighting best or near-best numbers.The datasets are CNN/DailyMail, Gigaword, and Reddit TIFU.
- Summarization: For CNN/DailyMail, Gigaword, and Reddit TIFU, R3F outperforms standard fine-tuning across all reported ROUGE metrics.Gigaword and Reddit TIFU improve by 1 and 4 points respectively on ROUGE-1.
4 REPRESENTATIONAL COLLAPSE
The paper defines representational collapse as degradation of generalizable pretrained representations during fine-tuning and measures it through probing across tasks and sequential training. Experiments find that standard fine-tuning degrades representations, whereas R3F/R4F retain them better and achieve stronger probing performance.
- Representational collapse is the degradation of generalizable pretrained representations during fine-tuning, independent of any single task.The paper contrasts this with catastrophic forgetting, which concerns forgetting examples on a fixed task.
- PROBING GENERALIZATION OF FINE-TUNED REPRESENTATIONS: The probing methodology freezes representations from one task and fine-tunes a linear layer for another task to measure representation quality.This evaluates how much representations collapse after fine-tuning on a sequence of tasks.
- PROBING GENERALIZATION OF FINE-TUNED REPRESENTATIONS: Across six GLUE probes after SST-2 fine-tuning, R3F or R4F performed best, while standard fine-tuning produced worse representations across the board.R3F/R4F also consistently outperformed the adversarial SMART method.
- PROBING REPRESENTATION DEGRADATION: Sequential probing on QNLI, QQP, and RTE showed that standard fine-tuning diverged from the source task, reducing source-task probe performance.The experiment began with SST-2 and used the best checkpoint from each prior iteration.
- PROBING REPRESENTATION RETENTION: R4F retained representation quality significantly better than standard fine-tuning across three cycles of cyclic sequential probing.The evaluated GLUE sequence was SST-2 → QNLI → QQP → RTE.
5 CONCLUSION
The conclusion presents R3F and R4F as efficient trust-region-based fine-tuning methods that address representational collapse. Across sentence-prediction and generation tasks, they improve on standard and adversarial fine-tuning and achieve several state-of-the-art results.
- R3F and R4F are trust-region-based fine-tuning approaches designed to retain generalizable representations while reducing computational cost.The paper attributes improved end-task performance to retaining representation generalizability.
- The methods improve monolingual and multilingual sentence prediction and generation tasks compared with standard and adversarial fine-tuning.Reported results include state of the art on DailyMail/CNN, Gigaword, Reddit TIFU, and zero-shot XNLI, plus improved RoBERTa GLUE results.
A.1 CONTROLLING CHANGE OF REPRESENTATION VIA CHANGE OF VARIABLE
This appendix analyzes representation change through a Markovian chain and change-of-variable reasoning. A spectral-norm assumption bounds the Jacobian determinant and motivates spectral normalization as a control on tightness.
- The analysis models representations as a Markovian chain x → y → z, with y = f(x; θ_f) and z = g(y; θ_g).
- Change-of-variable reasoning expresses the probability density after transforming x through f and g using the corresponding Jacobian determinants.
- Assuming g(y) = Wy with spectral norm ρ(W) = 1 gives |det W| ≤ 1, bounding the transformation's volume change.
- The tightness of the bound is controlled by |det W|, whose singular-value dependence provides intuition for using spectral normalization.
A.2 EXPERIMENT HYPER-PARAMETERS
The appendix lists implementation settings for GLUE, summarization, and XNLI experiments. Probing differs from full fine-tuning by freezing the RoBERTa encoder and removing encoder dropout.
- GLUE experiments include both full fine-tuning and probing configurations.For probing, the RoBERTa encoder is frozen and encoder dropout is removed.
- Table 5 contains task-specific hyperparameters for GLUE experiments.
- Table 6 contains R3F and R4F hyperparameters for GLUE experiments.
- Table 7 contains task-specific hyperparameters for summarization experiments.
- Table 8 contains R3F and R4F hyperparameters for summarization experiments, while Table 9 contains their XNLI hyperparameters.