Resumen Técnico: Memoria Condicional mediante Búsqueda Escalable (Engram)

1. Planteamiento del Problema

Los Modelos de Lenguaje Extensos (LLMs) actuales, particularmente aquellos que utilizan arquitecturas de Mezcla de Expertos (MoE), dependen de la computación condicional para escalar su capacidad. Sin embargo, los Transformers estándar carecen de un primitivo nativo para la búsqueda de conocimiento (knowledge lookup). En consecuencia, los modelos deben simular la recuperación de forma ineficiente mediante computación dinámica.

Los autores argumentan que las señales lingüísticas exhiben una heterogeneidad intrínseca:

• Razonamiento Composicional: Requiere computación profunda y dinámica.
• Recuperación de Conocimiento: Una parte sustancial del texto (por ejemplo, entidades con nombre, patrones formulaicos) es local, estática y altamente estereotipada.

Los Transformers estándar obligan al modelo a consumir múltiples capas tempranas de atención y redes de alimentación hacia adelante (FFN) para reconstruir búsquedas estáticas (por ejemplo, resolver "Diana, Princesa de Gales") que teóricamente podrían recuperarse mediante una simple búsqueda en tabla. Esto desperdicia una profundidad secuencial valiosa en operaciones triviales, impidiendo que el modelo asigne recursos a un razonamiento de mayor nivel.

2. Metodología: El Módulo Engram

Para abordar esto, el artículo introduce la Memoria Condicional como un eje de escasez complementario a la computación condicional. Esto se instancia mediante Engram, un módulo que moderniza las clásicas incrustaciones de $N$-gramas para permitir búsquedas de tiempo constante $O(1)$.

2.1. Arquitectura

Engram aumenta el núcleo del Transformer separando estructuralmente el almacenamiento de patrones estáticos de la computación dinámica. Opera en dos fases:

1. Recuperación Dispersa mediante $N$-gramas Hasheados:
   • Compresión de Tokenizadores: Una función suryectiva mapea los IDs de tokens brutos a identificadores canónicos (normalizando mayúsculas, puntuación, etc.) para maximizar la densidad semántica y reducir el tamaño del vocabulario.
   • Hasheo Multicabeza: Para manejar el espacio combinatorio de los $N$-gramas, el módulo utiliza $K$ cabezas de hasheo distintas por cada orden de $N$-gramas. Estas mapean contextos comprimidos a índices en tablas de incrustaciones mediante funciones deterministas, recuperando vectores estáticos.
   • Concatenación: Las incrustaciones recuperadas de varios órdenes de $N$-gramas y cabezas de hasheo se concatenan para formar el vector de memoria.
2. Puerta de Contexto y Fusión (Gating & Fusion):
   • Dado que las memorias estáticas carecen de contexto, un mecanismo de puerta modula la memoria recuperada utilizando el estado oculto actual (contexto global agregado) como una Consulta (Query).
   • Se computa una puerta escalar $\alpha_t$ mediante un producto punto entre el estado oculto normalizado y la Clave (Key) de la memoria. Si la memoria contradice el contexto, la puerta suprime la memoria (reducción de ruido).
   • Una convolución causal de profundidad ligera expande el campo receptivo antes de que la memoria con puerta sea fusionada al núcleo mediante una conexión residual.

2.2. Eficiencia del Sistema y Co-diseño de Infraestructura

A diferencia de MoE, que requiere enrutamiento dinámico y comunicación del estado oculto, Engram utiliza direccionamiento determinista basado únicamente en los IDs de los tokens de entrada.

• Pre-lectura (Prefetching): Los índices se conocen antes de la pasada hacia adelante (forward pass), lo que permite la pre-lecta asíncrona de incrustaciones desde la memoria del host (CPU/DRAM) a través de PCIe.
• Solapamiento (Overlap): Esto permite que la comunicación se solape con la computación de los bloques Transformer precedentes, incurriendo en un gasto mínimo.
• Descarga (Offloading): Tablas de incrustaciones masivas (por ejemplo, 100B de parámetros) pueden descargarse a la memoria del host, evitando las restricciones de memoria de la GPU.

2.3. Asignación de Escasez

Los autores formulan el problema de la Asignación de Escasez: dado un presupuesto de parámetros fijo, ¿cómo debe dividirse la capacidad entre los expertos de MoE (computación dinámica) y la memoria Engram (búsqueda estática)?

• Definen una relación de asignación $\rho$ que representa la fracción de parámetros inactivos asignados a MoE.
• Los experimentos revelan una ley de escala en forma de U: el MoE puro ($\rho=100\%$) y el Engram puro ($\rho=0\%$) son subóptimos. El mejor rendimiento se logra con una asignación híbrida (por ejemplo, $\rho \approx 75\%–80\%$), donde aproximadamente el 20–25% del presupuesto disperso se reasigna a Engram.

3. Resultados Clave

3.1. Rendimiento en Pre-entrenamiento

Los autores entrenaron cuatro modelos con 262B de tokens, igualando estrictamente los parámetros activados (3.8B) y los FLOPs de entrenamiento:

• Dense-4B: Línea base (Baseline).
• MoE-27B: Línea base estándar de MoE.
• Engram-27B: Iso-parámetro a MoE-27B, con 5.7B de parámetros reasignados a Engram (reduciendo los expertos de 72 a 55).
• Engram-40B: Escala la memoria Engram a 18.5B de parámetros (Total 39.5B).

Hallazgos:

• Superioridad sobre MoE: Engram-27B superó a la línea base iso-parámetro MoE-27B en diversos dominios.
• Ganancias en Razonamiento: Contrario a la intuición de que la memoria solo ayuda en la recuperación de conocimiento, Engram mostró ganancias significativas en razonamiento general (BBH +5.0, ARC-Challenge +3.7) y código/matemáticas (HumanEval +3.0, MATH +2.4).
• Ganancias en Conocimiento: Mejora el rendimiento en tareas intensivas en conocimiento (MMLU +3.0, CMMLU +4.0).
• Escalabilidad: Engram-40B redujo aún más la pérdida, sugiriendo que el eje de memoria escala efectivamente sin computación adicional.

3.2. Capacidades de Contexto Largo

Al descargar las dependencias locales hacia las búsquedas, Engram libera la capacidad de atención para el contexto global.

• Benchmark RULER: Engram-27B alcanzó una precisión del 97.0% en el Multi-Query Needle-in-a-Haystack (NIAH) comparado con el 84.2% de la línea base MoE.
• Eficiencia: Engram-27B igualó el rendimiento de LongPPL de la línea base con solo el 82% de los FLóps de pre-entrenamiento.

3.3. Análisis Mecanicista

• Profundidad Efectiva: Utilizando LogitLens y CKA (Alineación de Kernel Centrada), los autores encontraron que las capas tempranas de Engram producen representaciones semánticamente equivalentes a capas mucho más profundas en la línea base de MoE. Esto sugiere que Engram "profundiza" la red al evitar que las capas tempranas tengan que reconstruir características estáticas.
• Dicotomía Funcional: Los estudios de ablación mostraron que los benchmarks de conocimiento factual colapsaron (reteniendo solo el 29–44% del rendimiento) cuando se desactivó Engram, mientras que la comprensión lectora se mantuvo robusta (81–93%), confirmando que Engram actúa como el repositorio principal de conocimiento paramétrico.

3.4. Eficiencia del Sistema

• Rendimiento (Throughput): La descarga de una tabla Engram de 100B de parámetros al host resultó en una penalización de rendimiento insignificante de <3% (2.8% en un backbone de 8B), validando la estrategia de pre-lectra y solapamiento.

4. Significado y Reivindicaciones

El artículo afirma que la Memoria Condicional es una primitiva de modelado indispensable para la próxima generación de modelos dispersos. Su importancia radica en tres áreas principales:

1. Complementariedad Estructural: Establece que la búsqueda estática y la computación dinámica no son mutuamente excluyentes sino estructuralmente complementarias. Optimizar la asignación entre ambas produce una ley de escala en forma de U que supera al MoE puro.
2. Eficiencia mediante el Desacoplamiento: Al desacoplar el almacenamiento de parámetros de la computación mediante direccionamiento determinista, Engram permite que los modelos escalen a cantidades masivas de parámetros (por ejemplo, 40B+) sin verse limitados por la memoria de la GPU o los costos de computación, sorteando efectivamente los cuellos de botella del hardware.
3. Profundización de la Red: El análisis mecanicista sugiere que proporcionar búsqueda de conocimiento explícita libera a las capas tempranas del backbone de las tareas de reconstrucción estática. Esto aumenta efectivamente la profundidad disponible de la red para el razonamiento complejo, explicando las ganancias inesperadas en tareas de razonamiento y código.

Los autores posicionan a Engram no meramente como un aumento de recuperación (RAG), sino como un cambio arquitectónico fundamental que moderniza los conceptos de $N$-gramas para la era de los modelos dispersos a gran escala, ofreciendo un nuevo eje de escasez que complementa el paradigma predominante de MoE.