One Model, Two Roles: Emergent Specialization in a Shared Recurrent Transformer

Resumen Técnico: Un Modelo, Dos Roles: Especialización Emergente en un Transformer Recurrente Compartido

Enunciado del Problema
El artículo investiga si un único Transformer recurrente con pesos compartidos puede desarrollar roles computacionales internos distintos para dos estados latentes sin ser particionado en módulos de parámetros separados. Mientras que trabajos anteriores como HRM (Modelo de Razonamiento Jerárquico) utilizan redes separadas para diferentes estados latentes y TRM (Modelo de Razonamiento Transformer) utiliza una única red compartida, sigue sin estar claro si un modelo compartido puede especializarse espontáneamente en roles funcionales distintos (por ejemplo, un estado de "propuesta" frente a un estado de "refinamiento") cuando se le proporciona una señal estable que distingue los tipos de actualización. La pregunta central es si la especialización requiere modularidad estructural (parámetros separados) o si puede emerger de una diferenciación funcional impulsada por una señal clara de identidad de estado dentro de una arquitectura monolítica compartida.

Metodología: Recurrencia de Entrada Asimétrica (AIR)
Para abordar esto, los autores introducen la Recurrencia de Entrada Asimétrica (AIR), una plataforma de prueba mínima derivada de HRM y TRM.

• Arquitectura: AIR mantiene dos estados latentes, $z_H$ (de alto nivel) y $z_L$ (de bajo nivel), actualizados por un único modelo Transformer compartido $f(\cdot; \theta)$.
• Mecanismo de Actualización: La recurrencia sigue un patrón de subpasos intercalados (por ejemplo, $L, L, H, L, L, H$). La asimetría crítica reside en la inyección de entrada:
  • Actualizaciones L: La entrada codificada $\tilde{x}$ se inyecta en la regla de actualización ($z_L \leftarrow f(z_L + z_H + \alpha_L \tilde{x})$).
  • Actualizaciones H: La entrada codificada no se inyecta ($z_H \leftarrow f(z_H + z_L + \alpha_H \tilde{x})$ con $\alpha_H=0$).
• Entrenamiento: El modelo se entrena utilizando gradientes truncados de un paso (retropropagando solo a través de las actualizaciones finales L y H de un ciclo) en dos tareas de razonamiento sintéticas: Sudoku-Extreme (rompecabezas de 9x9) y Maze (rejillas de 30x30).
• Análisis: Los autores emplean varias técnicas para sondear el modelo:
  • Despliegues Decodificados: Visualización de estados intermedios para observar niveles de compromiso.
  • Experimentos de Congelación: Congelar un estado para medir cambios de contenido en el otro.
  • Ablaciones: Probar inyección de entrada simétrica (eliminando la asimetría) y señales alternativas de identidad de estado (tokens de nivel).
  • Análisis Mecanístico: Análisis de patrones de atención (localidad vs. globalidad) y sensibilidad a violaciones.

Contribuciones y Resultados Clave

1. Existencia de Especialización Funcional:
   Los despliegues decodificados revelan una división estable y emergente en el comportamiento del modelo compartido:

   • $z_H$ (Alto nivel): Se comporta como un "estado de propuesta totalmente comprometido". Se decodifica a una solución candidata completa en cada subpaso, incluso si es incorrecta.
   • $z_L$ (Bajo nivel): Se comporta como un "borrador de compromiso cambiante". Retiene incertidumbre local, con posiciones indecisas (BLANK en Sudoku, PAD en Maze) que cambian a través de los subpasos.
   • Control: En variantes simétricas (donde ambos estados reciben inyección de entrada idéntica), esta división de roles distinta casi desaparece, confirmando que la especialización no es inherente a tener dos estados, sino que depende de la capacidad de distinguir los tipos de actualización.

2. La Señal de Identidad de Estado Induce Especialización:
   El artículo demuestra que el modelo compartido requiere una señal fiable para distinguir los tipos de actualización e inducir especialización.

   • Asimetría de Entrada: La inyección de entrada asimétrica ($\Delta > 0$) mejora significativamente la precisión (por ejemplo, ~59% en Sudoku frente a ~51% para variantes simétricas).
   • Tokens de Nivel: Los autores probaron señales alternativas utilizando tokens de nivel aprendidos. Solo una estrategia de "prepending y stripping" (donde se añade un token, se atiende a él y se elimina) recuperó con éxito la brecha de rendimiento. La adición elemento a elemento simple o los tokens persistentes fallaron, lo que sugiere que la señal debe ser estructuralmente separable y no mezclarse con los tokens de contenido.
   • Ablación de Estado Único: Un modelo de estado único con inyección de entrada periódica no logró igualar el rendimiento del modelo asimétrico de dos estados, lo que indica que se necesitan búferes de memoria distintos para una especialización estable.

3. Dinámicas de Estado Interdependientes:
   Los experimentos de congelación revelan una relación específica entre los dos estados:

   • Sudoku: Congelar $z_H$ suprime la actividad de $z_L$ (reduciendo los cambios de contenido), mientras que congelar $z_L$ aumenta la actividad de $z_H$.
   • Maze: Congelar cualquiera de los estados aumenta los cambios de contenido en el otro.
   • Control Simétrico: Los modelos simétricos carecen de este patrón de interacción organizado y consistente con el rol, mostrando respuestas sensibles a perturbaciones pero no estructuradas.

4. Evidencia Mecanística (Análisis de Atención):
   El análisis de atención confirma la división funcional:

   • Localidad: Las actualizaciones $L$ atienden consistentemente más localmente (a vecindades inmediatas) que las actualizaciones $H$, que exhiben una atención más global.
   • Sensibilidad a Violaciones: En Sudoku, las capas más profundas de las actualizaciones $L$ muestran una enrutamiento de atención específico hacia las restricciones violadas, una señal no observada en las actualizaciones $H$ ni en la tarea Maze.

Significado y Afirmaciones
El artículo afirma que la modularidad explícita de parámetros (redes separadas) no es estrictamente necesaria para la especialización funcional en el razonamiento recurrente. En cambio, una señal clara de identidad de estado—suministrada ya sea mediante inyección de entrada asimétrica o un token de nivel estructuralmente separable—es suficiente para inducir roles funcionales estables y relacionados dentro de un Transformer recurrente con parámetros compartidos.

Los autores posicionan este trabajo como investigación fundamental sobre los sesgos inductivos de pequeños sistemas de razonamiento recurrente. Enfatizan que sus conclusiones se derivan de entornos sintéticos controlados (Sudoku y Maze) y no afirman universalidad a través del lenguaje natural o la observabilidad parcial. La contribución principal es demostrar que "un modelo, dos roles" es un paradigma viable para arquitecturas de razonamiento eficientes, siempre que el modelo pueda distinguir el contexto computacional de sus actualizaciones.