Confidence-based reward for group relative policy optimization in language models

Abstract

Certain aspects of the disclosure provide a method for training a language model (LM) including: generating, using an LM, one or more outputs; computing a confidence score of an output of the one or more outputs based on a perplexity value of the output; determining, by a group relative policy optimization (GRPO)-based model, that the output is: associated with a correct status based on a reference policy; and associated with an uncertain status based on the confidence score and a threshold; determining, by the GRPO-based model, an increased reward value for the output that is associated with the correct status and the uncertain status based at least in part on a base reward value and the confidence score; causing, by the GRPO-based model, a reinforcement of the output using the increased reward value; and training the LM in accordance with the reinforcement of the output.