Riemannian Gradient Descent for Low-Rank Architectures

Resumen Técnico: Descenso de Gradiente Riemanniano para Arquitecturas de Bajo Rango

Planteamiento del Problema
En el aprendizaje profundo moderno, los parámetros del modelo suelen ser matrices con rango restringido para reducir el coste computacional y la cantidad de parámetros (por ejemplo, en mecanismos de atención multicabezas donde $W = AB^T$). El entrenamiento estándar actualiza los factores $A$ y $B$ de forma independiente utilizando gradientes euclidianos, ignorando la estructura geométrica del conjunto de matrices de rango $r$. Este artículo investiga si tratar la matriz de rango restringido $W$ como un punto en una subvariedad suave y aplicar la optimización riemanniana puede producir una mejor convergencia o mejores soluciones. El desafío central es que la factorización $W=AB^T$ no es única (es invariante bajo $A \to AS, B \to BS^{-T}$), y el conjunto de matrices de rango $r$ es una variedad, no un espacio vectorial. Los autores pretenden restringir las modificaciones únicamente al optimizador, evitando cambios en la arquitectura del modelo, y manteniendo una complejidad por iteración comparable a los métodos estándar ($O(mr^2 + nr^2)$).

Metodología
Los autores proponen un marco de descenso de gradiente riemanniano adaptado para parámetros factorizados en bajo rango. La iteración general implica:

1. Recuperar el gradiente euclidiano ambiente $G_W$ a partir de los gradientes de los factores $G_A$ y $G_B$.
2. Proyectar este gradiente sobre el espacio tangente de la variedad específica para obtener el gradiente riemanniano $G_W^R$.
3. Actualizar el momento (momentum) y normalizar el tamaño del paso.
4. Moverse a lo largo de la geodésica (usando una retracción) hacia la siguiente iteración.
5. Transportar en paralelo el vector de momento al nuevo espacio tangente.

El artículo instancia este esquema a través de diez diseños algorítmicos específicos categorizados por geometría y restricciones:

1. Geometrías para Matrices de Rango Fijo ($R_{r}^{m \times n}$):

   • Geometría Embebida: Utiliza la métrica de Frobenius ambiente. La retracción se realiza mediante la SVD truncada de rango $r$ de $W + \Xi$. El transporte en paralelo utiliza la proyección métrica.
   • Geometría de Cociente: Trata la variedad como un cociente del espacio producto de los factores $(A, B)$. La métrica se define "arriba" en los factores, induciendo una métrica en el cociente. La retracción es aditiva en los factores; el transporte implica la proyección horizontal.

2. Geometrías para Semi-isometrías de Rango $r$ ($P_{r}^{m \times n}$):

   • Aquí, $W$ está restringida de tal manera que sus valores singulares no nulos son todos 1 (equivalente a que $A$ y $B$ sean ortonormales en sus columnas).
   • Geometría Embebida: Restringe la métrica de Frobenius. La retracción utiliza SVD truncada con restricciones de ortogonalidad específicas.
   • Geometría de Cociente: Utiliza el producto de variedades de Stiefel ($St_{m,r} \times St_{n,r}$) como el espacio "arriba". Las retracciones utilizan descomposición QR o polar.
   • Geometría Canónica: Utiliza la métrica canónica en los factores de Stiefel, induciendo una métrica de producto en el cociente.

3. Compartición de Pesos (Variedades de Rejilla/Grid Manifolds):

   • Extiende lo anterior a arquitecturas donde los factores se comparten entre bloques de filas y bloques de columnas (por ejemplo, Group-Query Attention).
   • Para Rango Fijo, los autores demuestran que apilar los factores en matrices más grandes permite el uso de algoritmos simples (singleton), siempre que se verifique la preservación del rango por bloques.
   • Para Semi-isometrías, el simple apilamiento falla debido a las restricciones de ortogonalidad. Los autores desarrollan algoritmos más complejos, incluyendo un enfoque "híbrido" para el caso embebido donde se elevan a la geometría de cociente para realizar retracciones y transporte, y luego proyectan de nuevo hacia abajo.

Contribuciones Clave

• Exploración Sistemática: El artículo proporciona una encuesta exhaustiva e implementación de diez variantes distintas de optimización riemanniana para parámetros de matrices de bajo rango, cubriendo tanto restricciones de rango fijo como de semi-isometría, así como escenarios de compartición de pesos.
• Detalles de Implementación: Detalla las retracciones específicas (SVD, QR, Polar) y los transportes de vectores requeridos para cada geometría, abordando la no unicidad de las factorizaciones y la preservación de las restricciones de la variedad (rango y ortogonalidad).
• Evaluación Empírica: Los métodos se aplican a las matrices Q, K, V, O en modelos de lenguaje de pequeña escala (6 capas, dimensión de incrustación 512) utilizando el conjunto de datos FineWeb.

Resultados Experimentales

• Rendimiento: Tras el ajuste de las tasas de aprendizaje, los métodos riemannianos no superaron de manera concluyente al baseline de AdamW. En los entornos de Multi-Head Attention (MHA) y Group-Query Attention (GQA), las diferencias de pérdida estuvieron a menudo dentro de la varianza de las semillas aleatorias.
• Hiperparámetros: La normalización de la métrica proporcionó mejoras leves (0.01–0.05 nats). Un momentum pesado ($\nu=7/8$) generalmente perjudicó el rendimiento o proporcionó ganancias insignificantes en comparación con el baseline.
• Estabilidad: Algunas ejecuciones exhibieron picos de pérdida (por ejemplo, alrededor del paso 6660) que se recuperaron, a menudo correlacionados con datos estructurados en el conjunto de entrenamiento.
• Complejidad: Los autores reconocen que, aunque los métodos son matemáticamente sólidos, su complejidad computacional por iteración es mayor que la de AdamW, lo que los hace menos atractivos para los experimentos de pequeña escala realizados.

Significancia y Afirmaciones
El artículo mantiene una postura modesta respecto a su novedad e impacto inmediato:

• Redescubrimiento: Los autores declaran explícitamente que cualquier aparente novedad es "casi seguramente un redescubrimiento", basándose fuertemente en libros de texto establecidos (Absil et al., Boumal) y literatura previa.
• Motivación Heurística: El trabajo está motivado por la heurística de que tener en cuenta mejor las relaciones entre parámetros (vía la geometría de la variedad) podría conducir a una convergencia o soluciones mejores, aunque esto no se demostró definitivamente en los experimentos.
• Potencial Futuro: Los autores argumentan que la falta de superación en modelos de pequeña escala no descarta su mérito en escalas mayores. Sugieren que el enfoque es matemáticamente sólido y que la "relación coste-precisión" podría desplazarse favorablemente con modelos más grandes o arquitecturas diferentes (por ejemplo, modelos basados en tensores como MLA, MoLAE).
• Limitaciones: El estudio admite limitaciones significativas, incluyendo el pequeño tamaño del modelo, la falta de un desvanecimiento (annealing) de la tasa de aprendizaje y el uso de un único baseline (AdamW) sin ajustar sus hiperparámetros.

En conclusión, el artículo sirve como una demostración técnica rigurosa de la aplicación de la optimización riemanniana a componentes de bajo rango en aprendizaje profundo, proporcionando una biblioteca de algoritmos y un punto de referencia para investigaciones futuras, mientras reconoce que los resultados empíricos actuales aún no justifican el reemplazo de optimizadores estándar como AdamW.