Adaptación eficiente

El problema

Tienes una red diminuta preentrenada para reconocer los dígitos 0 a 4 de load_digits, imágenes de 8×8 píxeles. Ahora necesitas reconocer los dígitos 5 a 9 y solo tienes 50 imágenes etiquetadas: 10 por dígito.

En Fine-tuning adaptaste un modelo de dos formas: con el cuerpo congelado, que solo entrena una cabeza nueva, y con fine-tuning completo, que entrena todos los pesos. Con fine-tuning completo cada tarea termina con su propia copia del modelo. El tamaño de esa copia depende del modelo:

Con diez tareas destino, el fine-tuning completo guarda diez copias. Durante el entrenamiento, cada parámetro entrenable necesita además un gradiente y el estado del optimizador.

Parámetros entrenables

Cada nn.Linear(entrada, salida) guarda una matriz weight de forma (salida, entrada) y un vector bias de longitud salida. p.numel() cuenta los números de un tensor. Un parámetro es entrenable si p.requires_grad es True: backward le calcula un gradiente y el optimizador lo actualiza.

La celda plegada carga load_digits, escala los píxeles al intervalo de 0 a 1 y separa los dígitos en dos grupos. También define exactitud, que devuelve la fracción de predicciones correctas.

import copy
import math

import numpy as np
import torch
import torch.nn.functional as F
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from torch import nn

X, y = load_digits(return_X_y=True)
X = torch.tensor(X / 16, dtype=torch.float32)  # píxeles de 0 a 16 → de 0 a 1
y = torch.tensor(y)

fuente = y < 5
Xf, yf = X[fuente], y[fuente]
Xd, yd = X[~fuente], y[~fuente] - 5  # etiquetas 5 a 9 → 0 a 4
ent, val = train_test_split(
    np.arange(len(yd)), train_size=50, stratify=yd.numpy(), random_state=0
)
Xd_ent, yd_ent = Xd[ent], yd[ent]
Xd_val, yd_val = Xd[val], yd[val]


def exactitud(modelo, X, y):
    with torch.no_grad():
        return (modelo(X).argmax(dim=1) == y).float().mean().item()

Hay 901 imágenes de los dígitos 0 a 4 para preentrenar. Los dígitos del 5 al 9 se dividen en 50 imágenes de entrenamiento y 846 de validación. Sus etiquetas se renumeran de 0 a 4.

nuevo_cuerpo construye el cuerpo 64 → 256 → 256 con ReLU. La cabeza es 256 → 5. entrenar le pasa a Adam solo los parámetros con requires_grad y hace pasos actualizaciones con todos los ejemplos a la vez. La red se preentrena con los dígitos 0 a 4.

def nuevo_cuerpo():
    return nn.Sequential(nn.Linear(64, 256), nn.ReLU(), nn.Linear(256, 256), nn.ReLU())


def entrenar(modelo, X, y, pasos, lr):
    entrenables = [p for p in modelo.parameters() if p.requires_grad]
    optimizador = torch.optim.Adam(entrenables, lr=lr)
    for _ in range(pasos):
        optimizador.zero_grad()
        F.cross_entropy(modelo(X), y).backward()
        optimizador.step()


torch.manual_seed(0)
preentrenado = nn.Sequential(nuevo_cuerpo(), nn.Linear(256, 5))
entrenar(preentrenado, Xf, yf, pasos=100, lr=1e-3)

named_parameters() recorre los parámetros junto con su nombre. En 0.2.weight, el primer 0 es el cuerpo dentro de la red y el 2 es la segunda nn.Linear dentro del cuerpo.

Pregunta. ¿Cuántos parámetros tiene cada capa? ¿Qué capa concentra la mayoría?

for nombre, p in preentrenado.named_parameters():
    print(nombre, tuple(p.shape), p.numel())
print(sum(p.numel() for p in preentrenado.parameters()))
print(round(exactitud(preentrenado, Xf, yf), 3))

0.0.weight (256, 64) 16384
0.0.bias (256,) 256
0.2.weight (256, 256) 65536
0.2.bias (256,) 256
1.weight (5, 256) 1280
1.bias (5,) 5
83717
1.0

Como en Fine-tuning, modulo.requires_grad_(False) congela todos los parámetros de un módulo: no reciben gradiente y el optimizador no los cambia. Adam guarda dos tensores de estado por cada parámetro entrenable, exp_avg y exp_avg_sq, con la misma forma que el parámetro.

con_cabeza_nueva pone una cabeza 256 → 5 nueva sobre un cuerpo, y contar_entrenables suma numel() solo de los parámetros con requires_grad. copy.deepcopy copia el cuerpo preentrenado, así que congelar o entrenar la copia no cambia preentrenado.

def con_cabeza_nueva(cuerpo):
    return nn.Sequential(cuerpo, nn.Linear(256, 5))


def contar_entrenables(modelo):
    return sum(p.numel() for p in modelo.parameters() if p.requires_grad)


torch.manual_seed(0)
todo = con_cabeza_nueva(copy.deepcopy(preentrenado[0]))
solo_cabeza = con_cabeza_nueva(copy.deepcopy(preentrenado[0]))
solo_cabeza[0].requires_grad_(False)

Para medir la memoria del optimizador, se crea Adam con los parámetros entrenables de cada red, se da un paso con los 50 ejemplos y se suman los números que quedan en optimizador.state, sin el contador step.

Pregunta. Con cabeza nueva y el cuerpo congelado, ¿cuántos parámetros entrenables quedan? Después de un paso de Adam, ¿cuántos números de estado guarda el optimizador con el cuerpo congelado y cuántos sin congelar?

Descomposición de rango bajo

El fine-tuning completo aprende un cambio ΔW con la misma forma que el peso W; el peso final es W + ΔW. LoRA (Low-Rank Adaptation) escribe ese cambio como el producto de dos matrices delgadas:

ΔW = B A, con A de forma (r, entrada) y B de forma (salida, r).

El número r es la dimensión interna del producto. El rango de una matriz es el número de columnas linealmente independientes que tiene, y torch.linalg.matrix_rank lo calcula.

Pregunta. Para la capa 256 → 256 con r = 4, ¿cuántos números tiene ΔW y cuántos tienen A y B juntos? ¿Qué rango tiene B A?

Para varios valores de r, en esta capa y en una 4096 → 4096, se cuentan los números de A y B frente a los entrada · salida de ΔW. En la tabla, la fila con 256 → 256 y r = 4 repite el caso anterior: 2 048 números, el 3.12 % de ΔW.

Capa LoRA

Una capa LoRA envuelve una nn.Linear congelada y le suma el cambio de rango bajo:

salida = W x + b + escala · B A x, con escala = alpha / r.

A empieza con valores aleatorios pequeños y B empieza en ceros. El artículo de LoRA divide entre r para no tener que reajustar el learning rate cada vez que cambia r. El código calcula (x Aᵀ) Bᵀ en ese orden, así nunca construye la matriz completa B A. nn.Parameter registra un tensor como parámetro entrenable del módulo, e in_features y out_features son las dimensiones de entrada y salida de la capa envuelta.

class LoRALinear(nn.Module):
    def __init__(self, base, r, alpha):
        super().__init__()
        self.base = base
        self.base.requires_grad_(False)
        # dividir entre √entrada mantiene pequeños los valores de A x
        self.A = nn.Parameter(torch.randn(r, base.in_features) / math.sqrt(base.in_features))
        self.B = nn.Parameter(torch.zeros(base.out_features, r))
        self.escala = alpha / r

    def forward(self, x):
        return self.base(x) + self.escala * (x @ self.A.T @ self.B.T)

Para probarla, h guarda la salida de preentrenado[0][:2], la primera capa lineal del cuerpo y su ReLU, con los 50 ejemplos: es la entrada que recibe la capa 0.2. torch.equal es True solo si dos tensores tienen la misma forma y los mismos valores exactos. pow(2).mean() hace de pérdida de prueba para que backward calcule gradientes.

Pregunta. Envuelves una copia de la capa 0.2 con r = 4 y alpha = 4. ¿La salida cambia respecto a la capa original? ¿Cuántos parámetros tiene la capa en total y cuántos son entrenables? Después de un backward, ¿qué tensores tienen gradiente distinto de cero: A, B o base.weight?

capa_base = copy.deepcopy(preentrenado[0][2])
torch.manual_seed(0)
capa_lora = LoRALinear(capa_base, r=4, alpha=4)

with torch.no_grad():
    h = preentrenado[0][:2](Xd_ent)

print(torch.equal(capa_lora(h), capa_base(h)))
print(sum(p.numel() for p in capa_lora.parameters()), contar_entrenables(capa_lora))

capa_lora(h).pow(2).mean().backward()
print(capa_lora.A.grad.abs().max().item())
print(capa_lora.B.grad.abs().max().item() > 0)
print(capa_lora.base.weight.grad is None)

True
67840 2048
0.0
True
True

Comparación con fine-tuning completo

Tres formas de adaptar el cuerpo preentrenado a los dígitos 5 a 9. Las tres usan una cabeza nueva 256 → 5, los mismos 50 ejemplos, 100 pasos de Adam y learning rate 0.003:

• cabeza: cuerpo congelado
• lora: cuerpo congelado con LoRALinear de r = 4 en sus dos capas lineales
• completo: todos los parámetros entrenables

Cada forma se repite con las semillas 0, 1 y 2, que cambian la cabeza nueva y las matrices A.

adaptar copia el cuerpo preentrenado y lo congela salvo en el modo completo. En el modo lora envuelve después las capas 0 y 2 del cuerpo.

def adaptar(modo, r=4, alpha=4, semilla=0):
    torch.manual_seed(semilla)
    cuerpo = copy.deepcopy(preentrenado[0])
    cuerpo.requires_grad_(modo == "completo")
    if modo == "lora":
        # se envuelve después de congelar, así A y B quedan entrenables
        cuerpo[0] = LoRALinear(cuerpo[0], r, alpha)
        cuerpo[2] = LoRALinear(cuerpo[2], r, alpha)
    modelo = con_cabeza_nueva(cuerpo)
    entrenar(modelo, Xd_ent, yd_ent, pasos=100, lr=3e-3)
    return modelo

modelos guarda cada red con la clave (modo, semilla) para usarla en la siguiente sección. Cada fila de la tabla muestra los parámetros entrenables, su fracción de los 83 717 y la media ± desviación estándar de la exactitud de validación en las tres semillas.

Pregunta. Ordena las tres formas por exactitud de validación. ¿Qué fracción de los 83 717 parámetros entrena LoRA?

modelos = {}
for modo in ["cabeza", "lora", "completo"]:
    exactitudes = []
    for semilla in [0, 1, 2]:
        modelos[modo, semilla] = adaptar(modo, semilla=semilla)
        exactitudes.append(exactitud(modelos[modo, semilla], Xd_val, yd_val))
    n = contar_entrenables(modelos[modo, 0])
    print(
        f"{modo:9} {n:6} {n / 83717:6.1%} "
        f"{np.mean(exactitudes):.3f} ± {np.std(exactitudes):.3f}"
    )

cabeza      1285   1.5% 0.723 ± 0.013
lora        4613   5.5% 0.898 ± 0.007
completo   83717 100.0% 0.917 ± 0.005

Fusión de pesos

Después de entrenar, W + escala · B A es una matriz normal de forma (salida, entrada). Se puede fusionar: escribirla en una nn.Linear y quitar A y B. La red fusionada tiene la misma arquitectura que la original.

fusionar crea esa nn.Linear y escribe en ella el peso combinado y el sesgo de la base con copy_. La copia va dentro de torch.no_grad() para que no quede registrada en el grafo de gradientes.

def fusionar(capa):
    fusionada = nn.Linear(capa.base.in_features, capa.base.out_features)
    with torch.no_grad():
        fusionada.weight.copy_(capa.base.weight + capa.escala * capa.B @ capa.A)
        fusionada.bias.copy_(capa.base.bias)
    return fusionada

modelo_fusionado es una copia de la red con LoRA de la semilla 0 con sus dos capas LoRA fusionadas. torch.allclose compara dos tensores con una tolerancia, aquí atol=1e-5, porque el producto fusionado suma en otro orden y puede cambiar los últimos decimales.

Pregunta. ¿La red fusionada da las mismas predicciones que la red con LoRA?

modelo_lora = modelos["lora", 0]
modelo_fusionado = copy.deepcopy(modelo_lora)
modelo_fusionado[0][0] = fusionar(modelo_lora[0][0])
modelo_fusionado[0][2] = fusionar(modelo_lora[0][2])

with torch.no_grad():
    iguales = torch.allclose(modelo_lora(Xd_val), modelo_fusionado(Xd_val), atol=1e-5)
print(iguales, sum(p.numel() for p in modelo_fusionado.parameters()))

True 83717

Dentro de cada LoRALinear, la capa original queda en el atributo base. Para comprobar si cambió, se toman los pesos y sesgos de esas capas base en la red con LoRA de la semilla 0 y torch.equal los compara con los de preentrenado[0].

Pregunta. ¿Cambió algún peso del cuerpo preentrenado al entrenar con LoRA?

adaptador copia con detach().clone() solo los parámetros entrenables de la red con LoRA, separados del modelo y de su grafo.

Pregunta. ¿Qué tensores de la red con LoRA forman adaptador?

adaptador = {
    nombre: p.detach().clone()
    for nombre, p in modelo_lora.named_parameters()
    if p.requires_grad
}
print(list(adaptador))
print(sum(t.numel() for t in adaptador.values()))

['0.0.A', '0.0.B', '0.2.A', '0.2.B', '1.weight', '1.bias']
4613

La biblioteca PEFT de Hugging Face aplica esta misma idea a modelos reales. LoraConfig fija r, lora_alpha y target_modules, las capas que se envuelven. get_peft_model congela el modelo base y añade A y B. merge_and_unload fusiona los pesos como fusionar. PEFT llama adapter al conjunto de pesos pequeños que se entrena y se guarda aparte del modelo base, como adaptador. Los adapters originales, anteriores a LoRA, insertan en cada bloque dos MLP pequeños d → m → d, uno después de la atención y otro después del MLP del bloque, con m mucho menor que d. Esos MLP tienen una activación no lineal, así que no se pueden fusionar con los pesos existentes y la red conserva sus capas extra al predecir.

La cuantización reduce la memoria del modelo base congelado: guarda cada peso con 8 o 4 bits en lugar de los 32 de float32. Un modelo de 7 000 millones de parámetros pasa de 28 GB a unos 3.5 GB con 4 bits. QLoRA entrena A y B en mayor precisión sobre un modelo base cuantizado a 4 bits. La guía de cuantización de PEFT muestra cómo combinar las dos técnicas.

Práctica

Traza guiada

Una capa nn.Linear(3, 2) tiene W = [[1, 0, 0], [0, 1, 0]] y b = [0, 0]. Su capa LoRA tiene r = 1, alpha = 1, A = [[1, 1, 0]] y B = [[2], [0]]. Antes de procesar, escribe ΔW = B A, su rango, la salida para x = [1, 2, 3] y el número de parámetros entrenables.

capa_traza = LoRALinear(nn.Linear(3, 2), r=1, alpha=1)
with torch.no_grad():
    capa_traza.base.weight.copy_(torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]))
    capa_traza.base.bias.zero_()
    capa_traza.A.copy_(torch.tensor([[1.0, 1.0, 0.0]]))
    capa_traza.B.copy_(torch.tensor([[2.0], [0.0]]))
x_traza = torch.tensor([1.0, 2.0, 3.0])

Comprueba con capa_traza(x_traza), capa_traza.B @ capa_traza.A y contar_entrenables(capa_traza). La salida debe ser [7, 2]. Si obtienes otro valor, calcula por separado W x y B (A x).

Rango de la adaptación

Repite adaptar("lora", r=r, alpha=r) con r igual a 1, 2, 8 y 16 y las semillas 0, 1 y 2. Antes de procesar, escribe una fórmula para el número de parámetros entrenables en función de r y predice si la exactitud de validación sube siempre con r. Comprueba la fórmula con assert contar_entrenables(modelo) == …. Compara cada exactitud con la del fine-tuning completo de la tabla de «Comparación con fine-tuning completo».

Reemplazo de capas

Escribe dos funciones para cualquier cuerpo nn.Sequential:

1. aplicar_lora(cuerpo, r, alpha) congela el cuerpo y reemplaza cada nn.Linear por una LoRALinear
2. fusionar_cuerpo(cuerpo) reemplaza cada LoRALinear por la nn.Linear que devuelve fusionar

Pruébalas con un cuerpo nuevo 64 → 128 → 128 → 128 con ReLU entre capas. Comprueba con assert que el número de entrenables es r · (64 + 128) + 2 · r · (128 + 128) y que las salidas antes y después de fusionar_cuerpo coinciden con torch.allclose. Adapta el cuerpo a los dígitos 5 a 9 con una cabeza nueva y reporta la exactitud de validación.

Reto opcional

Calcula el cambio que produjo el fine-tuning completo en la capa 0.2: modelos["completo", 0][0][2].weight - preentrenado[0][2].weight. Obtén sus valores singulares con torch.linalg.svdvals y la fracción de la suma de sus cuadrados que aportan los 4 mayores. Explica si el resultado apoya la suposición de LoRA para esta tarea.

Errores que revisar

• Si con LoRA contar_entrenables da solo los 1 285 parámetros de la cabeza, llamaste requires_grad_(False) sobre el cuerpo después de envolver las capas y también congelaste A y B. Congela primero y envuelve después.
• Si LoRA no mejora sobre entrenar solo la cabeza y A.grad y B.grad son cero, revisa la inicialización. Si A y B empiezan en cero, ninguna de las dos recibe gradiente. Solo B empieza en cero.
• Si LoRA da la misma exactitud que entrenar solo la cabeza, revisa cuándo creaste el optimizador. Si lo creaste antes de reemplazar las capas, A y B no están en su lista de parámetros.
• Si la red fusionada da predicciones distintas, revisa escala y el orden del producto. El cambio es B @ A, de forma (salida, entrada), igual que weight.
• Si el fine-tuning completo da lo mismo que entrenar solo la cabeza, revisa requires_grad del cuerpo copiado. copy.deepcopy conserva requires_grad=False si el original estaba congelado.

Cierre

Explica por qué la red con LoRA reconoce los dígitos 5 a 9 sin cambiar ningún peso del cuerpo preentrenado, y qué guardarías para usar la misma red base en las dos tareas.

Lista de comprobación

• ¿Cuántos parámetros añade LoRA de rango r a una nn.Linear(entrada, salida) y por qué?
• ¿Por qué una capa LoRA recién creada da la misma salida que su capa base?
• ¿Por qué en el primer paso solo B recibe gradiente distinto de cero?
• ¿Qué parte del modelo se guarda por tarea con LoRA y qué parte se comparte entre tareas?
• ¿Por qué fusionar los pesos no cambia las predicciones ni el tamaño de la red?

Documentación oficial

• PEFT: métodos de fine-tuning eficiente en parámetros, adapters e integración con Transformers.
• LoRA en PEFT: LoraConfig, r, lora_alpha, target_modules y fusión de pesos.
• Cuantización en PEFT: modelos base en 8 y 4 bits y entrenamiento de LoRA encima (QLoRA).
• LoRA: Low-Rank Adaptation of Large Language Models: artículo original del método, con la inicialización de A y B y la escala alpha / r.
• torch.nn.Module: parameters(), named_parameters() y requires_grad_().