5. LLM Architecture
Arquitectura LLM
El objetivo de esta quinta fase es muy simple: Desarrollar la arquitectura del LLM completo. Juntar todo, aplicar todas las capas y crear todas las funciones para generar texto o transformar texto a IDs y viceversa.
Esta arquitectura se utilizará tanto para entrenar como para predecir texto después de haber sido entrenada.
Ejemplo de arquitectura LLM de https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/ch04.ipynb:
Se puede observar una representación de alto nivel en:
Entrada (Texto Tokenizado): El proceso comienza con texto tokenizado, que se convierte en representaciones numéricas.
Capa de Embedding de Tokens y Capa de Embedding Posicional: El texto tokenizado pasa a través de una capa de embedding de tokens y una capa de embedding posicional, que captura la posición de los tokens en una secuencia, crítica para entender el orden de las palabras.
Bloques de Transformer: El modelo contiene 12 bloques de transformer, cada uno con múltiples capas. Estos bloques repiten la siguiente secuencia:
Atención Multi-Cabeza enmascarada: Permite que el modelo se enfoque en diferentes partes del texto de entrada a la vez.
Normalización de Capa: Un paso de normalización para estabilizar y mejorar el entrenamiento.
Capa Feed Forward: Responsable de procesar la información de la capa de atención y hacer predicciones sobre el siguiente token.
Capas de Dropout: Estas capas previenen el sobreajuste al eliminar aleatoriamente unidades durante el entrenamiento.
Capa de Salida Final: El modelo produce un tensor de 4x50,257 dimensiones, donde 50,257 representa el tamaño del vocabulario. Cada fila en este tensor corresponde a un vector que el modelo utiliza para predecir la siguiente palabra en la secuencia.
Objetivo: El objetivo es tomar estos embeddings y convertirlos de nuevo en texto. Específicamente, la última fila de la salida se utiliza para generar la siguiente palabra, representada como "forward" en este diagrama.
Representación de Código
Función de Activación GELU
Propósito y Funcionalidad
GELU (Unidad Lineal de Error Gaussiano): Una función de activación que introduce no linealidad en el modelo.
Activación Suave: A diferencia de ReLU, que anula las entradas negativas, GELU mapea suavemente las entradas a salidas, permitiendo valores pequeños y no nulos para entradas negativas.
Definición Matemática:
El objetivo del uso de esta función después de las capas lineales dentro de la capa FeedForward es cambiar los datos lineales a no lineales para permitir que el modelo aprenda relaciones complejas y no lineales.
Red Neuronal FeedForward
Las formas se han añadido como comentarios para entender mejor las formas de las matrices:
Propósito y Funcionalidad
Red FeedForward por Posición: Aplica una red completamente conectada de dos capas a cada posición de manera separada e idéntica.
Detalles de la Capa:
Primera Capa Lineal: Expande la dimensionalidad de
emb_dim
a4 * emb_dim
.Activación GELU: Aplica no linealidad.
Segunda Capa Lineal: Reduce la dimensionalidad de nuevo a
emb_dim
.
Como puedes ver, la red Feed Forward utiliza 3 capas. La primera es una capa lineal que multiplicará las dimensiones por 4 usando pesos lineales (parámetros a entrenar dentro del modelo). Luego, se utiliza la función GELU en todas esas dimensiones para aplicar variaciones no lineales y capturar representaciones más ricas y, finalmente, se utiliza otra capa lineal para volver al tamaño original de las dimensiones.
Mecanismo de Atención Multi-Cabeza
Esto ya fue explicado en una sección anterior.
Propósito y Funcionalidad
Autoatención Multi-Cabeza: Permite que el modelo se enfoque en diferentes posiciones dentro de la secuencia de entrada al codificar un token.
Componentes Clave:
Consultas, Claves, Valores: Proyecciones lineales de la entrada, utilizadas para calcular puntajes de atención.
Cabezas: Múltiples mecanismos de atención que funcionan en paralelo (
num_heads
), cada uno con una dimensión reducida (head_dim
).Puntajes de Atención: Calculados como el producto punto de consultas y claves, escalados y enmascarados.
Enmascaramiento: Se aplica una máscara causal para evitar que el modelo preste atención a tokens futuros (importante para modelos autorregresivos como GPT).
Pesos de Atención: Softmax de los puntajes de atención enmascarados y escalados.
Vector de Contexto: Suma ponderada de los valores, de acuerdo con los pesos de atención.
Proyección de Salida: Capa lineal para combinar las salidas de todas las cabezas.
El objetivo de esta red es encontrar las relaciones entre tokens en el mismo contexto. Además, los tokens se dividen en diferentes cabezas para prevenir el sobreajuste, aunque las relaciones finales encontradas por cabeza se combinan al final de esta red.
Además, durante el entrenamiento se aplica una máscara causal para que los tokens posteriores no se tengan en cuenta al buscar las relaciones específicas con un token y también se aplica algo de dropout para prevenir el sobreajuste.
Normalización de Capa
Propósito y Funcionalidad
Normalización por Capas: Una técnica utilizada para normalizar las entradas a través de las características (dimensiones de incrustación) para cada ejemplo individual en un lote.
Componentes:
eps
: Una constante pequeña (1e-5
) añadida a la varianza para prevenir la división por cero durante la normalización.scale
yshift
: Parámetros aprendibles (nn.Parameter
) que permiten al modelo escalar y desplazar la salida normalizada. Se inicializan en uno y cero, respectivamente.Proceso de Normalización:
Calcular Media (
mean
): Calcula la media de la entradax
a través de la dimensión de incrustación (dim=-1
), manteniendo la dimensión para la difusión (keepdim=True
).Calcular Varianza (
var
): Calcula la varianza dex
a través de la dimensión de incrustación, también manteniendo la dimensión. El parámetrounbiased=False
asegura que la varianza se calcule utilizando el estimador sesgado (dividiendo porN
en lugar deN-1
), lo cual es apropiado al normalizar sobre características en lugar de muestras.Normalizar (
norm_x
): Resta la media dex
y divide por la raíz cuadrada de la varianza máseps
.Escalar y Desplazar: Aplica los parámetros aprendibles
scale
yshift
a la salida normalizada.
El objetivo es asegurar una media de 0 con una varianza de 1 a través de todas las dimensiones del mismo token. El objetivo de esto es estabilizar el entrenamiento de redes neuronales profundas al reducir el cambio interno de covariables, que se refiere al cambio en la distribución de las activaciones de la red debido a la actualización de parámetros durante el entrenamiento.
Bloque Transformer
Las formas se han añadido como comentarios para entender mejor las formas de las matrices:
Propósito y Funcionalidad
Composición de Capas: Combina atención multi-cabeza, red de avance, normalización de capas y conexiones residuales.
Normalización de Capas: Aplicada antes de las capas de atención y avance para un entrenamiento estable.
Conexiones Residuales (Atajos): Agrega la entrada de una capa a su salida para mejorar el flujo de gradientes y permitir el entrenamiento de redes profundas.
Dropout: Aplicado después de las capas de atención y avance para regularización.
Funcionalidad Paso a Paso
Primer Camino Residual (Autoatención):
Entrada (
shortcut
): Guarda la entrada original para la conexión residual.Norma de Capa (
norm1
): Normaliza la entrada.Atención Multi-Cabeza (
att
): Aplica autoatención.Dropout (
drop_shortcut
): Aplica dropout para regularización.Agregar Residual (
x + shortcut
): Combina con la entrada original.
Segundo Camino Residual (FeedForward):
Entrada (
shortcut
): Guarda la entrada actualizada para la siguiente conexión residual.Norma de Capa (
norm2
): Normaliza la entrada.Red de Avance (
ff
): Aplica la transformación de avance.Dropout (
drop_shortcut
): Aplica dropout.Agregar Residual (
x + shortcut
): Combina con la entrada del primer camino residual.
El bloque transformer agrupa todas las redes y aplica alguna normalización y dropouts para mejorar la estabilidad y los resultados del entrenamiento. Nota cómo se realizan los dropouts después del uso de cada red mientras que la normalización se aplica antes.
Además, también utiliza atajos que consisten en agregar la salida de una red con su entrada. Esto ayuda a prevenir el problema del gradiente que se desvanece al asegurarse de que las capas iniciales contribuyan "tanto" como las últimas.
GPTModel
Las formas se han añadido como comentarios para entender mejor las formas de las matrices:
Propósito y Funcionalidad
Capas de Embedding:
Embeddings de Tokens (
tok_emb
): Convierte índices de tokens en embeddings. Como recordatorio, estos son los pesos dados a cada dimensión de cada token en el vocabulario.Embeddings Posicionales (
pos_emb
): Agrega información posicional a los embeddings para capturar el orden de los tokens. Como recordatorio, estos son los pesos dados a los tokens según su posición en el texto.Dropout (
drop_emb
): Aplicado a los embeddings para regularización.Bloques de Transformer (
trf_blocks
): Pila den_layers
bloques de transformer para procesar embeddings.Normalización Final (
final_norm
): Normalización de capa antes de la capa de salida.Capa de Salida (
out_head
): Proyecta los estados ocultos finales al tamaño del vocabulario para producir logits para la predicción.
El objetivo de esta clase es usar todas las otras redes mencionadas para predecir el siguiente token en una secuencia, lo cual es fundamental para tareas como la generación de texto.
Nota cómo usará tantos bloques de transformer como se indique y que cada bloque de transformer utiliza una red de atención multi-cabeza, una red de avance y varias normalizaciones. Así que si se utilizan 12 bloques de transformer, multiplica esto por 12.
Además, se agrega una capa de normalización antes de la salida y se aplica una capa lineal final al final para obtener los resultados con las dimensiones adecuadas. Nota cómo cada vector final tiene el tamaño del vocabulario utilizado. Esto se debe a que está tratando de obtener una probabilidad por cada token posible dentro del vocabulario.
Número de Parámetros a entrenar
Teniendo la estructura de GPT definida, es posible averiguar el número de parámetros a entrenar: