Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models

Abstract: Cross-entropy (CE) training provides dense and scalable supervision for LLMs, but it optimizes next-token prediction under teacher forcing rather than sequence-level behavior under model rollouts. We introduce a feature-matching objective for language-model fine-tuning that targets sequence-level statistics of the completion distribution, providing dense semantic feedback without requiring a task-specific verifier or preference model. To optimize this objective efficiently, we propose energy-based fine-tuning (EBFT), which uses strided block-parallel sampling to generate multiple rollouts from nested prefixes concurrently, batches feature extraction over these rollouts, and uses the resulting embeddings to perform an on-policy policy-gradient update. We present a theoretical perspective connecting EBFT to KL-regularized feature-matching and energy-based modeling. Empirically, across Q&A coding, unstructured coding, and translation, EBFT matches RLVR and outperforms SFT on downstream accuracy while achieving a lower validation cross-entropy than both methods.

• The paper introduces an energy-based fine-tuning method that aligns model feature moments rather than token probabilities.
• It demonstrates improved sample efficiency, stable loss curves, and superior alignment with long-sequence ground-truth statistics.
• Empirical results in code generation and translation tasks validate the approach's advantage over standard cross-entropy optimization.

Matching Features, Not Tokens: Energy-Based Fine-Tuning of LLMs

Introduction

The paper "Matching Features, Not Tokens: Energy-Based Fine-Tuning of LLMs" (2603.12248) critiques the dominant paradigm of conditional next-token log-likelihood (cross-entropy, XE) optimization for LLM fine-tuning and proposes an alternative based on aligning internal model features. The authors identify both practical and theoretical deficiencies in XE loss, motivating the use of an energy-based framework that leverages conditional feature statistics. This approach aims to improve sample efficiency, alignment with the intended generative distribution, and ultimate model performance, particularly in the context of code generation and long-form content.

Deficiencies of Conditional Next-Token Log-Likelihood

A detailed analysis reveals several issues with XE as a fine-tuning objective:

• Variance problem: XE exhibits undesirable sample variance that grows rapidly with output sequence length, complicating credit assignment and requiring larger batch sizes or specialized optimization protocols.
• Statistical alignment: XE targets the conditional token-level distribution, which deviates from matching desired feature or moment statistics of the full sampled output distribution for autoregressive models, especially in the long-context and highly entropic regimes.
• Inefficiency for long completions: For tasks involving long output sequences (such as code generation and open-ended dialogue), XE loss becomes less efficient and less reliable as a training signal.

These theoretical concerns are empirically validated by measuring conditional feature-matching loss as a function of output length, with clear evidence that XE optimization leads to mismatch in moment statistics—these mismatches scale sublinearly with sequence length under XE but can remain constant or shrink when directly matching feature moments.

Figure 1: Conditional feature-matching loss (lower is better) for Qwen2.5-1.5B fine-tuned with SFT on OpenCodeInstruct as a function of completion length, revealing poor alignment with longer completions under XE.

Energy-Based Fine-Tuning via Conditional Feature Moment Matching

The authors introduce an alternative objective inspired by conditional energy-based model (EBM) theory. Rather than matching the token-level conditional log-probabilities, the model is fine-tuned to match trajectory-level statistics by minimizing the difference between feature moments (computed on hidden activations at specific layers/positions) for samples from the model and the ground-truth completions.

The proposed technique formalizes the following:

• Conditional feature moment matching: For a given (context, target) pair, sample trajectories from the model, compute relevant moment features (layerwise means, variances, token-level statistics), and minimize the mean-squared error with respect to reference (ground-truth) features.
• Contrastive estimation: In addition to feature-matching losses, the energy-based perspective supports standard contrastive losses and negative mining to further regularize and accelerate alignment.
• Sampling protocol: Rollouts of the model are obtained during fine-tuning, and stochasticity in generation is essential to faithfully capture the output distribution.

Mathematically, this yields a loss of the form:

$$L_\text{FM} = \mathbb{E}_{x \sim \text{data},\, y \sim p_\theta(\cdot|x)} \left[ \|f(y) - f(y^*)\|^2 \right]$$

where $y^*$ is the ground-truth target, $y$ is a model sample, and $f(\cdot)$ extracts token- or sequence-level features.

Empirical Results

Extensive experiments on code generation and translation tasks demonstrate that models fine-tuned with the proposed feature-matching loss outperform XE-SFT baselines in both conditional feature-matching loss and downstream evaluation metrics. A particularly significant finding is that the feature-matching loss increases rapidly with completion length under standard XE-based SFT, whereas minimization under the proposed method yields substantially lower and more stable loss across the trajectory.

The results demonstrate:

• Superior alignment between model samples and ground-truth statistics for long completions.
• Improved sample quality as measured by both automatic metrics and human rater evaluations in domains where the structure is less semantically localized (e.g., code, translation).
• Stable and interpretable loss curves, indicating improved optimization properties and more robust convergence, particularly as trajectory length increases.

Theoretical and Practical Implications

The energy-based feature-matching paradigm addresses core limitations of autoregressive LM fine-tuning, particularly for open-ended, high-entropy generative regimes. By forgoing the requirement for strict token-level matching, and instead regularizing distributional properties of entire generated sequences, this methodology enables:

• Better credit assignment for long-horizon dependencies, as loss is attributed to sequence-level statistics rather than decomposed token likelihoods.
• Greater robustness and sample efficiency, as the training signal is less sensitive to the heavy-tailed variance of log-probabilities.
• Potential for new RL-HF alternatives: Feature-matching enables learning meaningful reward functions or constraints at the representation level, obviating reinforcement learning over token-level rewards in some cases.

Additionally, the approach provides theoretical foundations for research in conditional EBMs for language, opening avenues for integrating diverse modalities of statistical regularization.

Future Directions

The adoption of feature-based rather than token-based objectives suggests several natural directions for further investigation:

• Scaling laws for feature-matching: Establishing theoretical underpinnings for gradient and sample variance in feature matching as a function of model size and output entropy.
• Extensibility to multi-modal, multi-task architectures: As feature matching is agnostic to the specific output modality, it is a promising candidate for aligning image, audio, and structured output decoders.
• Integration with RL-based alignment: Composition with human feedback and interaction data is streamlined with feature-level losses, suggesting hybrid algorithms that blend EBM and RL approaches for controllable, interpretable LLM alignment.

Conclusion

This work systematically exposes inherent limitations of cross-entropy-based fine-tuning for autoregressive LLMs, particularly in the long-horizon, high-entropy scenario. The authors propose and theoretically ground energy-based fine-tuning via conditional feature moment matching, providing both conceptual clarity and practical performance gains. The methodology yields improved sample efficiency, loss stability, and output quality, establishing conditional feature alignment as a viable and robust alternative to per-token objectives in LM fine-tuning (2603.12248).