Learning Perturbations to Extrapolate Your LLM

TL;DR

This paper introduces a learnable perturbation framework for large language models that improves out-of-domain extrapolation by perturbing token prefixes in a continuous embedding space.

Contribution

It proposes a novel method of perturbing token prefixes with learnable transformations, overcoming fixed discrete perturbations and deriving unbiased estimators for model training.

Findings

Significant out-of-domain performance improvements over baseline methods.

The proposed estimator has desirable statistical properties in over-parameterized regimes.

Empirical results on synthetic and real datasets validate the approach.

Abstract

Recent advancements in large language models demonstrate that injecting perturbations can substantially enhance extrapolation performance. However, current approaches often rely on discrete perturbations with fixed designs, which limits their flexibility. In this work, we propose a framework where token prefixes are perturbed by a learnable transformation of a continuous latent vector within an embedding space. To overcome the challenge of an intractable marginal likelihood, we derive unbiased estimating equations for model parameters and optimize them via stochastic gradient descent. We establish the statistical properties of the resulting estimator in over-parameterized regimes. Empirical evaluations on both synthetic and real-world datasets demonstrate that our proposal yields significant gains in out-of-domain settings over a range of state-of-the-art baseline methods.