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

Dit paper introduceert een methode waarbij LLM-agenten native zoekopdrachten uitvoeren door hun eigen verborgen staten direct naar de embedding-ruimte te projecteren, waardoor een apart embedding-model overbodig wordt terwijl 97% van de zoekkwaliteit behouden blijft.

De kern

Stel je voor dat je een zeer slimme assistent hebt (een LLM, of "Groot Taalmodel") die voor je zoekt op internet. Op dit moment werkt die assistent in twee stappen, wat een beetje inefficiënt is:

1. Stap 1: De assistent denkt na en schrijft een zoekopdracht op een briefje (bijvoorbeeld: "Hoe maak ik een goede pizza?").
2. Stap 2: Hij geeft dat briefje aan een tweede persoon (een apart "inbeddingsmodel"). Die tweede persoon leest het briefje, denkt erover na, en maakt er een complexe code van om de beste antwoorden te vinden.

Het probleem? De assistent had het antwoord al in zijn hoofd! Hij had de volledige context al verwerkt toen hij het briefje schreef. De tweede persoon moet dus hetzelfde werk opnieuw doen, wat tijd kost en extra apparatuur vereist.

Dit artikel introduceert een slimme oplossing: "Eén model is genoeg."

De Nieuwe Methode: De "Inwendige Vertaler"

De auteurs van dit paper zeggen: "Waarom een tweede persoon inhuren? Laten we de assistent gewoon een klein, slim hulpmiddel geven dat direct zijn gedachten vertaalt."

In plaats van een briefje te schrijven en door te geven, krijgt de assistent een kleine projectie-kop (een soort vertaalbril) op zijn hoofd.

• Hoe het werkt: Terwijl de assistent denkt, heeft hij al "verborgen staten" (gedachten) in zijn hoofd. Deze nieuwe kop pakt die gedachten direct en vertaalt ze in de code die nodig is om te zoeken.
• Het resultaat: De assistent zoekt nu met zijn eigen gedachten, zonder een briefje te schrijven en zonder een tweede persoon nodig te hebben.

Hoe leren ze dit? (De Drie Leraars)

Om deze "vertaalbril" goed te laten werken, trainen ze het met drie soorten oefeningen (verliezen), alsof je een student voorbereidt op een examen:

1. De Spiegel (Alignement): De assistent kijkt in de spiegel van de oude, sterke tweede persoon. Hij probeert zijn gedachten zo dicht mogelijk bij die van de meester te houden. "Zie je hoe de meester dit denkt? Ik moet het net zo doen."
2. De Groepsdynamiek (Contrastief): De assistent leert welke vragen op elkaar lijken en welke verschillen. "Als vraag A lijkt op vraag B, dan moeten hun gedachten ook op elkaar lijken. Maar vraag C is heel anders, dus die moet ver weg staan."
3. De Ranglijst (Rank Distillation): De assistent leert niet alleen wat de gedachten zijn, maar ook welke antwoorden het beste zijn. Hij kijkt naar de ranglijst van de meester: "Voor deze vraag was antwoord 1 het beste, antwoord 2 was goed, en antwoord 3 slecht. Ik moet diezelfde volgorde kunnen maken."

Wat is het resultaat?

De tests tonen aan dat deze methode fantastisch werkt:

• Snelheid: Het is 21 keer sneller. Omdat ze de tweede persoon (de zware rekenmachine) niet meer nodig hebben, duurt het zoeken maar een fractie van de tijd.
• Kwaliteit: De kwaliteit van de zoekresultaten is bijna hetzelfde (97% van de originele kwaliteit). Het is alsof je een Formule 1-auto hebt, maar je rijdt er nu met een elektrische scooter op die net zo snel is, maar veel minder brandstof verbruikt.
• Kosten: Je hebt minder hardware nodig, omdat je geen extra zware modellen hoeft te draaien.

De Grootte van de "Vertaler"

De extra "vertaalbril" die ze toevoegen, is heel klein (ongeveer 25 miljoen parameters). Vergelijk het met het toevoegen van een klein navigatiesysteem aan een gigantische vrachtwagen (de LLM). De vrachtwagen is al groot en krachtig; het kleine navigatiesysteem maakt hem alleen nog slimmer en sneller, zonder dat hij zwaarder wordt.

Conclusie

Dit onderzoek laat zien dat we vaak onnodig complexe systemen bouwen. In plaats van een assistent te laten praten en dan iemand anders te laten luisteren, kunnen we de assistent gewoon laten denken in de taal die de zoekmachine begrijpt. Het is een stap naar slimmere, snellere en goedkopere AI-assistenten die direct kunnen handelen zonder onnodige tussenstappen.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "One Model Is Enough: Native Retrieval Embeddings from LLM Agent Hidden States" in het Nederlands.

Probleemstelling

In de huidige standaard voor Retrieval-Augmented Generation (RAG) werken LLM-agenten (Large Language Models) volgens een tweestaps-pijplijn:

1. De LLM genereert een zoekopdracht als natuurlijke tekst.
2. Een apart, gespecialiseerd embedding-model encodeert deze tekst naar een vector voor zoekopdrachten in een documentindex.

Het paper identificeert een fundamentele redundantie in deze architectuur. De LLM heeft de volledige conversatiestroom (gebruikersintentie, dialooggeschiedenis, taakeisen) al verwerkt en deze informatie is aanwezig in de verborgen staten (hidden states) van het model. Het genereren van tekst is slechts een verliesrijke, discrete projectie van deze rijke interne representatie. Een tweede model moet deze tekst vervolgens opnieuw verwerken om semantische informatie te herstellen die de LLM al bezat. Dit leidt tot onnodige infrastructuurcomplexiteit en latentie.

Methodologie

De auteurs stellen een nieuwe aanpak voor waarbij de LLM-agent wordt uitgerust met natieve zoekcapaciteit, waardoor een apart embedding-model tijdens de inferentie overbodig wordt.

Kerncomponenten:

• Projectie-Head: Een lichtgewicht leerbaar component dat wordt toegevoegd aan de LLM. Dit projecteert de verborgen staten van de LLM direct naar de bestaande embedding-ruimte.
• Extrahering: Tijdens de autoregeneratieve generatie van de zoekopdracht worden de verborgen staten ($h_i$) van de laatste laag van de LLM verzameld. Deze worden niet weggegooid, maar gebruikt als input.
• Architectuur van de Projectie-Head:
  1. Input Projectie: Een lineaire laag die de dimensie van de LLM ($d_h$) mapt naar een interne dimensie ($d_m$).
  2. Transformer Encoder: Een stack van encoder-lagen met zelf-attention om afhankelijkheden tussen tokens te aggregeren.
  3. Pooling: Mean pooling over de geldige posities om een vaste vector te krijgen.
  4. Output Projectie & Normalisatie: Een laatste lineaire laag gevolgd door L2-normalisatie, zodat de dot-product gelijkstaat aan cosijn-sequentie (compatibel met de bestaande index).

Trainingsdoelstellingen (Knowledge Distillation):
De projectie-head wordt getraind via kennisdistillatie van een leraar-model (het embedding-model) met een combinatie van drie verliesfuncties:

1. Alignment Loss: Minimaliseert de hoekafstand tussen de geprojecteerde vector en de leraar-embedding (directe supervisie).
2. Contrastive Loss (InfoNCE): Behoudt de relatieve structuur tussen verschillende zoekopdrachten binnen een batch.
3. Rank Distillation Loss: Transfereert de voorkeuren voor documentranking van de leraar via KL-divergentie, gebaseerd op de scores van de top-K kandidaat-documenten.

De totale loss is een gewogen som van deze drie componenten: $L = \lambda_a L_{align} + \lambda_c L_{contra} + \lambda_r L_{rank}$.

Belangrijkste Bijdragen

1. Formalisatie van Redundantie: Het paper identificeert en formaliseert de inefficiëntie van het huidige tweemodelsysteem voor RAG en stelt projectie van verborgen staten voor als een alternatief.
2. Nieuwe Trainingsobjectieven: Ontwerp van een drie-delige trainingsstrategie (alignment, contrastief, rank distillatie) die specifiek is afgestemd op het projecteren van LLM-verborgen staten naar een embedding-ruimte.
3. Uitgebreide Validatie: Biedt uitgebreide experimenten met 12 ablatie-configuraties, bootstrap-vertrouwensintervallen en statistische significantietests, wat aantoont dat de kwaliteit bijna gelijk is aan die van een apart embedding-model.

Resultaten

De methode is geëvalueerd op het QReCC-benchmark (conversational search) met een "same-family" setting (Qwen3-8B als agent, Qwen3-Embedding-8B als leraar).

• Kwaliteit: De beste configuratie behoudt 97% van de zoekkwaliteit van de baseline (standaard genereren-then-encoderen).
  • Recall@10: 0.607 (vs. 0.637 baseline, een daling van slechts 3,0%).
  • MRR@10: 0.293 (vs. 0.329 baseline).
• Latentie: Er is een drastische reductie in inferentie-tijd, van 43,5 ms naar 2,0 ms (een factor 21,8 sneller), omdat de dure forward pass van het embedding-model wordt verwijderd.
• Statistiek: McNemar's test bevestigt dat het verschil in kwaliteit statistisch significant is ($p=0.0005$), maar de methodes komen in 84,2% van de gevallen overeen.
• Ablatie-studies:
  • Alignment is de sterkste individuele component.
  • Rank distillation werkt niet alleen (het model "instort" zonder alignment), maar levert bijdragen wanneer gecombineerd.
  • Trainingsduur: Lange training (80 epochs) met een lagere leerfrequentie ($2 \times 10^{-4}$) bleek cruciaal voor de beste prestaties; te hoge leerfrequenties leiden tot instorting.

Betekenis en Impact

Dit werk is significant omdat het de noodzaak van een apart embedding-model tijdens de inferentie elimineert.

• Efficiëntie: Het vereenvoudigt de infrastructuur en verlaagt de latency aanzienlijk, wat essentieel is voor real-time applicaties.
• Architectuur: Het bewijst dat de interne representaties van een generatief LLM voldoende informatie bevatten voor effectieve zoekopdrachten, mits correct geprojecteerd.
• Toekomstperspectief: Hoewel de huidige resultaten veelbelovend zijn, zijn er beperkingen: het werkt momenteel alleen binnen dezelfde model-familie (shared pretraining) en vereist nog steeds het embedding-model tijdens de trainingsfase. Toekomstig werk moet generalisatie naar andere domeinen en cross-family settings onderzoeken.

Kortom, het paper toont aan dat "één model voldoende is" voor zowel generatie als retrieval, mits de juiste projectie-mechanismen worden toegepast op de reeds berekende verborgen staten.