Feature Matching for Language Model Fine-Tuning: Moving Beyond Teacher Forcing

Introduction

The dominant paradigm for training large language models rests on a subtle but consequential sleight of hand. During supervised fine-tuning, we optimize models to predict the next token while conditioning on ground-truth prefixes, a technique known as teacher forcing. Yet at deployment, these same models must condition on their own generated tokens, creating a feedback loop where early errors compound over long sequences. This discrepancy, long observed in the literature on exposure bias, means that a model can achieve near-optimal perplexity on held-out text while producing incoherent rollouts in practice.

In their recent work, Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models, Jelassi et al. identify the root cause as a mismatch between token-level supervision and sequence-level behavior. Rather than treating fine-tuning as a next-token prediction task, they frame it as a distributional alignment problem in an embedding space. Their proposed method, Energy-Based Fine-Tuning (EBFT), directly supervises the expected feature statistics of model-generated completions, offering a principled alternative to both standard supervised fine-tuning (SFT) and reinforcement learning with verifiable rewards (RLVR).

The Distribution Shift in Autoregressive Models

The technical core of the exposure bias problem lies in the divergence between training and inference distributions. Under teacher forcing, the model sees only valid context windows drawn from the data distribution. During generation, however, the model samples from its own predictive distribution, and any deviation from the data manifold shifts subsequent conditioning contexts into regions where the model has never trained. Braverman et al. (2019) quantified this phenomenon by measuring the expected conditional entropy of the k-th token as a function of its position k. For a perfectly calibrated model, this entropy should remain stable; empirically, it increases with sequence position even for models with low training loss.

Jelassi et al. operationalize this drift through a conditional feature-matching loss. For each prompt, they compute the mean feature embedding of the model's rollout distribution and compare it to the mean embedding of the ground-truth completion using squared error. A model is considered feature-calibrated when this loss vanishes, indicating that while individual tokens may vary, the expected semantic content of generations matches the data distribution. Their experiments with Qwen2.5-1.5B on OpenCodeInstruct reveal that SFT exhibits monotonically increasing feature-matching loss as completion length grows from 1 to 64 tokens. This degradation is not merely variance inflation; it reflects SFT's fundamental inability to calibrate rollout behavior over long horizons.

Standard RLVR methods attempt to solve this by optimizing scalar rewards under model rollouts, but this introduces its own pathologies. While RLVR improves downstream task accuracy, the paper demonstrates that it simultaneously worsens both validation cross-entropy and feature-matching loss. The model learns to exploit the verifier rather than align its generative distribution, effectively trading distributional calibration for task-specific optimization.

Energy-Based Fine-Tuning: Methodology and Theory

EBFT addresses these limitations by optimizing the feature-matching loss directly, treating it not as a diagnostic but as the training objective. The method relies on three technical components that merit detailed examination.

First, the authors employ a frozen feature network φ, initialized from the pretrained model, which embeds concatenated prompt-completion sequences into a semantic space. This embedding function serves as a sufficient statistic; by matching moments in this space, the model preserves high-level semantic coherence without requiring token-level identity.

Second, to make this optimization tractable, they introduce strided block-parallel sampling. Rather than generating full sequences serially, the method generates multiple rollouts from nested prefixes concurrently. This allows for efficient batched feature extraction over partial completions, addressing the computational bottleneck of sequential generation during training.

Third, optimization proceeds via an on-policy policy-gradient update using a REINFORCE-style estimator. The gradient updates push the generator p_θ toward feature embeddings that match the reference distribution, effectively performing moment matching in the embedding space.

Theoretical analysis reveals that EBFT implements KL-regularized feature-matching. Under this view, the feature-matching objective implicitly defines an energy function over sequences, with the optimal policy taking the form of an exponential tilt of the base model. This connection to energy-based modeling is significant; it suggests that EBFT is not merely a heuristic sampling strategy but a principled approach to density ratio estimation in high-dimensional sequence spaces.

Empirical Tradeoffs and Behavioral Calibration

The experimental results reveal a complex landscape of tradeoffs between fluency, calibration, and task performance. Across question-answering coding, unstructured coding, and translation tasks, EBFT matches RLVR in downstream accuracy while achieving lower validation cross-entropy than both RLVR and SFT. This is a notable finding. Typically, task-specific fine-tuning increases log-perplexity as the model sacrifices general language modeling capacity for specialized performance. EBFT appears to avoid this sacrifice by maintaining distributional alignment.

Particularly striking is the generalization behavior across completion lengths. EBFT models are trained with a rollout horizon of only 8 tokens, yet their feature-matching gains persist and grow at lengths of 64 tokens and beyond. This suggests that local calibration induces global coherence; by ensuring that short rollouts match feature statistics, the model learns internal representations that propagate stability through the autoregressive chain.

In contrast, RLVR exhibits deteriorating feature-matching loss relative even to the base model. This provides concrete evidence that scalar reward optimization can induce distributional shift away from natural language manifolds, even when it improves verifiable metrics. The model becomes confident in specific answer formats while losing calibration over the broader space of possible completions.

Original Insights and Limitations

The shift from token matching to feature matching represents a conceptual reorientation with implications extending beyond the specific algorithm proposed. Current alignment methods, including Direct Preference Optimization (DPO) and its variants, implicitly learn reward functions that are scalar-valued and typically trained on pairwise preferences. EBFT suggests an alternative pathway: instead of learning to rank completions, we could learn to match their statistical properties in a learned feature space. This connects to broader themes in generative modeling, particularly the success of energy-based models and flow matching in continuous domains such as image generation.

However, several limitations warrant careful consideration. The method's efficacy depends critically on the choice of feature network φ. If the frozen embeddings fail to capture task-relevant structure, moment matching in embedding space may not correlate with downstream utility. The paper uses the pretrained model's own representations, which provides a reasonable prior, but there is no guarantee that these features are optimal for all fine-tuning tasks.

Computational cost presents another constraint. While strided block-parallel sampling improves efficiency, EBFT still requires generating multiple rollouts per training example, making it more expensive than standard SFT. For large-scale applications, the tradeoff between training cost and improved calibration must be evaluated against the alternative of simply collecting more supervised data.

Furthermore, the theoretical guarantees assume sufficiently rich features. In practice, feature collapse, where the model learns to produce distinct sequences that map to identical embeddings, remains a risk. The paper does not extensively explore adversarial robustness of the feature-matching objective, and it remains unclear how the method behaves under distribution shift where the test-time feature statistics differ from training.

Conclusion and Open Questions

The work of Jelassi et al. provides a rigorous framework for thinking about fine-tuning as distributional alignment rather than next-token prediction. By supervising sequence-level feature statistics instead of individual token probabilities, EBFT offers a middle path between the distributional drift of SFT and the over-optimization risks of RLVR.

Looking forward, several questions remain open. Can the feature network φ be learned or adapted during fine-tuning without collapsing to trivial solutions? How does EBFT scale to very long horizon tasks, such as book-length generation or multi-step reasoning, where feature variance grows substantially? And perhaps most intriguingly, can feature-matching objectives be integrated with process-based reward models to provide both semantic calibration and step-level verification?

As language models increasingly serve as long-form generators and reasoning engines, the distinction between local token probability and global sequence coherence grows ever more consequential. Methods that supervise the statistical structure of generations, rather than just their token-level likelihood, may prove essential for building models that remain stable, calibrated, and controllable over arbitrary time horizons.