Linear-Time and Constant-Memory Text Embeddings Based on Recurrent Language Models

Probleemstelling

Huidige state-of-the-art tekst-embeddingmodellen zijn gebaseerd op Transformer-architecturen. Hoewel deze modellen uitstekende prestaties leveren, hebben ze twee fundamentele beperkingen die hun bruikbaarheid voor lange documenten en hulpbronnenbeperkte omgevingen in de weg staan:

1. Kwadratische computatiecomplexiteit: De aandachtsmechanismen (attention) vereisen $O(T^2)$ berekeningen voor een sequentie van lengte $T$.
2. Lineaire geheugencomplexiteit: De opslag van tussenliggende activaties (activations) groeit lineair met de sequentielengte ($O(T)$).

Dit maakt het onpraktisch om zeer lange documenten (bijv. duizenden tokens) te verwerken zonder dat het geheugen (VRAM) volloopt, zelfs met optimalisaties zoals FlashAttention. Recurrente modellen (zoals RNN's) bieden weliswaar lineaire schaalbaarheid en constant geheugengebruik, maar traditionele RNN's missen de parallelle verwerkingsefficiëntie die nodig is voor groot-schalig trainen en inferentie.

Methodologie

De auteurs stellen een nieuwe aanpak voor die recurrente architecturen (specifiek Mamba2, RWKV en xLSTM) combineert met een geavanceerde inferentiestrategie om de voordelen van beide werelden te verenigen.

1. Training van Recurrente Embedding Modellen

• De auteurs passen het trainingsrecept van E5mistral-7b toe op voorgeprogrammeerde recurrente taalmodellen (Mamba2).
• Ze gebruiken contrastief leren (InfoNCE loss) om de modellen te fine-tunen voor het genereren van embeddings.
• Instruction Tuning: Er wordt een instructie-template toegevoegd aan de query (niet aan het document) om het model aan te passen aan diverse downstream-taken zonder dat het document-corpus opnieuw hoeft te worden ingesloten.
• De embedding wordt gegenereerd door de output van het laatste token (EOS) van de laatste laag te gebruiken.

2. De "Vertically Chunked" Inferentiestrategie

Dit is de kerntechnische bijdrage. De auteurs benutten de dualiteit van moderne recurrente modellen (zoals Mamba2), die zowel als lineaire recurrente vergelijking als als een parallelle matrixtransformatie kunnen worden geformuleerd.

• Chunked Intra-Layer: De inputsequentie wordt opgedeeld in blokken (chunks). Binnen een blok wordt de berekening parallel uitgevoerd via matrixvermenigvuldiging (zoals bij Transformers), terwijl de recurrente staat tussen blokken sequentieel wordt bijgewerkt.
• Verticale Chunking (Cross-Layer): In plaats van de volledige sequentie door één laag te verwerken voordat de volgende laag begint (horizontale inferentie), verwerken ze een vast blok van tokens (verticale chunk) door alle lagen van het model voordat ze doorgaan met het volgende blok.
  • Geheugenvoordeel: Hierdoor hoeven ze niet de volledige sequentie-activaties voor elke laag in het geheugen te houden. Ze hoeven alleen de activaties van het huidige verticale blok en één recurrente staat per laag op te slaan.
  • Complexiteit: Zodra de inputlengte de grootte van het verticale blok overschrijdt, wordt het geheugenverbruik constant ($O(1)$ ten opzichte van de sequentielengte), in plaats van lineair. De rekentijd blijft lineair ($O(T)$).

Belangrijkste Bijdragen

1. Validatie van Recurrente Modellen als Embedders: Het paper toont aan dat recurrente modellen (Mamba2, RWKV, xLSTM) na fine-tuning concurrerende prestaties leveren ten opzichte van Transformer-baselines voor algemene tekstembeddingtaken.
2. Verticale Chunking Strategie: Een nieuwe inferentiemethode die de geheugencomplexiteit van recurrente modellen onafhankelijk maakt van de sequentielengte, terwijl de parallelle snelheid van moderne hardware behouden blijft.
3. Empirisch Bewijs: De strategie is getest op drie verschillende recurrente architecturen (Mamba2, RWKV7, xLSTM) en bevestigt dat het geheugenverbruik constant blijft voor lange sequenties, in tegenstelling tot Transformers.
4. Open Source: De fine-tuned checkpoints en de inferentie-implementatie zijn openbaar gemaakt om reproduceerbaarheid te garanderen.

Resultaten

De evaluatie is uitgevoerd op de MTEB (Massive Text Embedding Benchmark) en LongEmbed benchmarks.

• Kwaliteit (Performance):

  • De recurrente modellen (zoals Codestral Mamba2 7B) presteren concurrerend met Transformer-modellen (zoals Mistral 7B) op multilinguale taken en lange context-taken.
  • Op Engelse taken zijn de recurrente modellen iets minder goed (1-2% verschil), wat de auteurs toeschrijven aan het feit dat de onderliggende recurrente base-modellen minder data hebben gezien tijdens pre-training dan hun Transformer-tegenhangers.
  • Op de LongEmbed benchmark (tot 32k tokens) presteren ze vergelijkbaar met Transformers.

• Efficiëntie (Runtime & Geheugen):

  • Geheugen: Voor lange sequenties (boven de verticale chunk-grootte, bijv. 4096 tokens) blijft het geheugenverbruik van recurrente modellen constant, terwijl het bij Transformers lineair blijft stijgen. Dit maakt het mogelijk om zeer lange documenten te verwerken op hardware waar Transformers zouden crashen.
  • Runtime: De rekentijd is vergelijkbaar met of sneller dan Transformers voor lange sequenties. De parallelle verwerking binnen de chunks zorgt ervoor dat er geen significant runtime-overhead is ten opzichte van volledig parallelle inferentie, zolang de chunk-grootte goed is gekozen.
  • Hardware-onafhankelijkheid: De resultaten zijn consistent gebleken op verschillende GPU's (L40S, A100, H100).

Significantie

Dit paper markeert een belangrijke stap in de evolutie van NLP-modellen voor lange contexten:

• Schalbaarheid: Het biedt een praktische oplossing voor het probleem van "out-of-memory" fouten bij het verwerken van lange documenten, wat essentieel is voor toepassingen zoals juridische documentanalyse, medische dossiers of lange rapporten.
• Alternatief voor Transformers: Het bewijst dat recurrente architecturen, die eerder vooral werden gezien als een alternatief voor generatie, ook uitstekend werken voor representatieleer (embeddings) en dit met aanzienlijk minder geheugenkosten.
• Toekomstperspectief: De auteurs suggereren dat met beter voorgeprogrammeerde recurrente base-modellen de kwaliteitsgap met Transformers volledig kan worden gedicht, wat leidt tot een nieuwe klasse van efficiënte, lange-context embeddingmodellen.

Kortom, de auteurs hebben een brug geslagen tussen de efficiëntie van recurrente netwerken en de prestaties van moderne taalmodellen, waardoor tekstembedding voor lange documenten haalbaar wordt op beperkte hardware.