Transformers con Memoria Aumentada

¿Y si el modelo pudiera recordar más allá de su ventana?

La atención dispersa (el artículo anterior) limita cuáles tokens atienden a cuáles, reduciendo el costo cuadrático. Pero el contexto total sigue acotado: una ventana deslizante de tamaño $w$ solo puede ver $w$ tokens hacia el pasado, e incluso los patrones híbridos global + local siguen operando sobre una única secuencia de entrada finita. ¿Qué pasaría si el modelo tuviera una memoria separada que persiste a través de ventanas de contexto — un almacén de información del pasado que el modelo puede consultar en cada paso, independientemente de cuán atrás aparecieron los tokens originales?

Hay tres enfoques amplios para construir esta memoria:

• Recuperación exacta: almacenar todos los pares KV pasados en un gran caché externo, y recuperar los más relevantes mediante búsqueda aproximada de vecinos más cercanos (kNN) cuando sea necesario.
• Memoria compresiva: comprimir los pares KV antiguos en una matriz de memoria de tamaño fijo. La memoria nunca crece, sin importar cuánto contexto pasado se haya procesado, pero es con pérdida.
• Memoria neuronal aprendida: hacer que la memoria sea una pequeña red neuronal que se actualiza mediante descenso de gradiente en tiempo de prueba. Las nuevas entradas se "escriben" en la memoria entrenando la red con ellas.

Esto es diferente de RAG , que recupera de un corpus externo de documentos que el modelo nunca ha visto durante su paso hacia adelante. Los transformers con memoria aumentada recuerdan su propio contexto pasado — tokens que ya fluyeron a través del modelo en segmentos anteriores de la misma secuencia. RAG extiende lo que el modelo sabe ; las arquitecturas con memoria aumentada extienden lo que recuerda .

Memorizing Transformers: Recuperación exacta por kNN

La idea más simple: simplemente almacenar todos los pares key-value pasados en un gran caché externo, y en cada paso de atención, buscar en ese caché las entradas más relevantes. Este es el enfoque de Memorizing Transformers (Wu et al., 2022) .

Durante la atención en la capa $l$, el modelo ejecuta dos cálculos de atención en paralelo:

• Atención local: atención estándar de producto punto escalado sobre la ventana de contexto actual, produciendo $A_{ ext{local}}$. Esta es autoatención ordinaria — nada nuevo.
• Atención kNN: usar los vectores de query actuales para buscar en el caché KV externo mediante kNN aproximado. Recuperar los $k$ pares key-value pasados más similares para cada query, calcular puntuaciones de atención solo sobre esos $k$ pares recuperados, y producir $A_{ ext{kNN}}$.

Las dos salidas de atención se combinan con una compuerta aprendida:

La compuerta es un escalar aprendido (pasado por una sigmoide) que permite al modelo decidir, por cabeza, cuánto confiar en memorias recuperadas versus contexto local. Cuando la compuerta $ o 0$, la capa ignora el caché externo por completo y se comporta como un transformer estándar. Cuando la compuerta $ o 1$, la capa depende completamente del contexto pasado recuperado. En la práctica, el modelo aprende a usar el caché externo principalmente en las capas superiores, porque las capas inferiores tienden a manejar patrones sintácticos locales que no se benefician de la recuperación de largo alcance.

El caché externo en sí es un buffer FIFO : a medida que el modelo procesa nuevos segmentos, añade los pares KV del segmento actual al caché y expulsa las entradas más antiguas cuando el caché está lleno. Wu et al. muestran que el caché puede escalar a 262K tokens con búsqueda kNN aproximada eficiente usando cuantización de producto (la misma técnica de indexación usada en bases de datos vectoriales para RAG ). Como la búsqueda kNN recupera solo las $k$ entradas principales (típicamente $k = 32$ o $k = 64$) en lugar de atender sobre el caché completo, el costo de atención para la memoria externa es $O(k \cdot d)$ por query — constante con respecto al tamaño del caché.

Infini-Attention: Memoria compresiva

Almacenar pares KV exactos funciona, pero el caché crece linealmente con la cantidad de contexto pasado. ¿Podemos hacerlo mejor? Infini-attention (Munkhdalai et al., 2024) toma un enfoque diferente: en lugar de almacenar pares KV exactos, comprimir todo el contexto pasado en una matriz de memoria de tamaño fijo que nunca crece, sin importar cuán larga sea la secuencia.

Cada capa de atención mantiene dos objetos persistentes: una matriz de memoria $M \in \mathbb{R}^{d_k imes d_v}$ and a vector normalizador $z \in \mathbb{R}^{d_k}$. These are initialised to zeros and updated as the model processes successive segments of the input.

Después de procesar el segmento $s$ (con keys $K_s$ y values $V_s$ de la atención estándar), la memoria se actualiza:

where $\sigma$ is a nonlinearity (specifically ELU + 1, which ensures all values are positive). Let's unpack what this update does. The term $\sigma(K_s)^T V_s$ is a sum of outer products: each key vector $\sigma(k_{s,i})$ (of dimension $d_k$) is multiplied by its corresponding value vector $v_{s,i}$ (of dimension $d_v$), producing a $d_k imes d_v$ matrix. Summing over all tokens in the segment gives a matrix that encodes key-value associations. This gets suma to the existing memory $M_{s-1}$, accumulating information across segments. The normaliser $z_s$ tracks the total "mass" of the keys that have been written, which is needed for proper normalisation during retrieval.

Para recuperar de la memoria, el modelo usa las queries del segmento actual:

Esta es una consulta de atención lineal . Cada query $\sigma(q_{s,i})$ se multiplica por la matriz de memoria $M_s$ para producir una salida de dimensión $d_v$ (el value "recuperado"), luego se divide por $\sigma(q_{s,i})^T z_s$ para normalización (análogo a cómo softmax normaliza las puntuaciones de atención estándar). Esta recuperación cuesta $O(d_k \cdot d_v)$ por query — completamente independiente de cuántos tokens pasados se han comprimido en $M$. Esa es la ventaja clave: la memoria tiene tamaño fijo $d_k imes d_v$ regardless of sequence length.

The final output combines local attention and memory attention via a learned gate $\beta$:

Here $\sigma(\beta)$ is a sigmoid applied to a learned parameter $\beta$ (per head). Análisis de frontera: when $\beta o -\infty$, $\sigma(\beta) o 0$, and the output is pure local attention — the memory is completely ignored, and the layer behaves like standard attention. When $\beta o +\infty$, $\sigma(\beta) o 1$, and the output is pure memory — local context is ignored. In between, the model blends both sources. Each head can learn its own $\beta$, so some heads can specialise in local patterns while others focus on long-range memory.

Titans: Aprendiendo a memorizar en tiempo de prueba

Tanto Memorizing Transformers como Infini-attention almacenan el contexto pasado como datos (pares KV o matrices comprimidas). Titans (Behrouz et al., 2025) toma un enfoque más radical: la memoria es una red neuronal — un MLP pequeño — que se actualiza mediante descenso de gradiente en tiempo de prueba . New information is "written" to memory by training the MLP on it, and "read" by running a forward pass through it.

Titans organiza la información en tres sistemas de memoria:

• Memoria a corto plazo: standard atención estándar sobre una ventana local. Esta es la ventana de contexto normal — precisa pero acotada.
• Memoria a largo plazo: un MLP pequeño que se entrena mediante SGD durante la inferencia. Los tokens pasados se codifican en los pesos del MLP.
• Memoria persistente: parámetros aprendibles fijos (similares al ajuste de prompts) que codifican conocimiento a nivel de modelo compartido entre todas las entradas. Estos se entrenan durante el preentrenamiento y se congelan en tiempo de prueba.

La memoria a largo plazo es donde radica la novedad. En cada paso de tiempo $t$, el MLP de memoria $M$ se actualiza:

Este es un paso de descenso de gradiente estándar, pero ejecutándose durante la inferencia , no el entrenamiento. Así es como funciona:

• Recuperación: ejecutar un paso hacia adelante a través del MLP de memoria con la entrada actual $x_t$ para producir una predicción.
• Sorpresa: compute a reconstruction loss $\mathcal{L}(M_{t-1}, x_t)$ measuring how well the memory predicted the current input. If the input is surprising (high loss), the memory was missing relevant information.
• Almacenamiento: retropropagar a través del MLP y tomar un paso de gradiente para actualizar los pesos de la memoria. Las entradas sorprendentes causan actualizaciones de gradiente más grandes, escribiéndolas más fuertemente en la memoria.

Hay dos hiperparámetros clave. La tasa de aprendizaje de memoria $\eta$ controls how quickly new information is written. The tasa de olvido $\lambda$ applies weight decay after each update:

Este decaimiento de pesos causa que las memorias antiguas se desvanezcan gradualmente a menos que sean reforzadas por patrones recurrentes. Boundary analysis: when $\eta = 0$, the memory receives no gradient updates and remains frozen — the model has no long-term memory and relies only on its local attention window. When $\lambda = 0$, there is no forgetting: the MLP accumulates all past information without decay, which means its weights grow in magnitude unboundedly and can eventually become unstable. When $\lambda = 1$, the memory is completely erased at every step — the weight decay zeros out all parameters, giving the model no memory at all. The sweet spot is small $\eta$ (gentle writes) and small $\lambda$ (slow forgetting), allowing the MLP to build up a persistent representation of past context.

¿Por qué es esto novedoso? Reformula la memoria como una red neuronal entrenada continuamente durante la inferencia. La memoria tiene tamaño constante (los parámetros del MLP) independientemente de la longitud de la secuencia, igual que la matriz de memoria de Infini-attention. Pero a diferencia de una matriz de memoria lineal, un MLP puede representar asociaciones no lineales — en principio, puede aprender patrones complejos que una simple memoria de producto exterior no puede. El costo es cómputo adicional: cada actualización de memoria requiere un paso hacia adelante, cálculo de pérdida y paso hacia atrás a través del MLP. Esto conecta con la idea emergente más amplia de entrenamiento en tiempo de prueba / adaptación en tiempo de prueba, donde los modelos continúan aprendiendo durante la inferencia en lugar de usar pesos congelados.

Landmark Attention: Recuperación dentro de la atención

Los enfoques anteriores añaden módulos de memoria externa al transformer. Landmark attention (Mohtashami & Jaggi, 2023) toma un ángulo diferente: en lugar de una memoria separada, reestructurar el mecanismo de atención mismo para habilitar recuperación de acceso aleatorio sobre secuencias arbitrariamente largas.

La idea es insertar tokens landmark at regular intervals, every $k$ tokens. Each landmark is trained to summarise the block of $k$ tokens that precedes it — it acts as a learned "table of contents" entry for its block. During attention, the model uses a two-stage process:

• Etapa 1 — Atender a los landmarks: cada query primero atiende a todos los tokens landmark. Con una secuencia de $n$ tokens y un landmark cada $k$ tokens, hay solo $n / k$ landmarks. Esto es barato.
• Etapa 2 — Seleccionar y atender a bloques completos: basándose en las puntuaciones de atención de los landmarks, seleccionar los $b$ bloques más relevantes (aquellos cuyos landmarks recibieron la mayor atención). Cargar los pares KV completos para esos $b$ bloques y calcular atención de grano fino sobre ellos.

El costo total de atención por query es:

The first term $n/k$ is the cost of attending to all landmarks (coarse retrieval). The second term $b \cdot k$ is the cost of attending to the $b$ selected blocks of $k$ tokens each (fine-grained attention). Análisis de frontera: cuando $b = n/k$ (seleccionar todos los bloques), el costo se reduce a $n/k + n = O(n)$, que es esencialmente atención completa. Cuando $b = 1$ (seleccionar solo el bloque más relevante), el costo es $n/k + k$. Para un rendimiento equilibrado, elegir $k = \sqrt{n}$ y $b = O(1)$ da un costo de $O(\sqrt{n})$ por query — sublineal en la longitud de secuencia. Esto es mucho menos que los $O(n)$ de atención completa cuando solo unos pocos bloques son relevantes.

La ventaja clave sobre la atención de ventana deslizante es el acceso aleatorio . Una ventana deslizante solo puede ver los $w$ tokens más recientes — no tiene forma de saltar de vuelta a un pasaje importante de hace 50,000 tokens. La atención landmark puede alcanzar cualquier bloque en la secuencia, siempre que su landmark se considere relevante. La compensación es que los landmarks deben ser buenos resúmenes de sus bloques. Si un landmark falla en capturar la información relevante de su bloque, el bloque no será seleccionado y sus tokens serán invisibles para el query.

Elegir el enfoque de memoria adecuado

Estos cuatro enfoques forman un espectro desde la recuperación exacta hasta la abstracción aprendida. La siguiente tabla los compara a través de las dimensiones que importan en la práctica:

import json, js

rows = [
    ["Memorizing Transformers
(Wu et al., 2022)",
     "Exact KV cache + kNN",
     "O(n) memory",
     "Highest (exact recall)",
     "Upper layers only",
     "Cache eviction (FIFO)"],
    ["Infini-Attention
(Munkhdalai et al., 2024)",
     "Compressed matrix",
     "O(d_k * d_v) fixed",
     "Lossy (linear approx.)",
     "All layers",
     "Accumulation in matrix"],
    ["Titans
(Behrouz et al., 2025)",
     "Neural net (MLP) + SGD",
     "O(MLP params) fixed",
     "Lossy (nonlinear)",
     "Separate module",
     "Weight decay forgetting"],
    ["Landmark Attention
(Mohtashami & Jaggi, 2023)",
     "Landmark-guided retrieval",
     "O(n/k) landmarks",
     "Exact (for selected blocks)",
     "Within attention",
     "Block selection via scores"],
]

js.window.py_table_data = json.dumps({
    "headers": [
        "Method",
        "Memory Type",
        "Memory Cost",
        "Fidelity",
        "Integration",
        "Forgetting Mechanism"
    ],
    "rows": rows
})

print("Comparison of memory-augmented transformer approaches")
print()
print("Exact retrieval  = highest fidelity, highest memory cost")
print("Compressed       = fixed memory regardless of length, but lossy")
print("Learned (Titans) = most flexible, adds compute (test-time SGD)")
print("Landmark         = random-access retrieval, depends on landmark quality")

At one end, Memorizing Transformers almacenan pares KV exactos y recuperan mediante kNN — máxima fidelidad, pero el costo de memoria crece linealmente con el contexto pasado. En el otro extremo, Titans comprime todo en los pesos de un MLP pequeño — huella de memoria fija y la capacidad de capturar asociaciones no lineales, pero al costo de cómputo en tiempo de prueba (paso adelante + atrás a través del MLP en cada paso) y siendo el enfoque más nuevo y menos probado a escala. Infini-attention se sitúa en el medio: memoria de tamaño fijo con actualizaciones lineales simples, sin entrenamiento en tiempo de prueba, pero limitada por la expresividad de la compresión de producto exterior lineal. Landmark attention es ortogonal: no añade una memoria separada, sino que reestructura la atención para habilitar acceso aleatorio eficiente.

En la práctica, la mayoría de los modelos de producción hoy (GPT-4, Claude, Gemini) dependen de ventanas de contexto largas + FlashAttention en lugar de memoria externa explícita. Las ventanas de contexto han crecido lo suficiente (128K–1M+ tokens) para que muchos casos de uso quepan dentro de una sola ventana, y FlashAttention (Dao et al., 2022) hace que el costo de IO sea manejable incluso a esas longitudes. Pero el campo avanza rápido. A medida que se pide a los modelos procesar bases de código completas, documentos del tamaño de libros o conversaciones persistentes de múltiples turnos que abarcan miles de turnos, la necesidad de memoria explícita más allá de la ventana de contexto probablemente crecerá — y los enfoques en este artículo representan las direcciones de investigación líderes para satisfacer esa necesidad.

Quiz

Pon a prueba tu comprensión de las arquitecturas de transformers con memoria aumentada.