From Robustness to Improved Generalization and Calibration in Pre-trained Language Models (2025)

Abstract: Enforcing representation smoothness in pre-trained language models (PLMs) through Jacobian and Hessian regularization provides an effective approach for enhancing both robustness and generalization. Although such regularization methods have proven effective in computer vision, their application in natural language processing, where PLM inputs are derived from a discrete domain, poses unique challenges. We introduce JacHess, a regularization approach for PLMs that minimizes the norms of the Jacobian and Hessian matrices in intermediate representations, using embeddings as substitutes for discrete token inputs. JacHess supports dual-mode regularization, alternating between fine-tuning with labeled data and regularization with unlabeled data. We evaluate JacHess on the GLUE benchmark and demonstrate that it consistently and significantly improves in-distribution generalization and enhances performance under domain shift. Across diverse PLMs, JacHess outperforms comparable representation-based regularization methods and unregularized fine-tuning, while also improving model calibration. Our findings, coupled with a computationally efficient estimator for the Jacobian and Hessian norms, position JacHess as a robust and widely applicable solution for enhancing PLM performance.

## introduction effective generalization, broadly understood as the capability to transfer learned representations, knowledge, and strategies from familiar contexts to new ones @xcite , stands out as a key goal for models in natural language processing (nlp) and extends to the broader domain of machine learning. at the heart of machine learning optimization is the principle of minimizing empirical risk, which acts as a surrogate for the true risk. this approach inevitably leads to a generalization gap-the discrepancy between the empirical and true risks. the size of this gap can be gauged by considering certain model properties. notably, robustness, defined as the amount that the loss can vary with respect to changes in the inputs, offers an effective lens for understanding generalization @xcite @xcite . recently, @xcite provided tighter generalization bounds for robustness, a challenge that has persisted since robustness was first proposed as a tool for analyzing learning algorithms @xcite . however, their bounds are data-dependent and thus remain undetermined until the training data is specified, leaving room for further exploration. robustness in neural networks is intrinsically linked to the concept of representation smoothness @xcite . when representations are smooth, the outputs of the network vary minimally in response to small input changes, thereby maintaining a stable loss by controlling the network's geometric complexity @xcite . promoting representation smoothness not only aids in better generalization but also supports more reliable uncertainty quantification in the model's predictions @xcite , a critical area where neural networks often fail. typically, this failure manifests as overconfidence in the model's predictions, resulting in poor calibration with predicted probabilities misaligned with actual outcomes @xcite . promoting representation smoothness thus emerges as a promising strategy for boosting generalization through robustness and improving uncertainty quantification in neural networks. while traditional regularization methods such as weight decay @xcite and dropout @xcite enhance neural network generalization by promoting model simplicity and preventing overfitting, they do not ensure robustness against input variations. in contrast, more specialized approaches focus on manipulating the network's input-output jacobian and hessian matrices, involving first-order and second-order partial derivatives. minimizing the norms of these matrices increases robustness by promoting representation smoothness. this effect can be understood by interpreting a neural network as a function. the concept of lipschitz continuity reflects the smoothness of a function against input variations, representing the maximum rate of change in the function's output relative to a change in the input @xcite . reducing the network's lipschitz constant thus reduces its sensitivity to input perturbation. while computing the exact lipschitz constant is an np-hard problem @xcite , leveraging jacobian and hessian norms as proxies offers a practical way to emulate these effects. the norm-based regularization approach has proven highly successful in the field of computer vision, as evidenced by several studies @xcite @xcite . while regularization methods directly targeting robustness have effectively enhanced generalization in computer vision, none have been used in nlp, revealing a significant research gap. in particular, given the pivotal role of pre-trained language models (plms) in advancing nlp, there exists an untapped potential to leverage representation-based regularization methods to improve plms' generalization capabilities. however-unlike in computer vision-the discrete nature of the tokens processed by plms poses a significant barrier to this research direction in nlp. this issue echoes previous challenges encountered when methods proven effective in computer vision did not seamlessly transfer to nlp (e.g., the idea of generative adversarial networks @xcite that operate by applying slight continuous modifications to inputs). in this paper, we adapt and expand upon the representation-based regularization techniques used in computer vision and apply them to nlp. the discrete nature of data in plms makes it challenging to enhance robustness to input variations, but we find a workaround by leveraging the continuous embedding space as an effective alternative. we introduce jachess, a novel regularization approach that minimizes the norms of the jacobian and hessian matrices within plm representations relative to their inputs. by targeting both jacobian and hessian norms, jachess doubly enhances model robustness by reducing sensitivity to input changes and smoothing the curvature of representations. typically, calculating these norms in high-dimensional spaces is compute-intensive, but we address this using a computationally efficient estimator. our method applies regularization across the network's layers, promoting smoothness in intermediate representations. moreover, jachess employs a dual-mode strategy, cycling between iterative fine-tuning using labeled data and regularization with additional unlabeled data. while employing jachess with training data significantly improves generalization in most cases, utilizing a separate set of unlabeled data for regularization-which is typically more readily available than labeled data-enhances the model's generalization capabilities even further. additionally, we demonstrate that jachess proves highly beneficial under domain shift, where models need to generalize across different domains while performing the same underlying task. given their growing importance, our evaluation concentrates on decoder-based models, including the opt family and llama 2, where we fine-tune the plms and evaluate them on the glue benchmark @xcite . we validate jachess by examining plm robustness against the perturbations in embedding space and its impact on handling corrupted token inputs. subsequently, we assess the effectiveness of jachess in improving models' predictive accuracy and calibration. our results reveal that jachess markedly outperforms standard unregularized fine-tuning and other jacobian-and hessian-based methods, achieving a 2% to 4.5% absolute increase in average glue scores across various plms. this enhancement in generalization is complemented by improved calibration of model predictions. in summary, our work presents two significant contributions: (1) we empirically evaluate the enforcement of representation smoothness on plms through jacobian-and hessian-based regularization, an aspect that has thus far been overlooked in nlp, and (2) we introduce jachess, a novel regularization approach that not only improves model generalization beyond the capabilities of standard fine-tuning and other regularization methods in the representation space but also offers more reliable uncertainty quantification through improved calibration and enhanced robustness to domain shift. taken together, this work sets a new standard for improving the generalization and calibration of plms in the evolving nlp landscape. foot_0 2 related work generalization enhancers. applying regularization to plms is a common practice during training to improve generalization @xcite . although increasing the amount of labeled data can straightforwardly improve generalization, the challenge lies in efficiently utilizing existing data or capitalizing on the plentiful supply of unlabeled data. various strategies aim to address this challenge. notably, @xcite proposed task-adaptive pre-training, which involves additionally training the model on the task-specific unlabeled set via the language modeling task. in a different approach, @xcite leveraged unlabeled data for generalization enhancement through weak supervision. moreover, beyond the utilization of unlabeled data, the field has seen advancements in generalization through data augmentation methods @xcite @xcite . representation smoothness. examining the impact of input changes on a function's output is essential for enhancing the generalization capabilities of neural networks. @xcite explored the concept of model smoothness, or representation smoothness in the context of deep neural networks. they advocate for implementing smoothness constraints with respect to inputs to improve generalization and provide reliable uncertainty quantification. while numerous theoretical approaches predominantly focus on lipschitz-bounds to assess model sensitivity @xcite @xcite , it has been demonstrated that lipschitz continuity can be effectively leveraged for regularization purposes @xcite . representation-based regularization. @xcite were the first to explore the effects of the input-output jacobian matrix of neural networks. since then, several iterations of this technique have been proposed @xcite , primarily aimed at enhancing robustness against adversarial examples @xcite . extending this work, @xcite suggested that such robustness could also improve in-distribution generalization. further developments by @xcite introduced an estimator for the frobenius norm of the jacobian matrix of logits with respect to inputs, avoiding the need to compute the entire jacobian and significantly reducing resource demands. extending the scope of regularization, @xcite proposed regularizing the frobenius norm of the hessian matrix, focusing on the scalar outputs of binary classification, a technique they termed cross-hölder regularization. in a related approach, @xcite introduced an adversarial-style method by adding noise directly to the inputs during fine-tuning. known as robust fine-tuning with regularized representations (r3f), this approach draws on earlier works @xcite and leverages trust-region theory to enhance model stability. by reducing the divergence between outputs generated from noisy and clean inputs, r3f promotes smoothness within the representation space. sharpness-aware minimization. the proposed jachess method improves generalization by enhancing smoothness within the representation space of a plm, which is similar in spirit to sharpness-aware minimization (sam; @xcite . sam is an optimization technique designed to enhance model generalization by locating parameters that reside in flatter regions of the loss landscape. @xcite showed that sam improves plm generalization, particularly in transfer learning scenarios. extending this concept, @xcite leveraged sam for cross-lingual tasks, improving the transfer of pre-trained representations across languages. in our work, we focus on promoting smoothness with respect to inputs by applying regularization directly within the representation spaces of plms. specifically, we leverage jacobianand hessian-based regularization techniques to enforce smoothness across intermediate representations at all plm layers, as opposed to restricting regularization to the model's final outputs or logits. furthermore, we push beyond traditional regularization techniques by exploring the application of regularization on separate, unlabeled data. ## representation-based regularization in this section, we outline our approach, starting with the theoretical background of lipschitz continuity, a useful tool for assessing the stability and smoothness of neural networks. next, we address computational challenges and propose strategies to mitigate them. finally, we introduce our regularization technique, jachess, which leverages the continuous embedding space of plms to employ representation-based regularization with respect to the inputs. ## experimental setup this section describes the experimental framework, including the models, methods, and datasets used. ## enhancing generalization through robustness the goal of our experiments is to evaluate how well jachess improves robustness. we begin by testing our method against perturbations in the embedding space, followed by testing its impact on generalization for both in-distribution data and data under domain shift. ## calibration uncertainty quantification in nlp models has garnered significant attention lately, given its importance in various applications @xcite @xcite . motivated by research indicating a connection between smoothness in model representations and more reliable uncertainty quantification @xcite , we explore the effect of jachess on out-of-the-box calibration, using the plm's base cross-hölder jachess model brier/ece brier/ece brier/ece bert .213/.042 .202/.039 .184/.035 opt-125m .256/.057 .233/.045 .204/.040 opt-1.3b .184/.038 .193/.042 .157/.029 opt-6.7b .189/.037 .157/.032 .094/.024 llama-2-7b .167/.034 .163/.031 .089/.017 ## analysis several design choices behind jachess may impact its efficacy. we next analyze these factors. ## conclusion the success of plms depends significantly on their robustness and ability to generalize. in this work, we examined how enforcing representation smoothness affects these properties in transformer-based plms by leveraging jacobian and hessian norm minimization. we introduced jachess, a regularization method that enforces smooth representations with respect to the input embeddings, which serve as an alternative to discrete tokens. leveraging computationally efficient estimators for the jacobian and hessian norms, jachess is practical and effective, consistently surpassing unregularized fine-tuning and similar regularization approaches by improving both in-distribution and cross-domain generalization. across diverse plms and glue benchmark tasks, jachess enhances predictive accuracy and model calibration, leading to more reliable uncertainty quantification. these results highlight the importance of representation smoothness for model robustness, setting jachess as the new reference for improving the generalization and calibration of plms.