Energy-Based Fine-Tuning (EBFT)

• Energy-Based Fine-Tuning (EBFT) is a method that augments conventional fine-tuning by incorporating energy-based modeling and noise contrastive estimation to enhance sequence-level alignment and model calibration.
• It employs distinct energy function variants—scalar, hidden, and sharp-hidden—that offer different trade-offs between computational efficiency and calibration accuracy.
• The EBFT pipeline integrates MLM-trained noise generation and joint loss optimization, resulting in reduced calibration error while maintaining robust performance on various NLP tasks.

Energy-Based Fine-Tuning (EBFT) refers to a family of methods that incorporate energy-based modeling principles into the fine-tuning of LLMs and text encoders. In these approaches, training objectives augment conventional cross-entropy losses with additional terms that explicitly model input likelihoods or alignment with sequence-level feature statistics, leveraging energy-based models (EBMs) and noise-contrastive frameworks. The effect is to promote improved calibration, better sequence-level alignment, and robust downstream generalization on tasks such as natural language understanding, code generation, and machine translation (Jelassi et al., 12 Mar 2026, He et al., 2021).

1. Theoretical Foundations

EBFT grounds its methodology in energy-based modeling, where learning is guided by assigning lower energy to desirable data points (e.g., correctly classified inputs, plausible sequences) and higher energy to non-data or "noise" points. In conventional discriminative finetuning, models minimize expected negative log-likelihood of the correct label given input, $$L_{CE} = \mathbb{E}_{(x, y)\sim\mathcal{P}_D}[-\log P_\theta(y|x)]$$. EBFT augments this by also modeling the marginal likelihood of $x$ via an EBM head, typically parameterized as a scalar function $\hat{E}_\theta(x)$ acting on encoder representations.

Noise-Contrastive Estimation (NCE) is used to bypass intractable partition functions over discrete text, optimizing a loss $L_{NCE}$ that encourages $\tilde{P}_\theta(x) = \exp(-\hat{E}_\theta(x))$ to separate data samples $x^+$ from noise samples $x^-$. The joint loss is:

$L_{\text{joint}} = L_{\text{CE}} + L_{\text{NCE}},$

where $L_{\text{NCE}}$ takes the form (with $K$ noise samples per data sample):

$x$0

By parameterizing $x$1, the NCE loss can be simplified and the noise model’s effect analytically cancelled (He et al., 2021).

2. Energy Function Variants

EBFT deploys several parameterizations for the EBM head:

• Scalar variant: A linear projection $x$2 on the encoder output, i.e., $x$3.
• Hidden (LogSumExp) variant: The negative log-sum-exp of output logits, $x$4.
• Sharp-hidden (Max) variant: The negative maximum logit, $x$5.

The hidden variant is particularly notable for its tight coupling to the classifier decision surface, yielding strong calibration performance. The sharp-hidden max-over-logits variant is a computationally cheaper approximation, and the scalar variant provides implementation simplicity.

Variant	Definition	Calibration/Accuracy Trade-off
Scalar	$x$6	Simpler, still effective
Hidden	$x$7	Strongest calibration
Sharp-hidden	$x$8	Slightly higher accuracy

3. Noise Model Design and NCE Implementation

Effective contrastive learning in the text domain requires the negative samples (noise) to be neither trivial nor uninformative. EBFT commonly leverages an autoregressive GPT-2 small model as the noise generator, but crucially, this generator is finetuned with a masked language modeling (MLM) objective applied to the target task domain. For each training instance, a mask ratio $x$9 (typically 0.4) is used to randomly mask tokens, and the model is trained to reconstruct the full sequence, i.e., minimizing $\hat{E}_\theta(x)$0.

At NCE training time, noise samples are generated by re-masking fresh instances and using the MLM-trained GPT-2 with top-$\hat{E}_\theta(x)$1 sampling (commonly $\hat{E}_\theta(x)$2) to fill in masked spans. This procedure ensures $\hat{E}_\theta(x)$3 better matches the data distribution, resulting in a more stable and informative contrastive signal (He et al., 2021).

4. Practical EBFT Training Pipeline

A typical EBFT pipeline proceeds as follows:

• Initialize a pretrained encoder (e.g., RoBERTa) and attach both a classification head and an EBM head.
• For each minibatch:
  • Draw real data samples $\hat{E}_\theta(x)$4 and corresponding labels from the task distribution.
  • Generate $\hat{E}_\theta(x)$5 negative samples $\hat{E}_\theta(x)$6 per data point via the MLM-trained noise model.
  • Compute $\hat{E}_\theta(x)$7 over the batch.
  • Compute $\hat{E}_\theta(x)$8 contrasting real and noise samples.
  • Backpropagate the sum, $\hat{E}_\theta(x)$9.
• Use the Adam optimizer and maintain standard finetuning hyperparameters (batch sizes of $L_{NCE}$0–$L_{NCE}$1, learning rate $L_{NCE}$2).
• Recommended hyperparameters are $L_{NCE}$3, $L_{NCE}$4, and top-$L_{NCE}$5 for noise; increasing $L_{NCE}$6 beyond $L_{NCE}$7 yields diminishing calibration returns.

This pipeline adds a modest computational overhead (approximately $L_{NCE}$8 forward passes for noise), but the benefits in calibration are empirically robust.

5. Empirical Evaluation and Task Performance

EBFT has been evaluated primarily on the GLUE suite of NLU tasks (excluding STS-B), with accuracy (and Matthews Correlation Coefficient for CoLA) assessed alongside Expected Calibration Error (ECE, with $L_{NCE}$9 bins). Across multiple tasks, EBFT variants consistently reduce ECE by approximately half, with a decrease of less than $\tilde{P}_\theta(x) = \exp(-\hat{E}_\theta(x))$0 in accuracy—or a slight improvement in some settings.

Moreover, performance of the hidden energy variant matches or exceeds post-hoc calibration methods that employ separate dev sets for calibration (e.g., temperature scaling, scaling-binning calibrators, Posterior-Calibrated Training) (He et al., 2021). Calibration–accuracy trade-off analysis demonstrates that EBFT converges to optima with both high accuracy and strong calibration, while baseline finetuning models often over-optimize for accuracy at the cost of miscalibration.

In sequence-level tasks such as Q&A coding, unstructured code generation, and machine translation, feature-matching-based EBFT (as proposed by (Jelassi et al., 12 Mar 2026)) achieves downstream accuracies surpassing those of supervised finetuning (SFT), while achieving lower validation cross-entropy and performance competitive to reward learning methods such as RLVR.

6. Methodological Extensions and Theoretical Connections

Theoretical analyses motivate EBFT as an instantiation of KL-regularized feature-matching objectives, linking energy-based modeling with policy-gradient fine-tuning under on-policy rollouts. Strided block-parallel sampling is leveraged to efficiently perform sampling across multiple prefixes, which allows dense feature extraction and batched rollouts for efficient sequence-level inference (Jelassi et al., 12 Mar 2026). This approach contrasts explicit preference modeling or reliance on task-specific verifiers, aligning the model distribution with desirable sequence-level statistics through direct feature distribution matching.

A plausible implication is that EBFT provides a framework for sequence-level alignment in LLMs that is dense, stable, and does not depend on explicit human or automated preference signals.

7. Practical Insights, Pitfalls, and Recommendations

Variant selection depends on desired trade-offs: the hidden (LogSumExp) variant maximally leverages the relationship to classifier decisions for best calibration, while the sharp-hidden (max-based) variant offers a more computationally expedient option with marginal performance trade-off. The scalar variant remains an easy-to-implement choice that still achieves substantial improvements.

Noise generation using a non-MLM–trained LLM is discouraged due to insufficient challenge posed to the contrastive objective. Excessively high noise ratios ($\tilde{P}_\theta(x) = \exp(-\hat{E}_\theta(x))$1) may yield minimal additional calibration gains while unnecessarily increasing compute. It is also important to ensure the energy function possesses sufficient capacity to utilize the contrastive signal; under-parameterized scalar heads may be suboptimal.

EBFT represents a lightweight yet effective augmentation to classical encoder finetuning pipelines. Its joint objective and design yield robust improvements in model confidence calibration without sacrificing downstream accuracy.

(Jelassi et al., 12 Mar 2026, He et al., 2021)