Entropy regularization improves language models for verified correct outputs

Entropy Regularization: A Free Correction to Cross-Entropy for Verified Demonstrations

Machine Learning

Summary

Large language models are often trained to copy one expert answer exactly, but some problems have many correct answers. The authors found that training models to just copy expert examples can mistakenly put chance on wrong answers. They propose adding entropy regularization, which encourages models to be more confident and avoid spreading guesses over unsupported outputs. This change helps the models produce more verifier-accepted answers in tasks like math reasoning and coding.

What this means in practice

  • For software developers: Improve automated code generation so models produce outputs more likely to pass correctness verification tests.
  • For data science teams: Enhance mathematical reasoning tools by training language models to better align with verification-based correctness.

Authors

Mihir Dhanakshirur, Adam Ousherovitch, Ambuj Tewari

Abstract

Large language models are often post-trained on expert demonstrations using cross-entropy (CE), even when the downstream objective is not to imitate the demonstrated solution but to produce any output accepted by a verifier. This mismatch is seen in verifiable domains with multiple correct solutions, such as mathematical reasoning and code generation, where training data may contain only one expert solution per problem. We show that minimizing cross-entropy can be misaligned with minimizing verifier risk; two policies can assign identical likelihood to the observed demonstrations while placing different probability mass on incorrect outputs. This is formalized through a learning-theoretic counterexample in which CE minimization selects a suboptimal policy. We identify that controlling the support of the learned policy can solve this problem by preventing probability mass from spreading to unsupported outputs. Since support size is non-differentiable and computationally intractable, we propose entropy-regularized cross-entropy (ER-CE), using token-level Shannon entropy as a tractable proxy. Finally, across mathematical reasoning and code-generation benchmarks, we find that entropy-regularized training consistently improves verifier accuracy over standard cross-entropy. Our results identify a simple failure mode of imitation-based post-training in verifiable tasks and provide a practical objective that is better aligned with producing correct outputs.