cs.LGMar 12, 2026
SaveMatching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models
Organizations: Harvard University, Kempner Institute · MBZUAI · Microsoft Research New England
Abstract
Cross-entropy (CE) training provides dense and scalable supervision for language models, 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.
Figures & tables
Figure 1: Feature-matching loss grows with completion length. Conditional feature-matching loss (lower is better) as a function of completion length for Qwen2.5-1.5B fine-tuned with SFT on OpenCodeInstruct ( Ahmad et al., 2025 ) . Although this increase is expected even under a perfect model due to growing feature variance, part of the degradation reflects SFT’s inability to calibrate the model’s rollout distribution over long horizons.
Figure 2: EBFT achieves the lowest feature-matching loss across all completion lengths. Despite training with a rollout horizon of only 8 tokens, EBFT’s gains persist and grow at longer completions. RLVR worsens this loss relative to the base model.
Figure 3: EBFT improves downstream performance without sacrificing distributional calibration. From left to right, we plot HumanEval accuracy (greedy and pass@16), validation cross-entropy (CE), and conditional feature-matching (CFM) loss over training for Qwen2.5-1.5B fine-tuned on OpenCodeInstruct ( Ahmad et al., 2025 ) . SFT improves cross-entropy and CFM loss but lags on downstream accuracy. RLVR improves downstream accuracy but substantially degrades both calibration metrics relative to the base model (dashed line). EBFT achieves the best results across all four metrics, avoiding this tradeoff. CE and CFM losses are computed on a 1k-samples held-out subset of OpenCodeInstruct.
Figure 4: Overview of Energy-Based Fine-Tuning (EBFT). For each context , the generator samples completions. A frozen feature network embeds each prompt–completion pair, producing features for the sampled completions and for the ground truth. Each completion receives a feature-matching reward measuring alignment with the ground-truth feature moment, and the generator is updated via REINFORCE with an RLOO baseline.
Algorithm 1 One EBFT training iteration.
| Q&A Coding | Unstructured Coding | |||||||||||
| Method | CE | FM | greedy | pass@1 | pass@4 | pass@16 | CE | FM | greedy | pass@1 | pass@4 | pass@16 |
| Base | 0.338 | 0.361 | 0.484 | 0.424 | 0.606 | 0.715 | 0.631 | 0.369 | 0.473 | 0.419 | 0.596 | 0.702 |
| Warm start | 0.301 | 0.344 | 0.483 | 0.440 | 0.611 | 0.723 | 0.499 | 0.317 | 0.508 | 0.458 | 0.638 | 0.743 |
| SFT | 0.289 | 0.315 | 0.483 | 0.455 | 0.617 | 0.728 | 0.501 | 0.321 | 0.504 | 0.467 | 0.644 | 0.747 |
| EBFT | 0.207 | 0.258 | 0.548 | 0.510 | 0.659 | 0.771 | 0.499 | 0.320 | 0.548 | 0.524 | 0.664 | 0.769 |
| EBFT (ws.) | 0.190 | 0.255 | 0.534 | 0.508 | 0.658 | 0.756 | 0.481 | 0.312 | 0.536 | 0.514 | 0.659 | 0.769 |
Table 1: EBFT outperforms SFT and matches or exceeds RLVR on downstream metrics, while achieving the best distributional calibration across all tasks. Best results per method on Q&A coding, unstructured coding, and translation. CE: validation cross-entropy; FM: feature-matching loss (both lower is better). “ws.”: warm-started from an SFT checkpoint. RLVR is inapplicable to unstructured coding where no verifier exists; EBFT still yields substantial gains over SFT in this setting. See Table 6 for per-benchmark results and Section 3 for full experimental details.
Figure 5: On translation, EBFT outperforms both SFT and RLVR on downstream accuracy, cross-entropy, and feature-matching loss. From left to right, we plot COMET scores on WMT22 and MTNT, validation cross-entropy, and CFM loss over training for Llama-3.2-1B fine-tuned on ALMA ( Xu et al., 2023 ) . EBFT achieves the lowest CE and CFM losses and matches SFT on WMT22 while clearly outperforming it on MTNT. RLVR underperforms SFT on all four metrics, with cross-entropy rising well above the base model (dashed line).
Figure 6: The CE regularization weight controls cross-entropy reduction without affecting downstream performance or feature-matching loss. We ablate EBFT with on Qwen2.5-1.5B. Dashed and dotted lines indicate the base model and 2-epoch SFT, respectively. Larger accelerates CE reduction while downstream accuracy and CFM loss remain nearly identical across settings. Even pure feature matching ( ) surpasses SFT on cross-entropy, confirming that the two objectives are aligned rather than in tension.
Figure 7: Feature network ablations: whitening and last-token pooling matter most; scaling the feature network does not help. HumanEval accuracy, validation cross-entropy, and CFM loss over training for EBFT ( ) on Qwen2.5-1.5B with different feature network configurations. The default (last-token features with whitening from a frozen 1.5B copy) achieves the best downstream accuracy and CFM loss. Removing whitening and mean pooling cause the largest degradations. Random weights hurt only modestly, and replacing the 1.5B feature network with a frozen Qwen2.5-7B yields similar results, suggesting that pre-trained representations help but that naively scaling the feature network does not.
Figure 8: EBFT improvements are consistent across model scales. HumanEval accuracy, validation cross-entropy, and CFM loss over training for EBFT ( ) applied to Qwen2.5-1.5B, 3B, and 7B. Each model uses a frozen copy of itself as the feature network. Dashed lines indicate base model performance. All three scales show substantial and qualitatively similar improvements across all four metrics, with no sign of diminishing returns.
Appendix figures & tables17 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 9: Strided parallel rollouts for a sequence of length , stride and completion length .
| \displaystyle-\left[\begin{array}[]{cccccccc|cc|cc|cc}0&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty\\ 0&0&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty\\ 0&0&0&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty\\ 0&0&0&0&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty\\ 0&0&0&0&0&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty\\ 0&0&0&0&0&0&\infty&\infty&\infty&\infty&\infty&\infty&\infty&\infty\\ 0&0&0&0&0&0&0&\infty&\infty&\infty&\infty&\infty&\infty&\infty\\ 0&0&0&0&0&0&0&0&\infty&\infty&\infty&\infty&\infty&\infty\\ \hline\cr 0&0&0&0&\infty&\infty&\infty&\infty&0&\infty&\infty&\infty&\infty&\infty\\ 0&0&0&0&0&0&0&0&\infty&0&\infty&\infty&\infty&\infty\\ \hline\cr 0&0&0&0&\infty&\infty&\infty&\infty&0&\infty&0&\infty&\infty&\infty\\ 0&0&0&0&0&0&0&0&\infty&0&\infty&0&\infty&\infty\\ \hline\cr 0&0&0&0&\infty&\infty&\infty&\infty&0&\infty&0&\infty&0&\infty\\ 0&0&0&0&0&0&0&0&\infty&0&\infty&0&\infty&0\\ \end{array}\right] |
Figure 10: Custom attention matrix in for a sequence of length , stride and completion length . When the entry is 0, the token in position attends to the token in position , and when it is it does not.
| Parameter | Value | ||
|---|---|---|---|
| Q&A Code | Unstructured Code | Translation | |
| Rollout Batch Size | 16 | ||
| Sequence Length | 1024 | ||
| Completion Length | 8 | 8 | 4 |
| Stride | 8 | 8 | 2 |
| Actor Learning Rate | |||
Table 2: Hyperparameter Configuration for EBFT runs.
| Parameter | Value |
|---|---|
| Training Batch Size | 64 |
| Epochs | 1 |
| Max Length | 2048 |
| Learning Rate | |
| Scheduler | Warmup + Cosine Decay to 0.1 lr |
| Warmup | 0.03 |
Table 3: Hyperparameter Configuration for SFT warmstart runs.
| Parameter | Value | ||
|---|---|---|---|
| Q&A Code | Unstructured Code | Translation | |
| Training Batch Size | {64, 128} | ||
| Max Length | 2048 | ||
| (Learning Rate, Scheduler) | {( , Warmup + Constant), ( , Warmup + Cosine Decay)} | ||
| Warmup | 0.03 | ||
| Adam Betas | (0.9, 0.95) | ||
Table 4: Hyperparameter Configuration for SFT baseline runs.
| Parameter | Value | |
|---|---|---|
| Q&A Code | Translation | |
| Rollout Batch Size | {32, 64} | {128, 256} |
| Prompt Max Length | 1024 | |
| Generate Max Length | 1024 | |
| Actor Learning Rate | ||
| Temperature | 1.0 | |
Table 5: Hyperparameter Configuration for RLVR baseline runs.
Figure 11: On Q&A Coding, increasing the cross-entropy weight consistently lowers both validation cross-entropy and feature-matching losses, while SFT on full sequences yields faster initial downstream gains that quickly degrade. We sweep and for EBFT initialized from base Qwen2.5-1.5B. As a baseline, we compare against SFT trained on full sequences (solid red), whereas the main text reports SFT trained only on the answer. SFT on full sequences improves downstream performance faster early on but quickly deteriorates, and answer-level cross-entropy rises. Setting and (blue dotted) causes both cross-entropy and moment-matching losses to increase and leads to degraded pass@1 and pass@ scores, indicating that both terms are necessary for stable training. We also report feature-matching loss with and without whitening; the non-whitened variant tracks more closely with cross-entropy, which is why we use it for comparison. Overall, increasing has limited effect on downstream metrics but helps decrease both the cross-entropy and moment-matching objectives. The validation set is a 1k-sample held-out subset of OpenCodeInstruct.
Figure 12: On unstructured coding, EBFT achieves better downstream performance than SFT while reaching comparable levels on cross-entropy and feature-matching losses, even as downstream accuracy plateaus before these losses converge. We sweep and for EBFT initialized from base Qwen2.5-1.5B. Cross-entropy and feature-matching losses continue to decrease throughout training even after downstream accuracy has plateaued, which we attribute to distribution shift between the training data and the downstream evaluation benchmarks. Setting (blue dotted) again leads to diverging cross-entropy, confirming the importance of the cross-entropy term. SFT (solid red) achieves the lowest cross-entropy and feature-matching losses but at the cost of weaker downstream performance compared to EBFT configurations with . The validation set is a 1k-sample held-out subset of the unstructured coding data.
Figure 13: On translation, EBFT matches or exceeds RLVR on downstream COMET scores—outperforming it on MTNT—while consistently surpassing SFT across all benchmarks and evaluation settings. We sweep and for EBFT initialized from base Llama-3.2-1B. On WMT’22, EBFT performs comparably to or slightly below RLVR, while on MTNT it clearly outperforms RLVR, and on OpenSubtitles the two are similar. EBFT surpasses SFT across all benchmarks and decoding strategies (greedy, best-of-1/4/16). Unlike in the coding setting, using and (blue dotted) leads to degraded downstream performance, indicating that both terms are important for translation. SFT is trained only on answers, as training on full sequences led to worse downstream results. On cross-entropy and feature-matching losses, EBFT achieves lower values than SFT, with larger configurations (green) reaching the best feature-matching levels. The legend is shared with the previous figures. The validation set is a 1k-sample held-out subset of ALMA.
Figure 14: On translation BLEU scores, EBFT consistently outperforms both SFT and RLVR, while RLVR degrades at best-of-16 later in training—consistent with the distribution sharpening induced by RLVR. We report BLEU scores for the same and sweep as the previous figure, initialized from base Llama-3.2-1B. Although RLVR is more competitive on BLEU than on COMET, EBFT still achieves higher scores across most benchmarks and decoding strategies. Notably, at best-of-16 (temp. 0.6), RLVR performance drops in the second half of training, suggesting that at large the model suffers from reduced sample diversity, consistent with the distribution sharpening induced by RLVR ( Zhao et al., 2025 ) . EBFT configurations with do not exhibit this degradation and maintain stable or improving BLEU throughout training. The legend is shared with the previous figures.
Figure 15: HumanEval/18: counting overlapping substring occurrences. EBFT correctly advances by one character after each match and therefore counts overlapping occurrences. By contrast, the base model uses Python’s non-overlapping string.count , while both SFT and RLVR advance by len(substring) , which misses overlap-sensitive cases and causes test failures.
Figure 16: HumanEval/94: largest prime digit sum. EBFT produces a self-contained executable solution by defining is_prime and then computing the digit sum of the largest prime in the list. The base model and RLVR instead rely on an undefined helper function, while SFT truncates mid-function and yields syntactically invalid Python, illustrating a recurring executability gap between EBFT and the competing methods.
Figure 17: HumanEval/105: filtering, ordering, and lexical mapping. EBFT correctly implements the full prompt: it filters to digits 1 to 9, sorts the valid entries, reverses the order, and maps each digit to its English name. The other outputs each satisfy only part of the specification: the base model returns the input unchanged, SFT incorrectly returns the original array whenever any element is out of range instead of filtering, and RLVR outputs numeric strings rather than digit names.
Figure 18: HumanEval/69: selecting the greatest frequency-qualified integer. EBFT returns the maximum positive integer whose frequency is at least its value, matching the full prompt semantics. The base model fails to provide a usable implementation, SFT returns the first qualifying key rather than the greatest one, and RLVR follows the right high-level logic but is not executable under strict evaluation because the output includes formatting artifacts outside valid code.
Figure 19: MTNT EN FR: non-translation and multilingual runaway. EBFT remains on task and produces a direct French translation, whereas the base model and RLVR drift into multilingual template continuations that repeat the source sentence and append language-tagged variants. SFT avoids the runaway list structure but still fails to translate, remaining almost entirely in English.
Figure 20: MTNT EN FR: truncation after template drift. EBFT produces a complete on-task translation, while the base model drifts into unrelated multilingual continuation and SFT paraphrases too loosely away from the original meaning. RLVR begins a language-tagged template expansion and terminates mid-token, yielding an unfinished output that is incompatible with strict evaluation.
Figure 21: MTNT EN FR: semantic failure from dropped negation. Not all errors are formatting failures: this example shows that negation can be overlooked by the fine-tuned models. EBFT preserves the negative meaning of the source sentence, whereas RLVR drops the negation and changes “no more available rooms” into “more rooms.” The base model again exhibits multilingual repetition after an initially plausible translation, and SFT produces a weaker paraphrase that underspecifies the original statement.