vendredi, décembre 1, 2023

Comprendre LoRA avec un exemple minimal



LoRA (Low-Rank Adaptation) est une nouvelle technique permettant de peaufiner des modèles pré-entraînés à grande échelle. De tels modèles sont généralement formés sur des données de domaine général, afin de disposer du maximum de données. Afin d’obtenir de meilleurs résultats dans des tâches telles que le chat ou la réponse à des questions, ces modèles peuvent être davantage « affinés » ou adaptés sur des données spécifiques à un domaine.

Il est possible d’affiner un modèle simplement en initialisant le modèle avec les poids pré-entraînés et en effectuant une formation supplémentaire sur les données spécifiques au domaine. Avec la taille croissante des modèles pré-entraînés, un cycle complet avant et arrière nécessite une grande quantité de ressources informatiques. Un réglage fin par simple poursuite de la formation nécessite également une copie complète de tous les paramètres pour chaque tâche/domaine auquel le modèle est adapté.

LoRA : adaptation de bas rang de grands modèles de langage
propose une solution aux deux problèmes en utilisant une décomposition matricielle de bas rang. Il peut réduire de 10 000 fois le nombre de poids pouvant être entraînés et de 3 fois les besoins en mémoire GPU.

Méthode

Le problème du réglage fin d’un réseau neuronal peut s’exprimer en trouvant un \(\Delta \Thêta\)
qui minimise \(L(X, y; \Theta_0 + \Delta\Theta)\)\(L\) est une fonction de perte, \(X\) et \(y\)
sont les données et \(\Thêta_0\) les poids d’un modèle pré-entraîné.

Nous apprenons les paramètres \(\Delta \Thêta\) avec dimension \(|\Delta \Thêta|\)
est égal à \(|\Thêta_0|\). Quand \(|\Thêta_0|\) est très grand, comme dans les modèles pré-entraînés à grande échelle, trouvant \(\Delta \Thêta\) devient un défi informatique. De plus, pour chaque tâche, vous devez apprendre un nouveau \(\Delta \Thêta\) ensemble de paramètres, ce qui rend encore plus difficile le déploiement de modèles affinés si vous avez plusieurs tâches spécifiques.

LoRA propose d’utiliser une approximation \(\Delta \Phi \environ \Delta \Theta\) avec \(|\Delta \Phi| << |\Delta \Theta|\). L’observation est que les réseaux neuronaux comportent de nombreuses couches denses effectuant une multiplication matricielle, et bien qu’ils aient généralement un rang complet pendant le pré-entraînement, lors de l’adaptation à une tâche spécifique, les mises à jour de poids auront une faible « dimension intrinsèque ».

Une simple décomposition matricielle est appliquée pour chaque mise à jour de la matrice de poids \(\Delta \theta \in \Delta \Theta\). Considérant \(\Delta \theta_i \in \mathbb{R}^{d \times k}\) la mise à jour pour le \(je\)ème poids dans le réseau, LoRA l’approche avec :

\[\Delta \theta_i \approx \Delta \phi_i = BA\]

\(B \in \mathbb{R}^{d \times r}\), \(A \in \mathbb{R}^{r \times d}\) et le rang \(r << min(d, k)\). Ainsi, au lieu d’apprendre \(d \fois k\) paramètres que nous devons maintenant apprendre \((d + k) \times r\) ce qui est facilement beaucoup plus petit compte tenu de l’aspect multiplicatif. En pratique, \(\Delta \theta_i\) est mis à l’échelle par \(\frac{\alpha}{r}\) avant d’être ajouté à \(\theta_i\)qui peut être interprété comme un « taux d’apprentissage » pour la mise à jour LoRA.

LoRA n’augmente pas la latence d’inférence, car une fois le réglage fin effectué, vous pouvez simplement mettre à jour les poids dans \(\Thêta\) en ajoutant leurs respectifs \(\Delta \theta \approx \Delta \phi\). Cela simplifie également le déploiement de plusieurs modèles spécifiques à des tâches sur un seul grand modèle, comme \(|\Delta \Phi|\) est beaucoup plus petit que \(|\Delta \Thêta|\).

Mise en œuvre au flambeau

Maintenant que nous avons une idée du fonctionnement de LoRA, implémentons-le en utilisant Torch pour un problème minimal. Notre plan est le suivant :

  1. Simulez les données d’entraînement à l’aide d’un simple \(y = X \thêta\) modèle. \(\theta \in \mathbb{R}^{1001, 1000}\).
  2. Former un modèle linéaire de rang complet pour estimer \(\thêta\) – ce sera notre modèle « pré-entraîné ».
  3. Simulez une distribution différente en appliquant une transformation dans \(\thêta\).
  4. Entraînez un modèle de rang inférieur à l’aide des poids pré=entraînés.

Commençons par simuler les données d’entraînement :

library(torch)

n <- 10000
d_in <- 1001
d_out <- 1000

thetas <- torch_randn(d_in, d_out)

X <- torch_randn(n, d_in)
y <- torch_matmul(X, thetas)

Nous définissons maintenant notre modèle de base :

model <- nn_linear(d_in, d_out, bias = FALSE)

Nous définissons également une fonction d’entraînement d’un modèle, que nous réutiliserons également plus tard. La fonction effectue la boucle d’entraînement standard dans Torch à l’aide de l’optimiseur Adam. Les poids du modèle sont mis à jour sur place.

train <- function(model, X, y, batch_size = 128, epochs = 100) {
  opt <- optim_adam(model$parameters)

  for (epoch in 1:epochs) {
    for(i in seq_len(n/batch_size)) {
      idx <- sample.int(n, size = batch_size)
      loss <- nnf_mse_loss(model(X[idx,]), y[idx])
      
      with_no_grad({
        opt$zero_grad()
        loss$backward()
        opt$step()  
      })
    }
    
    if (epoch %% 10 == 0) {
      with_no_grad({
        loss <- nnf_mse_loss(model(X), y)
      })
      cat("[", epoch, "] Loss:", loss$item(), "\n")
    }
  }
}

Le modèle est ensuite entraîné :

train(model, 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, nous avons maintenant notre modèle de base pré-entraîné. Supposons que nous ayons des données provenant d’une distribution légèrement différente que nous simulons en utilisant :

thetas2 <- thetas + 1

X2 <- torch_randn(n, d_in)
y2 <- torch_matmul(X2, thetas2)

Si nous appliquons notre modèle de base à cette distribution, nous n’obtenons pas de bonnes performances :

nnf_mse_loss(model(X2), y2)
#> torch_tensor
#> 992.673
#> [ CPUFloatType{} ][ grad_fn = <MseLossBackward0> ]

Nous affinons maintenant notre modèle initial. La répartition des nouvelles données est légèrement différente de la distribution initiale. C’est juste une rotation des points de données, en ajoutant 1 à tous les thêtas. Cela signifie que les mises à jour de poids ne devraient pas être complexes et que nous ne devrions pas avoir besoin d’une mise à jour complète du classement pour obtenir de bons résultats.

Définissons un nouveau module torch qui implémente la logique LoRA :

lora_nn_linear <- nn_module(
  initialize = function(linear, r = 16, alpha = 1) {
    self$linear <- linear
    
    # parameters from the original linear module are 'freezed', so they are not
    # tracked by autograd. They are considered just constants.
    purrr::walk(self$linear$parameters, \(x) x$requires_grad_(FALSE))
    
    # the low rank parameters that will be trained
    self$A <- nn_parameter(torch_randn(linear$in_features, r))
    self$B <- nn_parameter(torch_zeros(r, linear$out_feature))
    
    # the scaling constant
    self$scaling <- alpha / r
  },
  forward = function(x) {
    # the modified forward, that just adds the result from the base model
    # and ABx.
    self$linear(x) + torch_matmul(x, torch_matmul(self$A, self$B)*self$scaling)
  }
)

Nous initialisons maintenant le modèle LoRA. Nous utiliserons \(r = 1\), ce qui signifie que A et B ne seront que des vecteurs. Le modèle de base comporte 1 001 x 1 000 paramètres entraînables. Le modèle LoRA que nous allons affiner n’a que (1001 + 1000), ce qui en fait 1/500 des paramètres du modèle de base.

lora <- lora_nn_linear(model, r = 1)

Entraîneons maintenant le modèle lora sur la nouvelle distribution :

train(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 

Si nous regardons \(\Delta \thêta\) nous verrons une matrice pleine de 1, la transformation exacte que nous avons appliquée aux poids :

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 = <SliceBackward0> ]

Pour éviter la latence d’inférence supplémentaire liée au calcul séparé des deltas, nous pourrions modifier le modèle d’origine en ajoutant les deltas estimés à ses paramètres. Nous utilisons le add_ méthode pour modifier le poids en place.

with_no_grad({
  model$weight$add_(delta_theta$t())  
})

Désormais, l’application du modèle de base aux données de la nouvelle distribution donne de bonnes performances, nous pouvons donc dire que le modèle est adapté à la nouvelle tâche.

nnf_mse_loss(model(X2), y2)
#> torch_tensor
#> 0.00130013
#> [ CPUFloatType{} ]

Final

Maintenant que nous avons appris comment LoRA fonctionne pour cet exemple simple, nous pouvons réfléchir à la manière dont elle pourrait fonctionner sur de grands modèles pré-entraînés.

Il s’avère que les modèles Transformers sont pour la plupart une organisation intelligente de ces multiplications matricielles, et appliquer LoRA uniquement à ces couches est suffisant pour réduire considérablement le coût de réglage fin tout en obtenant de bonnes performances. Vous pouvez voir les expériences dans l’article LoRA.

Bien entendu, l’idée de LoRA est suffisamment simple pour pouvoir être appliquée non seulement aux couches linéaires. Vous pouvez l’appliquer aux convolutions, en incorporant des calques et à tout autre calque.

Image de Hu et al sur le Papier LoRA

Related Articles

LAISSER UN COMMENTAIRE

S'il vous plaît entrez votre commentaire!
S'il vous plaît entrez votre nom ici

Latest Articles