GenARM: Generazione guidata da ricompensa con modello di ricompensa autoregressivo per l'allineamento al momento del test.

GenARM: Reward Guided Generation with Autoregressive Reward Model for Test-time Alignment

October 10, 2024

Autori: Yuancheng Xu, Udari Madhushani Sehwag, Alec Koppel, Sicheng Zhu, Bang An, Furong Huang, Sumitra Ganesh

cs.AI

Abstract

I Large Language Models (LLM) mostrano capacità impressionanti ma richiedono un'allineamento accurato con le preferenze umane. I metodi tradizionali di addestramento finetunano i LLM utilizzando set di dati sulle preferenze umane, ma comportano costi di addestramento significativi e richiedono addestramenti ripetuti per gestire diverse preferenze degli utenti. I metodi di allineamento al momento del test affrontano questo problema utilizzando modelli di ricompensa (RMs) per guidare i LLM congelati senza riallenamento. Tuttavia, gli approcci esistenti al momento del test si basano su RMs a livello di traiettoria progettati per valutare risposte complete, rendendoli inadatti alla generazione di testo autoregressivo che richiede il calcolo delle ricompense del token successivo dalle risposte parziali. Per affrontare questo problema, presentiamo GenARM, un approccio di allineamento al momento del test che sfrutta il Modello di Ricompensa Autoregressivo - una nuova parametrizzazione della ricompensa progettata per prevedere le ricompense del token successivo per una generazione autoregressiva efficiente ed efficace. Teoricamente, dimostriamo che questa parametrizzazione può guidare in modo provabile i LLM congelati verso qualsiasi distribuzione ottenibile dai RMs tradizionali all'interno del framework di apprendimento per rinforzo regolarizzato KL. I risultati sperimentali mostrano che GenARM supera significativamente i baselines di allineamento al momento del test precedenti e si allinea alle prestazioni dei metodi di addestramento. Inoltre, GenARM consente un efficiente allineamento da debole a forte, allineando LLM più grandi con RMs più piccoli senza i costi elevati dell'addestramento di modelli più grandi. Inoltre, GenARM supporta l'allineamento multi-obiettivo, consentendo scambi in tempo reale tra le dimensioni delle preferenze e soddisfacendo le diverse preferenze degli utenti senza riallenamento.

English

Large Language Models (LLMs) exhibit impressive capabilities but require careful alignment with human preferences. Traditional training-time methods finetune LLMs using human preference datasets but incur significant training costs and require repeated training to handle diverse user preferences. Test-time alignment methods address this by using reward models (RMs) to guide frozen LLMs without retraining. However, existing test-time approaches rely on trajectory-level RMs which are designed to evaluate complete responses, making them unsuitable for autoregressive text generation that requires computing next-token rewards from partial responses. To address this, we introduce GenARM, a test-time alignment approach that leverages the Autoregressive Reward Model--a novel reward parametrization designed to predict next-token rewards for efficient and effective autoregressive generation. Theoretically, we demonstrate that this parametrization can provably guide frozen LLMs toward any distribution achievable by traditional RMs within the KL-regularized reinforcement learning framework. Experimental results show that GenARM significantly outperforms prior test-time alignment baselines and matches the performance of training-time methods. Additionally, GenARM enables efficient weak-to-strong guidance, aligning larger LLMs with smaller RMs without the high costs of training larger models. Furthermore, GenARM supports multi-objective alignment, allowing real-time trade-offs between preference dimensions and catering to diverse user preferences without retraining.