Búsqueda de Memoria Reemplaza Atención Lineal en Prefijos Largos

Investigadores del Institute of Science Tokyo e Imperial College London han publicado un método que reemplaza atención lineal sobre prefijos largos con búsqueda de memoria de tiempo constante. En LLaMA 3.1-8B, el enfoque reduce latencia de atención 1.36x con presupuesto de memoria de 8K y supera RAG de atención completa en el benchmark NBA con 20% del footprint de memoria, sin actualizaciones de gradiente al modelo.

En inferencia, la atención sobre un prefijo (prompt de sistema, documentos recuperados, ejemplos in-context) escala linealmente con longitud del prefijo durante prefill y cada paso de decodificación. Prefix caching (enfoque de Anthropic en Claude Code) amortiza costo de prefill pero deja overhead de paso de decodificación y memoria KV-cache intactos. Compresión de atención aún lee el prefijo. Destilación de contexto y ajuste fino de hiperred requieren entrenamiento basado en gradiente, que es costoso y se rompe cuando el prefijo cambia. Este método evita ambos.

El método precomputa atención de prefijo en modo forward-only. Ejecuta queries representativas a través del modelo, recopila sus salidas de atención y las agrupa en centroides. En inferencia, una query entrante recupera el centroide más cercano y lo fusiona con self-attention usando una operación online-softmax—una reconstrucción sin pérdida que omite tokens de prefijo. El costo de búsqueda escala logarítmicamente con el número de centroides. Una vez construido, el costo por-decodificación permanece constante mientras el prefijo crece.

En ManyICLBench con LLaMA 3.1-8B, el método mejora precisión sobre aprendizaje in-context estándar en presupuestos de 1K a 8K centroides, reduciendo latencia 1.36x en 8K. En NBA (una tarea RAG), supera baselines de atención completa con 20% de la memoria KV. El paper también valida en RuleArena para rule-following en prompts de sistema largos. Especificidades de hardware y throughput no se divulgan.

Construir la memoria requiere un corpus de query representativo. La calidad de los centroides depende de cuán bien las queries de construcción cubren la distribución de query en tiempo de ejecución. Para RAG de propósito general con queries impredecibles, seleccionar ese corpus es no trivial. Segundo, la memoria se vincula a un prefijo específico. Cualquier actualización de prefijo requiere reconstrucción—un pase forward-only pero un paso operacional real. El paper no reporta costo de construcción en tiempo de wall-clock o GPU-hours, que los arquitectos necesitan para presupuestar cadencia de reconstrucción.

El paper señala que la influencia del prefijo decae conforme la generación procede, incluso bajo atención completa. Pero la memoria se construye a partir de una distribución de query estática. En secuencias generativas largas donde el estado del modelo cambia, el centroide recuperado puede divergir de lo que atención completa produciría. Las tareas de evaluación ejecutan generación de longitud moderada. Si la ganancia de latencia 1.36x se mantiene más allá de 32K pasos de decodificación permanece sin probar.

El código está disponible en github.com/yasu0001/AttentionMemory.

Conclusión del arquitecto: para prefijos estables y de alto reutilización (prompts de sistema fijos, corpus de documentos actualizados diariamente), precomputar centroides una vez y servir búsquedas en tiempo de decodificación es una ganancia de latencia plug-and-play sin reentrenamiento. Riesgo: un corpus de query mal construido en tiempo de compilación degrada la calidad de recuperación en ejecución.