Expected Reward Prediction, with Applications to Model Routing

Probleemstelling

Reward-modellen (beloningsmodellen) zijn een standaardtool om antwoorden van Large Language Models (LLMs) te scoren. Traditioneel worden deze modellen getraind om antwoorden op een vaste prompt te rangschikken (bijvoorbeeld om de beste van $n$ gegenereerde antwoorden te kiezen). Het artikel onderzoekt echter een fundamentele vraag: is het mogelijk om de verwachte beloning van een LLM voor een specifieke prompt te voorspellen voordat het model überhaupt een antwoord heeft gegenereerd?

De uitdaging ligt in het feit dat reward-modellen vaak worden getraind op relatieve voorkeuren (Bradley-Terry-model), waarbij alleen het verschil tussen twee antwoorden betekenisvol is, niet de absolute waarde. Dit maakt het theoretisch onzeker of de verwachte beloning (de gemiddelde score over alle mogelijke antwoorden van een model op een prompt) voorspelbaar is. Bovendien is het onduidelijk of een model zich consistent genoeg gedraagt om deze verwachting a priori te kunnen schatten.

Methodologie

1. Definitie van Verwachte Beloning
De auteurs definiëren de verwachte beloning van een taalmodel $\pi$ voor een prompt $x$ als:
$$ER_\pi(x) := \mathbb{E}_{Y \sim \pi(Y | x)}[r(x, Y)]$$
Dit is de verwachte waarde van de reward $r$ voor een willekeurige steekproef $Y$ uit het model.

2. Data en Experimenteel Opzet

• Dataset: Er wordt gebruik gemaakt van het open-perfectblend dataset, bestaande uit 4000 prompts verdeeld over vier categorieën: algemene chat, instructievolging, wiskundig redeneren en coderen.
• Modellen: Er worden diverse open-weight instructie-tuned modellen getest (Llama 3.1-IT 8B/70B, Gemma 2-IT 9B/27B, Gemma 1-IT 7B).
• Reward Modellen: Drie verschillende reward-modellen worden gebruikt (OpenAssistant-RM, GRM-2B-RM, InternLM-RM).
• Ground Truth: Voor elke prompt worden $K=32$ antwoorden gegenereerd per model om de empirische gemiddelde beloning te berekenen.

3. Voorspellingsmodel (Linear Probe)
In plaats van complexe generatieve modellen te gebruiken, trainen de auteurs een lineair model (ridge-regressie) om de verwachte beloning te voorspellen.

• Input: De prompt $x$ wordt omgezet in een vaste vectorrepresentatie (embedding) $v(x)$ met behulp van een kant-en-klaar model (gte-large-en-v1.5).
• Training: Het lineaire model leert de mapping van de embedding $v(x)$ naar de empirische gemiddelde beloning.
• Doel: Voorspel $ER_\pi(x)$ zonder dat er tijdens de inferentie herhaaldelijk gesampled hoeft te worden.

4. Toepassing: Model Routing
De voorspelde verwachte beloningen worden gebruikt voor Model Routing. Het doel is om prompts tijdens de inferentie naar het meest kosteneffectieve model te sturen dat de hoogste verwachte beloning oplevert.

• Doelfunctie: Maximaliseer de verwachte beloning onder een begrensde rekenkosten ($C$).
• Strategie: Een router $\rho$ kiest het model $i$ dat $ER_{\pi_i}(x) - \lambda \cdot c(i)$ maximaliseert, waarbij $c(i)$ de kosten van het model zijn en $\lambda$ de afweging tussen kosten en kwaliteit regelt.

Kernbijdragen

1. Voorspelbaarheid van Verwachte Beloning: Het artikel bewijst empirisch dat de verwachte beloning van een LLM voor een specifieke prompt met hoge precisie voorspelbaar is, zelfs met een zeer eenvoudig lineair model gebaseerd op prompt-embeddings.
2. Van Response-Level naar Model-Level: De methologie "lift" reward-scores van het niveau van individuele antwoorden naar het niveau van het model zelf. Dit stelt systemen in staat om de geschiktheid van een model voor een prompt te beoordelen zonder het model eerst te laten genereren.
3. Efficiënte Model Routing: De auteurs introduceren een "Expected Reward Prediction" (ERP) routing-protocol. Dit is eenvoudiger en schaalbaarder dan bestaande methoden (zoals Zooter of preference-based routing) omdat het geen paar-voor-paar vergelijkingen tussen modellen vereist.
4. Schaalbaarheid: Bij het toevoegen van een nieuw model aan de pool hoeft er slechts één voorspellingsmodel voor dat specifieke model te worden getraind, in plaats van een nieuw preference-model dat alle bestaande en nieuwe modellen met elkaar moet vergelijken (wat kwadratisch schaalt).

Resultaten

• Voorspellingsnauwkeurigheid: Lineaire modellen presteren verrassend goed. De $R^2$-waarden (variatie verklaard door het model) liggen vaak tussen 0,4 en 0,6, zelfs binnen specifieke categorieën zoals wiskunde en codering. Dit suggereert dat de variatie in modelkwaliteit per prompt lineair af te leiden is uit de prompt-embedding.
• Routing Performance:
  • De ERP-based router presteert aanzienlijk beter dan baselines zoals het gebruik van een vast model of het routeren naar het model met de beste gemiddelde prestatie per categorie.
  • In de Pareto-frontier (kosten vs. regret) domineert de ERP-router de meeste baselines.
  • De ERP-methode bereikt vergelijkbare prestaties als de complexere "Zooter"-methode (die een softmax-distributie voorspelt), maar is computerefficiënter.
• Pairwise Vergelijking: De voorspelde verwachte beloningen correleren sterk met de daadwerkelijke winstkansen in paar-voor-paar vergelijkingen. Het voorspellen van de winnaar op basis van het verschil in verwachte beloning is even effectief als het direct trainen van een classifier op paar-voor-paar data.

Betekenis en Conclusie

Dit onderzoek heeft belangrijke implicaties voor de efficiënte inzet van LLMs:

• Kostenefficiëntie: Het stelt systemen in staat om prompts dynamisch naar het juiste model te sturen, waarbij rekening wordt gehouden met zowel de kwaliteit van het antwoord als de rekenkosten (bijv. het kiezen van een kleiner model voor simpele vragen en een groter model voor complexe redenering).
• Simpelheid en Generalisatie: De bevinding dat een lineair model op embeddings voldoende is, vermindert de complexiteit van inferentie-systemen aanzienlijk.
• Toekomstige Toepassingen: Naast routing kunnen deze voorspellingen worden gebruikt voor andere inferentie-taken, zoals het aanpassen van systemprompts of het selecteren van het beste model voor een specifieke taak zonder uitgebreide testruns.

Kortom, het artikel toont aan dat we de "verwachte kwaliteit" van een model voor een specifieke vraag kunnen voorspellen op basis van de vraag zelf, wat een krachtige basis vormt voor slimme, kosteneffectieve AI-systemen.