A Single-Layer Model Can Do Language Modeling

Resumen Técnico: Un Modelo de Capa Única Puede Realizar Modelado de Lenguaje

Enunciado del Problema
Los modelos de lenguaje modernos suelen escalar su capacidad aumentando la profundidad, apilando múltiples capas donde cada capa mantiene su propio estado (por ejemplo, cachés KV por capa en Transformers o estados matriciales por capa en Mamba, GDN, RWKV y xLSTM). Esto distribuye la computación y la memoria a lo largo de la pila. En contraste, las redes neuronales biológicas dependen en gran medida de la recurrencia —revisitar el mismo circuito a lo largo del tiempo— en lugar de apilar capas. El artículo investiga la viabilidad de esta forma biológica para el modelado de lenguaje: específicamente, hasta qué punto puede llegar una única capa recurrente y poco profunda cuando revisita un único vector de estado en cada paso, en lugar de distribuir el estado a través de una pila profunda.

Metodología: Redes de Predicción Fundamentadas (GPN)
Los autores proponen las Redes de Predicción Fundamentadas (GPN), una arquitectura definida por un único vector de estado revisado en cada paso de tiempo. El bucle central consta de tres operaciones:

1. Fundamentación: El estado predicho anterior ($s^p_{t-1}$) se fusiona con el token observado actual ($x_t$) mediante un residuo con puerta para producir un estado fundamentado ($s^g_t$).
2. Predicción: El estado fundamentado se procesa para predecir el siguiente estado ($s^p_t$). Esta predicción se decodifica en una distribución de tokens para el siguiente paso.
3. Memoria (GPN+M): Los autores extienden la GPN base a GPN+M añadiendo una memoria matricial centralizada ($M$). Esta memoria se actualiza mediante una regla delta con puerta y se lee en paralelo con una Red de Alimentación Adelante (FFN). Tanto la FFN (parámetros fuera de línea) como la lectura de memoria (parámetros en línea) actúan como predictores aditivos para la actualización del estado.

A diferencia de los modelos recurrentes estándar que mantienen un estado por capa, GPN utiliza una sola capa ancha y poco profunda con una única memoria matricial compartida. La profundidad se logra a través del tiempo (revisitar el bucle) en lugar del espacio (apilar capas). El vector de estado cumple una doble función: actúa como una previsión para el siguiente paso y como una codificación de la historia observada hasta el momento, siendo la "fundamentación" la que completa la codificación para el siguiente ciclo.

Contribuciones Clave

1. Arquitectura: La introducción de GPN y GPN+M, que utilizan un único vector de estado, una FFN y una memoria matricial centralizada revisada en cada paso, en contraste con los estados apilados por capa de los Transformers y las RNN lineales.
2. Visión Conceptual: Un marco donde la predicción del estado del siguiente paso se considera parte de la codificación de ese paso, con un único vector desempeñando ambos roles.
3. Evidencia Empírica: Puntos de referencia de rendimiento con 130 millones de parámetros que muestran que una GPN+M de 1 capa logra una perplejidad competitiva en comparación con líneas base profundas, con una variante de 2 capas reduciendo aún más la brecha.
4. Observabilidad: Dado que todo el contexto de trabajo es un único vector, los autores pueden inspeccionar directamente su geometría sin necesidad de profundizar a través de flujos residuales. Identifican estructuras emergentes no diseñadas explícitamente en la arquitectura.

Resultados
Evaluados en el conjunto de datos FineWeb-Edu (130 millones de parámetros, ventana de contexto de 2048):

• GPN+M de 1 capa: Logra una perplejidad (PPL) de 18.06. Esto está dentro del 13% de un Transformer++ de 12 capas (PPL 16.05) y del 18% de un Gated DeltaNet (GDN) de 10 capas (PPL 15.34).
• GPN+M de 2 capas: Logra una PPL de 16.95, cerrando la brecha al 6% del Transformer++ y al 11% del GDN.
• Tarea LAMBADA: Los modelos poco profundos muestran una menor precisión en tareas dependientes del contexto en comparación con las líneas base profundas, lo que resalta un costo en la profundidad de computación por token.
• Sin Memoria (GPN): Una versión sin memoria matricial funciona significativamente peor (PPL 23.51), confirmando la necesidad del componente de memoria asociativa.

Estructuras Emergentes (Análisis)
El estado de vector único permite un análisis geométrico directo, revelando tres estructuras que surgieron espontáneamente sin regularización explícita:

1. Geometría de Alto Rango: El vector de estado es ancho ($d=2496$) y utiliza su ancho completo. Exhibe dimensiones casi no correlacionadas y ninguna dimensión "muerta", con un rango estable de ~40.
2. Dirección Persistente de Token Predeterminado: La media del vector de estado ($\bar{s}$) es un gran desplazamiento fijo que se decodifica en una distribución dominada por tokens funcionales frecuentes (por ejemplo, "el", "y", ","). Esta dirección persiste incluso cuando se añade un sesgo explícito al decodificador, lo que sugiere que cumple funciones estructurales más allá de la decodificación simple (por ejemplo, anclando RMSNorm).
3. División de Memoria Rápida/Lenta: Las 15 cabezas de memoria arquitectónicamente idénticas se dividen espontáneamente en dos grupos durante el entrenamiento: un grupo "rápido" (6 cabezas) con vidas medias de 1 a 3 tokens, y un grupo "lento" (9 cabezas) con vidas medias que oscilan entre 5 y 322 tokens. Esto crea un sistema de memoria de dos escalas temporales sin diseño arquitectónico explícito.

Significado y Limitaciones
El artículo afirma que una única capa recurrente poco profunda puede realizar el modelado de lenguaje de manera efectiva, alcanzando aproximadamente un 13% de las líneas base profundas con conteos de parámetros equivalentes. El significado radica en demostrar que la "profundidad" puede ser reemplazada por la recurrencia revisitada en el tiempo y que un único vector puede soportar estructuras de memoria complejas y emergentes (divisiones rápidas/lentas) y propiedades geométricas (anclajes persistentes) que normalmente están ocultas en flujos residuales profundos.

Sin embargo, los autores son modestos sobre el estado actual del trabajo:

• No Listo para Producción: El modelo se ejecuta secuencialmente en Python sin paralelización a lo largo de la dimensión de la secuencia, lo que hace que los costos de entrenamiento sean significativamente más altos que los de los Transformers.
• Escala: Los resultados se basan en una única semilla con 130 millones de parámetros, sin barridos de escalado.
• Naturaleza Descriptiva: El análisis de las estructuras emergentes (como la dirección del token predeterminado) es observacional; el artículo no afirma prueba causal de su utilidad mediante ablación.
• Costos de Contexto: El modelo poco profundo paga un costo medible en tareas que requieren contexto profundo (como LAMBADA), lo que indica que, aunque el enfoque es viable, las pilas profundas aún mantienen ventajas para predicciones específicas con mucho contexto.

El trabajo sirve como una prueba de concepto para el bucle de "predicción fundamentada" y la viabilidad de los modelos de estado de vector único, ofreciendo una nueva lente para observar cómo los modelos de lenguaje organizan la información.