Ottenere ricompense per token tramite l’attribuzione delle ricompense senza addestramento

Abstract degli autori

In questo articolo proponiamo un metodo estremamente efficiente, che non richiede addestramento, per estrarre segnali di ricompensa a livello di token direttamente da un modello di ricompensa profondo già esistente. L’idea centrale è attribuire la ricompensa complessiva di processo ai singoli token stimando l’influenza di ciascuno. Definiamo questa influenza come la variazione della ricompensa macroscopica finale, per esempio la ricompensa di processo, quando un token viene sostituito con un token semanticamente nullo. Calcolare l’influenza in modo diretto è computazionalmente impraticabile: per una sequenza di N token richiederebbe N passaggi in avanti attraverso il PRM. Superiamo questo ostacolo proponendo uno stimatore basato sul gradiente altamente efficiente. In particolare, usiamo un’approssimazione di Taylor del primo ordine, che riduce il calcolo dell’influenza al prodotto scalare tra la differenza degli embedding del token e del token nullo e il gradiente della ricompensa rispetto all’embedding del token. Basta così un solo passaggio in avanti e uno all’indietro. Le ricompense per token ottenute permettono agli algoritmi standard di RL di assegnare con precisione il merito dei risultati, senza richiedere un ulteriore addestramento del modello di ricompensa. Esperimenti su benchmark di ragionamento impegnativi mostrano che il nostro metodo migliora notevolmente l’efficienza dell’ottimizzazione della policy e la generalizzazione delle capacità di ragionamento degli LLM. Il nostro P2T supera la ricompensa basata sul risultato di +4,9% su MathVista per Qwen2.5-VL-7B-Instruct e di +11,5% su AIME24 per Qwen2.5-Math-7B. Il codice sarà pubblicato su https://github.com/stonewst/P2T.