Reward-Based Online LLM Routing via NeuralUCB

1. Définition du Problème

L'article aborde le défi du routage dynamique des grands modèles de langage (LLM). Les LLMs présentent des variations significatives en termes de capacités et de coûts d'inférence selon les tâches. Les modèles de pointe offrent souvent une meilleure qualité mais à un coût élevé, tandis que les modèles moins chers peuvent échouer sur des requêtes complexes (raisonnement, code).

L'objectif est de concevoir un système de routage en ligne qui, pour chaque requête utilisateur $x_t$, sélectionne le modèle $a_t$ optimal parmi un ensemble de candidats $A$ afin de maximiser l'utilité à long terme, en trouvant un équilibre entre la qualité de la réponse et le coût d'inférence.

Le problème est formulé comme un problème de bandit contextuel :

• Le système observe une requête et sélectionne un modèle.
• Il ne reçoit de retour (feedback) que pour le modèle sélectionné (feedback partiel), et non pour les autres modèles candidats.
• Le défi réside dans la nature non linéaire de la relation entre la requête et le choix optimal, ainsi que dans la rareté des données de feedback.

2. Méthodologie

La proposition des auteurs repose sur une approche combinant l'apprentissage profond et l'exploration par bandit, nommée NeuralUCB.

A. Fonction d'Utilité et Récompense

Au lieu de maximiser uniquement la qualité, l'objectif est de maximiser une récompense d'utilité qui pénalise le coût. La récompense $r(x, a)$ est définie comme suit :
$$r(x, a) = q(x, a) \cdot \exp(-\lambda \tilde{c}(x, a))$$
Où :

• $q(x, a)$ est la qualité du modèle.
• $\tilde{c}(x, a)$ est le coût normalisé (via une transformation logarithmique pour stabiliser l'échelle).
• $\lambda$ est un paramètre contrôlant la force de la pénalité de coût.

B. Architecture : UtilityNet

Le cœur du système est un réseau de neurones appelé UtilityNet, qui sert de prédicteur d'utilité et de source de représentations pour l'exploration.

• Entrées : Le réseau prend en entrée l'embedding du texte ($x_{emb}$), des caractéristiques auxiliaires ($x_{feat}$), l'identifiant du domaine ($d$) et l'identifiant de l'action (modèle) $a$.
• Branches :
  1. Branche d'utilité : Prédit la récompense attendue $\mu(x, a)$ et génère une représentation latente $h(x, a)$. Elle est entraînée pour régresser la récompense réelle via une fonction de perte Huber.
  2. Branche de "gating" (commutation) : Prédit une probabilité $p(x)$ indiquant si l'exploration (via le bonus UCB) est bénéfique pour cette requête spécifique.

C. Politique de Routage basée sur NeuralUCB

Pour sélectionner l'action, l'algorithme utilise le score UCB (Upper Confidence Bound) :
$$s(x, a) = \mu(x, a) + \beta \sqrt{g(x, a)^\top A^{-1} g(x, a)}$$

• $\mu(x, a)$ : Utilité prédite par UtilityNet.
• Le terme de racine carrée représente l'incertitude (bonus d'exploration) basée sur la représentation latente $h(x, a)$.
• $A^{-1}$ : Une matrice de covariance inverse partagée estimant l'incertitude dans l'espace des features de la dernière couche.

Mécanisme de sécurité :
La décision finale combine le score UCB et une action de repli conservatrice (basée uniquement sur l'utilité moyenne $\mu$) :

• Si la probabilité de la branche de gating $p(x)$ dépasse un seuil $\tau_g$, le système sélectionne l'action maximisant le score UCB (exploration active).
• Sinon, il choisit l'action "sûre" ($\text{arg max } \mu(x, a)$).

D. Protocole d'Évaluation

L'évaluation se fait dans un environnement simulé en ligne sur le benchmark RouterBench (36 497 échantillons, 86 domaines, 11 modèles). Les données sont divisées en 20 tranches (slices) traitées séquentiellement pour simuler un déploiement réel avec mise à jour continue du modèle.

3. Résultats Expérimentaux

Les expériences comparent la méthode NeuralUCB proposée avec plusieurs baselines :

• Random (sélection aléatoire).
• Min-cost (sélection du modèle le moins cher).
• RouteLLM-BERT (méthode supervisée existante).
• Référence Max-Quality (le meilleur modèle possible sans contrainte de coût).

Principales observations :

1. Performance de la Récompense : NeuralUCB surpasse systématiquement les baselines (Random, Min-cost, RouteLLM-BERT) en termes de récompense moyenne et cumulative. La récompense moyenne se stabilise autour de 0,59 - 0,61, contre ~0,35 pour RouteLLM-BERT et ~0,51 pour Min-cost.
2. Efficacité Coût-Qualité : Par rapport à la référence "Max-Quality", la méthode NeuralUCB réduit le coût d'inférence d'environ 67% (n'utilisant que ~33% du coût) tout en maintenant une qualité de sélection compétitive.
3. Impact des Encodeurs : L'ablation sur les encodeurs de texte montre que des modèles comme all-mpnet-base-v2 et all-MiniLM-L6-v2 fonctionnent mieux que des modèles plus grands ou spécifiques (comme multilingual-e5-large-instruct), suggérant que la taille de l'encodeur n'est pas le facteur déterminant dans ce contexte.
4. Stabilité : La méthode maintient une performance stable au fil des tranches de données, prouvant sa capacité à s'adapter en ligne.

4. Contributions Clés

• Cadre Unifié : Introduction d'une approche de routage LLM basée sur le bandit contextuel neuronal (NeuralUCB) qui gère nativement le compromis coût-qualité via une fonction d'utilité exponentielle.
• Architecture Hybride : Conception d'UtilityNet avec une branche de gating permettant de décider dynamiquement quand explorer (UCB) ou quand exploiter (action sûre), améliorant la robustesse face aux incertitudes.
• Validation Rigoureuse : Évaluation complète sur RouterBench dans un protocole simulé en ligne, démontrant une supériorité claire par rapport aux méthodes supervisées et aux heuristiques simples.

5. Signification et Perspectives

Cette étude démontre que NeuralUCB est une approche prometteuse pour le routage de LLMs sensible aux coûts, capable d'apprendre des politiques non linéaires complexes à partir de feedbacks partiels.

Cependant, l'article identifie des défis restants :

• Discrimination des actions : Pour certains échantillons, le choix du modèle de "meilleure qualité" coïncide déjà avec le choix optimal en termes de récompense, limitant le gain potentiel du routage.
• Exploration : L'amélioration de la discrimination des actions sous contraintes de coût nécessite davantage de données d'entraînement et un raffinement du modèle d'utilité.

En conclusion, ce travail ouvre la voie à des systèmes de routage plus intelligents et économiques pour les déploiements massifs de LLMs, réduisant les coûts opérationnels sans sacrifier significativement la qualité de service.