On the Convergence Rate of LoRA Gradient Descent

Resumen Técnico: Sobre la Tasa de Convergencia del Descenso de Gradiente LoRA

Enunciado del Problema
La Adaptación de Bajo Rango (LoRA) se ha convertido en un método estándar para el ajuste fino de modelos de lenguaje grandes debido a su eficiencia, actualizando únicamente las matrices adaptadoras de bajo rango $A$ y $B$ de modo que los pesos finales sean $W_0 + BA$. Sin embargo, la comprensión teórica de su convergencia es limitada. Si bien la función de pérdida original $\ell(W)$ se asume típicamente como Lipschitz suave, la función de pérdida reparametrizada $L(BA)$ no es Lipschitz suave con respecto a las matrices adaptadoras $A$ y $B$. Esta falta de suavidad impide la aplicación directa de técnicas clásicas de análisis de convergencia. Los trabajos teóricos existentes se dividen en tres categorías: (1) análisis asintóticos en regímenes infinitos (por ejemplo, Kernel de Tangente Neuronal o ancho infinito), que no proporcionan tasas no asintóticas para modelos finitos; (2) análisis de algoritmos modificados (por ejemplo, GaLore, LoRA-One) que no abordan estrictamente la formulación original de LoRA; y (3) análisis del algoritmo LoRA original que dependen de suposiciones fuertes y artificiales, como matrices adaptadoras uniformemente acotadas, para imponer suavidad Lipschitz. En consecuencia, la tasa de convergencia no asintótica del algoritmo original de descenso de gradiente LoRA sin suposiciones de acotación sigue siendo una pregunta abierta.

Metodología
Los autores proporcionan el primer análisis de convergencia no asintótico del algoritmo original de descenso de gradiente LoRA bajo suposiciones mínimas: la función de pérdida original es Lipschitz suave y está acotada inferiormente. La demostración procede en tres pasos clave:

1. Reformulación mediante Matrices Apiladas: El problema se reformula apilando las matrices adaptadoras $B \in \mathbb{R}^{m \times r}$ y $A^T \in \mathbb{R}^{r \times n}$ en una única variable $V \in \mathbb{R}^{(m+n) \times r}$. La actualización original $BA$ se recupera como un bloque específico del producto externo $VV^T$. Esto permite visualizar la optimización como descenso de gradiente sobre una nueva función $J(V) = L(E_1 V V^T E_2)$, donde $E_1$ y $E_2$ son matrices extractoras.
2. Lema de Descenso Modificado: Dado que $J(V)$ no es Lipschitz suave (debido al factor multiplicativo de $V$ en el gradiente), los autores derivan un lema de descenso modificado. Este lema acota la disminución de la función $J(V_{t+1}) - J(V_t)$ utilizando términos de orden superior que involucran las normas de los parámetros $\|V_t\|$ y del gradiente.
3. Control de la Tasa de Aprendizaje: Para garantizar el descenso en cada paso, la tasa de aprendizaje $\eta_t$ se controla dinámicamente en función de la norma actual del parámetro y la norma del gradiente. Específicamente, $\eta_t$ se establece inversamente proporcional a $\|V_t\|^2 + \|\nabla L(E_1 V_t V_t^T E_2)\|$. Los autores demuestran que incluso si la norma del parámetro $\|V_t\|$ crece sin límite (lo que obliga a que $\eta_t$ disminuya), la suma de las tasas de aprendizaje $\sum \eta_t$ diverge lo suficientemente lentamente (como una serie armónica) para garantizar la convergencia, aunque a una tasa más lenta que el descenso de gradiente estándar.

Contribuciones Clave

• Tasa de Convergencia No Asintótica: El artículo demuestra que la norma mínima del gradiente bajo el descenso de gradiente LoRA converge a cero a una tasa de $O(1/\log T)$, donde $T$ es el número de iteraciones. Este resultado se mantiene sin asumir matrices adaptadoras acotadas ni suavidad Lipschitz de la pérdida reparametrizada.
• Recuperación de Tasas Estándar: Los autores muestran que si las normas de las adaptadoras están acotadas (una condición a menudo asumida en trabajos anteriores), la tasa de convergencia mejora a la estándar $O(1/T)$, coincidiendo con la tasa para funciones Lipschitz suaves.
• Dependencia de la Posición: El análisis revela una "dependencia de la posición" en la dinámica de LoRA: la convergencia se ralentiza si las iteraciones se alejan del origen (aumentando $\|V_t\|$) y se acelera si se mueven hacia el origen. Esto se atribuye a la reestructuración geométrica del paisaje de pérdida por la parametrización LoRA, que crea un punto estacionario en el origen independientemente de la estructura de pérdida original.
• Esquemas Prácticos de Tasa de Aprendizaje: Motivados por el requisito teórico de controlar $\eta_t$ en función de las normas de los parámetros y del gradiente, los autores proponen y validan programas prácticos de tasas de aprendizaje adaptativas ($\eta_{adapt}$, $\eta_{adapt2}$ y $\eta_{norm}$) que se ajustan en función de la norma de las adaptadoras apiladas o del valor de la pérdida.

Resultados

• Teóricos: El Teorema 3.5 establece la tasa de convergencia $O(1/\log T)$ para el caso general y $O(1/T)$ para el caso acotado. La demostración muestra que la relación recursiva entre la tasa de aprendizaje, la norma del parámetro y la norma del gradiente introduce una ralentización en la convergencia en comparación con el descenso de gradiente estándar.
• Empíricos: Se realizaron experimentos en regresión logística con incrustaciones de CIFAR-10 y ResNet-18 entrenado en CIFAR-10.
  • Las tasas de aprendizaje adaptativas y normalizadas derivadas de la teoría demostraron una convergencia más rápida y estable en comparación con las tasas de aprendizaje constantes.
  • En los experimentos de regresión logística, se observó que la norma del parámetro $\|V_t\|$ crecía sin límite durante 1000 épocas, confirmando el escenario teórico donde la convergencia se ralentiza a $O(1/\log T)$.
  • En los experimentos con ResNet-18, $\|V_t\|$ convergió a un conjunto acotado, resultando en una convergencia más rápida.
  • En el ajuste fino de Modelos de Lenguaje Grandes (LLM) (TinyLlama-1.1B), las tasas adaptativas mostraron beneficios significativos cuando los parámetros se inicializaron con normas pequeñas, pero se comportaron de manera similar a las tasas constantes cuando se inicializaron con normas grandes (donde $\|V_t\|$ es grande y crece lentamente en relación con su tamaño).

Significado
El artículo afirma cerrar una brecha significativa en la comprensión teórica de LoRA al proporcionar el primer análisis de convergencia no asintótico para el algoritmo original sin suposiciones artificiales de acotación. Aclara que la tasa de convergencia de LoRA no depende del rango $r$ elegido, sino que está gobernada por la interacción entre la tasa de aprendizaje y la norma de las iteraciones. El trabajo destaca que la reparametrización de LoRA altera fundamentalmente la geometría de optimización, lo que potencialmente conduce a soluciones diferentes en comparación con el ajuste fino de rango completo. Además, las estrategias de tasa de aprendizaje propuestas ofrecen un método práctico para estabilizar y acelerar el entrenamiento de LoRA teniendo en cuenta la estructura única inducida por la reparametrización de bajo rango.