Expected Reward Prediction, with Applications to Model Routing

1. Il Problema

I modelli di ricompensa (Reward Models - RM) sono strumenti standard per valutare e classificare le risposte dei Large Language Models (LLM). Tradizionalmente, questi modelli operano a livello di risposta: data una prompt $x$ e una risposta $y$, restituiscono un punteggio $r(x, y)$ che quantifica la qualità della risposta.

Tuttavia, in scenari di inferenza come il routing dei modelli (decidere quale modello utilizzare per una specifica richiesta prima di generare la risposta), è necessario un approccio diverso. L'obiettivo è "elevare" (lift) la funzione di ricompensa dal livello di risposta a un livello di modello. Nello specifico, il paper si pone la domanda: è possibile prevedere a priori, basandosi solo sulla prompt $x$, il punteggio di ricompensa atteso ($ER$) che un modello specifico $\pi$ otterrebbe se generasse una risposta?

Il problema è complesso per due motivi principali:

1. Comportamento imprevedibile: Non è scontato che il comportamento di un modello su una prompt specifica sia prevedibile, data la sua natura stocastica e la sensibilità ai dettagli della prompt.
2. Ambiguità dei Reward Models: La maggior parte dei RM è addestrata su dati di preferenza (coppie di risposte $y^+, y^-$) per massimizzare la verosimiglianza di Bradley-Terry. Teoricamente, questo garantisce solo che la differenza di ricompensa sia significativa, lasciando un valore arbitrario dipendente dalla prompt. Di conseguenza, il valore assoluto della ricompensa attesa non dovrebbe essere teoricamente identificabile o apprendibile.

2. Metodologia

Gli autori propongono un approccio empirico per verificare se la ricompensa attesa è prevedibile e come utilizzarla.

Definizione Formale

La ricompensa attesa per un modello $\pi$ su una prompt $x$ è definita come:
$$ER_\pi(x) := \mathbb{E}_{Y \sim \pi(Y|x)}[r(x, Y)]$$
L'obiettivo è apprendere una funzione che predica direttamente $ER_\pi(x)$ partendo solo da $x$, senza dover campionare ripetutamente risposte durante l'inferenza.

Setup Sperimentale

• Dataset: Utilizzo di open-perfectblend, un dataset diversificato con 4000 prompt (1000 per categoria: Chat generale, Istruzioni, Matematica, Codice).
• Modelli Generativi (LLM): Un pool di modelli instruction-tuned di diverse dimensioni: Llama 3.1 (8B, 70B), Gemma 2 (9B, 27B) e Gemma 1 (7B).
• Modelli di Ricompensa (RM): Tre RM diversi (OpenAssistant-RM, GRM-2B-RM, InternLM-RM) addestrati con l'obiettivo di Bradley-Terry.
• Ground Truth: Per ogni prompt, vengono generate $K=32$ risposte da ogni modello. La ricompensa attesa empirica è calcolata come la media dei punteggi delle 32 risposte.

Modello di Predizione (Linear Probe)

Contrariamente alla complessità attesa, gli autori utilizzano un modello estremamente semplice:

1. Si utilizza un embedding pre-addestrato della prompt ($v(x)$), ottenuto tramite il modello gte-large-en-v1.5.
2. Si addestra un modello lineare regolarizzato (Ridge Regression) per mappare l'embedding $v(x)$ direttamente alla ricompensa attesa empirica $\hat{ER}_\pi(x)$.
3. L'addestramento avviene separatamente per ogni modello LLM e ogni Reward Model.

Applicazione al Routing

Il metodo proposto, chiamato Expected Reward Prediction (ERP), viene utilizzato per instradare le prompt al modello più adatto.

• Policy: Per una data prompt $x$, si sceglie il modello $i$ che massimizza la ricompensa attesa prevista meno un termine di costo: $\arg\max_i (\hat{ER}_{\pi_i}(x) - \lambda \cdot c(i))$, dove $c(i)$ è il costo computazionale (es. numero di parametri) e $\lambda$ è un parametro di bilanciamento.
• Vantaggio Scalabile: A differenza dei metodi basati su preferenze (che richiedono confronti a coppie tra tutti i modelli, scalando quadraticamente), l'ERP richiede un solo predittore per modello (scalabilità lineare). Aggiungere un nuovo modello richiede solo l'addestramento di un nuovo predittore lineare.

3. Risultati Chiave

Prevedibilità della Ricompensa Attesa

• Alta Accuratezza: È possibile prevedere la ricompensa attesa con alta fedeltà utilizzando un semplice modello lineare sugli embedding.
• Metriche: Il coefficiente di determinazione ($R^2$) mostra che i modelli lineari spiegano una frazione significativa della varianza nei dati di test (es. $R^2$ aggregato tra 0.46 e 0.59 a seconda del modello).
• Generalizzazione: La prevedibilità si mantiene sia all'interno delle categorie di prompt che tra diverse categorie, suggerendo che le capacità dei modelli variano in modo sistematico e prevedibile su prompt specifici.

Performance nel Routing

• Superiorità rispetto alle Baseline: La strategia ERP supera le baseline che instradano le prompt al modello con la media migliore per categoria (Per-category best oracle) o che usano modelli casuali/fissi.
• Frontiera di Pareto: L'ERP riesce a stabilire una frontiera di Pareto (trade-off costo-prestazione) che domina o eguaglia le baselines, inclusi metodi più complessi come Zooter (che prevede la distribuzione softmax del vincitore).
• Confronto con Zooter: Sebbene Zooter e ERP abbiano performance simili in termini di AUC per la previsione delle vittorie a coppie, ERP è molto più semplice, scalabile e non richiede dati di preferenza a coppie tra modelli diversi.
• Robustezza: I risultati sono coerenti su diversi Reward Models, anche se alcuni (come GRM-2B-RM) mostrano una minore prevedibilità in categorie specifiche come la matematica.

4. Contributi Principali

1. Dimostrazione Empirica: Si dimostra che, nonostante le limitazioni teoriche dei modelli di preferenza (Bradley-Terry), la ricompensa attesa di un LLM su una prompt è prevedibile con alta precisione.
2. Semplicità del Modello: Si evidenzia che un semplice probing lineare sugli embedding della prompt è sufficiente per catturare queste dinamiche, senza bisogno di architetture complesse.
3. Protocollo di Routing Scalabile: Si introduce un metodo di routing (ERP) che:
   • Non richiede dati di preferenza a coppie tra modelli.
   • È facilmente estendibile a nuovi modelli (aggiunta lineare, non quadratica).
   • Permette di ottimizzare il trade-off tra qualità della risposta e costo computazionale in tempo reale.
4. Spiegazione di Metodi Esistenti: Si suggerisce che il successo di metodi di routing complessi (come Zooter) deriva implicitamente dalla capacità di stimare la ricompensa attesa, che può essere ottenuta in modo molto più efficiente.

5. Significato e Implicazioni

Questo lavoro ha un impatto significativo sull'efficienza operativa degli LLM:

• Ottimizzazione dei Costi: Permette di instradare le richieste al modello "giusto" (es. un modello piccolo ed economico per task semplici, uno grande per task complessi) massimizzando la qualità e minimizzando i costi computazionali.
• Semplificazione dell'Infrastruttura: Riduce la necessità di raccogliere enormi dataset di preferenze a coppie per ogni nuova combinazione di modelli, semplificando la gestione di pool di modelli eterogenei.
• Nuove Applicazioni: Apre la strada ad altre applicazioni in tempo reale, come la modifica dinamica delle prompt o lo "hot-swapping" dei prompt di sistema, basandosi su stime predittive della qualità.

In sintesi, il paper stabilisce che la "bontà" di un modello su una specifica richiesta non è un mistero stocastico, ma una proprietà prevedibile che può essere sfruttata per costruire sistemi di inferenza più intelligenti ed economici.