Reward-Based Online LLM Routing via NeuralUCB

Titolo: Routing Online Basato su Ricompensa per LLM tramite NeuralUCB

Autori: Ming-Hua Tsai e Phat Tran (Oregon State University)

---

1. Definizione del Problema

Il lavoro affronta la sfida del routing dinamico dei Large Language Models (LLM). I diversi modelli LLM offrono un compromesso variabile tra capacità (qualità della risposta) e costo di inferenza. Modelli di fascia alta offrono prestazioni superiori ma a costi elevati, mentre modelli più economici possono fallire su compiti complessi (ragionamento, coding).

L'obiettivo è sviluppare una politica di routing che, per ogni query utente $x_t$, selezioni dinamicamente il modello appropriato $a_t$ da un insieme candidato $A = \{1, ..., K\}$, massimizzando l'utilità a lungo termine.
Il problema è formulato come un problema di bandito contestuale:

• Input: Una rappresentazione del contesto $x$ (embedding del testo, feature ausiliarie, dominio).
• Azione: Selezione di un modello $a$.
• Feedback: Si osserva solo la ricompensa del modello selezionato (feedback parziale), non quella degli altri modelli non scelti per la stessa query.
• Obiettivo: Massimizzare la ricompensa attesa cumulata $E[\sum r(x_t, a_t)]$, dove la ricompensa bilancia qualità e costo.

2. Metodologia

Funzione di Utilità e Ricompensa

Il sistema non massimizza solo la qualità, ma definisce una ricompensa di utilità che penalizza i costi:
$$r(x, a) = q(x, a) \cdot \exp(-\lambda \tilde{c}(x, a))$$
Dove:

• $q(x, a)$ è la qualità del modello.
• $\tilde{c}(x, a)$ è il costo normalizzato (tramite trasformazione logaritmica per stabilizzare la scala).
• $\lambda$ è un iperparametro che controlla la forza della penalità per il costo.

Architettura: UtilityNet

Il cuore del sistema è UtilityNet, una rete neurale che funge da predittore di utilità e fornisce rappresentazioni per la selezione delle azioni.

• Input: Combina embedding del testo ($x_{emb}$), feature ausiliarie ($x_{feat}$) e embedding del dominio.
• Branch di Utilità: Prevede la ricompensa $\mu(x, a)$ per ogni coppia contesto-azione. Utilizza un embedding specifico per l'azione (modello) e una rete MLP per regressione.
• Branch di Gating (Auxiliary): Predice una probabilità $p(x)$ che indica se l'esplorazione (basata sull'incertezza) è benefica per quel contesto specifico. Questo aiuta a evitare esplorazioni costose quando non necessarie.

Politica di Routing: NeuralUCB

Per la selezione delle azioni, il paper adotta l'algoritmo NeuralUCB (Upper Confidence Bound basato su reti neurali):

1. Stima dell'Incertezza: Utilizza l'ultimo strato nascosto di UtilityNet come rappresentazione delle feature. Calcola un bonus di esplorazione basato su una matrice di covarianza inversa condivisa ($A^{-1}$).
2. Punteggio UCB: $s(x, a) = \mu(x, a) + \beta \sqrt{g(x, a)^\top A^{-1} g(x, a)}$.
   • Favorisce azioni con alta utilità prevista e alta incertezza.
3. Logica di Gating: L'azione finale $a^*$ è determinata dalla probabilità di gating $p(x)$:
   • Se $p(x) \ge \tau_g$ (soglia): Si seleziona l'azione con il massimo punteggio UCB (esplorazione attiva).
   • Altrimenti: Si seleziona l'azione "sicura" basata solo sulla media prevista $\arg\max \mu(x, a)$ (sfruttamento conservativo).

Protocollo di Addestramento Online Simulato

Poiché l'esperimento online reale è costoso, i ricercatori utilizzano RouterBench per simulare un ambiente online:

• I dati sono divisi in 20 "fette" (slices).
• Per ogni fetta, la politica seleziona azioni, aggiorna il buffer di replay e la matrice $A^{-1}$, quindi addestra UtilityNet per 5 epoche.
• Questo ciclo permette di tracciare l'evoluzione della ricompensa cumulativa nel tempo.

3. Risultati Sperimentali

Gli esperimenti sono stati condotti su RouterBench (36.497 campioni, 86 domini, 11 modelli LLM candidati).

• Confronto con Baseline:

  • NeuralUCB supera costantemente le baseline: Random (scelta casuale), Min-Cost (sceglie sempre il modello più economico) e RouteLLM-BERT (un metodo supervisionato esistente).
  • Ricompensa Media: NeuralUCB mantiene una ricompensa stabile tra 0.59 e 0.61, contro lo 0.31-0.33 di Random e lo 0.35-0.36 di RouteLLM-BERT.
  • Ricompensa Cumulativa: Il divario tra NeuralUCB e le altre metodologie si amplia progressivamente man mano che vengono elaborate più fette di dati, dimostrando una migliore capacità di apprendimento adattivo.

• Analisi Costo-Qualità:

  • Rispetto al riferimento "Max-Quality" (che sceglie sempre il modello migliore in assoluto), la politica NeuralUCB utilizza solo circa il 33% del costo di inferenza mantenendo una qualità competitiva.
  • Sebbene esista ancora un divario di qualità rispetto all'ideale, il trade-off costo-prestazione è significativamente migliore rispetto alle strategie naive.

• Ablazione degli Encoder:

  • Sono stati testati diversi encoder di testo (es. multilingual-e5-large-instruct, all-mpnet-base-v2, Qwen3-Embedding).
  • Gli encoder più piccoli ma efficienti come all-MiniLM-L6-v2 e all-mpnet-base-v2 hanno mostrato le prestazioni migliori e più stabili, suggerendo che encoder più grandi non garantiscono necessariamente guadagni in questo contesto specifico.

4. Contributi Chiave

1. Framework Ibrido: Introduce un approccio che combina l'apprendimento supervisionato (UtilityNet) con l'esplorazione attiva (NeuralUCB) per il routing LLM in scenari a feedback parziale.
2. Meccanismo di Gating: Propone un ramo di gating neurale per attivare dinamicamente l'esplorazione solo quando utile, riducendo il rumore e i costi associati a esplorazioni non necessarie.
3. Validazione Pratica: Dimostra l'efficacia del metodo in un ambiente simulato realistico su RouterBench, superando sia approcci puramente casuali che metodi supervisionati statici.
4. Efficienza dei Costi: Fornisce evidenze empiriche che un routing intelligente basato su UCB può ridurre drasticamente i costi di inferenza (fino al 67% in meno rispetto al modello top) senza sacrificare eccessivamente la qualità.

5. Significato e Sfide Future

Il lavoro dimostra che NeuralUCB è un approccio promettente per il routing LLM "consapevole dei costi", offrendo un equilibrio superiore tra qualità e spesa rispetto alle metodologie attuali.

Tuttavia, il paper evidenzia alcune sfide residue:

• Discriminazione delle Azioni: Per molti campioni, la scelta del modello "max-quality" coincide già con la scelta ottimale per la ricompensa, limitando il potenziale guadagno del routing.
• Esplorazione: Rimane difficile discriminare perfettamente quando esplorare in contesti ad alta variabilità.
• Futuro: Il lavoro futuro si concentrerà sulla raccolta di più dati di addestramento e sul miglioramento del modello di utilità per affinare la discriminazione delle azioni sotto vincoli di costo.