One Model Is Enough: Native Retrieval Embeddings from LLM Agent Hidden States

Cette proposition de recherche introduit une méthode permettant aux agents LLM d'intégrer nativement la récupération d'informations via une projection légère de leurs états cachés, éliminant ainsi le besoin d'un modèle d'encodage séparé tout en préservant 97 % de la qualité de récupération.

L'essentiel

🧠 Le Problème : Deux Cerveaux pour une seule tâche

Imaginez que vous avez un assistant très intelligent (un LLM, comme un super-chatbot) qui vous aide à répondre à vos questions. Parfois, il ne connaît pas la réponse et doit aller chercher des informations dans une immense bibliothèque de documents.

Comment cela fonctionne-t-il aujourd'hui ?
C'est un peu comme si vous aviez deux personnes différentes dans la même équipe :

1. Le Chef (LLM) : Il écoute votre question, réfléchit, et écrit un résumé de ce qu'il faut chercher sur un bout de papier.
2. Le Bibliothécaire (Modèle d'encodage) : Il prend ce bout de papier, le relit, le traduit en un code secret (un vecteur) pour pouvoir le comparer aux livres de la bibliothèque.

Le problème ? C'est redondant !
Le Chef a déjà tout compris dans sa tête. Il a analysé votre question, le contexte de la conversation, et sait exactement ce qu'il faut chercher. Mais au lieu d'utiliser sa propre compréhension, il l'écrit sur un papier, et le Bibliothécaire doit recommencer tout le travail pour comprendre ce papier. C'est comme demander à quelqu'un de dessiner une carte, puis de demander à un autre de lire la carte pour savoir où aller. C'est lent et ça prend de la place.

💡 La Solution : "Un seul cerveau suffit"

Les auteurs de ce papier disent : "Pourquoi avoir deux modèles ?"

Ils proposent de donner au Chef (le LLM) un super-pouvoir : la capacité de se transformer directement en code secret, sans avoir besoin du Bibliothécaire.

Imaginez que le Chef a une "boîte à outils" cachée dans son cerveau (ses états cachés). Au lieu de sortir un papier, il ajoute simplement une petite loupe magique (une projection légère) qui transforme directement sa pensée en code de recherche.

• Avant : Pensée → Écriture → Traduction → Recherche.
• Maintenant : Pensée → Traduction directe → Recherche.

🛠️ Comment ils ont fait ? (L'entraînement)

Pour apprendre à cette "loupe magique" à bien fonctionner, ils ont utilisé une technique appelée distillation de connaissances. C'est comme un professeur (le modèle d'encodage actuel) qui enseigne à un élève (la loupe du LLM).

Ils ont utilisé trois méthodes d'enseignement :

1. L'Alignement : "Regarde, c'est comme ça que je traduis cette phrase. Fais pareil." (Pour que le code soit proche).
2. Le Contraste : "Si je compare deux questions différentes, mes codes doivent être très différents. Garde cette différence." (Pour ne pas tout mélanger).
3. Le Classement : "Quand je cherche un document, je mets le bon en premier. Toi aussi, essaie de mettre le bon document en premier." (Pour apprendre à bien trier).

🏆 Les Résultats : Rapide et Efficace

Ils ont testé leur méthode sur un jeu de données de conversations (QReCC). Voici ce qu'ils ont découvert :

• La Qualité : Le nouveau système est presque aussi bon que l'ancien (il garde 97% de la qualité). C'est comme si vous aviez perdu un tout petit peu de précision, mais que vous aviez gagné une vitesse incroyable.
• La Vitesse : C'est le vrai gagnant ! Le système est 21 fois plus rapide. Pourquoi ? Parce qu'il a supprimé l'étape où le "Bibliothécaire" devait relire le papier. Le Chef cherche directement avec sa propre pensée.
• La Complexité : Plus besoin de gérer deux modèles séparés. Juste un seul, avec une petite pièce ajoutée.

⚠️ Les Limites (Pour être honnête)

Ce n'est pas magique à 100% :

• Parfois, pour des questions très complexes ou avec des mots rares, le système fait encore quelques erreurs par rapport à l'ancien duo.
• Pour entraîner ce système, il faut quand même utiliser l'ancien "Bibliothécaire" au début (pour apprendre à la loupe), mais une fois entraîné, on peut le jeter !
• Pour l'instant, ça marche très bien quand le Chef et la Loupe viennent de la même "famille" de modèles (comme des frères).

🚀 En résumé

Ce papier nous dit qu'on peut rendre les assistants IA plus rapides et plus simples en leur apprenant à utiliser leur propre compréhension interne pour chercher des informations, au lieu de passer par un intermédiaire lent. C'est comme passer d'un système de messagerie postale (écrire, envoyer, lire) à un système de télépathie directe.

Résumé technique

1. Problématique : La redondance du pipeline à deux modèles

Les agents de langage (LLM) utilisant la génération augmentée par la récupération (RAG) suivent généralement un pipeline standard en deux étapes pour accéder à des connaissances externes :

1. Génération : L'agent LLM génère une requête de recherche sous forme de texte naturel.
2. Encodage : Un modèle d'embedding séparé encode ce texte en un vecteur dense pour la recherche dans un index documentaire.

L'auteur identifie une redondance fondamentale dans cette architecture :

• Le LLM a déjà traité l'intégralité du contexte conversationnel (intention de l'utilisateur, historique, tâches) et l'a encodé dans ses états cachés (hidden states).
• La génération du texte de la requête est une projection discrète et perteuse de cette représentation interne riche.
• Le modèle d'embedding séparé doit ensuite repartir de zéro pour extraire les informations sémantiques que le LLM possédait déjà, ce qui ajoute de la complexité infrastructurelle et de la latence.

Objectif : Éliminer le modèle d'embedding séparé en dotant l'agent LLM d'une capacité de récupération native, en projetant directement ses états cachés dans l'espace d'embedding.

---

2. Méthodologie : Projection native des états cachés

L'approche proposée consiste à attacher une tête de projection légère (projection head) au LLM pour transformer ses états cachés en vecteurs compatibles avec un espace d'embedding existant, sans nécessiter de nouvelle passe avant (forward pass) complète.

Architecture de la tête de projection

La fonction de projection $f$ transforme la séquence d'états cachés $H$ (variable en longueur) en un vecteur d'embedding fixe :

1. Projection d'entrée : Une couche linéaire mappe la dimension cachée du LLM ($d_h$) vers une dimension interne ($d_m$).
2. Encodeur Transformer : Une pile de couches Transformer (avec mécanisme d'attention auto) agrège les informations sur toute la séquence générée.
3. Pooling : Un pooling moyen (mean pooling) sur les positions non de remplissage (non-padding) pour compresser la séquence en un vecteur unique.
4. Projection de sortie et Normalisation : Une dernière couche linéaire projette vers la dimension cible ($d$), suivie d'une normalisation L2 pour garantir que la similarité par produit scalaire équivaut à la similarité cosinus.

Objectifs d'entraînement (Distillation de connaissances)

La tête est entraînée par distillation de connaissances à partir d'un modèle d'embedding "enseignant" (Teacher) en minimisant une combinaison de trois fonctions de perte :

1. Perte d'alignement ($L_{align}$) : Minimise la distance angulaire entre les vecteurs projetés et les embeddings de l'enseignant pour chaque requête. Cela fournit un signal de supervision direct.
2. Perte contrastive ($L_{contra}$) : Utilise une perte InfoNCE pour préserver la structure discriminative relative entre les requêtes au sein d'un batch.
3. Perte de distillation de rang ($L_{rank}$) : Transfère les préférences de classement de documents de l'enseignant vers l'élève via une divergence KL, en alignant les distributions de scores de similarité sur les documents candidats.

L'objectif total est une somme pondérée : $L = \lambda_a L_{align} + \lambda_c L_{contra} + \lambda_r L_{rank}$.

---

3. Contributions Clés

1. Identification et formalisation de la redondance : Démonstration que le pipeline standard à deux modèles gaspille les représentations sémantiques déjà présentes dans les états cachés du LLM.
2. Nouvelle architecture de projection : Conception d'une tête légère entraînée via une triple objective (alignement, contraste, rang) pour mapper les états cachés causaux du LLM vers un espace d'embedding bidirectionnel.
3. Validation empirique rigoureuse : Des expériences complètes avec 12 configurations d'ablation, des intervalles de confiance bootstrap et des tests de signification statistique, prouvant que la qualité de récupération est quasi équivalente à celle du pipeline standard, mais sans modèle d'embedding dédié.

---

4. Résultats Expérimentaux

Les expériences ont été menées sur le benchmark de recherche conversationnelle QReCC, dans un cadre "same-family" (LLM : Qwen3-8B, Enseignant : Qwen3-Embedding-8B).

Performance de Récupération

• Qualité : La méthode proposée conserve 97 % de la qualité de récupération de la ligne de base (Baseline) sur toutes les métriques.
  • Recall@10 : 0,607 (vs 0,637 pour la baseline, soit -3,0 %).
  • MRR@10 : 0,293 (vs 0,329, soit -3,6 %).
  • La différence est statistiquement significative (p = 0,0005), mais la méthode atteint une parité fonctionnelle.
• Latence : Réduction massive de la latence de 43,5 ms à 2,0 ms (accélération de 21,8x), car la passe avant du modèle d'embedding est supprimée.

Analyse des Ablations

• Composantes de la perte : L'alignement est la composante la plus critique. La distillation de rang seule échoue complètement (Recall@10 ≈ 0) car elle ne peut pas ancrer les embeddings dans l'espace sans l'alignement. La combinaison des trois pertes donne les meilleurs résultats.
• Recette d'entraînement : L'entraînement prolongé (80 époques) avec un taux d'apprentissage plus faible (2 × 10⁻⁴) est crucial. Un taux trop élevé (5 × 10⁻⁴) provoque un effondrement de l'entraînement.

---

5. Signification et Limites

Signification

Ce travail démontre qu'il n'est pas nécessaire de maintenir un modèle d'embedding séparé pour les agents LLM. En exploitant les états cachés déjà calculés lors de la génération, on peut obtenir une capacité de récupération native, simplifiant considérablement le déploiement (un seul modèle à l'inférence) et réduisant la latence, tout en maintenant une performance de recherche de haut niveau.

Limites

• Généralisation : Les résultats sont validés sur un seul dataset (QReCC) et dans un cadre "same-family". La généralisation à d'autres familles de modèles ou à des tâches de recherche open-domain reste à prouver.
• Écart statistique : Bien que faible, l'écart de performance (-3 % à -3,6 %) est statistiquement significatif, indiquant que la méthode n'est pas encore parfaitement équivalente à la baseline.
• Phase d'entraînement : Le modèle d'embedding reste nécessaire pendant la phase d'entraînement pour générer les cibles (enseignants), bien qu'il soit supprimé au moment du déploiement (inférence).
• Cas d'échec : La méthode échoue plus souvent qu'elle ne gagne sur des requêtes spécifiques (termes rares, chaînes de coréférence complexes), suggérant des difficultés dans la cartographie des états cachés causaux vers les représentations bidirectionnelles de l'enseignant pour les cas "long tail".

En conclusion, cette étude propose une voie prometteuse pour l'optimisation des systèmes RAG en intégrant la récupération directement dans le flux de génération du LLM.