Resumen Técnico: Estudio de la "Soupabilidad" de Documentos en Modelos de Espacio de Estados

Declaración del Problema

Muchas tareas de PLN en el mundo real, como la respuesta a preguntas multi-documento (QA), el resumen científico y el análisis legal, requieren el razonamiento sobre corpus enteros en lugar de documentos individuales largos. Los modelos actuales basados en Transformers tienen dificultades con esto debido a los costos prohibitivos de atención de $O(L^2)$ y la inflexibilidad de las estrategias de codificación monolíticas. En los enfoques estándar, todos los documentos se concatenan en una sola secuencia antes del procesamiento; por consiguiente, cualquier modificación al corpus (por ejemplo, añadir o eliminar un solo documento) requiere re-codificar todo el input desde cero. Esto impide la reutilización modular de las representaciones de los documentos y crea un cuello de botella para los sistemas de recuperación escalables y dinámicos.

Aunque los Modelos de Espacio de Estados (SSM) Estructurados como Mamba y Mamba2 ofrecen un procesamiento de secuencias en tiempo lineal y comprimen los inputs en estados ocultos de longitud fija, suelen desplegarse utilizando la misma estrategia de concatenación monolítica, heredando las mismas limitaciones de re-codificación.

Metodología: Document Souping (Sopa de Documentos)

Los autores proponen el Document Souping, una estrategia inspirada en el model souping (promedio de pesos de checkpoints ajustados) pero aplicada a los estados ocultos. La hipótesis central es que los mecanismos de recurrencia lineal en los SSM permiten codificar independientemente las representaciones de los documentos para luego fusionarlas post hoc mediante operaciones lineales simples (por ejemplo, el promedio) para formar un estado compuesto adecuado para el razonamiento posterior.

La Arquitectura

1. Codificación Independiente: Cada documento $d_i$ en un corpus se codifica independientemente mediante un codificador SSM compartido ($SSM_\theta$), produciendo un conjunto de estados ocultos por capa $\{h^{(l)}_i\}$.
2. Agrupación (Pooling): Estos estados por documento se agregan en cada capa utilizando un operador de agrupación conmutativo (por ejemplo, promedio elemento a elemento, suma o max-pooling) para crear un único estado "sopa" $h^{(l)}_{soup}$.
3. Decodificación: Los estados agrupados $\{h^{(1)}_{soup}, \dots, h^{(L)}_{soup}\}$ se inyectan en el decodificador junto con la consulta $q$ para generar la respuesta $\hat{y}$.

Entrenamiento y Optimización

• Ajuste Fino (Finetuning): El artículo demuestra que los SSM preentrenados no soportan de forma nativa el souping; se requiere un ajuste fino completo de codificador-decodificador para enseñar al modelo a interpretar eficazmente los estados agrupados.
• Eficiencia de Memoria: Para manejar la sobrecarga de memoria de la retropropagación a través de múltiples codificadores de documentos independientes, los autores emplean checkpointing de activación a nivel de documento. Esto recomputa las activaciones del forward durante el paso de backward, permitiendo que el uso de memoria permanezca constante independientemente del número de documentos ($k$).
• Estrategia de Agrupación: Experimentos extensos identifican el promedio elemento a elemento sin normalización como el método de agregación más robusto y estable, superando a la suma (que sufre de un crecimiento ilimitado de magnitud) y al max-pooling.

Contribuciones Clave

1. Estudio Empírico Sistemático: La primera investigación exhaustiva sobre la "soupabilidad": la capacidad de los SSM para fusionar representaciones de documentos codificadas independientemente para el razonamiento multi-documento.
2. Validación de Rendimiento: Demostración de que los modelos Mamba2 ajustados con representaciones "sopa" logran un rendimiento competitivo o superior en QA de saltos múltiples (HotpotQA), QA de documentos largos (RACE, QuALITY) y recuperación dispersa (NIAH) en comparación con los baselines de concatenación estándar.
3. Escalabilidad y Generalización: Evidencia de que el enfoque escala a cientos de documentos (probado hasta 256) y se generaliza a diversos tamaños de contexto, siempre que el modelo sea ajustado con un número suficiente de segmentos.
4. Especificidad Arquitectónica: Prueba de que este enfoque modular es único para los SSM. Un experimento análogo de "cache souping" en Transformers (LLaMA 3) falló significativamente, resaltando que el estado recurrente de tamaño fijo de los SSM es un habilitador crítico para esta técnica.
5. Eficiencia de Inferencia: Un nuevo flujo de trabajo de inferencia donde los documentos se codifican una vez, se cachean y se reensamblan dinámicamente, ofreciendo ahorros sustanciales de latencia (hasta 14.6x de aceleración para 32k tokens) en comparación con la re-codificación de contextos concatenados.

Resultados Experimentales

Los autores evaluaron modelos Mamba2-8B y Mamba2-2.7B a través de varios benchmarks:

• QA de Saltos Múltiples (HotpotQA): Los modelos con ajuste fino completo de codificador-decodificador y promedio de agrupación alcanzaron 47.8 EM con 5 documentos, superando el ajuste fino de solo decodificador y igualando o superando los baselines de concatenación. El método se mantuvo robusto a medida que aumentaba el número de distractores.
• QA de Documentos Largos (RACE/QuALITY): El souping superó sustancialmente a los enfoques tradicionales de concatenación en benchmarks de documentos largos, particularmente cuando la configuración de entrenamiento se alineaba con el tamaño de la sopa en tiempo de prueba.
• Recuperación Dispersa (NIAH): Los modelos entrenados con un mayor número de segmentos (por ejemplo, 8 segmentos) mantuvieron una alta precisión (71.45% EM) al ser probados con 16 segmentos, mientras que los modelos entrenados con menos segmentos fallaron en generalizar.
• Escalabilidad: Aunque un modelo entrenado en 64 documentos falló al extrapolar a 256, el ajuste fino continuo en 128 documentos permitió al modelo alcanzar 37.25 EM en 256 documentos, demostrando que la arquitectura puede adaptarse a contextos vastos.
• Latencia: Para un contexto de 32k tokens, el procesamiento de una consulta usando estados "sopa" cacheados tomó 65.8 ms, comparado con los 963.2 ms del método de concatenación estándar.

Significado y Reivindicaciones

El artículo afirma que el document souping desbloquea un nuevo paradigma para el razonamiento de corpus modular y escalable. Al desacoplar la codificación del documento del procesamiento de la consulta, el enfoque permite:

• Actualizaciones Eficientes: Los documentos pueden codificarse una vez y cachearse, permitiendo actualizaciones dinámicas del corpus y recuperación sin reprocesar todo el input.
• RAG Escalable: El método sirve como un reemplazo directo para el paso de "concatenar pasajes recuperados" en sistemas de Generación Aumentada por Recuperación (RAG), soportando un número variable pero acotado de documentos de manera eficiente.
• Adecuación Arquitectónica: El trabajo establece que las propiedades de recurrencia lineal de los SSM los hacen únicos para soportar la composición a nivel de estado, una capacidad que los Transformers estándar carecen debido a su dependencia de cachés KV de longitud variable.

Los autores concluyen que, si bien la extrapolación a tamaños de sopa no vistos puede ser un desafío, las estrategias de ajuste fino adecuadas permiten que los SSM manejen eficazmente conjuntos de documentos grandes y dinámicos, ofreciendo una solución ligera y efectiva para el razonamiento de contexto largo.