antonio leandro

ia generativa

LoRA: Low-Rank Adaptation of Large Language Models

paper · Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, et al. ·

a tese

o update do fine-tuning cabe num rank ridículo: duas matrizes finas ao lado do peso congelado igualam o ajuste completo, e como o produto delas soma no peso original, a inferência não paga nada por isso

o que fica

  1. O update de adaptação tem rank intrínseco baixo: no GPT-3 175B, r igual a 1 ou 2 em Wq e Wv já empata com o fine-tuning completo em WikiSQL e MNLI, com o rank cheio em 12.288.
  2. Como BA é linear e tem a mesma forma de W0, dá para somar os dois antes de servir — LoRA não cobra latência de inferência, enquanto adapters chegam a cobrar 30,3% com batch 1 e sequência curta.
  3. Com orçamento fixo de parâmetros, espalhar rank baixo por mais matrizes bate concentrar rank alto numa só: Wq mais Wv com r igual a 4 dá 73,7 no WikiSQL contra 70,4 de Wq sozinho com r igual a 8.
  4. O que economiza memória não é o forward, é o optimizer: sem gradiente para os pesos congelados o Adam para de guardar estado, e a VRAM de treino do GPT-3 175B cai de 1,2TB para 350GB.
  5. ΔW não repete as direções dominantes de W — ele amplifica direções que já estavam lá e não eram enfatizadas, com fator perto de 21,5 quando r é 4 e perto de 2 quando r é 64.
  6. Se os pesos estiverem merged, um mesmo batch não pode misturar tarefas diferentes: ou você mantém os módulos separados e paga o custo, ou separa os batches por adapter.

o problema

Fine-tuning atualiza todo parâmetro do modelo. Enquanto o modelo era GPT-2 ou RoBERTa large, isso era inconveniência: o checkpoint da tarefa tinha o tamanho do modelo, e você guardava. Com o GPT-3 175B virou outra categoria de problema. Cada tarefa produz um conjunto de 175 bilhões de parâmetros novos — 350GB de checkpoint — e servir dez tarefas quer dizer guardar e carregar dez desses. A barreira de hardware para treinar é a mesma do pré-treino, porque com Adam você carrega, além dos pesos, os estados do optimizer para cada parâmetro.

As alternativas da época tinham cada uma seu pedágio. Adapters inserem camadas novas entre os blocos, e camada nova é profundidade nova: o cálculo é sequencial e não desaparece no paralelismo do hardware. Os autores mediram isso no GPT-2 medium com batch 1 e sequência de 128 — o cenário de inferência online — e viram 20,7% e 30,3% de latência a mais em duas variantes de adapter, mesmo com bottleneck pequeno. Com o modelo shardado piora, porque a profundidade extra pede mais AllReduce e Broadcast. A outra linha, prefix tuning, não mexe em peso nenhum: reserva tokens especiais na entrada. O custo é o comprimento de sequência que some, e o paper registra que o método é difícil de otimizar e não melhora de forma monótona conforme ganha parâmetros.

a ideia

Trabalhos anteriores mostraram que modelos superparametrizados vivem numa dimensão intrínseca baixa. LoRA leva a observação um passo adiante e a aplica não ao modelo, mas à diferença: a hipótese é que o ΔW aprendido durante a adaptação também é pobre em rank. Se for, não faz sentido representar essa diferença como uma matriz densa d por k. Basta guardar o produto de duas matrizes finas.

É guardar um diff em vez do arquivo inteiro, com a restrição de rank fazendo o papel da compressão. E, diferente de um adapter, esse diff tem exatamente a forma do peso original — o que abre a porta do truque que vem a seguir.

como funciona

Para um peso pré-treinado W0, LoRA escreve o update como BA, com B de d por r, A de r por k e r muito menor que min(d, k). W0 fica congelado. O forward vira uma soma de dois caminhos que recebem o mesmo input:

h = W0 @ x + (alpha / r) * (B @ (A @ x))
# W0: d x k, congelado, sem gradiente
# A:  r x k, init gaussiano aleatório
# B:  d x r, init zero  ->  BA = 0 no passo 0

A inicialização importa: com B em zero, o treino começa exatamente no modelo pré-treinado. O fator alpha/r existe para você não ter que retunar learning rate toda vez que mexe em r; os autores fixam alpha no primeiro r que testam e não voltam nele.

No Transformer há quatro matrizes na self-attention (Wq, Wk, Wv, Wo) e duas no MLP. O paper adapta só as de atenção e congela o MLP; na maior parte dos experimentos, só Wq e Wv. A conta de parâmetros treináveis é 2 × número de matrizes adaptadas × d_model × r.

Na hora de servir, você calcula W = W0 + (alpha / r) * (B @ A) e roda inferência normal. Trocar de tarefa é subtrair BA e somar B’A’. Daí o “no additional inference latency by construction”.

Os números: no GPT-3 175B com r igual a 4 em Wq e Wv, o checkpoint da tarefa cai de 350GB para 35MB, a VRAM de treino cai de 1,2TB para 350GB, e o throughput sobe de 32,5 para 43,1 tokens/s por V100. Em qualidade, 73,4 no WikiSQL e 91,7 no MNLI-m com 4,7M de parâmetros, contra 73,8 e 89,5 do fine-tuning completo. No RoBERTa base, média 87,2 no GLUE com 0,3M de parâmetros contra 86,4 com 125M.

o que isso custou

O truque do merge tem preço. Se você absorve BA no peso, um mesmo batch não consegue atender tarefas diferentes num forward só. Dá para não fazer o merge e escolher os módulos por amostra, mas aí a latência que o paper vendeu como zero volta.

O rank baixo não é lei universal, e os autores dizem isso em nota de rodapé: se a tarefa de destino estiver num idioma diferente do pré-treino, retreinar o modelo inteiro provavelmente ganha de um r pequeno. No GPT-2 medium, o r ótimo em E2E fica entre 4 e 16, e a relação entre tamanho do modelo e rank ótimo continua em aberto.

A escolha de onde aplicar LoRA é heurística — palavra dos autores. MLP, LayerNorm e biases ficaram fora do estudo. E parte das baselines são números citados de trabalhos anteriores, não rodadas próprias; nos experimentos do GPT-3, o custo de treino permitiu reportar desvio padrão típico por tarefa, não por entrada de tabela.

onde isso aparece hoje

O paper já vinha com o pacote de integração com PyTorch e checkpoints para RoBERTa, DeBERTa e GPT-2, e o padrão que ele descreve virou o modo default de distribuir modelo ajustado: um base compartilhado na VRAM e um punhado de arquivos de dezenas de megabytes que entram e saem por cima. A conta que os autores fazem — 100 modelos adaptados ocupando 354GB em vez de 35TB — é a economia que sustenta serviço multi-tenant de modelo customizado.

A investigação da seção 7 teve vida própria. Mostrar que ΔW amplifica direções que já existiam em W, com fator perto de 21,5 para r igual a 4, deu a base empírica para a família inteira de métodos parameter-efficient que vieram depois escolher rank, alpha e conjunto de matrizes como os botões que importam.

lido na íntegra por pipeline de llm, revisado por antonio leandro antes de publicar ·