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

1. Planteamiento del Problema

Los modelos de incrustación (embeddings) de texto basados en Transformers son el estado del arte actual, pero presentan limitaciones críticas para documentos largos y entornos con recursos restringidos:

• Complejidad Computacional Cuadrática: El mecanismo de atención requiere $O(T^2)$ operaciones, donde $T$ es la longitud de la secuencia.
• Complejidad de Memoria Lineal: Aunque optimizaciones como FlashAttention reducen la memoria, el consumo sigue escalando linealmente con la longitud de la secuencia ($O(T)$), lo que hace inviable el procesamiento de documentos muy largos en hardware limitado.

El objetivo del trabajo es encontrar una alternativa eficiente que mantenga un rendimiento competitivo pero con una complejidad de tiempo lineal ($O(T)$) y un uso de memoria constante ($O(1)$) una vez superada una cierta longitud de entrada.

2. Metodología

Los autores proponen una solución basada en arquitecturas recurrentes modernas (Mamba2, RWKV, xLSTM) combinadas con una estrategia de inferencia novedosa.

A. Entrenamiento y Adaptación

• Base del Modelo: Se utilizan modelos de lenguaje preentrenados basados en Mamba2 (específicamente Mamba2 1.3B y Codestral Mamba2 7B).
• Procedimiento de Fine-tuning: Se adapta el procedimiento de entrenamiento de E5mistral-7b. Esto implica:
  • Eliminar la cabeza de modelado de lenguaje.
  • Usar un token de fin de secuencia (EOS) para generar la representación del texto completo.
  • Entrenamiento con pérdida InfoNCE (contrastiva) para distinguir pares positivos de negativos.
  • Instruction Tuning: Se añade una instrucción específica solo a la consulta (query) para permitir la adaptación a múltiples tareas sin reentrenamiento específico.

B. Estrategia de Inferencia: "Chunking Vertical"

Esta es la contribución técnica central. Los autores aprovechan la dualidad de modelos como Mamba2, que pueden formularse tanto como recurrencias lineales (eficientes en memoria) como transformaciones matriciales (paralelizables).

1. Inferencia Horizontal Estándar: Procesa toda la secuencia a través de una capa a la vez. Es paralela pero consume mucha memoria ($O(T)$).
2. Inferencia Vertical Chunked (Propuesta):
   • La secuencia de entrada se divide en bloques verticales de tamaño fijo $V$.
   • En lugar de procesar toda la secuencia a través de la capa 1, luego la capa 2, etc., el modelo procesa un bloque vertical completo a través de todas las capas antes de pasar al siguiente bloque.
   • Mecanismo de Estado: Se mantiene un estado recurrente (latente) entre los bloques verticales para cada capa.
   • Resultado de Memoria: Una vez que la longitud de la secuencia $T$ excede el tamaño del bloque vertical $V$, el consumo de memoria se vuelve constante ($O(L + Q \cdot V)$, donde $L$ es el número de capas y $Q$ el tamaño del bloque intra-capas), independientemente de cuán largo sea el documento.

3. Contribuciones Clave

1. Validación de Modelos Recurrentes como Embedders: Demuestran que los modelos recurrentes (Mamba2, RWKV, xLSTM) pueden ser ajustados finamente para ser generadores de embeddings de propósito general, logrando un rendimiento competitivo frente a los Transformers.
2. Estrategia de Inferencia de Memoria Constante: Introducen y formalizan la estrategia de "chunking vertical", que permite generar embeddings para secuencias arbitrariamente largas sin estallar la memoria de la GPU, manteniendo la complejidad temporal lineal.
3. Generalización Arquitectónica: Validan que esta estrategia de inferencia no es exclusiva de Mamba2, sino que funciona consistentemente en una familia de arquitecturas recurrentes modernas (RWKV y xLSTM), ofreciendo compromisos tiempo-memoria favorables en todas ellas.
4. Liberación de Recursos: Publican los checkpoints de los modelos ajustados y la implementación de inferencia para fomentar la reproducibilidad.

4. Resultados Experimentales

Los modelos fueron evaluados en los benchmarks MTEB (inglés y multilingüe) y LongEmbed (documentos largos).

• Calidad de Embeddings:
  • Los modelos basados en Mamba2 alcanzan un rendimiento competitivo con los basados en Transformers (como Mistral 7B y Qwen2) en la mayoría de las tareas.
  • En tareas multilingües, los modelos recurrentes a veces superan a los Transformers.
  • Existe una ligera brecha en tareas de recuperación en inglés (retrieval), atribuida a diferencias en la escala de preentrenamiento de los modelos base (Mamba2 1.3B fue entrenado con menos tokens que sus contrapartes Transformer).
• Eficiencia de Inferencia:
  • Memoria: Para secuencias largas (>4096 tokens), los modelos recurrentes con chunking vertical mantienen un uso de memoria constante, mientras que los modelos basados en Transformers (incluso con FlashAttention) muestran un crecimiento lineal, agotando rápidamente la memoria de la GPU.
  • Tiempo de Ejecución: Para secuencias largas, los modelos recurrentes son significativamente más rápidos que los Transformers debido a la ausencia de la complejidad cuadrática de la atención.
  • Saturación de Paralelismo: Se encontró que el tamaño del bloque vertical ($V$) necesario para saturar el hardware es relativamente pequeño (ej. 4096 tokens para un solo lote), lo que permite obtener la velocidad del paralelismo completo sin el costo de memoria.

5. Significado e Impacto

Este trabajo establece una nueva dirección para la generación de embeddings en escenarios de contexto largo y recursos limitados:

• Escalabilidad: Permite procesar documentos de decenas de miles de tokens en hardware estándar, algo que es prohibitivo para los Transformers actuales.
• Eficiencia de Recursos: Reduce drásticamente los requisitos de memoria, facilitando el despliegue de modelos de incrustación en entornos edge o con GPUs de menor capacidad.
• Alternativa Sólida: Demuestra que las arquitecturas recurrentes modernas no son solo para generación de texto, sino que son una alternativa viable y superior para la representación semántica de documentos largos, cerrando la brecha de calidad con los Transformers mientras ofrecen ventajas operativas significativas.

En resumen, el paper propone un cambio de paradigma hacia arquitecturas recurrentes para embeddings, resolviendo el cuello de botella de memoria de los Transformers mediante una ingeniosa estrategia de inferencia que mantiene la eficiencia sin sacrificar la calidad del modelo.