LoRA (Low-Rank Adaptation) é uma nova técnica para ajuste fino de modelos pré-treinados em larga escala. Esses modelos geralmente são treinados em dados de domínio geral, de modo a ter a quantidade máxima de dados. Para obter melhores resultados em tarefas como bate-papo ou resposta a perguntas, esses modelos podem ser ainda mais “ajustados” ou adaptados em dados específicos de domínio.
É possível fazer um ajuste fino de um modelo apenas inicializando-o com os pesos pré-treinados e treinando mais sobre os dados específicos do domínio. Com o aumento do tamanho dos modelos pré-treinados, um ciclo completo para frente e para trás requer uma grande quantidade de recursos de computação. O ajuste fino simplesmente continuando o treinamento também requer uma cópia completa de todos os parâmetros para cada tarefa/domínio ao qual o modelo é adaptado.
LoRA: Adaptação de baixo nível de modelos de linguagem grandes
propõe uma solução para ambos os problemas usando uma decomposição de matriz de baixa classificação. Ela pode reduzir o número de pesos treináveis em 10.000 vezes e os requisitos de memória da GPU em 3 vezes.
Método
O problema de ajuste fino de uma rede neural pode ser expresso pela descoberta de uma (Delta Teta)
que minimiza (L(X, y; Teta_0 + DeltaTeta)) onde (EU) é uma função de perda, (X) e (e)
são os dados e (Teta_0) os pesos de um modelo pré-treinado.
Nós aprendemos os parâmetros (Delta Teta) com dimensão (|Delta Teta|)
é igual a (|Teta_0|). Quando (|Teta_0|) é muito grande, como em modelos pré-treinados em larga escala, encontrando (Delta Teta) torna-se computacionalmente desafiador. Além disso, para cada tarefa você precisa aprender uma nova (Delta Teta) conjunto de parâmetros, tornando ainda mais desafiador implantar modelos ajustados se você tiver mais do que algumas tarefas específicas.
LoRA propõe usar uma aproximação (Delta Phi aprox Delta Theta) com (|Delta Phi| << |Delta Teta|). A observação é que as redes neurais têm muitas camadas densas realizando multiplicação de matrizes e, embora normalmente tenham classificação completa durante o pré-treinamento, ao se adaptar a uma tarefa específica, as atualizações de peso terão uma baixa “dimensão intrínseca”.
Uma decomposição de matriz simples é aplicada para cada atualização da matriz de peso (Delta teta em Delta Teta). Considerando (Delta theta_i em mathbb{R}^{d vezes okay}) a atualização para o (eu)º peso na rede, LoRA o aproxima com:
(Delta theta_i aprox Delta phi_i = BA)
onde (B em mathbb{R}^{d vezes r}), (A em mathbb{R}^{r vezes d}) e a classificação (r << min(d, okay)). Assim, em vez de aprender (d vezes okay) parâmetros que agora precisamos aprender ((d + okay) vezes r) que é facilmente muito menor dado o aspecto multiplicativo. Na prática, (Delta theta_i) é dimensionado por (frac{alfa}{r}) antes de ser adicionado a (teta_i)que pode ser interpretado como uma ‘taxa de aprendizado’ para a atualização do LoRA.
O LoRA não aumenta a latência de inferência, pois, uma vez feito o ajuste fino, você pode simplesmente atualizar os pesos em (Teta) adicionando seus respectivos (Delta theta aprox Delta phi). Também simplifica a implantação de vários modelos específicos de tarefas em cima de um modelo grande, como (|Delta Phi|) é muito menor que (|Delta Teta|).
Implementando na tocha
Agora que temos uma ideia de como LoRA funciona, vamos implementá-lo usando torch para um problema mínimo. Nosso plano é o seguinte:
- Simule dados de treinamento usando um simples (y = X teta) modelo. (teta em mathbb{R}^{1001, 1000}).
- Treine um modelo linear de classificação completa para estimar (teta) – este será nosso modelo ‘pré-treinado’.
- Simule uma distribuição diferente aplicando uma transformação em (teta).
- Treine um modelo de baixa classificação usando os pesos pré-treinados.
Vamos começar simulando os dados de treinamento:
Agora definimos nosso modelo base:
mannequin <- nn_linear(d_in, d_out, bias = FALSE)
Também definimos uma função para treinar um modelo, que também estamos reutilizando mais tarde. A função faz o loop de treinamento padrão no torch usando o otimizador Adam. Os pesos do modelo são atualizados no native.
practice <- operate(mannequin, X, y, batch_size = 128, epochs = 100) {
decide <- optim_adam(mannequin$parameters)
for (epoch in 1:epochs) {
for(i in seq_len(n/batch_size)) {
idx <- pattern.int(n, measurement = batch_size)
loss <- nnf_mse_loss(mannequin(X(idx,)), y(idx))
with_no_grad({
decide$zero_grad()
loss$backward()
decide$step()
})
}
if (epoch %% 10 == 0) {
with_no_grad({
loss <- nnf_mse_loss(mannequin(X), y)
})
cat("(", epoch, ") Loss:", loss$merchandise(), "n")
}
}
}
O modelo é então treinado:
practice(mannequin, X, y)
#> ( 10 ) Loss: 577.075
#> ( 20 ) Loss: 312.2
#> ( 30 ) Loss: 155.055
#> ( 40 ) Loss: 68.49202
#> ( 50 ) Loss: 25.68243
#> ( 60 ) Loss: 7.620944
#> ( 70 ) Loss: 1.607114
#> ( 80 ) Loss: 0.2077137
#> ( 90 ) Loss: 0.01392935
#> ( 100 ) Loss: 0.0004785107
OK, então agora temos nosso modelo base pré-treinado. Vamos supor que temos dados de uma distribuição ligeiramente diferente que simulamos usando:
thetas2 <- thetas + 1
X2 <- torch_randn(n, d_in)
y2 <- torch_matmul(X2, thetas2)
Se aplicarmos nosso modelo base a essa distribuição, não obteremos um bom desempenho:
nnf_mse_loss(mannequin(X2), y2)
#> torch_tensor
#> 992.673
#> ( CPUFloatType{} )( grad_fn = )
Agora, ajustamos nosso modelo inicial. A distribuição dos novos dados é apenas um pouco diferente da inicial. É apenas uma rotação dos pontos de dados, adicionando 1 a todos os thetas. Isso significa que não se espera que as atualizações de peso sejam complexas, e não devemos precisar de uma atualização de classificação completa para obter bons resultados.
Vamos definir um novo módulo de tocha que implementa a lógica LoRA:
lora_nn_linear <- nn_module(
initialize = operate(linear, r = 16, alpha = 1) {
self$linear <- linear
# parameters from the unique linear module are 'freezed', so they aren't
# tracked by autograd. They're thought of simply constants.
purrr::stroll(self$linear$parameters, (x) x$requires_grad_(FALSE))
# the low rank parameters that will probably be educated
self$A <- nn_parameter(torch_randn(linear$in_features, r))
self$B <- nn_parameter(torch_zeros(r, linear$out_feature))
# the scaling fixed
self$scaling <- alpha / r
},
ahead = operate(x) {
# the modified ahead, that simply provides the end result from the bottom mannequin
# and ABx.
self$linear(x) + torch_matmul(x, torch_matmul(self$A, self$B)*self$scaling)
}
)
Agora inicializamos o modelo LoRA. Usaremos (r = 1)significando que A e B serão apenas vetores. O modelo base tem 1001×1000 parâmetros treináveis. O modelo LoRA que vamos ajustar tem apenas (1001 + 1000), o que o torna 1/500 dos parâmetros do modelo base.
lora <- lora_nn_linear(mannequin, r = 1)
Agora vamos treinar o modelo lora na nova distribuição:
practice(lora, X2, Y2)
#> ( 10 ) Loss: 798.6073
#> ( 20 ) Loss: 485.8804
#> ( 30 ) Loss: 257.3518
#> ( 40 ) Loss: 118.4895
#> ( 50 ) Loss: 46.34769
#> ( 60 ) Loss: 14.46207
#> ( 70 ) Loss: 3.185689
#> ( 80 ) Loss: 0.4264134
#> ( 90 ) Loss: 0.02732975
#> ( 100 ) Loss: 0.001300132
Se olharmos para (Delta teta) veremos uma matriz cheia de 1s, a transformação exata que aplicamos aos pesos:
delta_theta <- torch_matmul(lora$A, lora$B)*lora$scaling
delta_theta(1:5, 1:5)
#> torch_tensor
#> 1.0002 1.0001 1.0001 1.0001 1.0001
#> 1.0011 1.0010 1.0011 1.0011 1.0011
#> 0.9999 0.9999 0.9999 0.9999 0.9999
#> 1.0015 1.0014 1.0014 1.0014 1.0014
#> 1.0008 1.0008 1.0008 1.0008 1.0008
#> ( CPUFloatType{5,5} )( grad_fn = )
Para evitar a latência de inferência adicional da computação separada dos deltas, poderíamos modificar o modelo unique adicionando os deltas estimados aos seus parâmetros. Usamos o add_
método para modificar o peso no native.
with_no_grad({
mannequin$weight$add_(delta_theta$t())
})
Agora, a aplicação do modelo base aos dados da nova distribuição produz um bom desempenho, então podemos dizer que o modelo está adaptado para a nova tarefa.
nnf_mse_loss(mannequin(X2), y2)
#> torch_tensor
#> 0.00130013
#> ( CPUFloatType{} )
Concluindo
Agora que aprendemos como o LoRA funciona neste exemplo simples, podemos pensar em como ele poderia funcionar em grandes modelos pré-treinados.
Acontece que os modelos Transformers são, em sua maioria, uma organização inteligente dessas multiplicações de matrizes, e aplicar LoRA somente a essas camadas é o suficiente para reduzir o custo do ajuste fino em uma grande quantidade, enquanto ainda obtém um bom desempenho. Você pode ver os experimentos no artigo LoRA.
Claro, a ideia de LoRA é simples o suficiente para que possa ser aplicada não apenas a camadas lineares. Você pode aplicá-la a convoluções, camadas de incorporação e, na verdade, a qualquer outra camada.
Imagem de Hu et al no Artigo LoRA