Papers en cristiano: LoRA
El paper de 2021 que hizo barata la adaptación de modelos gigantes cabe en una idea: en vez de tocar todos los pesos, aprender solo un atajo de bajo rango y sumarlo. Segunda entrega de la serie, con un experimento real: un GPT diminuto entrenado en Shakespeare y adaptado al Quijote de cinco formas distintas para medir cuánto cuesta cada una y cuánto olvida.
Segunda entrega de la serie en la que leo papers clásicos y los cuento como me habría gustado que me los contaran (la primera fue Attention is all you need). Esta vez toca LoRA: Low-Rank Adaptation of Large Language Models, de Hu y compañía, 2021. Lo he implementado desde cero, lo he probado con un experimento real y no con la palabra de un paper, y todo el código está en github.com/delcenjo/lora-from-scratch.
El problema que resuelve. Ajustar un modelo grande a una tarea nueva, del modo clásico, significa tocar todos sus pesos: si el modelo tiene siete mil millones de parámetros, guardar una copia adaptada cuesta siete mil millones de parámetros, aunque la tarea nueva sea pequeña. Y hay un segundo problema, menos hablado: al tocar todos los pesos, se sobrescribe también lo que el modelo sabía antes. Adaptar bien a la tarea nueva y conservar lo aprendido tiran en direcciones opuestas.
La idea del paper. La observación de partida es que el cambio que realmente aporta el ajuste fino, aunque se represente como una matriz del mismo tamaño que los pesos originales, tiende a tener rango bajo: casi toda la información de ese cambio cabe en muchas menos dimensiones. Así que en vez de aprender esa matriz de cambio entera, se aprende como el producto de dos matrices muy estrechas, B y A, con un rango r mucho menor que las dimensiones originales. Los pesos originales, W0, se congelan del todo. La capa adaptada calcula W0·x más (alpha/r)·B·A·x, en vez de W·x directamente.
Por qué B empieza en cero. A se inicializa con valores pequeños al azar; B se inicializa entera a cero. No es un detalle menor: con B en cero, el producto B·A es cero en el primer paso, sea lo que sea lo que contenga A, así que el modelo adaptado empieza siendo una copia exacta del modelo preentrenado. El entrenamiento se aleja de ahí poco a poco. Si ambas matrices arrancaran con ruido aleatorio, el modelo adaptado empezaría roto y las primeras iteraciones se irían en deshacer ese ruido antes de aprender nada útil.
El truco de fusionar pesos. Una vez entrenado, se puede sumar B·A (con su escala) directamente a W0 y quedarse con una única matriz de pesos, igual de rápida en inferencia que la original: la estructura de bajo rango ahorra memoria y cómputo durante el entrenamiento, no durante la inferencia una vez fusionada. Lo comprobé con un test: la salida del modelo con LoRA sin fusionar y la salida tras fusionar los pesos coinciden hasta la sexta cifra decimal.
El experimento. Un GPT de caracteres diminuto, dos capas, dos cabezas, 128 de dimensión, contexto de 128 caracteres, 438.016 parámetros. Lo preentrené 2.500 pasos en el Shakespeare de siempre (el de los tutoriales de Karpathy) y después lo adapté al Quijote —dominio público, en español, con un alfabeto visiblemente distinto— durante 800 pasos, de cinco formas: ajuste fino completo, LoRA con rango 1, 4 y 16, y ajustar solo la capa de salida. Todo en una CPU normal, 21 minutos en total para las seis fases.
La tabla, con números reales. El ajuste fino completo es el que mejor se adapta al Quijote (pérdida de validación 1,76) pero es también el que peor recuerda Shakespeare después: la pérdida ahí sube de 1,91 a 3,57, peor que un modelo que nunca hubiera visto inglés. LoRA rango 16 se adapta casi tan bien (2,29) tocando solo el 5,31% de los parámetros; rango 4 con el 1,38% llega a 2,40; rango 1, con un 0,35%, apenas se mueve (2,55). Ajustar solo la cabeza de salida, con un 2,83% de los parámetros, sorprendentemente iguala a LoRA rango 16 (2,24) — con un modelo de solo dos capas, buena parte de pasar de "predecir inglés" a "predecir español" está en qué caracteres favorece la última capa, no en cómo se mueve la atención.
El olvido catastrófico, en números. Aquí está el dato que de verdad me interesaba. Con los adaptadores de LoRA encendidos, la pérdida en Shakespeare empeora según crece el rango (2,52 en rango 1, hasta 2,90 en rango 16) — a primera vista parece que más capacidad adaptadora es más olvido. No lo es: los pesos originales de LoRA no se tocan en ningún momento de estos entrenamientos. Lo que sube es el propio adaptador, entrenado a fondo en español, que tira de las predicciones hacia el Quijote incluso cuando la entrada es inglesa. Apagando el adaptador —sin recargar nada, solo un booleano— la pérdida en Shakespeare vuelve a 1,906, 1,914 y 1,909 según el rango, contra un original de 1,907: diferencias de milésimas, dentro del ruido de muestreo de la validación. El modelo preentrenado sigue ahí, intacto, a un flag de distancia. El ajuste fino completo y el de solo-cabeza no tienen ese flag: sus pesos originales desaparecen en cuanto empieza el entrenamiento.
Lo que me llevo. Que "B empieza en cero" no es un truco de implementación sino la razón de que todo esto funcione sin sobresaltos, que fusionar pesos es gratis en inferencia y caro solo de explicar, y que la reversibilidad de LoRA no es una promesa de la introducción del paper sino algo que se puede medir con dos evaluaciones y un booleano. Dicho esto, este es un modelo de cientos de miles de parámetros, no de miles de millones, y dos corpus de un tamaño modesto: las conclusiones sobre qué rango gana aquí no deberían extrapolarse a un modelo real sin repetir el experimento a esa escala. Lo honesto es quedarme con la dirección de los resultados, no con el tercer decimal. El código, los tests y el notebook con todo ejecutado están en el repositorio.
