Addressing Distribution Shift at Test Time in Pre-trained Language Models

TL;DR

This paper introduces MEMO-CL, a test-time adaptation method for pre-trained language models that uses unsupervised data augmentation to improve performance under distribution shift without requiring labels or additional data.

Contribution

The paper proposes MEMO-CL, a novel unsupervised, domain-agnostic test-time adaptation technique leveraging data augmentation to enhance PLM robustness against distribution shifts.

Findings

Achieves 3% performance improvement over existing baselines.

Operates efficiently on a batch of augmented samples from a single test observation.

Requires no additional data or labels for adaptation.

Abstract

State-of-the-art pre-trained language models (PLMs) outperform other models when applied to the majority of language processing tasks. However, PLMs have been found to degrade in performance under distribution shift, a phenomenon that occurs when data at test-time does not come from the same distribution as the source training set. Equally as challenging is the task of obtaining labels in real-time due to issues like long-labeling feedback loops. The lack of adequate methods that address the aforementioned challenges constitutes the need for approaches that continuously adapt the PLM to a distinct distribution. Unsupervised domain adaptation adapts a source model to an unseen as well as unlabeled target domain. While some techniques such as data augmentation can adapt models in several scenarios, they have only been sparsely studied for addressing the distribution shift problem. In this work, we present an approach (MEMO-CL) that improves the performance of PLMs at test-time under distribution shift. Our approach takes advantage of the latest unsupervised techniques in data augmentation and adaptation to minimize the entropy of the PLM's output distribution. MEMO-CL operates on a batch of augmented samples from a single observation in the test set. The technique introduced is unsupervised, domain-agnostic, easy to implement, and requires no additional data. Our experiments result in a 3% improvement over current test-time adaptation baselines.