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

1. Il Problema

I modelli di embedding testuali basati su Transformer sono attualmente lo stato dell'arte per compiti come il recupero dell'informazione, il clustering e la stima della similarità semantica. Tuttavia, soffrono di due limitazioni fondamentali che ne riducono l'utilità per documenti lunghi o in ambienti con risorse limitate:

• Complessità computazionale quadratica: Il meccanismo di attenzione richiede $O(T^2)$ operazioni rispetto alla lunghezza della sequenza $T$.
• Complessità di memoria lineare: La necessità di memorizzare le attivazioni intermedie per l'attenzione porta a un uso della memoria che cresce linearmente con la lunghezza dell'input ($O(T)$).

Questi fattori rendono i Transformer inefficienti per documenti molto lunghi (es. >4k-8k token) e costosi da eseguire su hardware limitato. Sebbene le architetture ricorrenti moderne (come Mamba2, RWKV e xLSTM) offrano complessità computazionale lineare e uso di memoria costante per la generazione di testo, il loro potenziale come generatori di embedding generici non è stato sufficientemente esplorato.

2. Metodologia

Gli autori propongono un approccio basato su tre pilastri principali:

A. Addestramento dei Modelli Ricorrenti

Per valutare l'efficacia delle architetture ricorrenti, gli autori applicano la procedura di addestramento contrastivo standard utilizzata da E5mistral-7b (Wang et al., 2024) a modelli linguistici pre-addestrati basati su Mamba2.

• Architettura: Vengono rimossi i head di modellazione linguistica dai modelli pre-addestrati (Mamba2 1.3B e Codestral Mamba2 7B) per creare modelli di embedding.
• Input: Si aggiunge un token di fine sequenza (EOS) all'input. L'output dello strato finale in questa posizione viene utilizzato come embedding del testo.
• Loss: Viene utilizzata la funzione di perdita InfoNCE per l'addestramento contrastivo, con un tuning di istruzione (instruction tuning) per migliorare l'adattabilità ai diversi task a valle.

B. Strategia di Inferenza a Chunking Verticale

Il contributo tecnico principale è una strategia di inferenza ottimizzata che sfrutta la dualità dei modelli SSM (State Space Models) come Mamba2. Questi modelli possono essere visti sia come ricorrenze lineari (efficienti in memoria) sia come trasformazioni matriciali parallele (efficienti in velocità).

• Chunking Orizzontale (Standard): Divide la sequenza in blocchi per calcolare l'attenzione interna in parallelo, ma richiede la memorizzazione di tutte le attivazioni intermedie lungo la sequenza, mantenendo la memoria $O(T)$.
• Chunking Verticale (Proposta): Gli autori introducono un approccio che divide la sequenza di input in blocchi verticali di dimensione fissa $V$.
  • Il modello processa un intero blocco verticale attraverso tutti gli strati del modello prima di passare al blocco successivo.
  • Tra un blocco e l'altro, viene mantenuto solo lo stato ricorrente (latente) per ogni strato, invece di tutte le attivazioni della sequenza completa.
  • Risultato: Una volta che la lunghezza dell'input supera la dimensione del chunk verticale ($T > V$), l'uso della memoria diventa costante ($O(1)$ rispetto a $T$), dipendendo solo dalla profondità del modello e dalla dimensione del chunk.

C. Validazione su Multiple Architetture

La strategia di inferenza è stata validata empiricamente non solo su Mamba2, ma anche su RWKV e xLSTM, dimostrando che il trade-off runtime-memoria è coerente tra diverse architetture ricorrenti moderne.

3. Risultati Chiave

Gli esperimenti sono stati condotti sui benchmark MTEB (inglese e multilingue) e LongEmbed.

• Qualità degli Embedding:

  • I modelli basati su Mamba2 raggiungono prestazioni competitive rispetto ai baselines Transformer di dimensioni simili.
  • Su MTEB Multilingue, il modello Codestral Mamba2 7B supera il Transformer Mistral 7B v0.1.
  • Su MTEB Inglese, i modelli ricorrenti sono leggermente inferiori (1-2 punti percentuali), probabilmente a causa di differenze nella quantità di dati di pre-addestramento (Mamba2 1.3B è stato addestrato su meno dati rispetto a Qwen2/Mistral).
  • Su LongEmbed (documenti lunghi fino a 32k token), le prestazioni sono comparabili, con differenze minime (<0.2 punti).

• Efficienza di Inferenza (Runtime e Memoria):

  • Memoria: I modelli ricorrenti con chunking verticale mostrano un uso della memoria che rimane costante per sequenze lunghe, a differenza dei Transformer che crescono linearmente. Questo permette di processare documenti molto lunghi su GPU con memoria limitata (es. 80GB H100) dove i Transformer fallirebbero.
  • Runtime: Per sequenze singole, i modelli ricorrenti sono più veloci dei Transformer per lunghezze superiori a 4096 token. Per batch di grandi dimensioni, le prestazioni sono comparabili, ma i ricorrenti offrono un vantaggio significativo in termini di stabilità della memoria.
  • Saturazione: L'parallelizzazione si satura a dimensioni di chunk verticale relativamente piccole (es. 4096 token per sequenze singole), rendendo l'overhead della ricorrenza trascurabile.

4. Contributi Principali

1. Dimostrazione di fattibilità: Conferma che le architetture ricorrenti (Mamba2, RWKV, xLSTM) sono valide alternative generiche per la generazione di embedding, non solo per la generazione di testo.
2. Strategia di Inferenza Verticale: Introduzione di un metodo di scheduling che bilancia parallelizzazione e ricorrenza per ottenere complessità di memoria costante indipendentemente dalla lunghezza dell'input.
3. Open Source: Rilascio dei checkpoint dei modelli fine-tuned e dell'implementazione dell'inferenza per favorire la riproducibilità.

5. Significato e Impatto

Questo lavoro stabilisce le architetture ricorrenti come un'alternativa scalabile ed efficiente dal punto di vista delle risorse per la generazione di embedding.

• Impatto Pratico: Abilita l'uso di modelli di embedding su documenti estremamente lunghi (long-context) senza richiedere hardware di fascia alta o tecniche di quantizzazione aggressive.
• Futuro della Ricerca: Suggerisce che il divario di qualità tra Transformer e modelli ricorrenti sugli task in inglese può essere colmato con modelli base pre-addestrati su dataset più vasti.
• Accessibilità: Facilita l'accesso alle tecnologie NLP in ambienti con vincoli di risorse, rendendo possibile l'elaborazione di grandi corpora di documenti in modo economicamente sostenibile.

In sintesi, il paper dimostra che è possibile ottenere la qualità degli embedding dei Transformer mantenendo l'efficienza computazionale e di memoria dei modelli ricorrenti, aprendo la strada a una nuova classe di modelli di embedding ottimizzati per il contesto lungo.