Reward-Based Online LLM Routing via NeuralUCB

Titel: Reward-Based Online LLM Routing via NeuralUCB

Auteurs: Ming-Hua Tsai en Phat Tran (Oregon State University)

---

1. Probleemdefinitie

Grote Taalmodellen (LLM's) vertonen aanzienlijke variatie in zowel capaciteit als inferentiekosten. Hoogwaardige modellen bieden vaak betere antwoorden maar zijn duurder, terwijl goedkopere modellen kunnen falen bij complexe taken. Het doel van dit onderzoek is beloningsgebaseerde online LLM-routing: het dynamisch selecteren van het meest geschikte LLM voor elke gebruikersvraag om een evenwicht te vinden tussen antwoordkwaliteit en inferentiekosten.

Het probleem wordt geformuleerd als een contextueel bandit-probleem:

• Op elk tijdstip $t$ observeert het systeem een vraag $x_t$.
• Het systeem kiest een model $a_t$ uit een kandidaatset $A$.
• Er wordt alleen feedback ontvangen voor het gekozen model (gedeeltelijke feedback), niet voor de andere modellen.
• Het doel is het maximaliseren van de verwachte langdurige nuttigheid (utility), waarbij rekening wordt gehouden met zowel kwaliteit als kosten.

2. Methodologie

De auteurs stellen een nieuw kader voor dat NeuralUCB (Neural Upper Confidence Bound) combineert met een diep leernetwerk om een niet-lineair routingbeleid te leren.

A. Nuttigheidsfunctie (Utility Function)

In plaats van alleen de kwaliteit te maximaliseren, definiëren de auteurs een nuttigheidsbeloning ($r$) die een afweging maakt tussen kwaliteit ($q$) en genormaliseerde kosten ($\tilde{c}$):
$$r(x, a) = q(x, a) \cdot \exp(-\lambda \cdot \tilde{c}(x, a))$$
Hierbij is $\lambda$ een parameter die de strengheid van de kostenstraf bepaalt. De kosten worden logaritmisch genormaliseerd om instabiliteit door grote schaalverschillen te voorkomen.

B. Architectuur: UtilityNet

Het kerncomponent is UtilityNet, een neurale architectuur die twee taken uitvoert:

1. Voorspelling van de nuttigheid: Het schat de verwachte beloning voor een gegeven context (vraag) en actie (model).
   • Input: Prompt-embedding, hulpfeatures en domein-ID.
   • Output: Een geschatte nuttigheidswaarde ($\mu$).
2. Gating-mechanisme: Een extra tak voorspelt of exploratie (het proberen van onzekere opties) nuttig is voor een specifieke vraag. Dit helpt om te voorkomen dat het systeem onnodig risicovolle keuzes maakt wanneer een veilige, gemiddelde keuze beter is.

De architectuur gebruikt:

• Tekst-embeddings (via een voorgetrainde encoder).
• Domein- en feature-embeddings.
• Actie-embeddings (voor de specifieke LLM's).
• Deze worden samengevoegd en verwerkt door MLP's (Multi-Layer Perceptrons) om zowel de nuttigheid als de gate-probabiliteit te voorspellen.

C. Routingbeleid op basis van NeuralUCB

Voor de actiekeuze wordt de UCB-score gebruikt, die de voorspelde nuttigheid combineert met een exploratiebonus gebaseerd op onzekerheid:
$$s(x, a) = \mu(x, a) + \beta \sqrt{g(x, a)^\top A^{-1} g(x, a)}$$

• $\mu(x, a)$: De voorspelde nuttigheid van UtilityNet.
• $\beta$: Controleert de sterkte van de exploratie.
• $A^{-1}$: Een gedeelde inverse covariantiematrix die onzekerheid in de feature-ruimte schat.

Beslissingsregel:
Het beleid gebruikt het gating-mechanisme om te beslissen:

• Als de gate-probabiliteit $p(x)$ boven een drempel $\tau_g$ ligt, wordt de UCB-score gebruikt (exploratie + exploitatie).
• Anders wordt een veilige actie gekozen die puur gebaseerd is op de hoogste voorspelde gemiddelde nuttigheid ($\arg \max \mu(x, a)$).

D. Evaluatieprotocol

Het systeem wordt getest in een gesimuleerde online omgeving op de RouterBench-dataset. De data wordt in 20 "slices" verdeeld. Na elke slice wordt het model getraind op de opgeslagen ervaringen (replay buffer) en worden de UCB-statistieken herbouwd.

3. Belangrijkste Resultaten

De experimenten zijn uitgevoerd op RouterBench (36.497 samples, 86 domeinen, 11 kandidaat-LLM's).

• Prestaties ten opzichte van baselines:

  • De NeuralUCB-methode presteert consequent beter dan een random beleid, een min-cost beleid (altijd het goedkoopste model kiezen) en een gesuperviseerde baseline (RouteLLM-BERT).
  • De gemiddelde beloning van NeuralUCB stabiliseert rond 0,59 - 0,61, terwijl RouteLLM-BERT rond 0,35-0,36 blijft en min-cost rond 0,51-0,53.
  • De cumulatieve beloning van NeuralUCB blijft tijdens het hele proces boven die van de andere methoden, met een groeiende kloof naarmate meer data wordt verwerkt.

• Kosten-Kwaliteit Afweging:

  • In vergelijking met een referentie die altijd het model met de maximale kwaliteit kiest, bereikt de NeuralUCB-methode ongeveer 33% van de inferentiekosten terwijl het een concurrerende kwaliteit behoudt.
  • Hoewel er nog een kwaliteitskloof blijft ten opzichte van de "max-quality" referentie, is de kostenbesparing aanzienlijk.

• Ablatie Studies:

  • Verschillende tekst-encoders werden getest. De resultaten toonden aan dat grotere encoders (zoals multilingual-e5-large-instruct) niet per se beter presteerden; kleinere, efficiëntere encoders zoals all-MiniLM-L6-v2 en all-mpnet-base-v2 leverden de meest stabiele en sterke resultaten.

4. Bijdragen en Significantie

• Innovatie: Dit paper introduceert een effectieve toepassing van NeuralUCB voor LLM-routing, wat een oplossing biedt voor het probleem van gedeeltelijke feedback en niet-lineaire relaties tussen vraagcomplexiteit en modelprestaties.
• Efficiëntie: Het biedt een praktische methode om de kosten van LLM-inferentie drastisch te verlagen zonder significante kwaliteitsverlies, wat cruciaal is voor schaalbare AI-toepassingen.
• Adaptiviteit: In tegenstelling tot gesuperviseerde methoden die vaak opnieuw getraind moeten worden bij veranderingen in de modelpool, is de NeuralUCB-approach adaptief en leert het beleid online door feedback.
• Uitdagingen: De auteurs wijzen erop dat er nog uitdagingen zijn in het discrimineren van acties (het vinden van het perfecte model) en het optimaliseren van exploratie, vooral wanneer het "beste" model al vaak de optimale keuze is.

Conclusie:
De studie concludeert dat NeuralUCB een veelbelovende aanpak is voor kostenbewuste LLM-routing. Het slaagt erin om een superieure afweging te vinden tussen kosten en kwaliteit, wat het een sterke kandidaat maakt voor productieomgevingen waar efficiëntie en kostenbeheersing essentieel zijn.