🤖 AI

Looped State-Space Language Models with Adaptive Exit-State Selection

Este artículo demuestra que la aplicación de arquitecturas en bucle a los modelos de espacio de estados basados en Mamba mejora las capacidades de razonamiento y la eficiencia de parámetros mediante un aumento de la profundidad computacional, al tiempo que introduce una selección adaptativa de estado de salida para optimizar la profundidad de la predicción, aunque señala que el ahorro real en el tiempo de inferencia requiere mecanismos adicionales de gestión de estados.

Autores originales: Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

Publicado 2026-07-14

📖 6 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que tienes un robot chef brillante y superinteligente llamado Mamba. Este chef es famoso por ser increíblemente rápido y eficiente en el uso de la memoria, capaz de preparar un banquete sin necesidad de una cocina enorme llena de diferentes herramientas. Pero hay un inconveniente: cuando se enfrenta a una receta realmente difícil —como un rompecabezas matemático complejo o un juego de "encontrar el ingrediente oculto"— Mamba a veces se queda bloqueado. Necesita pensar más profundamente, pero no tiene suficientes "capas" de pensamiento para hacerlo sin construir una cocina nueva y gigante (lo que cuesta una fortuna).

Entran en escena los investigadores de la Universidad de Rikkyo y la Universidad de Tokio. Hicieron una pregunta simple y audaz: ¿Qué pasaría si no construyéramos una cocina más grande, sino que hiciéramos que Mamba cocinara el mismo plato una y otra vez, refinando el sabor con cada pasada?

La cocina "en bucle": Un chef, muchas pasadas

En el mundo de la IA, la mayoría de los modelos son como una línea de montaje de una fábrica. Un dato (como una palabra en una oración) se mueve a lo largo de la línea, pasando por 24 estaciones diferentes (capas), y luego termina. Si quieres que el robot piense más profundamente, normalmente solo tienes que añadir más estaciones.

Los autores intentaron algo diferente. Construyeron un Mamba en Bucle (Looped Mamba). Imagina que Mamba es una estación única y supertalentosa. En lugar de enviar el dato a lo largo de una larga línea de 24 estaciones diferentes, lo envían de vuelta a la misma estación, 24 veces seguidas.

La Magia: El robot utiliza el mismo cerebro (parámetros) para cada pasada. Es como un estudiante que vuelve a leer el mismo párrafo de un libro de texto cuatro veces para entenderlo, en lugar de leer cuatro párrafos diferentes y confusos.
El Resultado: En acertijos lógicos complicados llamados Mano (aritmética modular) y p-hop induction (encontrar patrones ocultos), esta estrategia de "relectura" funcionó de maravilla. Un modelo que solo tenía 4 capas pero las ejecutó en bucle 6 veces (4 ⊗6) aplastó a un modelo estándar con 24 capas únicas. Esto sugiere que para el razonamiento, la profundidad del pensamiento importa más que el número de herramientas únicas.

La "Puerta de Salida": Eligiendo la mejor versión

Pero, ¿qué pasa si el robot sigue releyendo el párrafo para siempre? Eso es un desperdicio de tiempo. Los investigadores añadieron una característica ingeniosa llamada Puerta de Salida (Exit Gate).

Piensa en esta puerta como un selector inteligente. A medida que Mamba procesa los datos, la puerta pregunta: "¿Se está aclarando la respuesta?".

Si la respuesta ya está clara después de 2 bucles, la puerta dice: "¡Detente! Usaremos el resultado de esta pasada".
Si el problema es difícil, la puerta dice: "¡Sigue adelante!" y selecciona el resultado de la 3ª o 4ª vuelta.

Esto se llama Selección de Estado de Salida Adaptativa (Adaptive Exit-State Selection). El artículo encontró que para modelos pequeños (alrededor de 140 millones de parámetros), esta puerta fue un factor decisivo. Permitió al modelo elegir el tiempo de pensamiento "justo" para cada palabra específica, a menudo funcionando mejor que forzar al robot a usar siempre el resultado de la cantidad máxima de bucles.

Sin embargo, hay un giro. Los investigadores fueron muy cuidadosos al señalar que esto aún no ahorra electricidad ni tiempo en la computadora. Aunque la puerta selecciona un resultado anterior, el robot todavía ejecuta físicamente todos los bucles en segundo plano. Esto se debe a que el estado de memoria de Mamba es continuo: el resultado del bucle 3 depende del estado dejado por el bucle 2, que depende del bucle 1. No puedes simplemente "soltar" al robot de los bucles posteriores sin romper la cadena para todas las palabras subsiguientes. Por lo tanto, la puerta selecciona la profundidad de predicción (qué versión de la respuesta usar), pero la computación de tiempo real (wall-clock computation) sigue siendo la misma. Para ahorrar tiempo realmente, se necesitaría una nueva forma de manejar el estado de memoria del robot, algo que los autores dicen que es un trabajo para el futuro.

El choque de realidad: Lo que no funcionó (y lo que aún es un "tal vez")

El artículo es refrescantemente honesto sobre lo que no resolvió.

Más grande no siempre es mejor (para este truco): Cuando compararon su Mamba en Bucle con un modelo estándar no bucleado que era igual de "costoso" de ejecutar (mismo número de cálculos, o "iso-FLOPs"), el modelo estándar seguía ganando en la prueba básica de "¿qué tan confundido está el modelo?" (perplejidad). El modelo en bucle era excelente en tareas de razonamiento, pero el modelo profundo estándar seguía siendo ligeramente mejor en la simple predicción de la siguiente palabra en una oración.
El tamaño importa: La "Puerta de Salida" funcionó de maravilla para los modelos de 140 millones de parámetros. Pero cuando lo probaron en los modelos más grandes de 370 millones de parámetros, la puerta no mejoró mucho las cosas. Los modelos más grandes parecían preferir simplemente ejecutar todos los bucles de todos modos. Los autores sugieren que los modelos más grandes podrían depender más de esa última y profunda pasada de pensamiento.
No es una solución mágica para todo: Los investigadores probaron esto en acertijos sintéticos (como árboles matemáticos y búsqueda de patrones). Aunque estos demuestran que el concepto funciona, admiten que aún no sabemos si este truco de "bucle" funcionará igual de bien en el razonamiento abierto y del mundo real, como escribir una novela o resolver un caso legal complejo.

La conclusión

El artículo sugiere que reutilizar el mismo cerebro inteligente una y otra vez es una forma poderosa de hacer que la IA sea más inteligente en el razonamiento sin construir un cerebro más grande y costoso. Es un eje de escalado "eficiente en parámetros".

Probado: En acertijos lógicos específicos, el bucle de un modelo Mamba pequeño supera a un modelo estándar de su mismo tamaño y se iguala a uno mucho más grande.
Sugerido: Este enfoque podría ser clave para hacer que la IA futura sea más eficiente, especialmente para tareas que requieren un pensamiento profundo y paso a paso.
No resuelto aún: Todavía tenemos que descubrir cómo hacer que la "Puerta de Salida" ahorre realmente tiempo de computación (en lugar de solo seleccionar la salida) y si esto funciona con los modelos masivos de miles de millones de parámetros que impulsan a las IA más grandes de hoy en día.

En resumen, los autores no solo construyeron un robot más grande; le enseñaron a un robot más pequeño a pensar más profundamente permitiéndole echar un segundo (o tercer, o cuarto) vistazo al problema. Y por ahora, ese segundo vistazo parece muy prometedor.

Resumen Técnico: Modelos de Lenguaje de Espacio de Estados en Bucle con Selección Adaptativa de Estado de Salida

Declaración del Problema

Los avances recientes en los Modelos de Lenguaje de Gran Escala (LLMs) han sido impulsados principalmente por el escalado del tamaño del modelo, los datos de entrenamiento y el cómputo intensivo en la inferencia. Sin embargo, este paradigma de escalado introduce desafíos significativos de despliegue, incluyendo el aumento del consumo de memoria y el crecimiento rápido de los costos de inferencia debido al mantenimiento de grandes cachés de clave-valor en los mecanismos de autoatención de los Transformers. Si bien los Modelos de Espacio de Estados (SSMs), particularmente Mamba-2, ofrecen una alternativa eficiente con un escalado favorable para contextos largos, aún no está claro si el principio de "profundidad computacional sobre parámetros independientes"—demostrado recientemente en los Transformers en bucle (Looped Transformers)—se aplica también a los SSMs. Además, la investigación existente sobre arquitecturas en bucle se ha centrado casi exclusivamente en backbones de Transformer, dejando un vacío en la comprensión de cómo funciona la compartición de parámetros recurrente dentro de las arquitecturas Mamba e híbridas de Mamba-Transformer.

Metodología

Los autores proponen las arquitecturas Looped Mamba y Looped Hybrid Mamba–Transformer. El mecanismo central consiste en compartir los parámetros de un único bloque Mamba (o híbrido) a través de las capas y aplicarlo repetidamente mediante bucles explícitos de profundidad finita.

Arquitectura

Notación: Un modelo se denota como $N \otimes R$ , donde $N$ es el número de capas con parámetros independientes y $R$ es el número de pasos recurrentes. La profundidad efectiva es $NR$.
Looped Mamba: Aplica repetidamente un backbone Mamba de $N$ $N$ capas con parámetros compartidos $M_\theta$ $M_{θ}$ .
- $h^{(r)} = M_\theta(h^{(r-1)})$ para $r = 1, \dots, R$ .
- La salida final se deriva del estado oculto en el paso $R$ .
Looped Hybrid: Extiende el bucle a un backbone que contiene tanto capas de Mamba-2 como un pequeño subconjunto de capas de autoatención causal. El bucle se aplica al bloque híbrido completo, permitiendo el refinamiento iterativo de las representaciones ocultas a través de actualizaciones de espacio de estados e interacciones basadas en atención.
Protocolo de Entrenamiento:
1. Pre-entrenamiento: Optimizado utilizando un objetivo de bucle final (la pérdida se calcula solo en el paso $R$ ).
2. Exit-Gate Etapa 1: Introduce un objetivo de computación adaptativa con regularización de entropía (adaptado de Ouro) para aprender una distribución por token sobre las profundidades de salida utilizando las salidas de todos los pasos recurrentes.
3. Exit-Gate Etapa 2: Congela el backbone y la cabeza, entrenando únicamente la puerta de salida (exit gate) para fomentar la continuación solo cuando el siguiente paso mejora la pérdida de predicción.

Selección Adaptativa de Estado de Salida

El artículo adapta el proceso de dos etapas de la puerta de salida de Ouro a Mamba. En la inferencia, un umbral $q$ controla la selección del paso de salida $z_i$ para cada token basándose en la probabilidad de salida acumulada.

Nota Semántica: Los autores aclaran que en los SSMs, la cadena de estados de todos los $R$ pasos debe preservarse debido a las restricciones de continuidad de estado. Por lo tanto, la "salida" selecciona qué estado oculto del paso recurrente se pasa a la cabeza de modelado de lenguaje; esto no reduce inherentemente el tiempo de ejecución (wall-clock computation) ni el número de pasos ejecutados sin mecanismos adicionales de manejo de estado.

Contribuciones Clave

Estudio Empírico Controlado: Los autores presentan una comparación rigurosa de los modelos Looped Mamba e Híbridos bajo protocolos de iso-parámetros (mismo $N$ , variando $R$ ) e iso-FLOPs (misma profundidad efectiva $NR$) frente a las líneas base de Mamba y Transformer no en bucle.
Evaluación de Tareas de Razonamiento: El estudio valida el enfoque en tareas de razonamiento sintético (Mano y $p$ -hop induction) y en el pre-entrenamiento de modelos de lenguaje con benchmarks de evaluación.
Mecanismo de Salida Adaptativa: La introducción de un mecanismo de selección de estado de salida controlado por umbral para modelos de espacio de estado en bucle, demostrando su capacidad para mejorar el rendimiento en tareas downstream en profundidades intermedias.

Resultados Experimentales

Tareas de Razonamiento Sintético

Tarea Mano (Aritmética Modular): Looped Mamba supera consistentemente a las líneas base no en bucle con parámetros emparejados. En el escenario más difícil ( $L=24$ ), un modelo $4 \otimes 6$ logró una precisión del 99.35% frente al 7.50% de la línea base $4 \otimes 1$ . Bajo condiciones de iso-FLOPs (profundidad efectiva 24), el modelo $4 \otimes 6$ superó significativamente a la línea base no en bucle $24 \otimes 1$ (99.35% vs. 37.45%), lo que sugiere que la compartición recurrente de parámetros proporciona un sesgo inductivo más fuerte para la computación composicional que el aumento de capas no compartidas.
$p$ -hop Induction (Recuperación Recursiva): Los modelos Mamba no en bucle de baja profundidad tuvieron dificultades (ej. $1 \otimes 1$ con ~~38-41%), mientras que añadir bucles mejoró drásticamente el rendimiento. El modelo $1 \otimes 12$ alcanzó una precisión casi perfecta (~~99.89%), indicando que la profundidad recursiva es más crítica que el recuento de parámetros distintos para la recuperación recursiva.

Modelado de Lenguaje y Benchmarks Downstream

Iso-Parámetro: La profundidad recurrente mejoró la perplejidad de validación (PPL) tanto para los backbones Mamba-2 como LLaMA-2 bajo presupuestos de parámetros fijos.
Iso-FLOPs: Los modelos no en bucle más profundos ( $24 \otimes 1$ ) retuvieron una ventaja en la perplejidad de validación sobre los modelos en bucle con profundidad efectiva igualada. Sin embargo, los modelos en bucle se mantuvieron competitivos en benchmarks downstream, superando en ocasiones a las líneas base no en bucle más profundas.
Salida Adaptativa (Adaptive Exit):
- Escala de 140M: La puerta de salida adaptativa mejoró el rendimiento downstream. Para el modelo Mamba-2 de 140M, la puerta de salida de la Etapa 2 logró un puntaje promedio de 38.87 con un paso de salida seleccionado promedio de 3.00, superando a la línea base de bucle fijo 4 (37.79).
- Escala de 370M: El mejor rendimiento downstream permaneció concentrado cerca de la configuración de bucle fijo 4. Las variantes con puerta de salida no superaron los puntajes máximos de los bucles fijos, lo que sugiere que los modelos más grandes pueden depender más fuertemente del último paso recurrente o requieren una mejor calibración de la puerta.
- Compromiso (Trade-off): La salida adaptativa proporciona una frontera controlable entre la profundidad seleccionada y el rendimiento, permitiendo al modelo utilizar estados recurrentes intermedios de manera efectiva en lugar de forzar una profundidad máxima fija.

Significancia y Reivindicaciones

El artículo sostiene que la profundidad recurrente ofrece un eje de escalado prometedor y eficiente en términos de parámetros para los modelos de lenguaje Mamba e híbridos. Los hallazgos clave incluyen:

Profundidad vs. Parámetros: Muchas tareas de razonamiento se benefician de una mayor profundidad computacional en lugar de parámetros independientes adicionales. El uso recurrente de capas Mamba mejora sustancialmente las capacidades de computación composicional y recursiva sin aumentar el recuento de parámetros.
Viabilidad Híbrida: El principio de profundidad en bucle se aplica eficazmente a las arquitecturas híbridas Mamba-Transformer, donde las capas de atención se integran en el bucle junto con las capas de Mamba.
Selección Adaptativa: Aunque las implementaciones actuales de SSM requieren la ejecución completa del estado, la selección adaptativa de estado de salida permite la selección de profundidades de predicción óptimas, mejorando el rendimiento downstream en profundidades intermedias (particularmente en la escala de 140M).

Los autores mantienen la modestia respecto a sus afirmaciones, señalando que los experimentos se limitan a modelos de hasta 370 millones de parámetros y que el ahorro real de cómputo durante la inferencia requiere trabajos futuros sobre mecanismos de manejo de estado para omitir completamente los pasos recurrentes. Los resultados se presentan como evidencia empírica controlada más que como una caracterización completa del rendimiento a gran escala.