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

TL;DR

UEGR addresses the dual challenge of generalization and robustness in language models by a bi-stage optimization framework that combines forward-stage adaptive dropout regularization with adversarial training and JS-divergence to stabilize outputs, and backward-stage selective parameter updates using saliency-based gradient masking. Theoretical analysis shows gradient regularization and loss landscape flattening, and empirical results on 13 datasets demonstrate state-of-the-art improvements across pretrained and non-pretrained models, along with robust performance under adversarial perturbations. Ablation studies identify key hyperparameters, including the number of adaptive dropout passes and the dropout range, validating the design choices. Overall, UEGR offers a practical, architecture-agnostic approach to achieving robust and generalizable language models with principled theoretical backing and broad empirical support.

Abstract

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.

Figures (3)

• Figure 1: Bi-stage optimization framework. This figure illustrates the employment of adaptive dropout regularization during the forward propagation stage and selective update during the backward propagation stage in the UEGR method, aimed at enhancing language model generalization and robustness. The solid lines denote the forward propagation process, whereas the dashed lines indicate the backward propagation process. Blue lines signify dropped connections within the language model, and red lines signify retained connections.
• Figure 2: A qualitative analysis of UEGR is provided as follows: (a) Examining the similarity between the original pre-trained model and the fine-tuned model, and (b) Assessing the model's robustness under varying levels of adversarial perturbation, along with the improvements achieved through UEGR.
• Figure 3: UEGR with two adaptive dropout rate combinations.

Theorems & Definitions (1)

• Theorem 1