De Muon al recorte de gradiente: la inestabilidad de QK en Transformers

Fuentes: From Muon to Gradient Clipping: Some Thoughts on QK Stability

El optimizador Muon, propuesto por Keller Jordan y Jeremy Bernstein, se diseñó desde una perspectiva de espacio de funciones y resuelve un problema con restricción de norma espectral cuya solución analítica es la función matrix sign. En la práctica, trabajos recientes como Kimi K2 han aplicado Muon con buenos resultados, pero los materiales públicos del proyecto señalan que, al usarlo directamente sobre las matrices de peso de consulta (Q) y clave (K) de un Transformer, el entrenamiento puede volverse muy inestable e incluso colapsar.

La explicación teórica, siguiendo el análisis de Su Jianlin, combina dos factores. Primero, Q y K forman un par acoplado: su producto QK^T define las puntuaciones de atención, por lo que limitar los cambios funcionales de cada matriz por separado no garantiza que el cambio acoplado esté acotado. Segundo, la actualización de Muon, al pasar por la función msign, tiene todos sus valores singulares iguales, lo que equivale a una actualización de rango efectivo completo. Esta característica hace que las direcciones singulares de la actualización choquen con más frecuencia con las de los pesos, inflando su norma espectral. Cuando este crecimiento ocurre simultáneamente en W_Q y W_K, los efectos se multiplican en el producto QK^T, generando una retroalimentación que puede llevar al MaxLogit explosion y al colapso.

Ante esta evidencia, una modificación fiel al espíritu de Muon debería dejar de acotar las matrices de peso por separado y pasar a acotar directamente el cambio en la matriz de puntuaciones de atención. La entrada expone con detalle la derivación teórica, el conflicto geométrico en la actualización de QK y la dirección de una nueva función objetivo, aunque el texto queda truncado antes de presentar la solución completa.