Preferencias y el modelo de recompensa

Con el ajuste supervisadoBloque 2 · Lección 3Ajuste supervisado: máxima verosimilitud sobre un formatoSeguir entrenando un modelo preentrenado sobre pares de instrucción y respuesta es la entropía cruzada de siempre, cobrada sólo en la respuesta: máxima verosimilitud sobre el conjunto de pares. La lección deriva su óptimo, que copia el conjunto, y su gradiente, que es el del preentrenamiento con las filas del prompt a cero. Después ajusta el mini-GPT con diez preguntas: aprende el formato, lee la pregunta, inventa el contenido y olvida su novela., el mini-GPT contesta en lugar de continuar el documento, y la forma de obedecer queda resuelta. El contenido no: un modelo ajustado llega hasta donde llegan los pares con que se ajustó, ni un punto más arriba. Mejorar lo que responde pide otro tipo de dato, y esta lección trabaja con comparaciones: un par de respuestas a la misma pregunta y la que alguien prefiere de las dos.

Los anotadores de InstructGPT hacían eso a escala. Por cada instrucción ordenaban de cuatro a nueve respuestas del modelo, y aun así dos de ellos coincidían en una comparación algo menos de tres veces de cada cuatro. Esta lección convierte datos así, ruidosos y sin una escala común, en un número por respuesta, la recompensa, y entrena una red que lo calcula. Con ese número, el bloque tendrá por fin algo que maximizar.

Pedir notas parece más directo, y tiene un defecto que ninguna cantidad de notas arregla. Una anotadora generosa le pone un 9 a una explicación de por qué el cielo es azul y un anotador severo le pone un 6 a otra; los dos podrían preferir la segunda. Cada nota se lee en la escala de quien la pone, y esa escala se mueve también a lo largo de una tarde de trabajo. Si una anotadora puntúa unas respuestas y otro anotador otras, juntar sus notas junta sus escalas.

Una comparación tira la escala y se queda con el orden. Entre dos explicaciones, la generosa y el severo dicen cuál prefieren, y lo generosos que sean no entra en la respuesta. Lo que se pierde es el tamaño de la diferencia, y se recupera contando: si tres de cada cuatro anotadores prefieren una, la ventaja es pequeña; si la prefieren todos, grande.

El ajedrez resolvió así el mismo problema. Nadie le pone nota a un jugador: juega partidas, gana o pierde, y su Elo sale de los resultados. El Elo es un número por jugador, y la probabilidad de que uno gane a otro depende sólo de la diferencia entre los dos: con 400 puntos más que su rival, se espera que un jugador gane diez partidas por cada una que pierda. Chatbot Arena clasifica hoy los modelos de lenguaje igual, con millones de votos entre dos respuestas anónimas. Con respuestas en lugar de jugadores, el número es la recompensa.

Una recompensa que reproduce las comparaciones

Fijemos los datos. Una comparación es una terna $(x, y_w, y_l)$: una instrucción, la respuesta que un anotador prefirió y la que rechazó, y se escribe $y_w \succ y_l$. Los subíndices vienen del inglés, win y lose. Un anotador (labeler) es la persona que compara, y su elección tiene algo de azar: los anotadores no siempre coinciden, y uno solo tampoco consigo mismo de un día a otro.

El modelo de Bradley–Terry (1952) le da a cada respuesta un número, su recompensa $r(x, y)$, y dice cómo salen las comparaciones. Para dos respuestas $y_A$ e $y_B$ a la misma instrucción, antes de que nadie elija, la probabilidad de que el anotador prefiera $y_A$ es la parte de $e^{r(x, y_A)}$ en la suma de las dos. Dividiendo arriba y abajo por ese término,

Es el softmax de las dos recompensas, el 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., y con dos entradas el softmax es una sigmoide de la diferencia. Llamemos margen a esa diferencia, $\text{margen} = r(x, y_A) - r(x, y_B)$. Despejando, $P/(1 - P) = e^{\text{margen}}$: el margen es el logaritmo de la ventaja. Uno de $\ln 3 \approx 1.1$ quiere decir que tres de cada cuatro anotadores prefieren $y_A$; uno de $\ln 9 \approx 2.2$, nueve de cada diez. El Elo es esta misma curva con otra unidad, y sus 400 puntos son un margen de $\ln 10$.

Un número por respuesta es una afirmación fuerte, y tiene consecuencias que se pueden poner a prueba. Los márgenes se suman a lo largo de una cadena: si $y_A$ gana a $y_B$ tres a uno y $y_B$ a $y_C$ tres a uno, el modelo dice que $y_A$ gana a $y_C$ nueve a uno, sin que nadie las haya comparado. Y no admite ciclos. Si los anotadores prefieren $y_A$ a $y_B$, $y_B$ a $y_C$ e $y_C$ a $y_A$, ninguna recompensa lo reproduce, y con las tres ventajas iguales el mejor ajuste es un empate. Es la concesión de Bradley–Terry: supone que detrás de cada elección hay un único eje de «calidad».

Lo que las comparaciones no fijan

Sumemos a todas las recompensas de una instrucción un mismo número, que puede ser distinto para cada instrucción, y llamémoslo $\kappa(x)$. Ninguna probabilidad cambia, porque el margen lo resta:

Las comparaciones enfrentan siempre dos respuestas a la misma instrucción, así que fijan los márgenes dentro de cada instrucción y nada más. La recompensa queda definida salvo una constante por instrucción, y eso tiene dos consecuencias. Comparar recompensas de instrucciones distintas no significa nada: que una respuesta sobre la nieve tenga más recompensa que otra sobre el mar no lo ha dicho ninguna comparación. Y quien aprende una recompensa de comparaciones tiene que fijar el nivel por su cuenta, o dejarlo donde caiga. InstructGPT, antes de usar su modelo de recompensa, le ajustaba a mano un sesgo para que las respuestas escritas por sus anotadores tuvieran media cero, y daba esta misma razón.

Lo que las comparaciones sí fijan es la escala. Multiplicar las recompensas por 2 dobla cada margen y cambia cada probabilidad. La unidad de la recompensa no es la de ningún anotador: es la frecuencia con que coinciden. Esta libertad de una constante por instrucción parece un defecto del modelo, y más adelante el bloque la aprovechará.

Aprender la recompensa de las comparaciones

Un modelo de recompensa (reward model) es una red $r_\phi(x, y)$, con sus parámetros $\phi$, que lee una instrucción y una respuesta y devuelve un número. Se entrena como todo en este bloque, maximizando la verosimilitud de sus datos. Llamemos $p_{\text{pref}}$ a la distribución que saca una comparación al azar del conjunto. Lo que el modelo dice de una comparación es la probabilidad de que el anotador eligiera la que eligió, y la pérdida es menos su logaritmo, en media:

Es la entropía cruzada binaria de aquella lección sobre funciones de pérdida, con el margen en el lugar de $z$ y la etiqueta siempre 1, porque cada comparación se escribe con la preferida delante. Su óptimo sale con el argumento de siempre. Si la red puede dar cualquier recompensa, agrupemos las comparaciones de un mismo par de respuestas: cada grupo es una entropía cruzada entre la proporción de veces que ganó cada una y la sigmoide del margen, 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 cuando son iguales. Un margen aprendido es una proporción de acuerdo escrita como logaritmo de una ventaja: tres anotadores de cuatro, $\ln 3$.

El gradiente es una línea de aquella lección y la regla de la cadena. La derivada de la entropía cruzada binaria respecto de $z$ era $\hat{y} - y$; con $\hat{y} = \sigma(\text{margen})$ e $y = 1$, es $\sigma(\text{margen}) - 1 = -\sigma(-\text{margen})$. Así que

Un paso de descenso sube la recompensa de la preferida y baja la de la rechazada, cada una en la dirección de su gradiente. El factor que lo multiplica, $\sigma(r_\phi(x, y_l) - r_\phi(x, y_w))$, es la probabilidad que el modelo le da todavía al orden contrario. Una comparación que ya ordena con un margen de 4 pesa $0.018$; una que tiene al revés por 1, $0.731$. El modelo aprende de lo que aún tiene mal, y lo que tiene bien sigue empujando, cada vez menos.

En la práctica, $r_\phi$ es el modelo ajustado con su última capa cambiada: en lugar de la que da un logit por entrada del vocabulario, una capa lineal que da un solo número. InstructGPT lo hizo con 6 000 millones de parámetros. Aquí lo haremos con el mini-GPT congelado y una capa lineal encima. Sea $\bar{\mathbf{h}}(x, y)$ la media de los vectores $\mathbf{h}$ que el último bloque deja en las posiciones de los tokens de $y$:

con $j$ por igual sobre los tokens de $y$, $\texttt{<|fin|>}$ incluido, y $\phi = (\mathbf{u}, b)$: 65 números. Es la neurona del curso anteriorDeep Learning para NLP: del Perceptrón al TransformerBloque 2 · Lección 1La neurona artificialUna suma ponderada y un umbral bastan para que un solo objeto convierta un vector en una decisión. Sus pesos son la perpendicular de una frontera que parte el espacio en dos, y el número que calcula antes de decidir mide a qué distancia de esa frontera ha caído la entrada. sin activación, con $\mathbf{u}$ para el peso porque $w$ ya marca la preferida.

El gradiente de $r_\phi$ es $\bar{\mathbf{h}}(x, y)$ respecto de $\mathbf{u}$ y 1 respecto de $b$, el mismo 1 en las dos respuestas. En la ecuación del gradiente, lo que las dos tienen en común se va en la resta, y el de $b$ es $1 - 1 = 0$. El sesgo no aprende nunca: es un $\kappa(x)$ como el de antes, el mismo para todas las instrucciones, hecho parámetro.

Un modelo de recompensa sobre el mini-GPT

Las dos celdas corren sobre el mini-GPT, tu mini-GPT del bloque 1 cargado desde el checkpoint con las cuatro filas de la plantilla, y no lo entrenan: lo congelan y entrenan encima los 65 números de $\phi$. He partido del modelo preentrenado y no del ajustado de la lección anterior, que olvidó la novela por aprender diez frases.

El anotador es un programa que sabe de qué color es cada cosa. Hay 36 preguntas como las del ajuste, seis cosas de cada color, y cada una tiene seis respuestas con la forma de los pares, una por color. El anotador prefiere la correcta, y cada pregunta da cinco comparaciones. En 24 de las preguntas, tres de sus comparaciones entran en el entrenamiento y las otras dos se guardan: son los rivales nuevos, respuestas que el modelo nunca ha visto enfrentadas a la correcta. Las 12 preguntas restantes, dos por color, son de cosas que la novela no nombra nunca, y ninguna de sus comparaciones entra. La primera celda calcula $\bar{\mathbf{h}}$ de las 216 respuestas en una sola ida.

Cada una de las 216 respuestas queda en un vector de 64 números, la media de los $\mathbf{h}$ en sus posiciones, y nada más del mini-GPT se usa en adelante. Las comparaciones son 72 de entrenamiento, 48 con rivales nuevos y 60 de preguntas nuevas. La segunda celda entrena $\phi$ con trescientos pasos de descenso por gradiente sobre $\mathcal{L}_{\text{R}}$ y mide el acierto en los tres grupos: la proporción de comparaciones en que la preferida recibe más recompensa. Después enseña las seis recompensas de una pregunta del entrenamiento y de una nueva.

La pérdida arranca en $0.693 = \ln 2$: con $\phi = 0$ todas las recompensas son iguales y cada comparación es una moneda al aire. Baja a $0.161$ en trescientos pasos, y el gradiente respecto de $b$ es cero en todos ellos, como decía la derivación. El modelo acierta 71 de sus 72 comparaciones de entrenamiento y el 85 % de las de rivales nuevos. Nunca vio La nieve es blanca. frente a La nieve es amarilla. y las ordena bien, porque no ha aprendido comparaciones sueltas sino un número por respuesta, y el de la correcta subió con cada comparación que ganó.

En las preguntas nuevas acierta el 52 %, lo que acertaría una moneda. A ¿De qué color es el plátano?, la respuesta El plátano es amarillo. recibe la menor recompensa de las seis. La capa lineal sólo puede pesar lo que $\bar{\mathbf{h}}$ trae, y el mini-GPT ha leído una novela en la que no sale ningún plátano: las comparaciones le enseñaron el color de 24 cosas, una a una. Un modelo de recompensa juzga con lo que su red sabe, y por eso los de verdad parten de un modelo grande. El de InstructGPT acertaba un 70 % de las comparaciones de anotadores que no había visto, cerca del 73 % en que los anotadores coincidían entre sí.

Mira también los niveles. La peor respuesta sobre la nieve tiene más recompensa, 18.1, que la mejor sobre el plátano, 14.2, y eso no significa nada: ninguna comparación ha puesto nunca una frente a la otra. Es la constante por instrucción, que aquí pone $\mathbf{u}^{\top}$ aplicado a lo que las seis respuestas de cada pregunta tienen en común.

Comprueba tu intuición

Cinco preguntas y un desafío: las recompensas de Bradley–Terry a partir de una tabla de votos, sin red.

Un modelo de recompensa da un número a cualquier respuesta, también a las que nadie ha comparado, y eso es lo que hacía falta para mejorar la política: cambiar $\pi_\theta$ para que sus respuestas reciban más. La celda dice también dónde se puede confiar en ese número. Cerca de lo que el modelo aprendió, ordena bien; lejos, adivina, y una política que lo maximice sin freno irá a buscar justo las respuestas en que se equivoca.

La lección siguiente, sobre el aprendizaje por refuerzo con retroalimentación humana (reinforcement learning from human feedback, RLHF), escribe ese freno. Maximiza la recompensa media de las respuestas de la política, menos lo que ésta se aleja del modelo ajustado, y deriva cómo se sube esa media con un gradiente aunque las respuestas salgan de un sorteo.