Optimizador NoamΒΆ

Vaswani et al. (2017) proponen un mecanismo de optimizaciΓ³n que, afirman, mejora el proceso de aprendizaje de los parΓ‘metros durante el entrenamiento del modelo de transformador. Este optimizador estΓ‘ basado en el optimizador Adam, que es un estΓ‘ndar en el entrenamiento de redes neuronales. Este optimizador, entonces, busca ajustar la tasa de aprendizaje para que Γ©ste se adapte a la convergencia hacia el mΓ­nimo de la funciΓ³n objetivo. Para entrar en detalle con el optimizador Noam, veamos antes los detalles de Adam.

AdamΒΆ

El optimizador Adam (Kingma & Ba, 2014) es un optimizador que a su vez se basa en otros dos optimizadores: 1) el optimizador adagrad (del inglΓ©s adaptative gradient (de donde viene "ada"; y 2) del mΓ©todo del momentum (de donde viene la "m"). Es decir, incorpora un gradiente adaptativo y un momentum a la regla de actualizaciΓ³n de los parΓ‘metros. Recordemos que el mΓ©todo de gradiente descendiente actualiza los parΓ‘metros $\theta$ (que en este caso son nuestros pesos en las matrices y los bias), dada una funciΓ³n objetivo $R(\theta)$, a partir de la siguiente regla:

$$\theta \leftarrow \theta - \eta \nabla R(\theta)$$

En el optimizador Adam, se proponen dos nuevos parΓ‘metros $\hat{\mu}$ y $\hat{\nu}$ que se estiman, al mismo, tiempo a partir de los valores de $\mu$ y $\nu$. Estos dos valores se calculan como:

$$\mu \leftarrow \beta_1 \mu + (1-\beta_1) \nabla R(\theta)$$$$\nu \leftarrow \beta_2 \nu + (1-\beta_2)|\nabla R(\theta)|^2$$

Donde $\beta_1, \beta_2 \in [0,1]$ son hiperparΓ‘metros, generalmente configurados como $\beta_1= 0.9$ y $\beta_2 = 0.98$. Como se puede observar, ambas ecuaciones son sumas convexas de el valor en el paso anterior y de un factor que tiene que ver con el gradiente de la funciΓ³n objetivo.

Como seΓ±alΓ‘bamos, Adam tiene una aplicaciΓ³n del momentum, que corresponde al valor de $\mu$. En este caso, los valores $\mu$ van acumulando los gradientes en los pasos anteriores, de tal forma que tenemos la suma de un momentum dado como $\nabla R(\theta) + \mu$, interpretando $\mu$ como el momentum. En este caso, $\beta_1$ simplemente pondera ambos factores de manera convexa. Por su parte, $\nu$ sirve para adaptar el gradiente; es decir, lo usaremos para dividir al gradiente y que este se adapte segΓΊn el tamaΓ±o del gradiente. Por tanto, este valor toma el gradiente al cuadrado que va a escalar la tasa de aprendizaje.

Antes de aplicar estos valores al optimizador, se realiza una normalizaciΓ³n de Γ©stos de la siguiente forma:

$$\hat{\mu} = \frac{\mu}{1-\beta_1}$$$$\hat{\nu} = \frac{\nu}{1-\beta_2}$$

De aquΓ­ se puede observar que si $\beta_1 = 0 = \beta_2$ el optimizador sΓ³lo adaptarΓ‘ el gradiente con respecto a $|\nabla R(\theta)|^2$. Haciendo los valores de $\beta_1$ y $\beta_2$ mΓ‘s cercanos a 1, lo que harΓ‘ es que la influencia estarΓ‘ mayormente en los gradientes de los pasos previos, pero la influencia de estos serΓ‘ pequeΓ±a (de orden $1-\beta$). Finalmente, el optimizador Adam utiliza la regla de actualizaciΓ³n:

$$\theta \leftarrow \theta - \frac{\eta}{\hat{\nu} + \epsilon} \hat{\mu}$$

Como seΓ±alΓ‘bamos $\hat{\nu}$ se utiliza para adapatar la tasa de aprendizaje, dividiendo entre este factor mΓ‘s un valor $\epsilon$ que evita divisiones entre 0. El factor $\hat{\mu}$ reemplaza al gradiente, por el valor del gradiente mΓ‘s el momento. El otpimizador Adam ha probado tener mejor desempeΓ±o que otros otpimizadores, ademΓ‘s de que puede simular a Adagrad cuando las betas son 0.

Detalles del optimizador NoamΒΆ

El optimizador Noam estΓ‘ basado en Adam, y lo incorpora dentro de este mismo optimizador. Sin embargo, hace una adaptaciΓ³n de la tasa de apendizaje de una forma distinta a Adam.

$$\eta = \frac{1}{\sqrt{d}} \min\{ \frac{1}{\sqrt{t}}, \frac{t}{\sqrt{\omega^3}}\}$$

En este caso, la adaptatibilidad de la tasa de aprendizaje depende de $t$ que refiere al paso del entrenamiento, de la dimensiΓ³n del modelo $d$ y finalmente del hiperparΓ‘metro $\omega$ que se conoce como warmup. Vaswani et al. (2017) proponen que $\omega = 4000$, esto lo determinan de manera empΓ­rica.

Para observar el comportamiento de la tasa de aprendizaje con respecto al optmizador Noam, podemos graficar los valores para un nΓΊmero de pasos creciente, tomamos la dimensiΓ³n del modelo $d=256$.

No description has been provided for this image

En este caso, las primeras iteraciones tienen un crecimiento casi lineal, mientras que despuΓ©s decrece aunque de manera lenta conforme se avanza en cada paso. De esta forma, la actualizaciΓ³n de los pesos con optimizador Noam incorpora el momentum de Adam, pero modifica la adaptatibilidad de la tasa de aprendizaje.

Para implementar el optimizador Noam, generamos una clase que se comporte como un optimizador de PyTorch. Definimos el valor de warump, asΓ­ como otros hiperparΓ‘metros. Utilizamos dentro de la clase al optimizador Adam. Definimos tambiΓ©n una funciΓ³n rate que estima la tasa de aprendizaje con respecto a la formulaciΓ³n arriba expuesta. Asimismo, definimos las funciones de step y zero_grad propias de un optimizador en PyTorch.

InΒ [1]:
import torch
import torch.nn as nn
import numpy as np

class NoamOptimizer:
    def __init__(self, parameters, d_model, warmup=40000, init_lr=0, eps=1e-9, decay=0.01):
        #optimizador
        self.optimizer = torch.optim.Adam(parameters, lr=init_lr, betas=(0.9, 0.98), eps=eps, weight_decay=decay)
        self._step = 0
        self.warmup = warmup
        self.model_size = d_model
        self._rate = 0
        
    def step(self):
        self._step += 1
        rate = self.rate()
        for p in self.optimizer.param_groups:
            p['lr'] = rate
        self._rate = rate
        self.optimizer.step()
        
    def rate(self):
        step = self._step
        lr_step = self.model_size**(-0.5) * min(step**(-0.5), step*self.warmup**(-1.5))
        return lr_step

    def zero_grad(self):
        self.optimizer.zero_grad()

Prueba del optimizadorΒΆ

Para probar el optimizador definimos el mismo problema de un modelo del lenguaje por medio de auto-atenciΓ³n con enmascaramiento en las palabras subsecuentes. De igual forma que en los casos pasados, definimos un dataset de entrenamiento y determinamos el modelo con los hiperparΓ‘metros necesarios:

InΒ [3]:
import pandas as pd
import matplotlib.pyplot as plt
from seaborn import heatmap as hm
from tqdm import tqdm
from transformers import *
from transformers import *
import copy

class MultiHeadMaskAttention(nn.Module):
    def __init__(self, in_size, d_model, hidden=128, heads=3, dropout=0.3):
        super(MultiHeadMaskAttention, self).__init__()
        self.d_model = d_model
        #Embedding y codificaciΓ³n posicional
        self.enc = Encoding(in_size, d_model)
        #Auto-atenciΓ³n enmascarada
        self.att = nn.ModuleList([copy.deepcopy(MaskAttention(d_model)) for _ in range(heads)])
        #Capa linear para multi cabezas
        self.lin = nn.Linear(heads*d_model, d_model, bias=True)
        #NormalizaciΓ³n
        self.norm = LayerNorm(d_model)
        #Red feedforward
        self.ffw = nn.Sequential(nn.Linear(d_model, hidden), nn.ReLU(),
                                nn.Linear(hidden, d_model))
        self.drop1 = nn.Dropout(p=dropout)
        self.drop2 = nn.Dropout(p=dropout)
        self.drop3 = nn.Dropout(p=dropout)
    
    def forward(self, x):
        x_e = self.enc(x)
        x_e = self.drop1(x_e)
        head_att = [head(x_e) for head in self.att]
        self.att_weights = [head[1] for head in head_att]
        heads = [head[0] for head in head_att]
        multi_heads = torch.cat(heads, dim=-1)
        h = self.lin(multi_heads)
        h_norm = x_e + self.norm(h)
        h_norm = self.drop2(h_norm)
        out = self.ffw(h)
        
        return self.drop3(h_norm + self.norm(out))
InΒ [4]:
#Corpus a utilizar
corpus = ['el perro come un hueso', 'un muchacho jugaba', 'el muchacho saltaba la cuerda',
          'un perro come croquetas', 'el perro come', 'el gato come croquetas', 
          'un gato come', 'un muchacho jugaba con la cuerda', 'el muchacho jugaba con la cuerda']
corpus = [w.split() for w in corpus]
#CreaciΓ³n del vocabulario
voc = vocab()
voc['[bos]'] = 0
voc['[eos]'] = 1
#IndexaciΓ³n de cadenas
sents = list(index(corpus, voc))

#Pares de entrenamiento
x = [torch.cat((torch.tensor([voc['[bos]']]),s), axis=0) for s in sents]
y = [torch.cat((s, torch.tensor([voc['[eos]']])), axis=0) for s in sents]
InΒ [5]:
len_voc = len(voc)
model = nn.Sequential(MultiHeadMaskAttention(len_voc, 128, heads=5), 
                      nn.Linear(128,len_voc), nn.Softmax(1)) 

#Carga del modelo
#model_heads.load_state_dict(torch.load('noam_model.model'))
#model_heads.eval()

Ahora procedemos a entrenar el modelo. En este caso, el procedimiento es similar a los casos anteriores y a cualquier entrenamiento de PyTorch, pero utilizamos el optimizador Noam. A este le damos los parΓ‘metros del modelo que debe optimizar, asΓ­ como la dimensiΓ³n del modelo que utiliza para la adaptaciΓ³n de la tasa de aprendizaje. En este caso, no usamos regularizaciΓ³n por weigth decay, por lo que le asignamos un valor de 0.

InΒ [6]:
criterion = nn.CrossEntropyLoss()
optimizer = NoamOptimizer(model.parameters(), model[0].d_model, decay=0)
epochs = range(300)

#Entrenamiento
model.train()
for t in tqdm(epochs):
    for i in torch.randperm(len(x)):
        prediction = model(x[i])
        optimizer.zero_grad()
        loss_value = criterion(prediction, y[i])
        loss_value.backward()
        optimizer.step()
/home/cienciasia/anaconda3/lib/python3.11/site-packages/torch/cuda/__init__.py:619: UserWarning: Can't initialize NVML
  warnings.warn("Can't initialize NVML")
100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 300/300 [00:43<00:00,  6.96it/s]

ExploraciΓ³n del modeloΒΆ

Al igual que en los casos anteriores, podemos ver quΓ© es lo que aprende el modelo a partir de las probabilidades que se obtienen y de sus matrices de atenciΓ³n. En este caso, esperarΓ­amos que con el optimizador Noam los modelos aprendidos fueran mΓ‘s adecuados, aunque hay que seΓ±alar que todavΓ­a falta una parte importante de los grandes modelos del lenguaje que son las grandes cantidades de datos, pues hasta ahora sΓ³lo estamos probando los modelos con datos de juguete.

InΒ [7]:
devoc = {i:t for t,i in voc.items()}
def result(text, model):
    #FunciΓ³n para predecir la siguiente palabra dado el contexto
    tokens = text.split()
    x = torch.tensor([voc[t] for t in tokens])
    pred = model(x)
    max_token = pred.argmax(axis=1).detach().numpy()
    
    return pred.detach().numpy(), ' '.join([devoc[i] for i in max_token])

p, pred_text = result('[bos]', model)
print('Palabra siguiente con mayor prob: {}'.format(pred_text))

#VisualizaciΓ³n de probabilidades mΓ‘s altas
args = np.argsort(p[-1])[::-1]
probs = np.sort(p[-1])[::-1]
pd.DataFrame(data=probs, columns=['prob. tΓ³ken'], index=[devoc[j] for j in args]).plot.bar()
plt.show()
Palabra siguiente con mayor prob: el
No description has been provided for this image
InΒ [8]:
text = '[bos] un gato come'
result(text, model)

for i, att_w in enumerate(model[0].att_weights):
    hm(att_w.detach().numpy(), xticklabels=text.split(), yticklabels=text.split(), vmin=0, vmax=1)
    plt.title('AtenciΓ³n en cabeza %i' %i)
    plt.show()
No description has been provided for this image
No description has been provided for this image
No description has been provided for this image
No description has been provided for this image
No description has been provided for this image

ReferenciasΒΆ

Kingma, D. P., & Ba, J. (2014). Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., ... & Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems, 30.


Principal