# Unified Enhancement of the Generalization and Robustness of Language Models via Bi-Stage Optimization

arXiv: https://arxiv.org/abs/2503.16550 · HTML: https://mlsanity.com/paper/2503.16550/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 31 of its 46 references that are in the map.
- First posted on arXiv 19 Mar 2025.

## Builds on (most cited first, 30 of 31)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (2017)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [DeBERTa: Decoding-enhanced BERT with Disentangled Attention](https://mlsanity.com/paper/2006.03654/) (2020)
- [Similarity of Neural Network Representations Revisited](https://mlsanity.com/paper/1905.00414/) (2019)
- [Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution](https://mlsanity.com/paper/2202.10054/) (2022)
- [Overfitting in adversarially robust deep learning](https://mlsanity.com/paper/2002.11569/) (2020)
- [Multi-Task Deep Neural Networks for Natural Language Understanding](https://mlsanity.com/paper/1901.11504/) (2019)
- [Pretrained Transformers Improve Out-of-Distribution Robustness](https://mlsanity.com/paper/2004.06100/) (2020)
- [CLUE: A Chinese Language Understanding Evaluation Benchmark](https://mlsanity.com/paper/2004.05986/) (2020)
- [FreeLB: Enhanced Adversarial Training for Natural Language Understanding](https://mlsanity.com/paper/1909.11764/) (2019)
- [Mind the Gap: Assessing Temporal Generalization in Neural Language Models](https://mlsanity.com/paper/2102.01951/) (2021)
- [Probing Neural Network Comprehension of Natural Language Arguments](https://mlsanity.com/paper/1907.07355/) (2019)
- [R-Drop: Regularized Dropout for Neural Networks](https://mlsanity.com/paper/2106.14448/) (2021)
- [Recall and Learn: Fine-tuning Deep Pretrained Language Models with Less Forgetting](https://mlsanity.com/paper/2004.12651/) (2020)
- [Better Fine-Tuning by Reducing Representational Collapse](https://mlsanity.com/paper/2008.03156/) (2020)
- [Raise a Child in Large Language Model: Towards Effective and Generalizable Fine-tuning](https://mlsanity.com/paper/2109.05687/) (2021)
- [Robustness Gym: Unifying the NLP Evaluation Landscape](https://mlsanity.com/paper/2101.04840/) (2021)
- [Sharpness-Aware Minimization Improves Language Model Generalization](https://mlsanity.com/paper/2110.08529/) (2021)
- [SSMBA: Self-Supervised Manifold Based Data Augmentation for Improving Out-of-Domain Robustness](https://mlsanity.com/paper/2009.10195/) (2020)
- [InfoBERT: Improving Robustness of Language Models from An Information Theoretic Perspective](https://mlsanity.com/paper/2010.02329/) (2020)
- [Contrastive Out-of-Distribution Detection for Pretrained Transformers](https://mlsanity.com/paper/2104.08812/) (2021)
- [HiddenCut: Simple Data Augmentation for Natural Language Understanding with Better Generalization](https://mlsanity.com/paper/2106.00149/) (2021)
- [Improving Generalization of Adversarial Training via Robust Critical Fine-Tuning](https://mlsanity.com/paper/2308.02533/) (2023)

## Abstract (arXiv)

Neural network language models (LMs) are confronted with significant challenges in generalization and robustness. Currently, many studies focus on improving either generalization or robustness in isolation, without methods addressing both aspects simultaneously, which presents a significant challenge in developing LMs that are both robust and generalized. In this paper, we propose a bi-stage optimization framework to uniformly enhance both the generalization and robustness of LMs, termed UEGR. Specifically, during the forward propagation stage, we enrich the output probability distributions of adversarial samples by adaptive dropout to generate diverse sub models, and incorporate JS divergence and adversarial losses of these output distributions to reinforce output stability. During backward propagation stage, we compute parameter saliency scores and selectively update only the most critical parameters to minimize unnecessary deviations and consolidate the model's resilience. Theoretical analysis shows that our framework includes gradient regularization to limit the model's sensitivity to input perturbations and selective parameter updates to flatten the loss landscape, thus improving both generalization and robustness. The experimental results show that our method significantly improves the generalization and robustness of LMs compared to other existing methods across 13 publicly available language datasets, achieving state-of-the-art (SOTA) performance.