Resumen Técnico: DeMuon - Un Muon Descentralizado para la Optimización de Matrices sobre Grafos

Declaración del Problema

Este artículo aborda el problema de la optimación de matrices totalmente descentralizada sobre grafos de comunicación. El objetivo es resolver el problema de optimización multiagente:
$$\min_{X \in \mathbb{R}^{m \times n}} f(X) := \frac{1}{N} \sum_{i=1}^N f_i(X)$$
donde $N$ agentes están conectados mediante un grafo dirigido $\mathcal{G}=(\mathcal{V}, \mathcal{E})$, y cada nodo $i$ posee una función objetivo local $f_i$. A diferencia de los métodos tradicionales de optimización descentralizada que suelen vectorizar las variables matriciales, este trabajo se centra en preservar la estructura de la matriz, aprovechando específicamente técnicas del algoritmo Muon centralizado (Jordan et al., 2024). El objetivo es lograr la convergencia global y el consenso entre los agentes sin un servidor central, mitigando al mismo tiempo la heterogeneidad en las funciones objetivo locales.

Metodología

El artículo propone dos algoritmos: DeMuon (una variante descentralizada de Muon) y DeMuon-A (una variante acelerada).

1. DeMuon: Muon Descentralizado

DeMuon integra actualizaciones basadas en la norma espectral dentro de un marco de seguimiento de gradientes (gradient-tracking). El algoritmo mantiene tres secuencias: estimadores de gradiente locales $\{M^{(k)}\}$, estimadores de gradiente globales $\{V^{(k)}\}$ e iterados locales $\{X^{(k)}\}$.

• Momento Local: Cada nodo actualiza su estimador local $M^{(k)}_i$ utilizando una media móvil exponencialmente ponderada de los gradientes estocásticos.
• Seguimiento de Gradiente: Los nodos actualizan $V^{(k)}_i$ agregando información de sus vecinos para rastrear el gradiente global, asegurando el consenso en la dirección del gradiente a pesar de la heterogeneidad local.
• Ortogonalización de Matrices: El iterado local $X^{(k+1)}_i$ se actualiza agregando las actualizaciones de los vecinos, aplicando específicamente una función de signo de matriz ($\text{msgn}(\cdot)$) al gradiente rastreado. Este paso resuelve un subproblema lineal con restricción de norma utilizando la norma espectral ($\|\cdot\|$), lo cual es central en la metodología de Muon.

Garantías Teóricas para DeMuon:
Bajo supuestos de ruido acotado y continuidad de Lipschitz del gradiente con respecto a las normas espectral y nuclear, los autores establecen:

• Error de Consenso: El error $\|X^{(k)} - \mathbf{1} \otimes \bar{X}^{(k)}\|$ está acotado por $O(\frac{\sqrt{N}\lambda \eta}{1-\lambda})$, donde $\lambda$ caracteriza la tasa de mezcla del grafo.
• Error de Estacionariedad: El error de estacionariedad promedio converge a una tasa de $O(K^{-1/4})$ con un tamaño de paso $\eta$ y un momento $\theta$ debidamente ajustados. Esto coincide con la tasa de convergencia óptima de SGD normalizado con momento en entornos centralizados.

2. DeMuon-A: Muon Descentralizado Acelerado

DeMuon-A extiende la técnica de multi-extrapolación (He et al., 2026b) al entorno de matrices descentralizadas.

• Mecanismo: Realiza $q = p-1$ pasos de extrapolación local por iteración para generar puntos $Z^{(k,s)}_i$, explotando la suavidad de orden superior de la función objetivo.
• Supuestos: Requiere una condición de Lipschitz de orden superior en la $p$-ésima derivada de las funciones $f_i$, inducida por las normas espectral y nuclear.
• Convergencia: Bajo estas condiciones, DeMuon-A logra una tasa de error de estacionariedad mejorada de $O(K^{-p/(3p+1)})$ y una tasa de error de consenso de $O(K^{-(2p+1)/(3p+1)})$.

Principales Contribuciones

1. Propuesta Algorítmica: Los autores introducen DeMuon, el primer método de optimización de matrices totalmente descentralizado que incorpora la ortogonalización basada en la norma espectral de Muon y el seguimiento de gradientes.
2. Análisis de Convergencia: Proporcionan garantías rigurosas de convergencia global tanto para DeMuon como para DeMuon-A en términos de errores de estacionariedad y de consenso. Notablemente, el análisis se basa en desigualdades de momento de martingalas de matrices para establecer una convergencia independiente de la dimensión (independiente de las dimensiones de la matriz $m, n$).
3. Aceleración: El artículo propone DeMuon-A, que integra la multi-extrapolación para aprovechar la suavidad de orden superior, ofreciendo tasas de convergencia mejoradas sobre el algoritmo base DeMuon.
4. Validación Empírica: Experimentos numéricos en aprendizaje profundo distribuido (entrenamiento de Transformers) y regresión de matrices distribuida demuestran las ventajas prácticas de los métodos propuestos sobre las líneas base descentralizadas existentes.

Resultados Experimentales

Los autores evaluaron DeMuon frente a tres líneas base: DSGD (SGD descentralizado vanilla), DSGD-C (SGD con recorte o clipping) y DSGD-N (SGD normalizado con seguimiento de gradiente).

• Configuración: Los experimentos se realizaron en un modelo GPT de 7.4M de parámetros entrenado en el conjunto de datos Multi30k a través de tres topologías de grafo: Completa, Exponencial Dirigida y de Anillo.
• Rendimiento: DeMuon superó consistentemente a todas las líneas base en términos de pérdida de entrenamiento, pérdida de validación y perplejidad en todas las topologías.
• Consenso: El error de consenso para DeMuon se mantuvo competitivo o inferior al de las líneas base, lo que indica un acuerdo efectivo entre los agentes.
• Especificaciones: En la topología de Grafo Completo, DeMuon alcanzó una pérdida de validación final de 0.240 (perplejidad 1.27), significativamente menor que la de DSGD-N (0.302) y DSGD (7.031).

Significado y Reivindicaciones

El artículo afirma que DeMuon logra cerrar la brecha entre el éxito empírico de Muon en el aprendizaje profundo centralizado y los requisitos de la optimización descentralizada.

• Convergencia Independiente de la Dimensión: A diferencia de algunos trabajos recientes (por ejemplo, Zhang et al., 2026a) que pueden depender explícitamente de la dimensionalidad del problema, los autores afirman que su uso de desigualdades de momento de martingalas de matrices permite garantías de convergencia que no dependen explícitamente de las dimensiones de la matriz $(m, n)$.
• Ventaja Práctica: Los resultados sugieren que la estructura específica de Muon (restricciones de norma espectral y ortogonalización de matrices) proporciona claras ventajas prácticas en entornos descentralizados, particularmente para el entrenamiento de grandes modelos de lenguaje donde las estructuras de matrices son inherentes.
• Rigor Teórico: El trabajo establece que los beneficios del momento y la normalización en la optimización de matrices pueden preservarse en un entorno totalmente descentralizado sin un coordinador central, siempre que se utilicen el seguimiento de gradientes y las actualizaciones apropiadas basadas en normas.

Los autores posicionan este trabajo como un paso hacia la llegada de técnicas avanzadas de optimización de matrices al mundo descentralizado, motivado por aplicaciones en aprendizaje profundo distribuido y problemas clásicos como PCA distribuido y completación de matrices.