Energy-Based Reward Models (EBRM)

• Energy-Based Reward Models (EBRM) are probabilistic frameworks that use energy functions to model reward distributions, supporting reinforcement learning, imitation learning, and language tuning.
• They employ contrastive training methods to filter noisy data and capture multi-modal human preferences, enhancing robustness and reward calibration.
• EBRMs offer unique maximum likelihood estimation properties, efficient post-hoc integration, and empirical gains across benchmarks in diverse RL and sequence modeling tasks.

Energy-Based Reward Models (EBRM) are a class of probabilistic models designed to capture the distributional structure of reward signals in machine learning, with particular application to alignment, reinforcement learning, imitation learning, and LLM tuning. Unlike scalar reward models that map input-output pairs to a single score, EBRMs introduce an explicit energy function that induces a conditional distribution over reward or value assignments. This approach allows for robust handling of uncertainty, label noise, multi-modality in human annotations, and improved generalization properties. EBRMs have found utility as post-hoc refinements for reward model calibration, as theoretically optimal structures in KL-regularized policy tuning, and as practical engines for preference modeling in RLHF pipelines, imitation learning, and sequence reranking.

1. Mathematical Formulation and Conditional Distribution

At the core of an EBRM is a parameterized energy function

$E_\theta(x, y, r)$

that acts over a prompt $x$, response $y$, and reward $r$. For a fixed embedding $e(x, y)$ from a base model, the conditional reward distribution is given by

$$p_\theta(r \mid x, y) = \frac{\exp(-E_\theta(x, y, r))}{Z_\theta(x, y)}$$

where the normalizer (partition function) is

$$Z_\theta(x, y) = \int_{\tilde r} \exp(-E_\theta(x, y, \tilde r))\, d\tilde r.$$

Alternately, one may use $f_\theta(e, r) = -E_\theta(e, r)$ as an unnormalized log-density. This construction generalizes beyond deterministic RM outputs, allowing EBRM to "push up" the energy surface around implausible rewards and "push down" near likely ones, thereby representing multi-modal or uncertain human preferences and reducing overconfident reward exploitation (Lochab et al., 17 Apr 2025). The same energy-based mechanism appears in EBM parameterizations for reward-conditioned policies, imitation learning, and preference modeling (Ding et al., 2023, Hong et al., 2024).

2. Contrastive Training and Data Filtering

EBRM training typically employs contrastive objectives that distinguish positive (preferred) samples from negative (less preferred or noise-induced) samples. For LLM alignment, conflict-aware data filtering is implemented: any pair where the pretrained RM contradicts the human label ($r^+ < r^-$) is excluded, yielding a cleaner proxy-reward set (Lochab et al., 17 Apr 2025). Noise-contrastive estimation (NCE) is adapted to handle label uncertainty by sampling noisy positives $r_i^{(0)}=r_i+\nu_i$ and multiple negatives $r_i^{(m)}\sim\mathcal{N}(r_i,\sigma^2)$, optimizing a log-softmax InfoNCE loss over a batch.

In preference alignment, the Energy Preference Alignment (EPA) loss generalizes the standard pairwise Bradley–Terry (BT) setting by contrasting each positive example against multiple strong and weak negatives (Hong et al., 2024). The EPA loss converges to the ideal energy discrepancy minimizer as the number of negatives increases, and uniquely identifies the correct slope-1 linear relationship between learned and true rewards.

3. Theoretical Properties and Inductive Bias

EBRMs provide rigorous guarantees not available to scalar or pairwise frameworks. For RLHF, the EBM maximum likelihood estimator (MLE) exists uniquely and enforces the required linearity $r_\theta(x, y) = r_{\mathrm{true}}(x, y) + C(x)$ under mild support conditions (Hong et al., 2024). In contrast, the BT model may admit multiple minimizers, failing the linearity criterion when certain responses are absent from paired comparisons. This theoretical distinction informs best practices for reward model selection in RLHF pipelines.

In KL-regularized RL, the optimal policy adopts a Gibbs (Boltzmann) structure

$\pi^*(a|x) \propto \exp(-E_\theta(x,a))$

with $E_\theta(x,a) = -R(x,a) - \log p_0(a|x)$, where $R$ is the reward and $p_0$ is the reference distribution (Tan et al., 21 Dec 2025). This energy-based construction yields detailed balance, monotonic KL convergence to stationarity, and explicit entropy–accuracy trade-offs.

4. Algorithmic Structures and Computational Efficiency

EBRM architectures are lightweight, typically adding only a small percentage (~3%) of parameters relative to the underlying RM (Lochab et al., 17 Apr 2025). Training is post-hoc, requiring ~465 seconds for 5 epochs over 70M embeddings on a single GPU; inference involves ≤50 gradient steps per example, with batch inference times comparable to or lower than ensemble or retrained models. Scalability is established across RMs with parameter counts from 70M to 8B (Lochab et al., 17 Apr 2025).

For reward-conditioned RL, energy-based Bayesian reparameterization explicitly decomposes the conditional

$$E_\theta(s, a, R) = -\log \bar\beta_\theta(a|s) - \log \bar\beta_\theta(R|s,a)$$

which enables adaptive inference under fixed data distributions, guards against out-of-distribution (OOD) RTG conditioning, and yields state-of-the-art scores on Gym-MuJoCo and Atari tasks (Ding et al., 2023).

5. Empirical Performance and Benchmark Gains

EBRM methods consistently outperform classical scalar or pairwise reward models in alignment and RL tasks:

Benchmark	Metric	Base RM/DPO	EBRM/EPA	Gain
RMB Harmlessness	Pairwise Acc	52.14	58.11	+5.97%
RMB Total	Mean Acc	42.66	45.98	+3.32
RewardBench	Avg Win Rate	55.10	56.16	+1.06
MT-Bench	GPT-4 Score	7.55	7.71	+0.16
Alpaca-Eval 2.0	Win Rate (%)	15.24	19.26	+4.02

Empirical studies highlight improved robustness to noisy preference annotations, delayed reward hacking in RL, and superior alignment at matched KL budget compared to DPO (Lochab et al., 17 Apr 2025, Hong et al., 2024).

6. Extensions to Sequence Modeling and Imitation Learning

EBRM structures generalize beyond RLHF and RL. In energy-based imitation learning (EBIL), the expert occupancy measure is represented as $\rho_E(s,a) \propto \exp(-E_\phi(s,a))$, the energy is estimated by score matching, and the learned reward $r_\phi(s,a) = -E_\phi(s,a)$ is directly used in policy optimization, bypassing adversarial IRL loops (Liu et al., 2020). This method is theoretically equivalent to Maximum Entropy IRL and delivers stable, interpretable solutions.

In neural sequence modeling, energy-based reranking uses an energy model $E_\theta(x, y)$ to directly score candidate translations, encouraging better task metrics (e.g., BLEU score) relative to MLE-trained NMT models. Joint-EBR models yield consistent +1–4 BLEU gains by ranking low-energy (high-quality) outputs (Bhattacharyya et al., 2020).

7. Practical Integration and Guidelines

EBRM implementations are compatible with standard RLHF and RL pipelines, requiring only extraction of embeddings/scores and training the top-layer EBM. Hyperparameters for post-hoc refinement include negative sample count (e.g., $M=768$), noise scales ($\sigma=3.5$, $\beta=0.1$), and optimization steps ($\leq 50$). EPA contrastive alignment recommends $K_{\mathrm{strong}}\geq 3$, $N_{\mathrm{weak}}\approx 2$–$10$, and batch size suitable for available memory (e.g., 64 for 7B models on 8×A100) (Lochab et al., 17 Apr 2025, Hong et al., 2024).

Best practices suggest preferring EBM/EPA models over BT/DPO when unique MLE and guaranteed linearity are required, or when abundant off-policy negatives can provide regularization for more stable alignment. Weak negatives and margin tricks further increase alignment quality, especially under high-KL regimes.

---

In summary, Energy-Based Reward Models offer principled, empirically validated solutions for representing, aligning, and leveraging reward uncertainty across diverse domains such as RLHF, reinforcement learning, sequence modeling, and offline alignment. EBRMs ensure unique MLE recovery, robust handling of noisy data, and straightforward integration into modern pipelines, establishing them as foundational tools for preference-based learning and control.