El fine-tuning es mala herramienta para añadir conocimiento: tres fases de entrenamiento y el atajo de bajo rango

En la parte 9 entrenamos un modelo desde cero, pero lo normal en este trabajo es partir de unos pesos que ha entrenado otro y modificarlos. Preentrenamiento, fine-tuning, RLHF y LoRA son cuatro palabras que describen esa tarea como si fueran cuatro máquinas distintas.

He leído la capa LoRA de peft y la pérdida DPO de trl, y después le he enseñado a SmolLM2-135M quince capitales, tres que ya sabía y doce de países que no existen. Se las aprendió todas, y ahí está el problema.

Tres fases y un modelo en entrenamiento

El preentrenamiento consiste en predecir el siguiente token sobre todos los datos que tengas, con la entropía cruzada de la parte 8 aplicada a un corpus. Es en esta fase donde el modelo adquiere casi todo lo que sabe.

El fine-tuning supervisado (SFT) usa la misma pérdida sobre pares de instrucción y respuesta. Lo habitual es cambiar qué posiciones cuentan y dejar fuera el prompt con una máscara, como hace trl por defecto con los datos de tipo prompt-completion. En un ejemplo con cinco tokens de prompt y siete de respuesta, la pérdida fue de 2,9883 sobre la secuencia entera y de 2,4845 sobre la respuesta sola, los mismos dos números que he sacado a mano de los logits. Lo único distinto es un array de etiquetas, igual que cuando se ejecuta el mismo binario con otro fichero de configuración.

En la tercera fase, el ajuste por preferencias, no hay respuesta correcta. Cada ejemplo son dos respuestas y una nota que indica cuál prefirió una persona.

Por defecto, las tres fases actualizan todos los pesos del modelo que se entrena, al que en el ajuste por preferencias se llama política, un término heredado del aprendizaje por refuerzo. Lo que distingue a la tercera es cuántos modelos más mantiene a su lado.

El ajuste por preferencias optimiza una diferencia

El aprendizaje por refuerzo a partir de retroalimentación humana (RLHF) entrena un modelo de recompensa para puntuar respuestas y luego optimiza la política según esas puntuaciones. Su algoritmo clásico, PPO (optimización proximal de políticas), añade un crítico que predice qué puntuación acabará teniendo una respuesta a medio escribir. También guarda una referencia congelada, el modelo tal como estaba antes de esta fase, para penalizar que la política se aleje de ella. Son cuatro modelos para una sola actualización. La optimización directa de preferencias (DPO) prescinde del modelo de recompensa y del crítico, pero conserva la referencia. Su núcleo, condensado de dpo_trainer.py:1445-1483 de trl, es este:

chosen_logratios = chosen_logps - ref_chosen_logps        # how far the policy moved, on the good answer
rejected_logratios = rejected_logps - ref_rejected_logps  # and on the bad one
chosen_scores, rejected_scores = chosen_logratios, rejected_logratios   # the default divergence
delta_score = chosen_scores - rejected_scores
per_sequence_loss = -F.logsigmoid(self.beta * delta_score)

La pérdida busca que la diferencia entre esos dos movimientos sea grande. beta, que vale 0,1 por defecto, fija lo que cuesta alejarse de la referencia. Cuanto más alta, menos se mueve la política.

La referencia parece una segunda copia de los pesos, pero con una política LoRA (la explico más abajo) no lo es. La parte entrenable es un pequeño desvío al lado de los pesos congelados, así que basta con apagarlo para tener la referencia. Con el desvío apagado, las log-probabilidades de la política coincidieron hasta el último bit con las de un modelo congelado cargado aparte, tanto antes de entrenar como tras veinte pasos. En el paso cero, la política y la referencia son el mismo modelo y la diferencia es cero. La pérdida vale entonces menos el logaritmo de un medio, es decir, log 2, que con seis decimales es 0,693147.

Esas cuentas las he hecho en un bucle propio, fuera de DPOTrainer. Con una tasa de aprendizaje brutal de 0,001 y veinte pasos sobre dos pares de preferencias, la diferencia pasó de 0 a 95,4. La respuesta buena se volvió 2,9 nats más probable y la mala, 92,5 nats menos.

En parte, eso viene de serie. Como una log-probabilidad no puede pasar de cero y las respuestas buenas partían de −2,3 y −3,5, esos 2,9 nats eran casi todo el margen que tenían; las malas, en cambio, medían nueve y doce tokens y no tenían suelo. Por eso lo he relanzado con pares de la misma longitud, París contra Lyon y 4 contra 5, tres veces en cada sentido. Con la tasa brutal, el lado rechazado cayó entre 31 y 47 nats en todas las ejecuciones, mientras que el preferido se movió entre −1,2 y +4,2 y en una de ellas acabó siendo menos probable que al principio.

Con una tasa realista de 0,00001, el reparto depende de dónde partía cada respuesta. Cuando las preferidas eran París y 4, ganaron entre 0,01 y 0,02 nats y las otras perdieron entre 0,16 y 0,17; cuando lo eran Lyon y 5, ganaron 0,15 y las otras perdieron 0,03. En estos pares, el movimiento mayor se lo llevó la respuesta que al principio era menos probable. La pérdida pide una diferencia y la saca del lado más barato. Nada en ella obliga a que la respuesta preferida se vuelva más probable.

LoRA entrena otro objeto

LoRA (adaptación de bajo rango) es una decisión sobre en qué consiste la actualización, válida para cualquiera de las tres fases. Se deja intacta una matriz de pesos y se le ponen al lado dos matrices pequeñas, que juntas forman el adaptador. Así lo hace peft, en una versión condensada de lora/layer.py:263-1102.

self.lora_A[adapter_name] = nn.Linear(self.in_features, r, bias=False)  # d_in to r
self.lora_B[adapter_name] = nn.Linear(r, self.out_features, bias=lora_bias)  # r to d_out
self.scaling[adapter_name] = lora_alpha / r
nn.init.zeros_(self.lora_B[adapter_name].weight)
result = result + lora_B(lora_A(dropout(x))) * scaling

El rango r marca lo estrecho que es el tramo intermedio. Antes de sumarse al resultado, el desvío se multiplica por lora_alpha / r, que con alpha 16 y rango 8 vale 2,0.

El truco está en los ceros de la última línea de la inicialización. En el paso cero el producto de las dos matrices es cero, así que el modelo es idéntico, bit a bit, al que tenías al empezar. Después el entrenamiento mueve A y B. Lo que aprenden es una diferencia que, al multiplicarse, da un cambio de rango 8 en una de las matrices de 576 por 576 de SmolLM2 (rango en el sentido del álgebra lineal, es decir, el número de direcciones independientes que tiene el cambio). Fusionado con la matriz, el adaptador dio los mismos logits que cuando iba al lado, con una discrepancia máxima de 0,00008.

De ese rango sale el ahorro. Si se adaptan las matrices query y value de Llama-3.1-8B con rango 8, se entrenan 3 407 872 números de 8 030 261 248, un 0,04 %.

Si subes el rango, el total apenas se mueve, porque casi todo es la base congelada. Ahora pasa la base a NF4 con doble cuantización, el formato de 4 bits de QLoRA (la receta para entrenar LoRA sobre una base de 4 bits), y desaparecen casi dos tercios de la barra de LoRA. Solo las capas de dentro de los bloques bajan a algo más de cuatro bits por peso; los embeddings y la capa de salida se quedan en dieciséis.

Los hechos entran despacio y con ellos llega la confianza

Hasta aquí la mecánica. Ahora viene la parte que se suele entender mal en las reuniones de planificación.

El punto de partida son doce capitales de países inventados y doce reales, todas con la forma «The capital of Vantressa is Molduin». De las reales, el modelo completó tres palabra por palabra. En seis de las otras nueve contestó cosas como «the city of Rome», que no cuentan como coincidencia exacta. Las nueve quedaron fuera. Luego he ajustado todos los pesos sobre las quince a la vez, con una tasa de 0,00001.

Las tres que ya sabía se afianzaron enseguida y en veinte pasos la confianza en el primer token de la respuesta correcta pasó del 22 % al 97 %. Para entonces las doce nuevas estaban en el 2 % y el modelo acertaba una. En el paso cuarenta ya respondía las doce. En parte, la lentitud se debe al punto de partida. En ese primer token, los hechos nuevos estaban a 8,6 nats de distancia, frente a los 1,5 de los conocidos, más de cinco veces más lejos. Además, cada uno necesitaba acertar dos o tres tokens cuando a los conocidos les bastaba uno. Gekhman y colaboradores describen el mismo patrón a gran escala: el conocimiento nuevo se ajusta significativamente más despacio y, a medida que entra, aumentan las alucinaciones.

Eso sí, «respondía» quiere decir que completaba la frase con la que lo había entrenado. Con «What is the capital of Vantressa?» no acertó ninguna de las doce, ni al preguntarle de qué país es capital Molduin; con «Vantressa’s capital city is», en cambio, acertó seis. Lo que había aprendido era una frase.

Algo parecido a ese aumento de las alucinaciones apareció también aquí. Otros seis nombres de país, inventados o sacados de lugares poco conocidos, se quedaron fuera del entrenamiento. Antes de entrenar, el modelo los completaba igual que las capitales reales de más arriba, con alrededor de un 24 % de confianza. Cinco respuestas empezaban por «the town of» o «the city of», cuatro de ellas seguidas del propio nombre del país. Tras cien pasos, cuatro de las seis acababan en una ciudad real seguida de un punto. Según el modelo, la capital de Volnaya es Moscú, con un 82 %. La de Quenby es Shillong, la de Dalmerie, Edimburgo, y la de Rhuddlan, Birmingham sin la B.

En el entrenamiento no había nada sobre esos lugares. Moscú estaba en el batch como capital de Rusia, pero Volnaya también suena a ruso, y en el paso 10 la respuesta era Vorónezh. Otras respuestas se guiaban por el sonido, como Maribor para Marovska o Cardiff para Rhuddlan. Aprendió que una frase de esta forma acaba en una ciudad y a menudo escogió una que encajaba por cómo sonaba.

Arrastra hasta el paso 20 y verás que los hechos conocidos están en el 97,0 %, los nuevos en el 1,9 % y el modelo recuerda uno de los doce. Después recorre todo el entrenamiento fijándote en la línea punteada, la de los seis países que nadie le enseñó. Baja al 5,7 % en el paso 10, cuando el modelo deja de decir «the», y termina en el 41,1 %, frente al 23,7 % del principio. Esa subida se debe a dos respuestas: Moscú, con un 82 %, y «Th», el comienzo del propio nombre de Threndale, con un 89 % (el modelo lo completa como «Thwapth»). Las otras cuatro se quedan en un 19 % de media, cuando empezaron en un 24 %.

El parche que nadie puede leer

Un fine-tuning completo es un diff, así que lo he contado. De 134 515 008 parámetros cambiaron 134 513 152, aunque ninguno lo bastante como para llamar la atención. La mediana del cambio fue 0,0002 y el máximo llegó a 0,0012, que es más o menos la tasa multiplicada por los cien pasos. Con el Adam de la parte 9, un peso con un gradiente estable se mueve en cada paso aproximadamente lo que vale la tasa. El cambio en una matriz de atención ocupó 572 de sus 576 direcciones posibles. Si se mide por la suma de los cuadrados de los cambios, las ocho primeras direcciones recogen el 71 % y hacen falta 27 para llegar al 90 %. Las ocho de LoRA lo recogen todo por construcción. Lo único que el diff comparte con un parche que se pueda revisar es que, al restarlo, se recuperan los pesos originales.

Un parche ilegible solo es un problema si ha cambiado algo que te importaba, y por eso lo he medido. He hecho ocho versiones del mismo fine-tuning, con tres tasas de aprendizaje para la actualización completa y cinco para un adaptador LoRA de rango 8 sobre las matrices query y value. A cada una la he puntuado por cuánto subía su pérdida en seis entradas de este blog con las que no se entrenó ninguna.

El fine-tuning completo aprendió los doce hechos con todas las tasas. LoRA no aprendió ninguno hasta la tasa de 0,0003, que es treinta veces la que ya le bastaba al completo. Si se toma para cada uno la tasa más baja de las probadas con la que aprendió, LoRA subió la pérdida en esas entradas más de veinte veces más que la actualización completa (0,111 nats frente a 0,005), una factura curiosa para la opción barata.

A igualdad de tasa, con 0,001 la actualización completa subió la pérdida quince veces más que LoRA (4,419 frente a 0,287). Este adaptador sí limita el daño, solo que necesita una tasa a la que haya daño que limitar.

Para ser justos, todo esto son ocho ejecuciones con una sola semilla sobre un modelo de 135M, con quince hechos, cien pasos, una única configuración de adaptador y una rejilla de tasas poco densa. Un adaptador en todas las capas lineales quizá aprendería con una tasa más baja, y con un modelo más grande o un dataset real el orden podría salir distinto. Lo que sí esperaría que se transfiriera es el método: comparar tanto a igualdad de resultado como de configuración y medir el daño colateral en algo ajeno al fine-tuning.

Qué vigila el ingeniero de ML

Son cuatro hábitos, y el primero es una decisión previa a todo lo demás.

Pregúntate qué quieres cambiar. Los hechos que faltan deberían llegar por recuperación a la ventana de contexto, donde puedes verlos y quitarlos. El comportamiento, el formato y las negativas a responder sí se le dan bien al fine-tuning.

Puntúa el daño colateral además del objetivo. Todas las ejecuciones que aprendieron los doce hechos acabaron por encima del 99 % en ellos, incluida la actualización completa a 0,001, que subió más de cuatro nats la pérdida en las seis entradas. La parte 8 pedía una sola puntuación sobre el conjunto de test, pero en esta fase hacen falta dos, una sobre texto general y otra sobre lo que de verdad vendes.

Trata la tasa como la palanca del daño. En las ocho ejecuciones, la tasa movió el daño colateral tres órdenes de magnitud, frente al factor de quince que hubo entre la actualización completa y la de bajo rango en la única tasa con la que aprendieron las dos. Haz un barrido, como en la parte 9, y elige la tasa por la cifra de daño colateral.

Mantén separados la base, el adaptador y el modelo fusionado. El adaptador de Llama con rango 8 ocupa 6,8 MB al lado de los 16 GB de la base congelada y se apaga con una sola llamada. Por eso el modelo de referencia de antes salía gratis en memoria. El cómputo sí se paga, porque en cada paso hay que volver a ejecutar el modelo con el adaptador apagado. Fusiona cuando vayas a servir un solo adaptador sobre una base sin cuantizar. Mientras decides, o cuando una misma base sirva muchos adaptadores, mantén las piezas separadas.

Esto tampoco es trabajo de modelado: todo se reduce a una máscara sobre las etiquetas, una sigmoide sobre una diferencia y, en la parte de bajo rango, dos matrices estrechas, una de ellas a cero. Las cuentas se leen en una tarde, y aun así acaban en un diff de cien millones de números que nadie puede revisar.