Muestreo: temperatura, top-k y top-p

La lección sobre el entrenamientoBloque 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. dejó el mini-GPT del bloque cargado en el navegador y aprendiendo, y todo lo que se le ha pedido hasta aquí cabe en un verbo: repartir. Lee un prefijo y devuelve 512 logits, y con ellos una distribución para el token siguiente. Escribir con él es repetir cuatro cosas: pasarle el texto, leer esa distribución, quedarse con un token y pegarlo al final. Tres de ellas son mecánicas; la otra, quedarse con un token, es una decisión.

¿Con cuál se queda? Parece un detalle y decide si lo que sale se puede leer. Las dos respuestas que caben en una línea fallan, y por lados opuestos: la favorita en cada posición acaba repitiéndose para siempre, y sortear tal cual, con las probabilidades del modelo, cuela palabras que no existen. Esta lección demuestra por qué, con una cuenta para cada fallo, y construye lo que se usa en su lugar: una temperatura entre los dos extremos y un corte que tira lo que el sorteo no debería tocar. Todo corre sobre el mini-GPT.

Mira una distribución de verdad: la del mini-GPT tras La Nela, el nombre de la protagonista, que reparte la apuesta entre muchas entradas (la favorita, no, no llega al 12 %), y tras La Nela bajó la cabe, que casi no deja dudas. En cada fila, la barra clara es la probabilidad tras la temperatura; la verde intenso, la que de verdad se sortea; lo rayado, lo que un corte tira.

Haz tres cosas. Lleva $\tau$ hasta el cero y súbela hasta 2: las barras crecen y encogen, y ninguna cambia de sitio. Pon top-k con $k = 10$: en el primer contexto tira la mitad de la masa, y en el segundo conserva nueve entradas más, que juntas no llegan al 2 %. Y pon top-p con $p = 0.9$: 44 entradas en el primero, una en el segundo.

La temperatura estira la distribución sin cambiar su orden

Fijemos una posición y llamemos $\mathbf{z} \in \mathbb{R}^{\lvert V \rvert}$ a sus logits: el $\mathbf{z}_t$ de la lección sobre el modelo de lenguaje causalBloque 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., sin el subíndice, porque aquí sólo hay una posición en juego. Para un número $\tau > 0$, la temperatura, definimos

el softmax de $\mathbf{z}/\tau$. (En el curso anterior, $\tau$ era el tokenizador.) Con $\tau = 1$, $\mathbf{q}$ es la distribución del modelo; con cualquier otra es una distribución distinta, fabricada con sus logits, y por eso no se llama $p_\theta$. Muestrear es elegir la entrada siguiente al azar, cada $v$ con probabilidad $q_v$: un sorteo por posición.

Todo lo que hace $\tau$ se lee en un cociente. Entre dos entradas $u$ y $v$ el denominador se va,

y $e^{z_u - z_v}$ es el cociente que da el propio modelo, a $\tau = 1$. La temperatura eleva cada cociente a $1/\tau$. Tras La Nela bajó la cabe, el modelo da a za unas 195 veces lo que da a da; a $\tau = 2$, la raíz, unas 14 veces, y a $\tau = 0.5$, el cuadrado, más de treinta mil. Y elevar a una potencia positiva no cambia qué número es mayor.

El orden de las entradas no depende de $\tau$. La favorita sigue siéndolo a cualquier temperatura, y las filas del explorable no se mueven.

Los dos extremos salen de esa misma cuenta. Llamemos $v^{\star}$ a la favorita, la entrada de logit mayor, y dividamos arriba y abajo entre $e^{z_{v^{\star}}/\tau}$:

Cuando $\tau \to 0$, cada exponente con $z_u < z_{v^{\star}}$ es un número negativo dividido entre algo que tiende a cero, y se va a $-\infty$: su término tiende a $0$. Sólo sobrevive el de $v^{\star}$, que vale $e^{0} = 1$ arriba y abajo, así que

(Si dos entradas empatan en el máximo, se reparten el $1$.) Sortear de ahí ya no es sortear: sale $v^{\star}$ siempre. Quedarse con la favorita en cada posición, la generación voraz (greedy decoding; no la llamamos decodificación porque decodificar ya es pasar de ids a texto), no es otra regla. Es el muestreo a temperatura cero.

En el otro extremo, cuando $\tau \to \infty$ cada exponente tiende a $0$, cada término a $1$, y $q_v \to 1/\lvert V \rvert$: el sorteo uniforme, que ignora todo lo que el modelo aprendió. Entre los dos está la familia entera, con $\tau$ como único mando: por debajo de $1$ la distribución se afila hacia la favorita, por encima se aplana hacia la uniforme.

La generación voraz entra en un ciclo y no sale

La voraz es determinista: elige en función de lo que el mini-GPT lee, que son como mucho los últimos $T_{\text{ctx}} = 64$ tokens, y cada posición pega la elegida al final y, cuando ya hay 64, suelta el más viejo. Lo que leerá después depende sólo de lo que lee ahora, y el texto entero, sólo del prompt.

De ahí sale el argumento. Si en algún momento lo que lee se repite, los mismos 64 tokens que $L$ posiciones antes, elige lo mismo que entonces, y en la posición siguiente otra vez, y así para siempre: el texto se vuelve periódico, con periodo $L$, y en la regla no hay nada que pueda romperlo. Es un ciclo. Y se cae en uno en cuanto la favorita continúa un fragmento que se repite, como no,, durante 64 posiciones seguidas: a partir de ahí, lo que lee es ese fragmento repetido, idéntico cada $L$ posiciones.

Lo que la regla no mira es cuánto. $\arg\max$ sólo ve qué entrada tiene el logit mayor, y una favorita del 90 % y una del 30 % dan la misma elección. Un sorteo sí lo ve: si la entrada que continúa el ciclo tiene probabilidad $0.3$, sale de él siete veces de cada diez, en cada posición. La voraz no tiene puerta. (La búsqueda en haz, beam search, cae en lo mismo, y queda fuera de este curso.)

Sin corte, el sorteo acaba en la cola

El otro extremo falla por el lado contrario. Sortear a $\tau = 1$ sin tocar nada no se atasca, pero reparte masa entre cientos de entradas que por separado son improbables: tras La Nela, las 468 que quedan fuera de las 44 más probables suman casi un 10 %. Ninguna saldrá en un sorteo concreto; alguna saldrá, casi seguro, en treinta.

La regla de la cadena lo dice exacto: la probabilidad de escribir $n$ tokens sin tocar nunca esa cola es el producto de la masa que queda dentro en cada posición, sabido lo anterior. Si cada una dejara fuera un 10 %, para 32 tokens sería

Treinta y dos tokens sin corte traen, casi seguro, alguno de la cola, y lo que se escribe se queda: el modelo lo lee en todas las posiciones siguientes.

La respuesta es tirar la cola antes de sortear. Ordenemos las probabilidades de mayor a menor, $q_{(1)} \ge q_{(2)} \ge \dots \ge q_{(\lvert V \rvert)}$ (el paréntesis del subíndice dice «la $i$-ésima mayor», no «la entrada $i$»), y llamemos masa de las $j$ primeras a su suma, $M_j = \sum_{i=1}^{j} q_{(i)}$. Un corte se queda con las $j$ más probables, pone a cero las demás y renormaliza, es decir, divide lo que queda entre su masa para que vuelva a sumar $1$:

Lo que se sortea es $\tilde{\mathbf{q}}$, y $1 - M_j$ es la masa que el corte tira, la parte rayada del explorable.

Top-p corta a la medida de cada posición

Los dos cortes que se usan sólo difieren en cómo eligen $j$. Top-k fija $j = k$ en todas las posiciones. (La $k$ del curso anterior cortaba el vocabulario; ésta corta una distribución.) Top-p fija la masa: se queda con las más probables hasta reunir al menos $p$, que son

entradas, el núcleo (el nombre original es nucleus sampling). La $p$ de top-p es un número, el umbral; la distribución del modelo sigue siendo $p_\theta$, con su subíndice. Y la entrada que cruza el umbral se queda.

Top-p es, entonces, top-k con una $k$ que cada posición elige según la forma de su distribución, y eso es lo que top-k no puede hacer. Tras La Nela, el núcleo de $0.9$ tiene 44 entradas y top-10 tira la mitad de la masa. Tras La Nela bajó la cabe, el núcleo es za sola y top-p no sortea, mientras top-10 conserva nueve entradas que juntas no llegan al 2 % y de vez en cuando escriben cabeda. Una $k$ fija es pequeña en una posición y grande en la otra.

La temperatura y el corte se combinan, en ese orden: primero $\tau$, después el corte. Como $\tau$ no cambia el orden, las entradas que conserva top-k no dependen de ella. El núcleo sí, porque $\tau$ mueve la masa.

Voraz, sorteo y núcleo, sobre el mini-GPT

Las dos celdas corren sobre el mini-GPT: minigpt.py y el checkpoint del bloque, sin entrenar nada. Su función muestrear(z, temperatura, top_k, top_p, rng) es lo de arriba: con temperatura=0 devuelve la favorita y, si no, divide, aplica el softmax, corta y sortea. generar la llama una vez por token.

La primera mira una sola posición con lupa: los dos contextos del explorable a tres temperaturas, con sus dos favoritas, la masa que guarda top-10 y el núcleo de $0.9$. Después genera 32 tokens de forma voraz tras La Nela y tras —¿Qué, apuntando cuánto le daba el modelo a cada favorita.

Las cuentas de arriba están ahí. A $\tau = 0.5$, za se queda con todo, y a $\tau = 2$ baja a $0.387$, catorce veces lo de da. El núcleo de La Nela pasa de 15 a 44 y a 92 entradas mientras top-10 guarda $0.852$, $0.502$ y $0.252$, y el de cabe salta de una a 133, porque $\tau = 2$ levanta a la vez las 511 entradas de la cola.

Y el ciclo. El texto voraz dice no debía decirle., abre un diálogo y ya no sale de no y la coma. El modelo apostaba $0.30$ por no y $0.24$ por la coma en cada vuelta: nunca estuvo seguro, y un sorteo habría salido siete veces de cada diez. Que siguiera dentro veinte posiciones tiene probabilidad $1.1 \times 10^{-11}$, el producto de las veinte apuestas; para la voraz, ese número no existe.

La línea del coste dice algo más. Esos veinte tokens le cuestan al modelo $1.26$ por token, la media de sus $-\log$, y las páginas de Marianela que no vio le cuestan $3.2$: el ciclo le parece más probable que la novela. Holtzman y sus coautores vieron lo mismo en GPT-2, un modelo miles de veces mayor. El texto que escribe una persona no es el que el modelo encuentra más probable, y maximizar la probabilidad posición a posición lleva a sitios como éste. Por eso se sortea.

La segunda continuación llega al mismo ciclo por otro camino, y en ella la apuesta sí sube con cada vuelta: de $0.18$ a $0.31$ por no, y ahí se queda. Es la realimentación que Holtzman midió en GPT-2, donde cada repetición hace más probable la siguiente hasta rozar el $1$. En el mini-GPT es débil, y el argumento de arriba no la necesita: con $0.18$ la voraz ya estaba atrapada.

La segunda compara tres maneras de elegir, con 32 tokens tras cada uno de tres prompts y dos semillas (una sola para la voraz, que no sortea). Cuenta cuántas de las palabras generadas existen en la novela y cuántas continuaciones repiten algún grupo de cuatro tokens. En un dispositivo lento se detiene antes de que el navegador la corte, y dice dónde.

Los dos fallos, medidos. La voraz escribe casi sólo palabras de la novela, 38 de 39, y sus tres continuaciones repiten, dos de ellas ya en el ciclo de no y la coma. Sortear a $\tau = 1$ no repite ninguna y lo paga: 54 de 70, con estreglando y ficill entre las otras. El núcleo de $0.9$ recupera parte, 59 de 73, sin repetir. El mini-GPT es pequeño y todas las muestras se leen mal; lo que la celda mide es hacia dónde empuja cada regla.

Ahora cambia la última línea de la lista por ("tau = 1.5", dict(temperatura=1.5), [0, 1]): las palabras de la novela bajan a 49 de 70, y la muestra dice lapacuma. Con temperatura=0.5 suben a 78 de 83, y dos de las seis continuaciones vuelven a repetirse: una temperatura baja es una voraz con algo de ruido. Seis continuaciones son pocas para separar dos reglas parecidas; los extremos, en cambio, no dejan duda.

Comprueba tu intuición

Cinco preguntas y un reto: escribir la distribución de la que sortea top-p.

El mini-GPT ya escribe: cada token sale de un sorteo sobre el núcleo de su distribución, a la temperatura que se elija, y la voraz queda como el caso $\tau = 0$, que conviene conocer y no usar. Pero mira lo que hace generar por cada token. Llama a adelante con el texto entero, hasta 64 tokens, se queda con la última fila de logits y tira las demás; en la posición siguiente vuelve a pasar los mismos tokens, más uno.

La máscara de la lección sobre el modelo de lenguaje causalBloque 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. garantiza que, en cada capa, las claves y los valores de esas posiciones no han cambiado: la fila $t$ sólo depende de lo que hay a su izquierda, y ahí no se ha movido nada. Ese trabajo se repite en cada token y crece con el texto. Guardarlo en lugar de rehacerlo, y contar cuánto cuesta entonces cada token, es la lección siguiente, sobre la caché de claves y valores.