Muon falla en QK: la solución que propone un ingeniero

From Muon to Gradient Clipping: Some Thoughts on QK Stability

El optimizador Muon, que ha mostrado resultados prometedores en modelos como Kimi K2, se vuelve inestable cuando se aplica directamente a las matrices de Query y Key en Transformers, pudiendo colapsar el entrenamiento. Este artículo explora la causa geométrica de este fallo, ligada al acoplamiento funcional de Q y K y al comportamiento de rango completo de Muon. Propone una modificación basada en los primeros principios de Muon, que restringe el cambio en la matriz de atención en lugar de los pesos individuales, y detalla el proceso de derivación, los desafíos de ingeniería y las soluciones heurísticas finales.