Leyes de escala: Kaplan y Chinchilla

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. midió al mini-GPT y le encontró una grieta: sobre las páginas de Marianela que no ha leído paga unos 3.2 nats por token, bastante más que sobre las que le sirvieron para entrenar. Detrás de esa medida hay una factura que el bloque todavía no ha sumado. El checkpoint es el paso 1 750 de un entrenamiento de 32 ventanas de 64 tokens por paso: $3\,584\,000$ tokens leídos por $136\,448$ parámetros, cuando la parte de entrenamiento de la novela sólo tiene $125\,789$. El mini-GPT leyó cada uno de ellos, de media, 28 veces y media. Con ese mismo gasto se podía haber pagado otra cosa, un mini-GPT mayor que leyera menos o uno menor que leyera más, y falta saber qué habría salido mejor.

La misma pregunta, con presupuestos once órdenes de magnitud mayores, decide cómo se entrena un modelo grande. En 2020, Kaplan y sus coautores, en OpenAI, encontraron que multiplicar por diez los parámetros, los tokens leídos o el cálculo multiplica la pérdida, cada vez, por un mismo factor, y concluyeron que un presupuesto mayor debía ir sobre todo a parámetros. GPT-3 y Gopher se entrenaron en esa línea: 175 000 y 280 000 millones de parámetros, y 300 000 millones de tokens cada uno, menos de dos por parámetro. En 2022, Hoffmann y sus coautores, en DeepMind, repitieron la medida con más de 400 modelos y llegaron a otro reparto: a partes iguales, unos 20 tokens por parámetro. Para comprobarlo entrenaron Chinchilla, 70 000 millones de parámetros y 1.4 billones (millones de millones) de tokens, con el cálculo que había costado Gopher, y ganó a Gopher y a GPT-3. Esta lección deriva de dónde sale un reparto así y le hace la cuenta al mini-GPT.

Antes de las fórmulas, las formas. La primera vista del explorable pone la pérdida frente a los parámetros, los tokens o el cálculo, con los dos ejes logarítmicos y los ajustes de los dos artículos. La de Kaplan es una recta: mueve el punto y verás el mismo 0.839 por cada multiplicación por diez de los parámetros, estés donde estés. La de Chinchilla se curva hacia un suelo, y cada multiplicación compra menos que la anterior.

La segunda vista fija el presupuesto. Con el cálculo fijo, un modelo mayor lee menos tokens, y la pérdida frente al tamaño es un valle: a la izquierda, modelos pequeños que leen mucho; a la derecha, grandes que apenas leen. Su fondo es el mejor reparto, y se desplaza por la recta de abajo cuando cambias el presupuesto. Con el botón del presupuesto de Gopher, un modelo del tamaño de Chinchilla queda en el fondo, y uno del de Gopher, en la pared de la derecha.

Una ley de potencia es una recta en escala logarítmica

Llamemos $N$ a los parámetros de un modelo, todos, embeddings incluidos; $D$ a los tokens que lee durante el entrenamiento, contando dos veces el que lee dos; y $C$ al cálculo que gasta en ello, en operaciones de coma flotante (floating-point operations, FLOPs), cada multiplicación y cada suma una. Son las letras de Kaplan, y en el curso anterior eran otras cosas (documentos, corpus, caracteres) que aquí no vuelven a aparecer.

Kaplan y sus coautores entrenaron modelos de entre 768 y 1 500 millones de parámetros (sin los embeddings, que ellos dejan fuera de $N$) sobre páginas web en inglés, y encontraron que, cuando una sola de las cantidades limita y las otras sobran, la pérdida sobre texto no visto es

con $\alpha_N \approx 0.076$ y $\alpha_D \approx 0.095$, y una tercera igual para $C$, de exponente $0.050$. (Los artículos escriben la pérdida $L$; aquí $L$ es el número de capas.) Una cantidad que es una constante por una potencia de otra sigue una ley de potencia, y una ley de potencia de la pérdida es una ley de escala (scaling law). La ley de Heaps del curso anteriorDeep Learning para NLP: del Perceptrón al TransformerBloque 1 · Lección 3Vocabulario, frecuencia y el problema OOVLas palabras de un corpus no se reparten por igual: unas pocas lo copan y la mayoría aparece una sola vez. De esa forma salen el tamaño del vocabulario, lo que cuesta recortarlo y el token que recoge todo lo que queda fuera., la de los tipos que siguen apareciendo al alargar el corpus, era de la misma familia. Tomemos logaritmos en la primera:

que en los ejes $(\log N, \log \mathcal{L})$ es una recta de pendiente $-\alpha_N$. Por eso estas leyes se dibujan con los dos ejes logarítmicos, y por eso se miden así: con la pérdida de unos cuantos modelos de tamaños distintos se ajusta una recta a los logaritmos, y su pendiente es el exponente. Y por eso lo que compra cada multiplicación de $N$ es siempre lo mismo. Multiplicarlo por $k$ multiplica la pérdida por $k^{-\alpha_N}$, en cualquier punto: por 0.949 al duplicarlo, por 0.839 al multiplicarlo por diez. Para bajar la pérdida a la mitad habría que multiplicar $N$ por $2^{1/0.076}$, unas 9 000 veces.

Un suelo, y lo que miden las constantes

Una recta así no puede seguir para siempre. Si $N$ creciera sin límite, $\mathcal{L}(N)$ tendería a cero, y ningún modelo predice sin error un texto que tiene algo de impredecible. Hoffmann y sus coautores le ponen suelo:

$\mathcal{L}_{\infty}$ es la pérdida irreducible, el suelo del explorable: lo que quedaría con $N$ y $D$ infinitos. Por encima hay dos leyes de potencia, una que se paga por tener pocos parámetros y otra por haber leído pocos tokens. (El artículo escribe $E$, $A$, $B$, $\alpha$ y $\beta$; aquí $B$ es el tamaño del batch, y $\beta$ la reserva el bloque siguiente.) Sus exponentes son del orden de 0.3, y no se comparan con los 0.076 de Kaplan, porque miden sólo lo que queda por encima del suelo. Y la suma ya no es una recta en ejes logarítmicos: es la curva del explorable, que se tuerce hacia $\mathcal{L}_{\infty}$.

Queda una advertencia, la de la perplejidad: una pérdida por token mide también al tokenizador. $N_c$, $D_c$ y $\mathcal{L}_{\infty}$ cambian con el corpus y el tokenizador, y Kaplan lo avisa: otro tokenizador multiplica la pérdida por un factor, y ese factor se lo quedan las constantes, no los exponentes. Por eso el mini-GPT no tiene punto en la primera vista: sus 3.2 nats son por token de su tokenizador sobre una novela en español, y la recta de Kaplan, por token del de GPT-2 sobre páginas web en inglés.

Lo que cuesta entrenar, y cómo repartirlo

El cálculo se puede contar. La lección sobre la cachéBloque 1 · Lección 5La caché de claves y valores: cuánto cuesta cada tokenPara escribir cada token, el mini-GPT volvía a pasar el texto entero por la red. La lección demuestra que las claves y los valores de las posiciones ya leídas no cambian, los guarda en una caché y cuenta lo que cuesta cada token con ella y sin ella: sin ella, exactamente t veces más; con ella, una memoria que crece con el texto y que en un modelo grande llega a ocupar más que sus pesos. Todo medido sobre el mini-GPT, y contra el reloj. contó lo que cuesta pasar una fila por la red, $c_{\text{fila}}$: una multiplicación por cada peso de cada matriz, que en el mini-GPT son $131\,072$ de sus $136\,448$ parámetros. Con su suma son dos operaciones por peso, así que la ida cuesta unas $2N$ por token. La vuelta cuesta el doble, y 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. dice por qué: por cada producto de la ida hace dos, uno que lleva el error hacia atrás con la traspuesta de la matriz y otro que forma el gradiente de sus pesos, los dos del tamaño del de la ida. Tres productos por matriz y dos operaciones por multiplicación dan

La cuenta deja fuera la atención, las $2Lt \cdot d_{\text{model}}$ multiplicaciones por token que dependen de la posición: en las ventanas de 64 del mini-GPT suman a lo sumo un 12 % más, y en un modelo grande, con miles de dimensiones, unos pocos por ciento. Chinchilla, que la cuenta entera, no se aparta de $6ND$ más de un 10 %. Para el mini-GPT, $C = 6 \cdot 136\,448 \cdot 3\,584\,000 \approx 2.9 \times 10^{12}$ FLOPs; GPT-3 gastó unos $3 \times 10^{23}$.

El fondo del valle

Fijemos un presupuesto, un $C$ decidido de antemano. Con $C$ dado, $D = C/(6N)$, y la pérdida de Chinchilla pasa a depender sólo de $N$: es el valle del explorable, que el artículo llama perfil IsoFLOP. Busquemos su fondo moviendo $N$ en escala logarítmica, que es como se mueve en el explorable. Como $\log N + \log D = \log(C/6)$, subir $\log N$ un poco baja $\log D$ lo mismo. Y la derivada de $A_N N^{-\alpha_N}$ respecto a $\log N$ es $-\alpha_N A_N N^{-\alpha_N}$: el exponente por la potencia. Así que el término de los parámetros baja a ritmo $\alpha_N A_N/N^{\alpha_N}$, el de los datos sube a ritmo $\alpha_D A_D/D^{\alpha_D}$, y en el fondo los dos ritmos se igualan:

Se lee sin despejar nada. En el fondo del valle, lo que baja el término de los parámetros al darle a $N$ un 1 % más es lo que sube el de los datos al quitarle ese 1 % a $D$. Despejando con $D = C/(6N)$ (el detalle, abajo),

donde el par $(N^{\star}, D^{\star})$ es el fondo del valle, el reparto óptimo (compute-optimal) de ese presupuesto. La estrella marca el óptimo, como en las lecciones anteriores; los artículos escriben $N_{\text{opt}}$. Los dos exponentes suman 1, como tiene que ser si $N^{\star}D^{\star} = C/6$. El cociente es la cuenta que importa:

Si los dos términos bajan igual de deprisa, $\alpha_N = \alpha_D$, el exponente es cero y la proporción de tokens por parámetro es la misma para cualquier presupuesto. Ésa es la afirmación de Chinchilla dicha con exponentes: $N^{\star}$ y $D^{\star}$ crecen los dos como $C^{1/2}$. En el ajuste del explorable (el del artículo, rehecho en 2024 porque el publicado se detuvo antes de tiempo), $\alpha_N = 0.348$ y $\alpha_D = 0.366$. Cuánto vale la proporción no sale de la derivación, porque depende de $A_N$ y $A_D$: hay que medirla.

Kaplan midió otra cosa: $N^{\star} \propto C^{0.73}$ y $D^{\star} \propto C^{0.27}$, así que por cada diez veces más de cálculo, 5.4 veces más parámetros y 1.9 veces más tokens. El artículo de Chinchilla lo achaca a que Kaplan entrenó todos sus modelos con el mismo calendario y leyó la pérdida de los entrenamientos cortos a mitad del coseno de 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., antes de que la tasa bajara: salían peores de lo que eran, y con poco presupuesto los modelos grandes, que llegan a él con pocos tokens leídos, parecían peor negocio. El óptimo de los presupuestos pequeños salía menor de lo que era, y la recta de $N^{\star}$, más empinada. Un trabajo de 2024 señala además los embeddings: Kaplan no los contaba, y en un modelo pequeño son mucho. En el mini-GPT, $36\,864$ de sus $136\,448$ parámetros.

Las cuentas del mini-GPT

Falta el número. La celda no carga el mini-GPT: sólo necesita sus parámetros y los tokens que había leído el checkpoint. Lo que trae escrito es la tabla 3 del artículo de Chinchilla: nueve tamaños de modelo, de 400 millones a diez billones de parámetros, con los tokens que harían óptimo a cada uno según el primero de sus tres métodos, un ajuste a cientos de entrenamientos. La celda les ajusta en logaritmos las dos rectas de la derivación y las prolonga hasta el presupuesto del mini-GPT.

Los exponentes salen $0.498$ y $0.502$, un medio los dos. Que sumen 1 no comprueba nada: cada fila cumple $\log N + \log D = \log C - \log 6$, y una recta de mínimos cuadrados respeta esa suma. La proporción sí es un resultado. A lo largo de nueve órdenes de magnitud de cálculo, las dos rectas dan entre $20.2$ y $21.8$ tokens por parámetro: los 20 de Chinchilla, con un 10 % de margen.

El mini-GPT queda casi siete órdenes de magnitud por debajo de la tabla, donde la recta ya no se mide sino que se prolonga (la parte discontinua del explorable, sobre la que cae su punto): la segunda cifra de lo que da no está garantizada, y el orden de magnitud sí. Para su presupuesto pide unos $160\,000$ parámetros y $3.1$ millones de tokens, y el checkpoint tiene $136\,448$ y leyó $3.6$ millones. En cálculo, el mini-GPT es un modelo óptimo.

Pero la recta cuenta tokens leídos, y en los experimentos de Chinchilla casi todo lo leído era texto nuevo para el modelo: su $D$ son, en la práctica, tokens distintos. Los $3.6$ millones del mini-GPT son $125\,789$ leídos 28 veces y media. Ésa es la respuesta a la pregunta del principio: ni mayor ni menor. Con el cálculo que gastó, su tamaño es el bueno, y lo que le falta es texto, unas 24 veces más del que tiene la novela. Y releer no es leer. Una época es una pasada entera por el texto de entrenamiento, como en el descenso de gradiente del curso anteriorDeep Learning para NLP: del Perceptrón al TransformerBloque 2 · Lección 6Descenso de gradienteEl gradiente reúne una derivada parcial por parámetro y apunta, entre todas las direcciones posibles, a la que más sube la pérdida; restarlo un poco cada vez es todo el método. Cuánto es «un poco» decide entre converger, rebotar y divergir, y la cuenta sale entera sobre una parábola., y un estudio de 2023 encontró que hasta unas cuatro valen casi como texto nuevo y que, a partir de ahí, cada una vale menos. Las 28 del mini-GPT están muy lejos de ese margen, y la distancia entre su pérdida de entrenamiento y la del texto reservado es lo que se ve de ello.

Comprueba tu intuición

Cinco preguntas: una ley de potencia a mano, el 6 de $6ND$, un reparto con exponentes desiguales, lo que dicen las cuentas del mini-GPT y un punto que no está en la gráfica.

Las leyes de esta lección predicen un número, la pérdida sobre texto no visto, y lo predicen a lo largo de muchos órdenes de magnitud: baja sin saltos, un poco con cada multiplicación del presupuesto. Lo que no dicen es qué sabe hacer un modelo con una pérdida dada. Entre el mini-GPT y GPT-3 hay once órdenes de magnitud de cálculo, y la diferencia no se agota en que uno escriba frases de Galdós con faltas y el otro no.

El artículo de GPT-3 lleva esa diferencia en el título: Language Models are Few-Shot Learners. Con unos cuantos ejemplos de una tarea en el prompt, una para la que nadie lo entrenó (traducir, ordenar las letras de una palabra, sumar), hace la tarea con el ejemplo siguiente, sin que cambie un solo peso. Qué significa eso en la notación de este bloque, y por qué el mini-GPT no puede hacerlo, es la lección siguiente, sobre el aprendizaje en contexto.