Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers

Resumen Técnico: Colapso Post-Grokking en la Interfaz de Representación–Lectura en Transformers Entrenados con Muon

Planteamiento del Problema
El artículo investiga una inestabilidad específica que ocurre en modelos Transformer entrenados en aritmética modular (específicamente suma y resta modular) después de haber alcanzado el "grokking"—el fenómeno donde la generalización emerge mucho después de que la precisión de entrenamiento se ha saturado. Si bien trabajos previos establecieron que Muon (un optimizador ortogonalizante) acelera el grokking en comparación con AdamW, este estudio reporta que los modelos entrenados con Muon sufren frecuentemente un colapso recurrente y abrupto en el rendimiento de la generalización post-grokking. Crucialmente, este fallo no se limita a Muon; los modelos AdamW también exhiben inestabilidad. El estudio demuestra que ni el optimizador ni el baseline es estrictamente estable a través de cinco semillas: el baseline de AdamW seleccionado registra evaluaciones sub-umbral en cuatro de cinco semillas (alcanzando un mínimo del 27.59%), mientras que Muon las registra en las cinco. El problema central es que estos modelos pierden la capacidad de generalizar a pesar de que el conjunto de entrenamiento permanece resuelto, y las medidas de progreso mecanístico estándar (soporte espectral y distribución de potencia) no logran detectar el fallo, reportando un circuito intacto incluso cuando la precisión cae a niveles cercanos al azar.

Metodología
Los autores emplean un régimen de entrenamiento de optimizador dividido donde las matrices de pesos ocultos son actualizadas por Muon, mientras que los embeddings de tokens/posiciones y el unembedding de salida (lectura o readout) son actualizados por AdamW. El estudio utiliza tareas de aritmética modular ($a+b \pmod p$ y $a-b \pmod p$) con $p=113$ (y variantes con $p=97$) en Transformers de solo decodificador de diversas profundidades (1, 2 y 4).

Componentes metodológicos clave incluyen:

1. Ramificación Emparejada (Matched Branching): Las trayectorias se ramifican desde estados de parámetros idénticos en bits para aislar factores causales. Una rama continúa con la optimización estándar, mientras que otras congelan grupos específicos de parámetros (matrices ocultas, embeddings o lectura) para probar su necesidad.
2. Descomposición de Fourier: Los autores descomponen los estados residuales internos del modelo en modos de Fourier 2D sobre la rejilla de los operandos. Definen "familias" específicas de modos (por ejemplo, la familia de la suma $(k, k)$, la familia de la resta $(k, -k)$) y utilizan intervenciones de filtrado para medir la suficiencia y necesidad de estas familias para el rendimiento de la tarea.
3. Suite de Intervención: El estudio aplica intervenciones como el desorden de fase (phase scrambling), la reubicación de frecuencia y la sustitución de lectura cruzada para determinar la naturaleza exacta del algoritmo aprendido y la fuente de desalineación.
4. Ablación: Se ablaciona la trayectoria de ortogonalización normalizada de Muon para probar su papel en la dispersión espectral e inestabilidad observadas.

Hallazgos Clave y Resultados

• La Naturaleza del Colapso: El fallo es una "desalineación oculto–lectura" (hidden–readout misalignment). A medida que la pérdida de entrenamiento se acerca a cero, el gradiente desaparece, pero los dos optimizadores responden de manera diferente. El paso de actualización de Muon se vuelve mayormente independiente de la magnitud del gradiente (elasticidad $\approx -0.03$), mientras que el de AdamW permanece proporcional (elasticidad $\approx +1.5$). Esto causa que la representación oculta y la lectura se distancien a una tasa de 8.0 veces por parámetro. El conjunto de entrenamiento permanece resuelto, pero la lectura ya no puede decodificar la representación con la que fue entrenada junto a ella.
• Enmascaramiento vs. Fallo del Circuito: El estudio distingue dos modos de fallo. En un fallo del circuito, la familia de Fourier deja de computar la tarea. En un enmascaramiento del circuito, la familia alineada con la tarea (por ejemplo, la familia de la suma) continúa resolviendo la tarea perfectamente de forma aislada (100% de precisión), pero su contribución es superada por un remanente adversarial en el modelo completo. Los autores muestran que reescalar únicamente el componente alineado con la tarea restaura la precisión total del modelo, probando que el circuito estaba intacto pero fue "vencido por votos".
• Diferencias entre Optimizadores: Muon alcanza el grokking significativamente más rápido (mediana de 13,011 pasos frente a 30,486 para AdamW) pero es universalmente inestable en todas las configuraciones probadas. AdamW también es inestable a través de cinco semillas, aunque la severidad del colapso (medida por el conteo de evaluaciones sub-umbral) es dos órdenes de magnitud menor que la de Muon. En un barrido de una sola semilla, las configuraciones de AdamW mostraron una inestabilidad condicional donde la severidad escalaba con la tasa de aprendizaje, pero la replicación a través de cinco semillas confirmó que el baseline "estable" seleccionado no era estrictamente estable entre semillas. La inestabilidad en Muon se atribuye a su regla de actualización, la cual dispersa la representación a través de un amplio espectro (326 pares conjugados efectivos vs. ~5 para AdamW), dejando a la familia alineada con la tarea con una parte menor del poder total.
• Selección de la Familia por la Tarea: La tarea, no el optimizador, selecciona la familia de Fourier. Los modelos entrenados en suma usan los modos $(k, k)$; aquellos entrenados en resta usan $(k, -k)$. La base dentro de la familia es exacta y no compartida; dos modelos entrenados desde la misma semilla divergen para usar diferentes frecuencias y fases, y sus lecturas se vuelcian casi ortogonales, fallando al decodificar las representaciones del otro.
• Profundidad y Alineación: En modelos más profundos, el código de Fourier es escrito por un MLP en el bloque penúltimo, pero solo se vuelve "legible" para el unembedding en el bloque final. La desalineación ocurre específicamente en la interfaz entre el estado residual del bloque final y el unembedding.
• Prevención: Congelar las matrices ocultas o los parámetros no ocultos (embeddings y lectura) después de que el circuito se forma previene completamente el colapso. El "Muon Estable", donde los embeddings y la lectura se congelan post-grokking, logra una estabilidad estricta a través de 451,400 pasos post-grokking y cinco semillas, mientras que el Muon estándar falla repetidamente.

Significancia y Reivindicaciones
El artículo afirma identificar una limitación fundamental en el entrenamiento con optimizadores divididos y en la dinámica post-grokking: el flujo residual no tiene una base privilegiada y, cuando la superficie de pérdida se aplana, la pérdida no proporciona presión para mantener la alineación específica entre la representación y la lectura. Los autores argumentan que:

1. Las medidas estándar son insuficientes: Las medidas espectrales (soporte y distribución de potencia) pueden permanecer invariantes durante un colapso porque el fallo es un cambio de base (desalineación) en lugar de una pérdida del algoritmo en sí mismo.
2. El grokking es un enmascaramiento reversible: La transición de la generalización al colapso es la inversa de la transición de la memorización a la generalización; ambas implican que la participación de la familia alineada con la tarea aumenta o disminuye respecto al remanente adversarial.
3. La interfaz es el locus del fallo: La inestabilidad no es una propiedad de la representación oculta o de la lectura por sí solas, sino de su identificación conjunta. Fijar un lado de la interfaz estabiliza el sistema.
4. La velocidad de Muon viene con un costo: Aunque Muon acelera el descubrimiento del circuito, su regla de actualización crea una representación que es más susceptible a la desalineación con la lectura bajo condiciones de entrenamiento estándar, lo que requiere intervención (congelación) para la estabilidad a largo plazo.

El estudio no propone nuevas aplicaciones, sino que ofrece un relato mecanístico de un modo de fallo específico en la interpretabilidad mecanística, resaltando que un modelo puede poseer un circuito correcto mientras falla en utilizarlo debido a la deriva representacional.