On the Convergence Rate of LoRA Gradient Descent

Resumo Técnico: Sobre a Taxa de Convergência do Descenso de Gradiente LoRA

Declaração do Problema
A Adaptação de Baixo Rango (LoRA) tornou-se um método padrão para o ajuste fino de modelos de linguagem grandes devido à sua eficiência, atualizando apenas as matrizes adaptadoras de baixo rango $A$ e $B$ de modo que os pesos finais sejam $W_0 + BA$. No entanto, a compreensão teórica de sua convergência é limitada. Embora a função de perda original $\ell(W)$ seja tipicamente assumida como Lipschitz suave, a função de perda reparametrizada $L(BA)$ não é Lipschitz suave em relação às matrizes adaptadoras $A$ e $B$. Essa falta de suavidade impede a aplicação direta de técnicas clássicas de análise de convergência. Trabalhos teóricos existentes caem em três categorias: (1) análises assintóticas em regimes infinitos (por exemplo, Kernel Tangente Neural ou largura infinita), que não fornecem taxas não assintóticas para modelos finitos; (2) análises de algoritmos modificados (por exemplo, GaLore, LoRA-One) que não abordam estritamente a formulação original do LoRA; e (3) análises do algoritmo LoRA original que dependem de suposições fortes e artificiais, como matrizes adaptadoras uniformemente limitadas, para impor suavidade Lipschitz. Consequentemente, a taxa de convergência não assintótica do algoritmo original de descenso de gradiente LoRA sem suposições de limitação permanece uma questão em aberto.

Metodologia
Os autores fornecem a primeira análise de convergência não assintótica do algoritmo original de descenso de gradiente LoRA sob suposições mínimas: a função de perda original é Lipschitz suave e limitada inferiormente. A prova procede em três etapas-chave:

1. Reformulação via Matrizes Empilhadas: O problema é reformulado empilhando as matrizes adaptadoras $B \in \mathbb{R}^{m \times r}$ e $A^T \in \mathbb{R}^{r \times n}$ em uma única variável $V \in \mathbb{R}^{(m+n) \times r}$. A atualização original $BA$ é recuperada como um bloco específico do produto externo $VV^T$. Isso permite que a otimização seja vista como descenso de gradiente em uma nova função $J(V) = L(E_1 V V^T E_2)$, onde $E_1$ e $E_2$ são matrizes extratoras.
2. Lema de Descenso Modificado: Como $J(V)$ não é Lipschitz suave (devido ao fator multiplicativo de $V$ no gradiente), os autores derivam um lema de descenso modificado. Este lema limita a diminuição da função $J(V_{t+1}) - J(V_t)$ usando termos de ordem superior envolvendo as normas dos parâmetros $\|V_t\|$ e do gradiente.
3. Controle da Taxa de Aprendizado: Para garantir a descida em cada etapa, a taxa de aprendizado $\eta_t$ é controlada dinamicamente com base na norma atual do parâmetro e na norma do gradiente. Especificamente, $\eta_t$ é definida inversamente proporcional a $\|V_t\|^2 + \|\nabla L(E_1 V_t V_t^T E_2)\|$. Os autores provam que, mesmo que a norma do parâmetro $\|V_t\|$ cresça ilimitadamente (o que força $\eta_t$ a diminuir), a soma das taxas de aprendizado $\sum \eta_t$ diverge suficientemente lentamente (como uma série harmônica) para garantir a convergência, embora a uma taxa mais lenta que o descenso de gradiente padrão.

Principais Contribuições

• Taxa de Convergência Não Assintótica: O artigo prova que a norma mínima do gradiente sob o descenso de gradiente LoRA converge para zero a uma taxa de $O(1/\log T)$, onde $T$ é o número de iterações. Este resultado vale sem assumir matrizes adaptadoras limitadas ou suavidade Lipschitz da perda reparametrizada.
• Recuperação de Taxas Padrão: Os autores mostram que, se as normas das adaptadoras forem limitadas (uma condição frequentemente assumida em trabalhos anteriores), a taxa de convergência melhora para o padrão $O(1/T)$, igualando a taxa para funções Lipschitz suaves.
• Dependência de Posição: A análise revela uma "dependência de posição" na dinâmica do LoRA: a convergência é desacelerada se as iterações se afastam da origem (aumentando $\|V_t\|$) e acelerada se se movem em direção à origem. Isso é atribuído à reestruturação geométrica da paisagem de perda pela parametrização LoRA, que cria um ponto estacionário na origem independentemente da estrutura de perda original.
• Esquemas Práticos de Taxa de Aprendizado: Motivados pela exigência teórica de controlar $\eta_t$ com base nas normas dos parâmetros e do gradiente, os autores propõem e validam cronogramas práticos de taxas de aprendizado adaptativas ($\eta_{adapt}$, $\eta_{adapt2}$ e $\eta_{norm}$) que se ajustam com base na norma das adaptadoras empilhadas ou no valor da perda.

Resultados

• Teórico: O Teorema 3.5 estabelece a taxa de convergência $O(1/\log T)$ para o caso geral e $O(1/T)$ para o caso limitado. A prova demonstra que a relação recursiva entre a taxa de aprendizado, a norma do parâmetro e a norma do gradiente introduz um desaceleração na convergência em comparação com o descenso de gradiente padrão.
• Empírico: Experimentos foram conduzidos em regressão logística com embeddings do CIFAR-10 e ResNet-18 treinado no CIFAR-10.
  • Taxas de aprendizado adaptativas e normalizadas derivadas da teoria demonstraram convergência mais rápida e estável em comparação com taxas de aprendizado constantes.
  • Nos experimentos de regressão logística, a norma do parâmetro $\|V_t\|$ foi observada crescendo ilimitadamente ao longo de 1000 épocas, confirmando o cenário teórico onde a convergência desacelera para $O(1/\log T)$.
  • Nos experimentos com ResNet-18, $\|V_t\|$ convergiu para um conjunto limitado, resultando em convergência mais rápida.
  • No ajuste fino de Modelos de Linguagem Grandes (LLM) (TinyLlama-1.1B), taxas adaptativas mostraram benefícios significativos quando os parâmetros foram inicializados com normas pequenas, mas comportaram-se de maneira semelhante às taxas constantes quando inicializados com normas grandes (onde $\|V_t\|$ é grande e cresce lentamente em relação ao seu tamanho).

Significado
O artigo afirma fechar uma lacuna significativa na compreensão teórica do LoRA ao fornecer a primeira análise de convergência não assintótica para o algoritmo original sem suposições artificiais de limitação. Ele esclarece que a taxa de convergência do LoRA não depende do rango escolhido $r$, mas é governada pela interação entre a taxa de aprendizado e a norma das iterações. O trabalho destaca que a reparametrização do LoRA altera fundamentalmente a geometria de otimização, potencialmente levando a soluções diferentes em comparação com o ajuste fino de rango completo. Além disso, as estratégias de taxa de aprendizado propostas oferecem um método prático para estabilizar e acelerar o treinamento do LoRA, levando em conta a estrutura única induzida pela reparametrização de baixo rango.