Résumé Technique : MixLM

Problématique

Les grands modèles de langage (LLM) ont démontré des capacités de classement de pertinence supérieures dans les systèmes de recherche et de recommandation en capturant des signaux sémantiques subtils que les méthodes lexicales traditionnelles ignorent. Cependant, le déploiement des LLM en tant que classificateurs de type « cross-encoder » dans des environnements industriels se heurte à un goulot d'étranglement critique : le coût computationnel.

Les approches standard de cross-encoder construisent des prompts combinant les requêtes utilisateur avec de longues descriptions d'articles candidates (souvent des milliers de jetons). Cela entraîne :

1. Une inférence de préremplissage (prefill) lourde : Le coût computationnel est dominé par la phase de préremplissage, qui croît de manière quadratique avec la longueur du contexte en raison des mécanismes d'attention.
2. Des contraintes de débit (throughput) : Sous des budgets de latence stricts (par exemple, 500 ms), les systèmes sont contraints d'utiliser soit des modèles plus petits et moins capables, soit de supprimer des caractéristiques informatives (via la résumé de texte ou l'élagage), sacrifiant ainsi la qualité du classement.

Le défi central est de réduire considérablement la longueur du contexte d'entrée tout en préservant l'interaction sémantique riche entre la requête et le document complet.

Méthodologie : Le cadre MixLM

Les auteurs proposent MixLM, un cadre de classement basé sur les LLM qui remplace le long texte de l'article par un ensemble compact de jetons d'embedding appris. Cette architecture de « mélange texte-embedding d'interaction » découple le traitement des descriptions d'articles de la phase de classement en ligne.

1. Architecture

MixLM se compose de deux composants principaux :

• L'encodeur LLM : Traite la description complète de l'article hors ligne (ou en quasi-temps réel) pour générer une représentation compacte. Au lieu de produire un vecteur unique, il produit une séquence de jetons d'embedding. En production, les auteurs utilisent l'échantillonnage du dernier-$N$ (last-$N$ sampling) pour sélectionner un petit nombre de jetons (par exemple, $T_S=1$) parmi les états cachés de l'encodeur pour représenter l'article.
• Le classificateur (Ranker) LLM : Opère en ligne. Il reçoit la requête utilisateur, le texte auxiliaire et les jetons d'embedding précalculés de l'encodeur. Ceux-ci sont concaténés pour former un prompt d'« entrée mixte », qui est nettement plus court que le prompt de texte intégral.

Le classificateur prédit la pertinence ($p_{yes}$) sur la base de cette entrée mixte. Comme les embeddings d'articles sont précalculés et mis en cache, l'inférence en ligne ne traite que le préfixe de la requête et quelques jetons d'embedding, réduisant ainsi drastiquement la longueur de la séquence.

2. Pipeline d'entraînement

Les auteurs emploient une recette d'entraînement en trois étapes pour aligner l'encodeur et le classificateur :

• Étape I (Fine-tuning de raisonnement de domaine) : Un LLM pré-entraîné de 0,6B est affiné sur un ensemble de données de raisonnement spécifique au domaine (180k échantillons) en utilisant des traces de chaîne de pensée (chain-of-thought) issues d'un juge de pertinence de 7B. Cela établit un modèle de base capable de comprendre les rationales de pertinence.
• Étape II (Entraînement du professeur de classement) : Le modèle de l'étape I est ensuite affiné en tant que classificateur de texte pur utilisant des descriptions d'articles complètes (10,9M d'échantillons). Ce modèle sert de « professeur » de haute qualité mais est trop coûteux pour un service en ligne.
• Étape III (Entraînement conjoint Encodeur-Classificateur) : L'encodeur et le classificateur sont co-entraînés en utilisant l'architecture d'entrée mixte. L'objectif d'entraînement combine :
  • Perte SFT de labels mous (Soft-Label SFT Loss) : Aligne les prédictions avec les étiquettes de vérité terrain.
  • Perte de distillation (Distillation Loss) : Aligne la distribution de sortie du modèle à entrée mixte avec celle du modèle professeur à texte pur (Étape II).
  • Perte d'auto-alignement (Self-Alignment Loss) : Garantit que les sorties d'embedding de l'encodeur sont compatibles avec les états cachés internes du classificateur. Cela inclut l'alignement de l'état caché (faire correspondre le dernier état caché de l'entrée mixte à celui de l'entrée de texte complet) et l'alignement de la prédiction.

Une stratégie d'apprentissage par curriculum est utilisée dans l'étape III, privilégiant initialement l'alignement de l'espace (compatibilité encodeur-classificateur) avant de se concentrer sur la distillation (précision de la prédiction de pertinence).

3. Optimisations de service

Pour maximifier le débit, les auteurs ont implémenté plusieurs optimisations au niveau du système :

• Préremplissage de préfixe partagé (Shared-Prefix Prefill) : Puisque le préfixe de la requête est identique pour tous les articles d'une liste de classement, le cache KV (Key-Value) pour le préfixe est calculé une seule fois et réutilisé pour tous les articles d'un batch.
• Mise en cache de préfixe en batch (In-Batch Prefix Caching) : Le moteur d'inférence (SGLang) est modifié pour fusionner les calculs d'attention, permettant aux jetons de suffixe (embeddings) d'être attentifs au préfixe partagé lors d'un seul passage direct (forward pass).
• Architecture multi-processus : La pile de service découple le frontend gRPC du moteur d'inférence en utilisant le multiprocessus pour contourner les goulots d'étranglement du verrou global de l'interpréteur (GIL) de Python.
• Gestion du cache KV : Les opérations de cache redondantes sont éliminées pour les charges de travail de préremplissage uniquement afin de réduire la surcharge CPU et d'augmenter la capacité effective de la mémoire GPU.

Contributions clés

1. Cadre MixLM : Une nouvelle architecture de classement à entrée mixte qui remplace les longs documents textuels par des jetons d'embedding compacts, réduisant la longueur des prompts de plusieurs ordres de grandeur tout en maintenant l'expressivité sémantique.
2. Pipeline d'entraînement de bout en bout : Une stratégie de co-entraînement avec des pertes personnalisées de distillation et d'auto-alignement qui transfère avec succès les connaissances d'un professeur à texte complet de haute qualité vers un modèle étudiant à entrée mixte.
3. Pile de service de production : Une optimisation complète du moteur d'inférence (incluant la mise en cache de préfixe partagé et la gestion multi-processus) qui permet un service LLM à haut débit.
4. Déploiement industriel : Le premier déploiement en plein trafic d'un système de classement sémantique basé sur les LLM dans la recherche d'emploi de LinkedIn, validé par des tests A/B à grande échelle.

Résultats expérimentaux

Évaluation hors ligne

• Qualité de la pertinence : MixLM a atteint un NDCG@10 de 0,9239, comparable au modèle professeur à texte pur (0,9432) et nettement supérieur à la récupération par embedding seul (0,8380) et aux baselines de texte résumé (0,9218).
• Débit (Throughput) : Sous un budget de latence fixe de 500 ms :
  • MixLM a atteint 22 000 articles/s/GPU.
  • Cela représente une amélioration de 10,0× par rapport au classement LLM par texte résumé (2 200 articles/s/GPU).
  • Cela représente une amélioration de 75,9× par rapport au classement LLM par texte complet (290 articles/s/GPU).

Tests A/B en ligne

• MixLM a été déployé sur 100 % du trafic pour la recherche d'emploi pilotée par l'IA de LinkedIn.
• Le système a maintenu des métriques de pertinence comparables à la baseline de production existante à échelle limitée.
• L'augmentation du débit a permis le déploiement complet du classement basé sur les LLM, entraînant une augmentation de 0,47 % des utilisateurs actifs quotidiens (DAU).

Signification et affirmations

L'article affirme que MixLM fournit une voie pratique et évolutive pour déployer des systèmes de classement basés sur les LLM dans des environnements réels. En comblant l'écart entre la richesse sémantique des LLM à texte complet et les exigences d'efficacité de la recherche industrielle, MixLM démontre que :

1. Compression de contexte via les embeddings : Il est possible de compresser des milliers de jetons en quelques jetons d'embedding sans perte significative de la qualité du classement.
2. Co-conception Système-Modèle : Un haut débit est obtenu non seulement par des changements d'architecture de modèle, mais aussi par une intégration étroite du design du modèle (entrée mixte) avec les optimisations d'infrastructure (mise en cache KV, batching).
3. Impact métier : Les gains d'efficacité permettent la transition de l'expérimentation à échelle limitée ou par texte résumé vers un déploiement à plein trafic des LLM, stimulant directement les mesures d'engagement des utilisateurs (DAU) dans un réseau professionnel à grande échelle.

Les auteurs positionnent ce travail comme une solution au compromis « latence-débit » qui a précédemment empêché l'adoption généralisée des cross-encoders LLM dans les systèmes de recherche en production.