Energy-Based Fine-Tuning: Teaching LLMs to Match Meaning, Not Just Tokens

Paper: 2603.12248 Authors: Samy Jelassi, Mujin Kwun, Rosie Zhao, Yuanzhi Li, Nicolo Fusi, Yilun Du, Sham M. Kakade, Carles Domingo-Enrich Categories: cs.LG

The Gap

Language model fine-tuning has two dominant paradigms, and both have real problems.

SFT (supervised fine-tuning) trains on cross-entropy loss token by token. It’s cheap and stable, but it optimizes for copying the teacher’s exact token sequence under teacher forcing — meaning the model never sees its own mistakes compound. It’s great at mimicry, bad at generalization.

RLVR (reinforcement learning with verifiable rewards) fixes this by letting the model roll out full sequences and scoring them with a reward signal. It actually trains on the distribution the model produces at inference time. But it requires a verifier — something that can judge whether the output is correct. For math and code with unit tests, fine. For translation, summarization, open-ended reasoning? You’re stuck building a task-specific judge, which is expensive and brittle.

The gap: is there a way to get sequence-level, on-policy training signal without needing a task-specific verifier?

Problem: SFT = token-level, off-policy | RLVR = sequence-level, on-policy but needs verifier
         |
         v
Assumption: Sequence-level *feature statistics* (embeddings) carry semantic signal
            without requiring explicit correctness labels
         |
         v
Method: EBFT -- match feature distributions via energy-based modeling
        + strided block-parallel rollouts for efficiency
        + on-policy policy gradient update
         |
         v
Evidence: Q&A coding, unstructured coding, translation benchmarks
          EBFT ~= RLVR accuracy, > SFT accuracy
          EBFT < RLVR and SFT in validation cross-entropy (better generalization)
         |
         v
Conclusion: Feature-matching is a viable verifier-free alternative to RLVR

The Increment

One sentence: Before this paper, getting on-policy sequence-level training signal required a task-specific verifier; after it, you can use the model’s own embedding space as a verifier-free semantic oracle.

Core Mechanism

EBFT has three interlocking pieces. First, the objective: instead of minimizing cross-entropy token by token, you define a feature-matching loss that asks “does the distribution of embeddings produced by my model’s rollouts match the distribution of embeddings from reference completions?” This is operationalized as an energy-based model where the energy function scores how well a generated sequence’s features align with the reference feature distribution.

Second, the sampling strategy: to make this tractable, you need many rollouts from many prefixes simultaneously. EBFT uses strided block-parallel sampling — you take a batch of prompts, generate completions in parallel, but also sample from nested prefixes (prefix of length k, k+stride, k+2*stride, …) so you get dense coverage of the sequence space without sequential bottlenecks. This is what makes the gradient estimates low-variance enough to be useful.

Third, the update: once you have rollouts and their embeddings, you compute a policy gradient update that pushes the model toward rollouts whose features match the reference, weighted by the energy score. The theoretical connection is to KL-regularized feature matching — you’re minimizing KL divergence between the model’s feature distribution and the reference’s, with an energy-based density ratio as the importance weight.

Reference completions                 Model rollouts
       |                                    |
       v                                    v
  [Embed with f(.)]              [Strided block-parallel sample]
       |                                    |
       v                                    v
  Reference feature               Rollout feature
  distribution mu_ref              distribution mu_model
       |                                    |
       +-----------> Energy score <---------+
                     E(x) = f(x)^T W f_ref
                          |
                          v
                  Policy gradient update
                  grad = E_rollout[ E(x) * grad log pi(x) ]
                          |
                          v
                  Updated model weights

Think of it like a wine sommelier training program. In SFT, you’re memorizing tasting notes word for word from a textbook — you learn to reproduce the exact description, but you never actually taste the wine. In RLVR, you taste the wine and a judge tells you “correct” or “wrong” — great feedback, but you need an expert judge for every wine variety. EBFT is different: you taste the wine, and instead of a judge, you compare the flavor profile (the embedding) of your description against a reference profile. No judge needed — the semantic fingerprint of the text is the signal. The sommelier learns to produce descriptions whose flavor profile matches the reference, not whose exact words match. The strided block-parallel sampling is like tasting the wine at multiple stages of aeration simultaneously — you get richer signal per training step.

Key Concepts

• Feature matching: Instead of asking “did you predict the right next token?”, you ask “does the statistical fingerprint of your output match the fingerprint of good outputs?” Concretely: run both the reference completion and your model’s completion through an encoder, get embedding vectors, and measure how close those vectors are in distribution. If your model writes “the capital of France is Paris” and the reference says “Paris is France’s capital city”, the token sequences differ but the embeddings are close — feature matching rewards this, cross-entropy doesn’t.

• Energy-based modeling: An energy function assigns a scalar score to each possible output — low energy = good, high energy = bad. You don’t need it to be a proper probability (normalized over all sequences), which is what makes it tractable. Here, the energy is defined as the inner product between the rollout’s embedding and the reference feature distribution. It’s a soft measure of “how much does this output look like the reference in feature space?” — no binary correct/incorrect label required.

• On-policy vs off-policy training: Off-policy (SFT) means you train on data generated by someone else (the teacher), under teacher forcing where the model never sees its own errors. On-policy means you train on data generated by your current model. The difference matters because a model that’s slightly wrong at step 5 will compound that error through step 50 — but SFT never shows it what step 50 looks like when step 5 was wrong. EBFT is on-policy: it rolls out the model’s own completions and trains on those, so the model learns to recover from its own mistakes.

Framework Shift

Before (mainstream approach):        After (this paper):

Prompt                               Prompt
  |                                    |
  v                                    v
Teacher forcing                      Model rollouts (on-policy)
  |                                    |  (strided block-parallel)
  v                                    v
Token-level CE loss                  Feature extraction
  |                                    |
  v                                    v
Gradient update                      Energy score vs reference features
                                       |
  [No rollouts, no sequence-level]     v
  [signal, no semantic feedback]     Policy gradient update
                                       |
  [Needs verifier for RLVR]            v
                                     [Sequence-level, on-policy,]
                                     [no verifier needed         ]

From token prediction to feature distribution matching, the core shift is: the training signal moves from “did you pick the right word?” to “does your output live in the right region of semantic space?”

Expert Assessment

Problem choice: This is a real gap. The verifier bottleneck in RLVR is a genuine pain point — the field has been quietly aware that RLVR’s success on math/code doesn’t transfer cleanly to tasks without cheap verifiers. Framing feature matching as the solution is a natural move, and the timing is right given the maturity of embedding models. Not a manufactured gap.

Method maturity: The core insight is clever — using embedding-space statistics as a verifier-free reward signal is genuinely novel in this framing. The strided block-parallel sampling is a solid engineering contribution. That said, the energy function design (inner product with reference features) is fairly simple, and it’s not obvious why this specific form is optimal. The theoretical connection to KL-regularized feature matching is reassuring but not tight enough to fully justify the design choices.

Experimental integrity: The benchmarks (Q&A coding, unstructured coding, translation) are reasonable choices that span verifiable and non-verifiable tasks. Comparing against SFT and RLVR is the right baseline set. One flag: “matches RLVR” is doing a lot of work — the margin matters, and without seeing the full numbers it’s hard to know if this is “within noise” or “genuinely competitive.” The lower validation cross-entropy than both SFT and RLVR is the most interesting empirical claim and deserves more scrutiny — it suggests better generalization, which would be a strong result if it holds up.

Writing quality: The theoretical section connecting EBFT to energy-based modeling is likely the weakest link — these connections are often stated more cleanly than they’re proven, and the gap between “we can write this as an EBM” and “therefore our method is principled” is where reviewers will push back. If that section were rewritten to either tighten the theory or honestly scope its claims, the paper would be significantly stronger.

Verdict: weak accept — the core idea is sound and the empirical results are promising, but the theoretical justification needs tightening and the “matches RLVR” claim needs more careful quantification.

Takeaways

The most transferable idea here is using pre-trained embeddings as a task-agnostic reward signal. If you’re fine-tuning a model on a task where you have reference outputs but no verifier, you can compute embedding-space distances between your model’s rollouts and the references and use that as a training signal — no reward model training required. This is immediately applicable to translation, summarization, and any generation task with reference outputs.

The strided block-parallel sampling trick is also worth stealing for any on-policy training setup where you need diverse rollouts efficiently. Sampling from nested prefixes gives you much denser coverage of the sequence space per GPU step than sampling full sequences independently.

Finally, the framing of “feature matching as a verifier substitute” is a useful mental model for thinking about when RLVR is and isn’t necessary. If your task has cheap verifiers, use RLVR. If it doesn’t, feature matching might be your best option before investing in a reward model.