Online Pandora's Box for Contextual LLM Cascading

Riepilogo Tecnico: Pandora's Box Online per il Cascading Contestuale di LLM

1. Formulazione del Problema

Il documento affronta la sfida operativa del cascading di LLM, in cui un'organizzazione deve interrogare sequenzialmente un portafoglio di API di Large Language Model (LLM) per generare contenuti di alta qualità in modo efficiente dal punto di vista dei costi. Diversamente dal routing standard dei modelli (selezione di una singola API per richiesta), il cascading implica un processo di ricerca sequenziale: il decisore (DM) interroga le API una alla volta, sostenendo costi e osservando gli output, finché non viene soddisfatta una condizione di arresto, al termine della quale viene selezionato e distribuito un singolo output.

Gli autori formalizzano questo scenario come un problema di Pandora's Box Contestuale Online con una distinta struttura a due fasi:

1. Fase di Query: Il DM osserva un contesto di richiesta $x_t$ e interroga sequenzialmente le scatole (API). Interrogare la scatola $a$ rivela un output stocastico $\omega_{at}$ e comporta un costo $c_{at}$. Il DM decide se continuare a interrogare o fermarsi sulla base degli output e dei costi osservati.
2. Fase di Selezione: Al momento dell'arresto, il DM seleziona uno degli output interrogati da distribuire. Il DM osserva solo la ricompensa a valle (downstream reward) $r_t$ associata all'output selezionato, non le ricompense degli output interrogati ma non selezionati.

Distinzione Chiave: Questo differisce dai classici modelli di Pandora's Box in cui l'apertura di una scatola rivela direttamente la sua ricompensa. Qui, la ricompensa è "mediata dall'output": la ricompensa dipende dalla coppia richiesta-output $(x_t, \omega_{at})$ tramite una funzione di ricompensa condivisa $\mu^*$, mentre l'eterogeneità specifica dell'API emerge solo attraverso le distribuzioni di output e costi. Il feedback è parziale: il DM apprende la funzione di ricompensa solo dagli output selezionati, ma apprende gli indici di prenotazione (reservation indices) da tutti gli output interrogati.

L'obiettivo è minimizzare il regret cumulativo su un orizzonte $T$ rispetto a una policy oracle che conosce le distribuzioni condizionate degli output, le funzioni di costo e la funzione di ricompensa.

2. Metodologia

Gli autori propongono COSMOS (COst-aware Sequential MOdel query and Selection), una policy di apprendimento online che evita di stimare le intere distribuzioni condizionate. Invece, modella e apprende direttamente gli indici di prenotazione e la funzione di ricompensa utilizzando strutture parametriche.

Componenti Core:

• Modellazione Parametrica:
  • Funzione di Ricompensa: Modellata come una funzione lineare generalizzata: $\mu^*(x, \omega) = G(\theta^*{}^\top \phi(x, \omega))$.
  • Indici di Prenotazione: Modellati come funzioni lineari generalizzate del contesto: $\sigma^*_a(x) = \Lambda(\rho_a^\top \psi(x))$. Questi indici sono derivati dalla policy ottimale di Weitzman, soddisfacendo l'equazione del momento $E[(\mu^*(x, \omega) - \sigma^*_a(x))^+ | x] = E[c_a | x]$.
• Apprendimento Basato sull'Ottimismo: La policy costruisce limiti superiori ottimistici sia per la funzione di ricompensa che per gli indici di prenotazione.
  • Stima della Ricompensa: Utilizza un Estimatore di Massima Verosimiglianza (MLE) penalizzato per $\theta^*$ basato sugli output distribuiti, costruendo Upper Confidence Bounds (UCB) tramite un'ineguaglianza di martingala auto-normalizzata.
  • Stima dell'Indice: Utilizza un approccio basato sul Metodo dei Momenti Generalizzato (GMM) (specificamente una formulazione di $M$-stima) per stimare $\rho_a$. La funzione di perdita è costruita in modo che la sua condizione del primo ordine recuperi l'analogo campionario dell'equazione del momento di Weitzman. Ciò consente di apprendere l'indice direttamente dalle coppie output-costo interrogate senza stimare l'intera distribuzione.
• Decomposizione del Regret: Un insight teorico chiave è che, sotto l'ipotesi di ottimismo, il regret per periodo si decompone in due termini di errore indipendenti:
  1. L'errore nella stima della ricompensa per l'output selezionato.
  2. La somma degli errori nella stima degli indici di prenotazione per tutte le scatole interrogate.
     Questa decomposizione permette agli autori di applicare l'analisi standard di tipo UCB (argomenti di potenziale ellittico) a entrambi i componenti separatamente.

Flusso dell'Algoritmo (COSMOS):

1. Osserva il contesto $x_t$.
2. Costruisce stime ottimistiche $\hat{\sigma}_{at}$ e $\hat{\mu}_t$.
3. Ordina le scatole per indice $\hat{\sigma}_{at}$ decrescente.
4. Interroga sequenzialmente le scatole. Interrompe l'interrogazione della scatola $k$ se la ricompensa ottimistica massima degli output interrogati supera l'indice ottimistico della successiva scatola.
5. Seleziona la scatola interrogata con la stima della ricenza ottimistica più elevata.
6. Osserva la ricompensa realizzata e aggiorna gli stimatori.

3. Contributi Chiave

1. Nuova Formulazione del Modello: Il documento introduce un modello di Pandora's Box contestuale online specificamente progettato per il cascading di LLM. Cattura la critica struttura di feedback "mediata dall'output" dove le ricompense sono osservate solo per gli output distribuiti, distinguendosi dai classici setting in cui le ricompense vengono rivelate all'ispezione.
2. Modellazione Diretta dell'Indice: Inveve di apprendere complesse distribuzioni condizionate di output e costi, gli autori propongono di modellare direttamente gli indici di prenotazione utilizzando strutture parametriche. Ciò è ottenuto tramite un approccio di $M$-stima basato su GMM che sfrutta la condizione del momento che definisce l'indice ottimale.
3. Garanzie di Regret: Il documento fornisce limiti formali di regret per la policy COSMOS.
   • Regime con Ricompensa Nota: Se la funzione di ricompensa è nota (ad esempio, tramite dati offline), la policy raggiunge un limite di regret di $\tilde{O}(A m \sqrt{T})$, dove $A$ è il numero di API e $m$ è la dimensione della mappa delle feature dell'indice.
   • Regime con Ricompensa Sconosciuta: Quando la funzione di ricompensa deve essere appresa online, la policy raggiunge un limite di regret di $\tilde{O}([d + A(m + \sqrt{dm})]\sqrt{T})$, dove $d$ è la dimensione della mappa delle feature della ricompensa.
     Questi limiti migliorano i risultati esistenti di Pandora's Box contestuale (es. $\tilde{O}(T^{5/6})$) sfruttando assunzioni di curvatura locale per raggiungere il tasso ottimale $\tilde{O}(\sqrt{T})$.

4. Risultati e Assunzioni

I risultati teorici si basano su diverse condizioni di regolarità:

• Assunzione 1 & 2 (Indici): La funzione di legame $\Lambda$ è smooth e strettamente crescente; la funzione di perdita esibisce curvatura locale attorno al vero indice (garantendo che l'equazione del momento attraversi lo zero con una pendenza limitata).
• Asszione 3 & 4 (Ricompensa): La funzione di ricompensa segue un modello lineare generalizzato con un legame $G$ strettamente crescente; le ricompense appartengono a una famiglia esponenziale canonica.
• Assunzione 5 (Anti-concentrazione): Garantisce che le feature degli output selezionati (usati per l'apprendimento della ricompensa) siano sufficientemente informative per controllare l'errore di stima degli indici (che dipendono dagli output interrogati). Ciò previene un "disallineamento" tra i dati usati per apprendere la ricompensa e i dati usati per apprendere gli indici.

Il documento dimostra che, in queste condizioni, la policy COSMOS raggiunge un regret cumulativo $\tilde{O}(\sqrt{T})$ dipendente dalla dimensione.

5. Significato e Rivendicazioni

Il documento sostiene di essere il primo a formulare il cascading di LLM come un problema di Pandora's Box contestuale online con query e selezione congiunte, fornendo garanzie formali di regret. Colma il divario tra i metodi euristici di cascading di LLM (come FrugalGPT) e la rigorosa teoria dell'apprendimento online.

La significanza risiede nel:

• Rilevanza Pratica: Affronta il trade-off reale tra costo di inferenza e qualità dell'output nell'implementazione di LLM, offrendo un modo principato per gestire le query sequenziali alle API.
• Avanzamento Teorico: Estende la letteratura su Pandora's Box a setting con feedback mediato dall'output e sviluppa un nuovo algoritmo di apprendimento che combina tecniche GMM e UCB per gestire la specifica struttura di feedback del cascading di LLM.
• Efficienza: Modellando direttamente gli indici di prenotazione invece delle intere distribuzioni, l'approccio risulta più trattabile dal punto di vista computazionale ed efficiente dal punto di vista statistico per contesti ad alta dimensionalità.

Gli autori notano che il loro lavoro assume una standard struttura di query sequenziale e modelli parametrici, lasciando l'estensione a setting non parametrici, querying a lotti (batched) ed ensembling di output come direzioni future.