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

Josip Jukić, Jan Šnajder


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.
Anthology ID:
2025.tacl-1.13
Volume:
Transactions of the Association for Computational Linguistics, Volume 13
Month:
Year:
2025
Address:
Cambridge, MA
Venue:
TACL
SIG:
Publisher:
MIT Press
Note:
Pages:
264–280
Language:
URL:
https://preview.aclanthology.org/corrections-2025-07/2025.tacl-1.13/
DOI:
10.1162/tacl_a_00739
Bibkey:
Cite (ACL):
Josip Jukić and Jan Šnajder. 2025. From Robustness to Improved Generalization and Calibration in Pre-trained Language Models. Transactions of the Association for Computational Linguistics, 13:264–280.
Cite (Informal):
From Robustness to Improved Generalization and Calibration in Pre-trained Language Models (Jukić & Šnajder, TACL 2025)
Copy Citation:
PDF:
https://preview.aclanthology.org/corrections-2025-07/2025.tacl-1.13.pdf