Learning Perturbations to Extrapolate Your LLM

Abstract: Recent advancements in LLMs 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.

• The paper introduces a continuous, trainable perturbation technique in the embedding space that enhances out-of-distribution performance in LLMs.
• It employs a neural perturbation module with unbiased estimating equations to achieve statistical consistency, resulting in lower perplexity and improved Mauve scores.
• Empirical evaluations across synthetic and real-world datasets demonstrate that jointly applying perturbations during training and inference maximizes model extrapolation.

Continuous, Trainable Perturbation for LLM Extrapolation

Overview and Motivation

The paper "Learning Perturbations to Extrapolate Your LLM" (2605.13284) addresses the persistent challenge of improving out-of-distribution (OOD) generalization in autoregressive LLMs. Existing data augmentation and perturbation techniques for language modeling are predominately heuristic, discrete, and fixed in design—operating at the lexical or syntactic level and thus limited in finer-grained adaptivity. These methods inadequately capture the complex, context-dependent regularities underlying generalization. In contrast, this work proposes a unified framework for continuous, trainable perturbation in the embedding space, designed to be adaptive and jointly learned with the core model parameters.

Methodological Framework

Learnable, Latent Embedding-Space Perturbations

At the core of the proposed method is the introduction of a continuous, low-dimensional latent variable $w \in \mathbb{R}^r$ drawn from a standard distribution (typically Gaussian). A neural perturbation module $T_\beta(w|X_{<t})$ transforms $w$ conditioned on the embedded sequence prefix $X_{<t}$, generating a context-sensitive matrix perturbation in embedding space. This perturbed representation, $\widetilde{X}_{<t} = X_{<t} + T_\beta(w|X_{<t})$, serves as the input context for next-token prediction by an autoregressive LLM $P_\theta$.

This structure enables both:

• Context-aware adaptation: The perturbation is a learned function of both context and a latent variable, facilitating distributional shifts tailored to the local content.
• Smooth, continuous augmentation: Embedding-space operations afford a greater degree of flexibility and coverage than discrete perturbations.

Optimization via Unbiased Estimating Equations

Because the overall marginal likelihood of observed data under this model structure requires integration over latent perturbations, direct maximum likelihood is intractable. Following score-based estimation theory, the paper derives an unbiased estimating equation framework where the gradient vanishes at the true parameter set. Stochastic gradient descent is employed on an empirical loss formed as a difference of log-likelihoods between observed and synthetically generated tokens under sampled perturbations.

The theoretical analysis addresses overparameterization, relaxing the usual identifiability assumption and allowing for a set of stationary points $\Gamma^*$. The results establish:

• Statistical consistency: Under compactness, differentiability, and mild regularity assumptions, the estimators converge (in probability) to the set of true parameters in the large-sample limit.
• Convergence rate: Explicit separation of statistical estimation error ($O_P(n^{-1/2})$) and optimization error ($O_P(\|\Psi(\widehat{\gamma})\|)$), the latter accounting for finite-iteration SGD effects.

Training and Inference Procedure

The procedure involves, at each training step, sampling perturbation latents, applying the perturbation module, and constructing both real and model-sampled tokens for unbiased estimation. At inference, generative decoding is performed autoregressively, continually sampling new perturbations at each generation step.

Empirical Evaluation

Synthetic Language Modeling

Simulated bigram LLMs with embedding perturbations parameterized by neural networks evaluate the methodology under known, controllable OOD shifts. Out-of-sample mean absolute error (MAE) on transition matrices is the primary metric.

• Results: Continuous, trainable perturbations consistently outperform discrete baseline augmentations across all vocabulary sizes.
• Sensitivity analysis: The performance gap narrows with increasing vocabulary size, likely due to diminishing informativeness of fixed-dimensional embeddings for smooth generalization.

Real-World Datasets

Experiments employ WikiText-2 for training, with evaluation on both in-domain (WikiText-2 test, WikiText-103) and OOD domains (WebText, WritingPrompts, CodeParrot, GermanQuAD). Representative decoder-only Transformers (OPT, Qwen3, GPT-Neo) are used as base models, with comparison against standard MLE, discrete perturbations [Cen et al., 2026], and embedding-level perturbation NEFTune.

Key metrics:

• Perplexity (PPL): Lower is better.
• Mauve: Higher indicates better distributional similarity to human text.

Highlights:

• Strong OOD performance: The continuous perturbation framework achieves substantially lower PPL and often higher Mauve on all OOD settings, for all model architectures. On in-domain test sets, performance is mostly competitive, occasionally ceding Mauve to fixed baselines.
• Debiasing effect: Incorporating a debiasing correction in estimation provides improved stability and further reduction in PPL for certain models.
• Comparison with strong baselines: NEFTune and discrete approaches are consistently outperformed in OOD generalization, with marked margin.

Ablation Studies

Ablations varying the presence of perturbations during training versus inference demonstrate that only jointly applying perturbation during both phases maximizes extrapolative benefit. Applying perturbations solely in training or inference degrades both PPL and Mauve.

Qualitative Analysis

Generated outputs on highly OOD tasks (e.g., GermanQuAD evaluated with a model trained on English Wikipedia) reveal:

• Baselines exhibit semantic drift, degenerate token repetitions, and mixed-language contamination.
• Continuous perturbation methods maintain topic relevance, better structure, and more coherent outputs, mitigating OOD degeneration.

Theoretical and Practical Implications

The proposed method establishes a principled, theoretically-justified design for context-dependent, continuous data perturbation integrated with LLM training. Unlike fixed or discrete augmentations, the learned perturbation process is adaptable, optimally leveraging both the structure of the embedding space and the dynamics of the base model for robust OOD generalization.

Key implications include:

• Statistical efficiency: The unbiased estimator guarantees optimal utilization of the full empirical distribution, avoiding the biases and inefficiencies of heuristic augmentations.
• Generalizability: The perturbation module serves as a plug-in adaptation layer and can, in principle, be retrofitted onto a variety of base LLM architectures with minimal additional computational overhead.
• Reduction of OOD collapse: By encouraging exploration in semantically meaningful directions during both training and inference, the method mitigates the memorization and mode-collapse pathologies typical in base LLMs under domain shift.

Future Directions

Potential avenues:

• Extension to instruction-finetuning and alignment: Incorporation of the perturbation framework in human preference-based post-training and RLHF.
• Latent structure exploration: Adaptive, hierarchical, or more structured latent perturbation processes could yield further gains in complex generalization.
• Robustness to adversarial or distributional attacks: Augmented resilience in safety-critical or multilingual settings.

Conclusion

This work formalizes and tests a continuous, data-driven perturbation paradigm for LLMs, bridging the gap between heuristic augmentation and fully-learned domain adaptation. The method is supported through rigorous theoretical guarantees and demonstrates pronounced empirical effectiveness in improving LLM extrapolation, particularly under distributional shift.