Ajuste supervisado: máxima verosimilitud sobre un formato

Con la plantilla de chatBloque 2 · Lección 2La plantilla de chat es un formato entrenadoUn modelo lee una sola secuencia de tokens, y una conversación es una lista de mensajes con autor. La plantilla de chat escribe la una como la otra: cada mensaje entre el token especial de su rol y el de fin, que ningún texto puede producir, y eso la hace inyectiva, cosa que unas marcas de texto no consiguen. Dentro de ella la respuesta tiene un sitio exacto, su probabilidad es un producto y la máscara dice qué tokens se cobran al entrenar. Y un token especial que nadie ha entrenado no significa nada: la lección acota cuánto puede valer y lo mide en el mini-GPT, que no sabe terminar su respuesta., una conversación cabe ya en lo único que un modelo lee, una secuencia de tokens, y la respuesta tiene dentro de ella un principio y un final marcados por tokens que ningún texto puede falsificar. Que el modelo los reconozca es otra cosa. Para el mini-GPT son cuatro entradas más de su vocabulario, con el mismo vector las cuatro, que no ha leído nunca en ningún texto: el token que le cede la palabra no le anuncia nada, y eso sólo lo cambia entrenarlo.

En 2023, un modelo de 65 000 millones de parámetros aprendió a comportarse como un asistente con mil pares de instrucción y respuesta, sin ningún otro entrenamiento detrás. Sus autores lo leyeron así: casi todo lo que el modelo sabía venía del preentrenamiento, y los mil pares le enseñaron sobre todo la forma de decirlo. Esta lección escribe qué le pide a un modelo un entrenamiento sobre pares y, sobre todo, qué deja sin pedir, que es lo que explica aquella cifra. Después lo prueba en el mini-GPT.

Mira los pares con que vas a ajustarlo. Son diez preguntas por el color de algo, cada una con su respuesta de una frase: a ¿De qué color es la nieve? le toca La nieve es blanca., y así las diez. Comparten casi todo: la forma de la pregunta y la de la respuesta, el artículo de la una repetido en la otra, el punto y el token de fin. De un par a otro cambian dos palabras, el sustantivo y el color, y cada paso de entrenamiento empuja a la vez lo común y lo de cada uno.

Lo común se aprende diez veces, y es lo único que una pregunta nueva puede tomar prestado. Lo particular sólo vale para su par: el color de la nieve no dice nada del de la tierra, y la novela con que se preentrenó el mini-GPT tampoco. Diez pares son suficientes para cambiar cómo responde y no para decidir qué responde.

Ajustar es maximizar la verosimilitud de los pares

Fijemos el objeto. Hay un modelo preentrenado y un conjunto de pares $(x, y)$, cada uno una conversación escrita con la plantilla: $x$ es la conversación hasta el token del asistente, incluido, e $y$ la respuesta con su $\texttt{<|fin|>}$, escrita por alguien y que obedece a $x$. El ajuste supervisado (supervised fine-tuning, SFT) es seguir entrenando el modelo sobre ese conjunto, desde los pesos del preentrenamiento. Es supervisado porque cada instrucción llega con la respuesta que se quiere.

Lo que se maximiza es la verosimilitud del conjunto, con la palabra de la lección sobre funciones de pérdida del curso anteriorDeep Learning para NLP: del Perceptrón al TransformerBloque 2 · Lección 5Funciones de pérdida: MSE y entropía cruzadaUna pérdida resume en un número lo mal que lo hace la red, y entre dos candidatas decide su derivada. El error cuadrático medio entrega su corrección más pequeña justo donde la red se equivoca con más seguridad; la entropía cruzada, que sale de leer la salida como probabilidad, deja esa derivada en la diferencia entre lo predicho y la etiqueta.: la probabilidad que el modelo le da entero, el producto de las $\pi_\theta(y \mid x)$ de sus pares. Maximizarla es minimizar menos su logaritmo, una suma, que dividida entre el número de pares es la media de $-\log \pi_\theta(y \mid x)$ sobre un par sacado al azar del conjunto. Llamemos $p_{\text{SFT}}$ a la distribución de ese sorteo. Con el producto de la lección anterior dentro, la pérdida del ajuste es

Es la entropía cruzada del bloque anterior con una sola diferencia: la suma recorre los tokens de $y$, que es lo que hace la máscara. Y tiene la forma exacta del preentrenamiento de la lección sobre predecir y obedecerBloque 2 · Lección 1Predecir no es obedecerPor qué un modelo entrenado a predecir texto no hace lo que se le pide. En el óptimo del preentrenamiento copia lo que su texto hace después de una instrucción, y obedecer es sólo una de esas cosas. La lección lo deriva, lo mide en el mini-GPT y separa lo que el prompt puede arreglar de lo que exige cambiar el texto con que se entrena., con $p_{\text{SFT}}$ en el lugar de $p_{\text{texto}}$. Aquella lección terminó proponiendo cambiar el texto con que se entrena. Esto es cambiarlo.

Su mínimo sale, por tanto, con el argumento de aquella lección. Agrupemos los pares por su instrucción: uno sacado al azar tiene la instrucción $x$ con probabilidad $p_{\text{SFT}}(x)$ y, dada ésa, la respuesta $y$ con probabilidad $p_{\text{SFT}}(y \mid x)$, así que

donde la suma de fuera recorre las instrucciones del conjunto. Cada suma de dentro es la entropía cruzada entre dos distribuciones de respuestas a una misma instrucción, la del conjunto y la del modelo, y la desigualdad de la lección sobre la perplejidadBloque 1 · Lección 6Perplejidad: medir un modelo de lenguajeUn modelo que sólo cuenta bytes paga, sobre las páginas de Marianela que el mini-GPT no vio, la misma pérdida por token que el mini-GPT. La lección explica por qué eso no dice nada: convierte la pérdida en perplejidad, el número de candidatas entre las que duda el modelo, y en bits por byte, la medida que no depende del tokenizador; le pone al lado dos varas de medir, el modelo uniforme y el de unigramas; y demuestra por qué la pérdida sobre el texto de entrenamiento engaña y la del texto reservado no. la hace mínima en un solo sitio: cuando son iguales. Con un modelo capaz de dar a cada instrucción la distribución que quiera, el mínimo de la suma es el de cada sumando,

La igualdad dice tres cosas. Con una sola respuesta por instrucción, como en los diez pares de esta lección, $p_{\text{SFT}}(y \mid x)$ vale 1 en ella, y el óptimo es aprenderse cada par de memoria. Si todas las respuestas del conjunto obedecen, el modelo obedece siempre en sus instrucciones: es el primer factor de la lección sobre predecir y obedecer, resuelto. Y el segundo factor, lo buena que es la respuesta cuando obedece, queda en el del conjunto. Un modelo ajustado contesta tan bien como quien escribió sus pares, con sus mismos errores y en la misma proporción.

El gradiente es el del preentrenamiento, con filas a cero

Del gradiente sólo cambia el principio. Cada factor de $\pi_\theta(y \mid x)$ es un softmax, el de los logits de la posición que ha leído $x$ y los $j - 1$ primeros tokens de $y$. Escribámoslos $\mathbf{z}(x, y_{<j})$, con lo que leyeron por argumento, como $\mathbf{z}_t(x_{\le t})$ en la primera lección del bloque anteriorBloque 1 · Lección 1Una sola columna: el modelo de lenguaje causalEl curso anterior eligió una columna del Transformer, le puso la máscara en todos los bloques y le pidió el token siguiente, posición a posición. Esta lección hace la pregunta que aquél no hizo, qué probabilidad le da ese modelo a un texto entero, y de la respuesta sale todo lo demás: la regla de la cadena convierte esa probabilidad, sin perder nada, en un producto de condicionales; la máscara causal es lo que permite calcular las T de una ventana en una sola pasada; y la pérdida es el logaritmo de ese producto, cambiado de signo y repartido por token. Sin código a propósito: la cuenta se hace a mano sobre una ventana de cuatro palabras, y sale un dieciseisavo por los dos caminos.. La pérdida de un par es

y cada término depende de una sola fila de logits. Su gradiente respecto de ella es el que dejó la lección sobre funciones de pérdida: el softmax menos el one-hot $\mathbf{o}_{y_j}$ del token que de verdad vino,

Ésas son las filas que apuestan por un token de $y$. La primera, $\mathbf{z}(x, y_{<1})$, lee $x$ entera, hasta $\texttt{<|asistente|>}$: el token del asistente no se cobra, pero la fila que lo lee sí. Las demás filas leen sólo un trozo de $x$ y apuestan por el token siguiente de $x$. Se calculan en la misma ida y no aparecen en ningún término, así que su gradiente es cero.

Esas dos líneas son todo el cambio. Lo que queda es la vuelta de la lección sobre backpropagationDeep Learning para NLP: del Perceptrón al TransformerBloque 2 · Lección 8Backpropagation: la derivación completaDerivar peso a peso obliga a recorrer la red una vez por parámetro, y esos recorridos comparten casi todo. Backpropagation calcula en su lugar un solo número por preactivación (el error de esa neurona), con una única pasada del final al principio, y deja que cada peso lea su derivada como el producto de ese error por lo que tenía delante., que es lineal en esos vectores y lleva cada uno a los pesos que calcularon su fila; un vector cero no lleva nada. En el preentrenamiento todas las filas tenían el suyo. En el ajuste, sólo las de la respuesta.

Que una fila no tenga gradiente no quiere decir que su posición no aprenda. Cada fila de $y$ lee, por la atención, las claves y los valores de todas las posiciones anteriores, las de $x$ incluidas, y la vuelta pasa por ellos hasta los pesos y las filas de $\mathbf{E}$ de los tokens de la pregunta. El ajuste le enseña al modelo a leer la instrucción, en la medida en que eso le ayuda a escribir la respuesta, y nunca a escribirla. Así se separan las filas de rol, que empezaron siendo la misma media: cada una aprende de los sitios en que se lee.

Por el otro lado de la tabla atada, el de la salida, cada fila cobrada empuja hacia abajo el logit de toda entrada que no es $y_j$, en proporción a lo que apostó por ella. $\texttt{<|fin|>}$ sube en la fila que cierra cada respuesta y baja en las demás, y los tokens de rol, que no son $y_j$ en ninguna, sólo bajan: el modelo aprende a no escribirlos nunca, que es lo que la plantilla necesita.

Lo que la pérdida no dice

La suma de $\mathcal{L}_{\text{SFT}}$ recorre las instrucciones del conjunto y ninguna más. Para una instrucción que no está en él no hay término, y la pérdida no dice nada de lo que $\pi_\theta$ haga ahí: ni que conteste bien, ni que conteste. Lo que pase lo deciden los pesos, que son los mismos para todas las instrucciones, y el sitio del que partieron.

Eso corta en dos direcciones. Hacia delante, una instrucción nueva recibe lo que los pares tienen en común y lo que el preentrenamiento dejó: con mil pares variados y un modelo que lo ha leído casi todo, una respuesta útil; con diez pares que sólo se distinguen en dos palabras y un modelo que ha leído una novela, el formato de los diez. Hacia atrás, el texto del preentrenamiento tampoco tiene término. Nada en $\mathcal{L}_{\text{SFT}}$ le pide al modelo que lo siga prediciendo, y cada paso de Adam mueve casi todos los pesos casi $\eta$, tengan o no que ver con los pares, como vio la lección sobre entrenar el mini-GPTBloque 1 · Lección 3Entrenar un mini-GPT en NumPyEl checkpoint del bloque es el resultado de 2 000 pasos de un bucle que esta lección escribe entero. Cada paso sigue la pérdida media de unas pocas ventanas del corpus sacadas al azar, y la lección deriva por qué eso apunta, de media, adonde apuntaría el corpus entero, y cuánto ruido cuesta ahorrarse el resto. El paso es AdamW, con una tasa de aprendizaje que se calienta y luego decae. Con todo ello el mini-GPT sigue entrenando en el navegador, sobre las últimas páginas de Marianela, que no había visto..

Perder lo que un modelo sabía por ajustarlo a otra cosa se llama olvido catastrófico (catastrophic forgetting). En el ajuste supervisado se frena con muchos pares variados, un $\eta$ pequeño y pocas épocas. El bloque vuelve a él más adelante, cuando haga falta mejorar un modelo sin dejar que se aleje del de partida.

El mini-GPT, ajustado con diez pares

Las tres celdas corren sobre el mini-GPT, tu mini-GPT del bloque 1 cargado desde el checkpoint con las cuatro filas de la plantilla, y esta vez lo entrenan, en el navegador. La primera escribe los diez pares con la plantilla de la lección anterior, los apila en un batch y mide el modelo antes de tocarlo.

Las diez conversaciones ocupan 277 posiciones de un batch de $10 \times 30$, y se cobran 116, las de las respuestas con su <|fin|>, como enseña la primera fila. Esa fila es la ecuación del gradiente en código, porque peso es la máscara corrida una posición: la fila $t$ apuesta por el token $t + 1$. Antes del ajuste, el mini-GPT paga $5.25$ nats por cada token de las respuestas, un nat menos que los $\ln 516 = 6.25$ de repartir por igual, y detrás de <|asistente|> escribe novela.

perdida divide entre la suma de peso, los 116 tokens cobrados, y no entre los diez pares. Con el conjunto entero en cada paso, eso es una constante: no mueve el mínimo, y Adam, que divide cada coordenada del gradiente por su propia escala, ni siquiera la nota.

La segunda da doce pasos de Adam sobre el batch entero. Reanuda un modelo ya entrenado con un Adam nuevo, así que empieza con cinco pasos de calentamiento hasta $\eta = 5 \times 10^{-3}$.

La pérdida en los pares baja de $5.25$ a $0.33$ nats por token en doce pasos, y seguiría bajando hacia el cero del óptimo, que es aprenderse los diez. La tercera celda le hace seis preguntas que no están entre ellos, mide la novela antes y después del ajuste y cuenta los pesos que se han movido. Al final prueba, sin plantilla, el marco Pregunta: … Respuesta: de la lección sobre predecir y obedecer.

Las seis respuestas son frases con la forma de los pares, y las seis terminan en <|fin|>: el mini-GPT ha aprendido a acabar. En las cuatro preguntas de colores, el artículo de la respuesta es el de la pregunta: la tierra recibe La leche es blanca., y el vino, El cielo es azul.. Ninguna fila cobrada apuesta por un token de la pregunta, y aun así la respuesta depende de ella. Es el gradiente que entró por las filas de la respuesta y llegó a lo que esas filas leen.

El contenido, en cambio, sale de los pares tal cual, sin relación con lo que se pregunta. La Nela, la protagonista de la novela que el mini-GPT leyó entera, es La nieve es blanca., y un número es El cielo es azul.. Diez pares han enseñado el formato y nada más.

Tampoco ha salido gratis. La novela le cuesta ahora $6.69$ nats por token, frente a $2.96$ antes del ajuste: más que los $6.25$ de repartir por igual entre las 516 entradas. Doce pasos han movido más de una milésima $131\,833$ de sus $136\,704$ pesos. Y el marco de la lección sobre predecir y obedecer, al que el bloque prometió volver, recibe ahora una frase de colores con su <|fin|>. No es que el mini-GPT haya aprendido el marco: es que ya no sabe escribir otra cosa.

Comprueba tu intuición

Cuatro preguntas y un desafío: la pérdida del ajuste y su gradiente, para una conversación.

El ajuste supervisado resuelve el primer factor de la lección sobre predecir y obedecer: en las instrucciones de su conjunto, el modelo obedece. El segundo lo deja donde lo tenía el conjunto, porque su óptimo es copiar a quien escribió los pares, aciertos y errores en la misma proporción. Y escribir buenos pares es caro: cada uno es una respuesta entera, redactada por alguien que sepa hacerlo.

Juzgar es más barato que escribir. Entre dos respuestas a una misma instrucción, casi cualquiera sabe decir cuál es mejor sin saber escribir ninguna de las dos, y ése es un dato que la pérdida de esta lección no sabe usar. La lección siguiente, sobre preferencias y el modelo de recompensa, lo convierte en un número: entrena un modelo que puntúa respuestas a partir de comparaciones, y explica por qué comparaciones y no notas.