Resumen Técnico: Memoria Tensorial: Estado Recurrente de Tamaño Fijo para Transformadores de Largo Horizonte

1. Declaración del Problema

Los transformadores procesan imágenes y videos aplanando las dimensiones espaciales y temporales en largas secuencias de tokens. Aunque los mecanismos de atención y la caché KV preservan las características pasadas, sufren dos limitaciones principales a medida que crece el horizonte de la secuencia:

1. Escalado de Memoria: La inferencia práctica depende del almacenamiento de claves y valores pasados, lo que escala linealmente con el número de tokens. Esto se vuelve computacionalmente prohibitivo y consume demasiada memoria para videos y documentos largos.
2. Falta de Estado Persistente: Los transformadores estándar carecen de un estado explícito y persistente. Aunque atienden a características pasadas, no mantienen una representación estructurada que pueda actualizarse y consultarse a medida que cambia la entrada. Esto conduce a fragilidad ante oclusiones, observaciones faltantes y a la necesidad de reconstruir repetidamente el contexto global a partir de una historia en constante crecimiento, en lugar de mantener un estado compacto y persistente.

Los enfoques existentes, como la atención dispersa, la compresión de historial o los contextos aumentados por recuperación, a menudo extienden la longitud del contexto pero almacenan la información como colecciones no estructuradas de vectores de tokens o ranuras. No preservan explícitamente la estructura espacial ni mantienen un "estado del mundo" compacto adecuado para tareas basadas en estados, como la comprensión de video.

2. Metodología: Memoria Tensorial

Los autores proponen Memoria Tensorial, un módulo ligero que aumenta los bloques estándar de transformadores con un tensor de memoria recurrente 3D de tamaño fijo. Este módulo desacopla la capacidad del estado de la longitud de la entrada, preservando al mismo tiempo un sesgo inductivo espacial.

Arquitectura Central

La memoria se define como un par de volúmenes (estado oculto $h_t$ y estado de celda $c_t$) con dimensiones $C \times D \times H \times W$, donde las dimensiones espaciales son constantes independientemente de la longitud de la entrada. El módulo opera en un ciclo por fragmento de tokens:

1. Fragmentación: Los tokens de entrada se agrupan en fragmentos. Se deriva una "consulta de lectura" del primer token del fragmento, y un "resumen de escritura" es una proyección aprendida de todo el fragmento.
2. Direccionamiento Continuo: Un predictor de coordenadas compartido (usando pesos atados) genera coordenadas 3D continuas para la lectura ($\mu_{read}$) y la escritura ($\mu_{write}$) dentro de la cuadrícula de voxels.
3. Lectura (Muestreo Continuo): El modelo lee del estado oculto $h_t$ en $\mu_{read}$ utilizando interpolación trilineal. La característica recuperada se proyecta de nuevo a la dimensión del token y se fusiona en la corriente de tokens mediante una conexión residual con puerta. Una puerta escalar aprendida $\gamma$ controla la fuerza de esta fusión, permitiendo al modelo suprimir la ruta de memoria si no es beneficiosa para una tarea específica.
4. Escritura (Escritura Suave Diferenciable): El modelo escribe información prediciendo un vector de contenido y un parámetro de dispersión $\sigma$. Deposita este contenido como un volumen ponderado por Gaussiana alrededor de la ubicación 3D predicha $\mu_{write}$. Esta "escritura suave" evita la asignación rígida de voxels, permitiendo que los gradientes moldeen dónde se coloca la información.
5. Actualización Recurrente: El volumen de escritura se concatena con el estado oculto actual. Un operador 3D ligero y factorizado (convoluciones 3D profundas seguidas de una mezcla punto a punto) calcula puertas al estilo ConvLSTM ($i_t, f_t, o_t, g_t$). Estas puertas actualizan los estados de celda y oculto, permitiendo que la información persista, se propague localmente y permanezca disponible incluso cuando faltan observaciones.

Decisiones Clave de Diseño

• Estado de Tamaño Fijo: El tamaño del tensor de memoria es constante, lo que previene el crecimiento lineal de la memoria asociado con la caché KV.
• Estructura Espacial: La cuadrícula de voxels 3D preserva explícitamente las relaciones espaciales, a diferencia de las ranuras de tokens no estructuradas.
• Escrituras Suaves Diferenciables: El uso de volúmenes ponderados por Gaussiana permite actualizaciones suaves y continuas en lugar de asignaciones discretas y rígidas.
• Fusión con Puerta: La puerta aprendida $\gamma$ asegura que el módulo pueda adjuntarse a modelos preentrenados sin desestabilizar el entrenamiento; si la memoria es innecesaria, la puerta suprime la ruta.

3. Contribuciones Clave

• Mecanismo de Memoria Novel: Introducción de un estado de voxel 3D recurrente de tamaño fijo al que los tokens pueden escribir y desde el cual pueden leer mediante coordenadas continuas, actuando como un estado espacial explícito.
• Integración: El módulo se integra limpiamente en las tuberías de entrenamiento estándar de transformadores y puede adjuntarse o eliminarse de bloques existentes sin cambios arquitectónicos.
• Suite Diagnóstica: Desarrollo de una suite diagnóstica controlada de "juguete" diseñada para aislar escenarios donde el estado persistente es beneficioso (por ejemplo, vinculación de coordenadas, oclusión, construcción de mapas) frente a escenarios donde no lo es.

4. Resultados Experimentales

Los autores evaluaron la Memoria Tensorial en benchmarks estándar en los dominios de lenguaje, imagen y video, así como en la suite diagnóstica.

• Modelado de Lenguaje (WikiText-2): La Memoria Tensorial logró las mayores ganancias absolutas, reduciendo la perplejidad de prueba de 130.30 (Transformador Base) a 100.07. Superó a líneas base que incluyen Transformer-XL, atención local y atención lineal.
• Reconstrucción de Parches de Imagen (CUB-200-2011): En una tarea de reconstrucción de parches enmascarados, la variante de Memoria Tensorial (Spatial ViT) mostró mejoras consistentes en PSNR y SSIM a través de ratios de máscara (25%, 50%, 75%) en comparación con Base ViT, Local ViT y Registers ViT.
• Reconocimiento de Acciones en Video (UCF-101): El modelo superó tanto a una línea base de transformador espaciotemporal conjunta como a una fuerte línea base de "token de registro" en precisión Top-1 (89.21% vs 82.12%) y Top-5. Los autores notan que esto vino a costa de un tiempo de entrenamiento en reloj significativamente mayor debido a la recurrencia de voxels por paso.
• Diagnósticos de Juguete:
  • Vinculación de Coordenadas: La Memoria Tensorial mantuvo una alta precisión (75–91%) a medida que aumentaba el número de escrituras, mientras que las líneas base basadas en atención colapsaron a un rendimiento cercano al aleatorio (7–14%).
  • Control Sin Daño: En tareas que no requieren memoria, la puerta aprendida $\sigma(\gamma)$ permaneció cerca de su inicialización (suprimiendo la ruta de memoria), y todos los modelos alcanzaron una precisión del 100%.
  • Construcción de Mapas: La Memoria Tensorial superó significativamente a las líneas base en tareas de navegación en cuadrícula de largo horizonte donde la información debe retenerse a lo largo del tiempo.

5. Significado y Afirmaciones

El artículo afirma que la Memoria Tensorial ofrece una forma práctica de manejar el razonamiento de largo horizonte al proporcionar a los transformadores un estado explícito, compacto y estructurado.

• Desacoplamiento de Capacidad y Longitud: Al usar un tensor de tamaño fijo, el modelo evita la escalada lineal de memoria de las cachés KV, haciéndolo teóricamente más escalable para secuencias largas.
• Robustez ante Oclusión: El estado persistente permite al modelo mantener un "estado de creencia" sobre la escena, haciéndolo más robusto ante observaciones faltantes u oclusiones en comparación con modelos que dependen únicamente de reconstruir el contexto a partir de una corriente de tokens.
• Versatilidad: Aunque motivado por video de largo horizonte, se demuestra que el mecanismo es beneficioso para el lenguaje (acumulando contexto global) y para imágenes (agregación espacial estructurada).
• Fundamento para Trabajo Futuro: Los autores ven la implementación actual como un fundamento. Sugieren que las direcciones futuras incluyen memorias de mayor dimensión (por ejemplo, 4D para memoria espacial consciente del tiempo), dinámicas más ricas (por ejemplo, deformación de flujo óptico) y actualizaciones dispersas para mejorar la eficiencia.

Los autores se mantienen modestos sobre las limitaciones actuales, reconociendo que las actualizaciones densas de voxels son un cuello de botella en el rendimiento y que el estado de tamaño fijo puede saturarse si una tarea requiere almacenar una gran cantidad de detalles distintos. Sin embargo, afirman que agregar un pequeño estado estructurado es un camino viable hacia sistemas de percepción y razonamiento de largo horizonte más capaces.