Resumen Técnico: "Sin subespacio para rastrear: No identificabilidad y estado del optimizador en el entrenamiento de bajo rango"

Planteamiento del problema

Los optimizadores eficientes en memoria como GaLore (Zhao et al., 2024) reducen la huella de memoria del entrenamiento de modelos de lenguaje extensos proyectando los gradientes sobre un subespacio de bajo rango (rango $r$) computado periódicamente y ejecutando Adam dentro de dicho subespio. La premisa geométrica subyacente de este enfoque es que el subespacio de los top-$r$ gradientes es un objeto significativo y de deriva lenta que puede rastrearse a través de los intervalos de actualización ($T$ pasos). Si esta premisa fuera cierta, las estimaciones consecutivas del subespacio serían similares, y llevar el estado del optimizador (momentos) a través de una actualización requeriría solo correcciones menores.

Este artículo desafía esa premisa, argumentando que, más allá de un pequeño núcleo reproducible, el subespacio top-$r$ no es un objeto estable sino que es estadísticamente no identificable. En consecuencia, la aparente "rotación" del subespacio en cada actualización está dominada por el ruido del estimador más que por el movimiento de una estructura geométrica estable.

Metodología

Los autores emplean un marco de medición riguroso para cuantificar la estabilidad del subespacio y el impacto de la no identificabilidad en el rendimiento del optimizador:

1. Medición de la rotación del subespacio: El artículo introduce una métrica de distancia de chordal con corrección de signo ($\|\Delta U\|_F$) para medir la rotación entre la base antigua ($U_{old}$) y la nueva base ($U_{new}$) en cada actualización. Esta métrica oscila entre 0 (idéntica) y $\sqrt{2r}$ (ortogonal).
2. Control de mismo paso (La prueba de "sin tiempo"): Para distinguir entre la deriva temporal genuina y el ruido del estimador, los autores dividen un único minibatch en dos mitades disjuntas en el mismo paso de entrenamiento. Computan el subespacio top-$r$ para cada mitad y miden el desacuerdo. Si el subespacio fuera una propiedad estable del gradiente, las estimaciones del mismo paso deberían estar de acuerdo. Si discrepan tanto como las estimaciones tomadas con $T$ pasos de diferencia, la rotación es ruido.
3. Análisis espectral: Los autores analizan el espectro de valores singulares de los gradientes para determinar si existe una brecha espectral en el rango $r$ que justifique un corte de subespacio estable.
4. Experimentos de promediado: El artículo pone a prueba si el promediar gradientes durante $N$ pasos (o promediar la base a lo largo del tiempo) puede "rescatar" la identificabilidad al abrir una brecha espectral.
5. Evaluación de remedios: Los autores prueban tres estrategias para mitigar los costos de los rediseños de subespacio:
   • Promedio de la base: Utilizar un promedio móvil exponencial (EMA) de la matriz de covarianza para estabilizar la base.
   • Transporte de estado: Transportar explícitamente los momentos del optimizador (primero y segundo) a través de la matriz de rotación $R = U_{new}^\top U_{old}$ (como se hace en LDAdam).
   • Acortamiento de la memoria: Reducir el parámetro de decaimiento de Adam $\beta_2$ para acortar la memoria del segundo momento, evitando que integre estadísticas obsoletas de subespacios previos.

Los experimentos se realizaron en cuatro familias de modelos (Pythia, GPT-2, Qwen2.5, Llama-3) que van desde 70M hasta 6.9B de parámetros, y un Vision Transformer (ViT-B).

Hallazgos clave

1. El subespacio es no identificable

• Rotación máxima: En cada actualización, la rotación medida $\|\Delta U\|_F$ se satura entre el 96.7% y el 99.6% del máximo geomético $\sqrt{2r}$. La nueva base es casi ortogonal a la anterior.
• Dominancia del ruido: El control de mismo paso revela que dos estimaciones de mitades disjuntas del mismo lote discrepan por $0.725\sqrt{2r}$, lo cual es casi idéntico a la rotación a través del tiempo de $0.742\sqrt{2r}$. La "rotación" no es movimiento; es ruido del estimador.
• Núcleo reproducible ($k^\star$): Solo una pequeña fracción de las direcciones es reproducible. Para $r=128$, solo $k^\star \approx 39$ direcciones son consistentes a través de los minibatches. Las $r - k^\star$ direcciones restantes son esencialmente muestras aleatorias de una cola espectral sin brechas.

2. Causas espectrales y el fallo del promediado

• Espectro sin brechas: El espectro del gradiente consiste en un pequeño pico dominante (unas pocas a ~20 direcciones) seguido de una cola suave y casi isotrópica sin brecha en el rango $r$. La relación de valores singulares adyacentes en $r=128$ es $\approx 1.005$.
• Señal, no ruido: La cola espectral profunda no es ruido de muestreo; es señal con un espectro de ley de potencia. Bajo un promediado de $N$ veces, esta cola se reduce solo como $N^{-1/4}$ en lugar del $N^{-1/2}$ esperado para ruido puro. Por lo tanto, ningún presupuesto de promediado finito puede abrir una brecha o estabilizar el subespacio.
• El promediado falla: Promediar la base durante 10,000 pasos reduce la rotación por actualización solo ligeramente (del 97% al 84% del máximo) y no tiene un efecto medible en la perplejidad final (un rango de solo 0.36).

3. Por qué GaLore sigue funcionando

A pesar de la falta de un subespacio estable, GaLore funciona porque el rol principal de la proyección es la captura de energía, no el rastreo. Dado que la cola del gradiente es casi isotrópica, cualquier marco de $r$ dimensiones que incluya el pico dominante captura una fracción aproximadamente constante de la energía del gradiente ($\gamma_g \approx 0.67–0.73$). La elección específica de las direcciones importa poco para la dirección del descenso, siempre que se preserve la fracción de energía.

4. Remedios efectivos

El costo del rediseño es la obsolescencia del estado del optimizador transportado, particularmente el segundo momento (varianza).

• Transporte de estado: Transportar el primer momento vía $R$ y el segundo momento vía las entradas al cuadrado de $R$ es óptimo. Transportar el estado ciegamente es demostrablemente subóptimo por un factor de $\approx (r - k^\star)/2$. En pruebas controladas, transportar el estado con el $\beta_2=0.999$ por defecto logró una perplejidad de 18.73, superando al GaLore sin transporte (19.28).
• Acortamiento de la memoria: Reducir $\beta_2$ de 0.999 a 0.99 ayuda a los optimizadores de actualización al prevenir la integración de estadísticas obsoletas de subespacios previos, ahora inválidos. Esto mejoró la perplejidad entre 1.8 y 26 puntos dependiendo del optimizador y la receta.
• Ineficacia del promedio de la base: Como se señaló, estabilizar la base mediante el promediado no funciona porque el subespacio es fundamentalmente no identificable.

Significado y Reivindicaciones

El artículo afirma proporcionar la primera medición directa de la premisa geométrica que subyace a los optimizadores de gradiente de bajo rango, encontrando que es falsa en los rangos utilizados habitualmente.

• Clarificación del mecanismo: El artículo aclifica que GaLore funciona no porque rastree una estructura estable, sino porque captura una fracción constante de la energía del gradiente en un marco que es efectivamente redibujado.
• Herramienta de diagnóstico: Los autores proponen una sonda práctica: medir el rango reproducible $k^\star$ dividiendo un lote. Si el rango elegido $r > k^\star$, las direcciones excedentes son ruido.
• Prescripción de ingeniería: El artículo prescribe dos soluciones accionables para el entrenamiento de bajo rango:
  1. Transportar el estado a través de la rotación (como en LDAdam) en lugar de transportarlo ciegamente.
  2. Reducir $\beta_2$ a 0.99 para reducir la memoria del segundo momento, mitigando la obsolescencia causada por los rediseños periódicos.
• Alcance: Los hallazgos se aplican específicamente a optimizadores que computan y refrescan explícitamente un subespacio de gradiente (familia GaLore). No se aplican a parametrizaciones como LoRA que no computan un subespacio de gradiente.

Los autores concluyen que el "objeto" (el subespacio estable de bajo rango) no existe más allá de un pequeño núcleo, e intentar rastrearlo es perseguir un fantasma. El camino a seguir consiste en reconocer el rediseño y adaptar el estado del optimizador en consecuencia.