LightRetriever: A LLM-based Text Retrieval Architecture with Extremely Faster Query Inference

Technisches Resümee: LightRetriever

Problemstellung

Aktuelle auf Large Language Models (LLM) basierende Textabrufsysteme verwenden typischerweise symmetrische Dual-Encoder-Architekturen, bei denen sowohl Dokumente als auch Abfragen durch tief parametrisierte LLMs verarbeitet werden, um dichte oder dünne (sparse) Repräsentationen zu generieren. Während dieser Ansatz die Abrufkapazitäten durch tiefes kontextuelles Modellierung erheblich verbessert, führt er einen kritischen Engpass beim Online-Einsatz ein. Dokumente können offline vorcodiert und indiziert werden, aber Suchanfragen kommen in Echtzeit an und erfordern eine sofortige Kodierung. Das Bereitstellen eines vollwertigen LLMs für die Query-Inferenz auf Online-Servern (z. B. A800 GPUs) führt zu hoher Latenz, niedrigem Durchsatz und übermäßigem Ressourcenverbrauch. Das Paper postuliert ein grundlegendes Spannungsverhältnis: Während Dokumente von schwerer Präkomputation profitieren, müssen Query-Encoder unter strengen Online-Beschränkungen operieren, doch bestehende Methoden zwingen Abfragen dazu, dieselbe rechenintensive tiefe Modellierung zu durchlaufen wie Dokumente.

Methodik: LightRetriever

Die Autoren schlagen LightRetriever vor, eine neuartige hybride Retrieval-Architektur, die die Symmetrie zwischen Dokument- und Query-Encodern explizit aufbricht. Die zentrale Erkenntik ist, dass Dokumente zwar die volle Modellierungskraft von LLMs benötigen, Abfragen jedoch keine gleichermaßen schwere Verarbeitung benötigen, um eine starke Retrieval-Leistung zu erzielen. Die Architektur verlagert den Großteil der Modellierungskosten auf die Dokumentenseite und reduziert die Query-Seite auf eine extrem leichte Operation.

1. Asymmetrische Architektur

• Dokumentenseite: Behält einen vollwertigen LLM-Encoder ($Enc_d$) für sowohl dichte als auch dünne Kodierung bei. Dokumente werden offline präkomputiert.
• Query-Seite: Eliminiert die tiefe Transformer-Inferenz während des Online-Servings. Der Query-Encoder ($Enc_q$) wird auf einen einfachen Embedding-Lookup-Mechanismus reduziert.

2. Dichter Retrieval-Mechanismus (Dense Retrieval)

Für das dichte Retrieval verwendet LightRetriever einen dreistufigen Prozess:

• Training: Ein vollwertiges LLM wird verwendet, um das System zu trainieren. Für ein gegebenes Query-Token $t_i$ (vorangestellt mit einer Aufgabeninstruktion) generiert das Modell eine Token-Level-Repräsentation $v^{den}_{t_i} = Enc_q(Inst; t_i)$. Der finale dichte Query-Vektor $v^{den}_q$ wird durch Mittelung dieser Token-Vektoren berechnet:
  $$v^{den}_q = \frac{1}{n} \sum_{i=0}^{n-1} v^{den}_{t_i}$$
  Das System wird mittels Contrastive Loss (Gleichung 2) optimiert, um Query- und relevante Dokumentenvektoren aufeinander abzustimmen.
• Caching: Da Query-Token während der Inferenz nicht tief miteinander interagieren, sind die Token-Vektoren cachebar. Vor dem Deployment präkomputiert das System die Token-Embeddings für das gesamte Vokabular ($V$) des LLMs und speichert sie in einer einzelnen Embedding-Lookup-Matrix $E$ (Form $[V, H]$).
• Online-Serving: Zur Inferenzzeit wird die Query tokenisiert, und der dichte Vektor wird generiert, indem die gecachten Embeddings für jedes Token nachgeschlagen und gemittelt werden. Dies ersetzt den Forward-Pass eines großen Transformers durch eine einzige Embedding-Lookup- und Mittelungsoperation, was lediglich eine im RAM befindliche Embedding-Schicht erfordert.

3. Dünner Retrieval-Mechanismus (Sparse Retrieval)

Für das dünne Retrieval ist die Query-Seite noch leichter:

• Query-Kodierung: Der dünne Query-Vektor $v^{spr}_q$ ist ein termbasierter, ungelernter Vektor, der direkt aus der Tokenisierung abgeleitet wird, wobei der Wert für ein Token $t$ einfach dessen Anzahl in der Query ist. Ein Encoder ist nicht erforderlich.
• Dokument-Kodierung: Das Dokument wird durch das volle LLM verarbeitet. Die Hidden States der letzten Schicht werden über den Language Modeling Head des Modells in den Vokabularraum projiziert. Um Sparsamkeit (Sparsity) zu induzieren, werden eine ReLU-Aktivierung und eine Log-Sättigungsfunktion angewendet, gefolgt von Max-Pooling über die Sequenzlänge.
• Hybrides Scoring: Dichte und dünne Scores werden linear interpoliert, um das finale Retrieval-Ranking zu erzeugen.

Zentrale Beiträge

1. Extreme Query-Inferenzgeschwindigkeit: Durch den Ersatz der tiefen Query-Modellierung durch einen Embedding-Lookup erreicht LightRetriever eine 1000-fache Beschleunigung der Query-Kodierung im Vergleich zum Servieren eines vollen Llama-8b-Modells auf einer A800 GPU.
2. Durchsatzverbesserung: Die Architektur ermöglicht eine 10-fache Steigerung des End-to-End-Retrieval-Durchsatzes (Queries Per Second).
3. Leistungserhalt: Trotz der drastischen Reduktion der Rechenlast auf der Query-Seite behält die Methode etwa 95 % der Retrieval-Leistung voller symmetrischer LLM-basierter Retriever über verschiedene Benchmarks hinweg bei.
4. Generalisierung: Der Ansatz wird über verschiedene Foundation-Modelle (Llama-1B/3B/8B, Qwen-1.5B/3B/7B) und Retrieval-Aufgaben (BeIR, CMTEB) hinweg validiert, was die Robustheit ohne die Notwendigkeit aufgabenspezifischer architektonischer Änderungen demonstriert.

Experimentelle Ergebnisse

Das Paper evaluiert LightRetriever auf 15 englischen Aufgaben aus BeIR und 8 chinesischen Aufgaben aus CMTEB Retrieval.

• Geschwindigkeit: Auf einem Benchmark von 65.536 Queries gegen 1 Mio. Passagen reduzierte LightRetriever die gesamte End-to-End-Retrieval-Zeit von ~119s (Full-Llama8b) auf ~9,36s und erreichte einen Durchsatz von 6.999 QPS im Vergleich zu 549 QPS des vollen Modells.
• Effektivität:
  • LightRetriever-Llama3.1-8b erreichte ein nDCG@10 von 54,4 auf BeIR und 63,0 auf CMTEB-R.
  • Dies stellt eine moderate Verschlechterung von -2,4 bzw. -4,6 Punkten gegenüber der vollen symmetrischen Llama3.1-8b Baseline (56,8 und 67,6) dar, vermeidet jedoch die hohen Kosten der Query-Inferenz.
  • Die Methode übertrifft andere effiziente Baselines wie BGE-m3 (49,6 auf BeIR) und nähert sich der Performance von E5-Mistral und LLM2Vec an.
• Ablationsstudien:
  • Das Entfernen des vollen Dokument-Encoders (symmetrisch leichtgewichtig) verursachte schwere Leistungseinbußen (z. B. -13,8 auf BeIR), was die Notwendigkeit der tiefen Dokument-Modellierung bestätigt.
  • Die Verwendung eines reinen Embedding Bags für das Query-Training (ohne den vollen Transformer) führte ebenfalls zu signifikanten Leistungsverlusten (-11,2 auf BeIR), was darauf hindeutet, dass tiefe Modellierung während der Trainingsphase essenziell ist, um effektive Repräsentationen zu lernen, selbst wenn sie während der Inferenz verworfen wird.

Bedeutung und Ansprüche

Das Paper behauptet, dass LightRetriever ein skalierbares Paradigma für latenzsensitive und hochdurchsatzfähige Retrieval-Szenarien bietet. Es adresst die praktische Ineffizienz beim Einsatz riesiger LLMs für die Echtzeit-Query-Kodierung, ohne die Vorteile des semantischen Verständnisses der LLM-basierten Retrieval zu opfern. Die Autoren betonen, dass ihre Arbeit die Rechenlast von der Online-Query-Seite auf die Offline-Dokument-Seite überträgt, wodurch Systeme die Kraft großer Modelle nutzen können, während sie gleichzeitig die für Produktions-Suchmaschinen erforderliche Reaktionsfähigkeit beibehalten. Die Methode wird als die erste präsentiert, die explizit den Online-Effizienz-Engpass auf der Query-Seite von LLM-basierten Retrivern durch diese extreme Entkopplung adressiert.