Causal Fine-Tuning of Pre-trained Language Models for Robust Test Time Adaptation

Jialin Yu, Yuxiang Zhou, Yulan He, Nevin L. Zhang, Junchi Yu, Philip Torr, Ricardo Silva

Keywords: distribution shift, causality, test time generalisation, causal fine-tuning

Abstract: Fine-tuning pre-trained language models (PLMs) with supervised data improves performance, but often fails to generalise under unknown distribution shifts, as models tend to rely on spurious, non-causal features. Existing approaches typically make restrictive assumptions or require multi-domain data, limiting their applicability in real-world, single-domain settings. We propose a novel causal adjustment framework that improves out-of-distribution generalisation by decomposing the representation of PLMs into causal and spurious components, and recombine them for testing time adaptation. Extensive experiments on semi-synthetic datasets demonstrate that our causal fine-tuning method consistently outperforms state-of-the-art domain generalisation baselines.

Submission Number: 25