Rango intrínseco bajo

LoRA, low-rank adaptation. Episodio dos: rango intrínseco bajo. En el episodio uno vimos que ajustar todos los pesos cuesta cuatro copias del modelo. LoRA entrena muchísimo menos, y apuesta a que alcanza. Hoy ponemos esa apuesta a prueba: hacemos un fine-tuning completo en un modelo de juguete, miramos el cambio que dejó en los pesos, y medimos en cuántas direcciones vive.

Parte 1 · Qué es el cambio

Un fine-tuning arranca de unos pesos y termina en otros. La resta es una matriz, y las matrices tienen rango.

Un cambio de rango 1, en el plano

Primero, en dos dimensiones. La grilla muestra cómo una matriz de dos por dos mueve el plano. Delta W, la letra griega delta mayúscula delante de W, es el cambio: acá, cero coma cinco, cero coma cinco; cero coma veinticinco, cero coma veinticinco. Es de rango uno: todo el plano cae sobre una sola recta, la dirección uno, cero coma cinco. Ahora W sub cero, los pesos de antes, que acá son la identidad, más ese cambio. El plano se deforma, pero solo empujado a lo largo de esa recta. Un cambio de rango uno lee una dirección de la entrada y escribe en una de la salida.

El cambio, como matriz

$W_0$: antes del ajuste · $W_1$: después · misma forma

SVD: capas de rango 1 ordenadas por $\sigma_i$

qué fracción del cambio carga cada capa

W sub cero son los pesos antes del ajuste, W sub uno los de después, de la misma forma. Delta W es la resta: en el dibujo, una columna de cada una. Con la SVD, singular value decomposition, de Álgebra lineal, episodio seis: cualquier matriz es una suma de capas de rango uno: una columna u sub i por una fila v sub i transpuesta, pesada por sigma sub i, su valor singular. En el dibujo, los seis sigma más grandes del cambio que vamos a medir hoy. La energía de cada capa es su sigma al cuadrado sobre la suma de todos. Rango intrínseco bajo quiere decir esto: pocas capas cargan casi toda la energía. En el dibujo, las ocho primeras suman cero coma noventa y dos.

Dos hipótesis parecidas

• Dimensión intrínseca: el ajuste entero cabe en pocos números
• Se mide sobre todos los parámetros a la vez
• Rango intrínseco: el cambio de cada matriz tiene pocas direcciones
• La segunda es la apuesta de LoRA

Dos ideas con nombres parecidos. La primera es de Aghajanyan y otros, en dos mil veinte: la dimensión intrínseca. Dice que todo el ajuste se puede hacer moviendo pocos números. Se mide sobre el modelo entero, todos los parámetros como un solo vector. La segunda es de Hu y otros, el paper de LoRA: el rango intrínseco. Dice que el cambio de cada matriz de pesos, por separado, tiene rango bajo. LoRA se apoya en la segunda. Hoy medimos las dos.

Parte 2 · Un fine-tuning de juguete

Un modelo chico que aprende una tarea, y después se ajusta a otra parecida. Corre en segundos.

El modelo base: dígitos escritos a mano

Los datos: dígitos escritos a mano de scikit-learn, imágenes de ocho por ocho. Una red con sesenta y cuatro entradas, una por píxel, y doscientas cincuenta y seis neuronas. Una capa del medio, de doscientos cincuenta y seis por doscientos cincuenta y seis: la que más vamos a mirar. Y diez salidas, un logit por dígito. Mil quinientos pasos de Adam, adaptive moment estimation: nuestro preentrenamiento. La tarea nueva: los mismos dígitos en espejo, dados vuelta de izquierda a derecha. Ochenta y cinco mil dos parámetros. Con dígitos normales acierta el noventa y siete coma cinco por ciento. En espejo, el cuarenta y tres coma tres: hay que ajustarla.

red = nn.Sequential(nn.Linear(64, 256), nn.ReLU(),     # 8×8 píxeles
                    nn.Linear(256, 256), nn.ReLU(),    # la capa del medio
                    nn.Linear(256, 10))                # un logit por dígito
print(f"parámetros: {sum(p.numel() for p in red.parameters()):,}")

entrenar(red, Xtr, ytr, pasos=1500)                    # "preentrenamiento"
print("dígitos normales: ", f"{precision(red, Xte, yte):.1%}")
print("dígitos en espejo:", f"{precision(red, espejo(Xte), yte):.1%}")

parámetros: 85,002
dígitos normales:  97.5%
dígitos en espejo: 43.3%

El ajuste completo, y la resta

El fine-tuning completo. Guardamos una copia de las tres matrices de pesos antes de tocar nada. Adam recibe todos los parámetros. Trescientos pasos con batches de sesenta y cuatro dígitos en espejo. Delta W: después menos antes, para cada matriz. Y qué tan grande es el cambio, con la norma: la raíz de la suma de todos los números al cuadrado. En espejo ahora acierta el noventa y tres coma cinco por ciento. Y el cambio no es chico: en la capa del medio, la norma de delta W es el treinta y cinco por ciento de la de W sub cero. ¿En cuántas direcciones?

red = cargar_base()                                    # la red preentrenada
W0 = [red[i].weight.detach().clone() for i in (0, 2, 4)]

opt = torch.optim.Adam(red.parameters(), lr=1e-3)      # ajuste completo
for paso in range(300):
    i = torch.randint(0, len(Xtr), (64,))              # batch de 64
    perdida = F.cross_entropy(red(espejo(Xtr[i])), ytr[i])
    opt.zero_grad(); perdida.backward(); opt.step()

dW = [red[i].weight.detach() - w0 for i, w0 in zip((0, 2, 4), W0)]
print("en espejo, después:", f"{precision(red, espejo(Xte), yte):.1%}")
for d, w0 in zip(dW, W0):
    print(tuple(d.shape), f"|ΔW|/|W0| = {d.norm() / w0.norm():.2f}")

en espejo, después: 93.5%
(256, 64) |ΔW|/|W0| = 0.30
(256, 256) |ΔW|/|W0| = 0.35
(10, 256) |ΔW|/|W0| = 0.18

Parte 3 · El espectro del cambio

La SVD ordena las direcciones del cambio de mayor a menor. Si la hipótesis es cierta, pocas cargan casi todo.

Los valores singulares de ΔW

Tomamos el delta W de la capa del medio, de doscientos cincuenta y seis por doscientos cincuenta y seis. svdvals da sus doscientos cincuenta y seis valores singulares, de mayor a menor. La energía de cada uno: su cuadrado, sobre la suma. Para comparar, una matriz al azar de la misma forma, sin ninguna estructura. Y cuánta energía juntan los primeros k. Sigma uno, tres coma ocho dos; sigma dos, dos coma cinco cuatro. El primero solo carga el treinta y ocho por ciento; los cuatro primeros, setenta y ocho; los ocho primeros, noventa y dos. La matriz al azar, con ocho, junta el once por ciento.

dW = c["dW"][1]                                        # capa del medio: 256 × 256
S = torch.linalg.svdvals(dW)                           # 256 valores singulares
energia = S**2 / (S**2).sum()                          # fracción de cada uno

azar = torch.randn(256, 256)                           # misma forma, sin estructura
Sa = torch.linalg.svdvals(azar)
energia_azar = Sa**2 / (Sa**2).sum()

print("sigma 1..6:", [round(s, 2) for s in S[:6].tolist()])
for k in [1, 4, 8, 32]:
    e, ea = energia[:k].sum(), energia_azar[:k].sum()
    print(f"top {k:2}: ΔW {e:.0%}   al azar {ea:.0%}")

sigma 1..6: [3.82, 2.54, 2.32, 1.97, 1.39, 1.22]
top  1: ΔW 38%   al azar 2%
top  4: ΔW 78%   al azar 6%
top  8: ΔW 92%   al azar 11%
top 32: ΔW 99%   al azar 38%

La energía de cada dirección

En barras, la energía de cada valor singular de delta W. El primero, treinta y ocho por ciento. El segundo, diecisiete. El tercero, catorce. El cuarto, diez. Y después caen rápido: cinco, cuatro, dos, dos. Las otras doscientas cuarenta y ocho direcciones, todas juntas, apenas el ocho por ciento. El cambio usa una matriz de doscientos cincuenta y seis por doscientos cincuenta y seis, pero casi todo pasa por ocho direcciones.

• σ₁: 38 %
• σ₂: 17 %
• σ₃: 14 %
• σ₄: 10 %
• σ₅: 5 %
• σ₆: 4 %
• σ₇: 2 %
• σ₈: 2 %
• los otros 248: 8 %

Energía acumulada: ΔW contra una matriz al azar

Acumulado: en horizontal, cuántas direcciones guardamos; en vertical, cuánta energía juntan. La punteada es la matriz al azar: sube casi en línea recta, porque todas sus direcciones pesan parecido. Con sesenta y cuatro, ni dos tercios. La del fine-tuning se dispara: con ocho, noventa y dos por ciento; con treinta y dos, noventa y nueve. Eso es rango intrínseco bajo, medido. Pero energía no es precisión: falta ver si alcanza para la tarea.

Quedarse con k direcciones

las k capas más grandes; el resto, apagado

dos matrices flacas: $256 \times k$ y $k \times 256$

Eckart–Young: el mejor rango k posible

La prueba de fuego: tirar todo menos k direcciones. Delta W sub k suma solo las primeras k capas. En el dibujo, los sigma reales con k igual a cuatro. Escrito como producto, son dos matrices flacas: U sub k, de doscientos cincuenta y seis por k, por Sigma V transpuesta, de k por doscientos cincuenta y seis. Esta forma, alta por ancha, es exactamente la de B por A en LoRA. Y la red usa W sub cero más ese cambio recortado. Por Eckart y Young, de Álgebra lineal, episodio seis, es la mejor aproximación de rango k. En el dibujo, lo que vamos a medir.

Recortar ΔW y medir

Una función que recorta: nos quedamos con las k primeras columnas de U, los k primeros sigma y las k primeras filas de V transpuesta. Probamos siete valores de k. En las tres matrices de pesos ponemos W sub cero más el cambio recortado; los sesgos quedan los del ajuste. Y medimos en espejo. Con k igual a cero, sin cambio en los pesos, cuarenta y tres coma seis por ciento. Con una dirección, cuarenta y seis. Con dos, sesenta y tres. Con cuatro, ochenta y seis coma seis. Y con ocho, noventa y tres coma cinco: lo mismo que el ajuste completo.

def rango_k(d, k):                                     # lo mejor de rango k
    U, S, Vt = torch.linalg.svd(d, full_matrices=False)
    return U[:, :k] @ torch.diag(S[:k]) @ Vt[:k]

res = {}
for k in [0, 1, 2, 4, 8, 16, 32]:
    red = cargar_base()
    red.load_state_dict({n: p for n, p in c["red"].items() if "bias" in n},
                        strict=False)                  # sesgos del ajuste
    for j, i in enumerate((0, 2, 4)):
        red[i].weight.data = W0[j] + rango_k(dW[j], k)
    res[k] = precision(red, espejo(Xte), yte)
    print(f"ΔW de rango {k:2}: {res[k]:.1%}")

ΔW de rango  0: 43.6%
ΔW de rango  1: 46.3%
ΔW de rango  2: 63.2%
ΔW de rango  4: 86.6%
ΔW de rango  8: 93.5%
ΔW de rango 16: 93.5%
ΔW de rango 32: 93.5%

Precisión según el rango del cambio

La línea verde es el ajuste completo: noventa y tres coma cinco. La azul es el cambio recortado a rango k. Sube rápido y en ocho la alcanza. En este juguete, un cambio de rango ocho en cada matriz hace todo lo que hizo el fine-tuning completo. En la capa del medio, ese cambio son cuatro mil noventa y seis números, contra sesenta y cinco mil quinientos treinta y seis.

verde punteada: el ajuste completo, 93,5 %

Hu et al., 2021 — LoRA: Low-Rank Adaptation of Large Language Models · arXiv:2106.09685

«intrinsic rank» · «Inspired by this, we hypothesize the updates to the weights also have a low “intrinsic rank” during adaptation.»

Es la hipótesis del paper de LoRA, sección cuatro. Inspirados por el trabajo de Aghajanyan, escriben: suponemos que los cambios de los pesos durante la adaptación también tienen un rango intrínseco bajo. Es una hipótesis, no un teorema.

Parte 4 · Dimensión intrínseca

La otra hipótesis: no por matriz, sino para todo el modelo a la vez. ¿Cuántos números hacen falta para el ajuste entero?

Entrenar en un subespacio al azar

$\theta$: los $D$ parámetros en fila · $z$: $d$ números · $P$: $D \times d$, al azar y fija

el gradiente solo puede moverse en d direcciones

la dimensión intrínseca

Theta es un vector con todos los parámetros del modelo uno detrás del otro: D números, ochenta y cinco mil dos en nuestro juguete. Theta sub cero, los preentrenados. En vez de entrenar theta, entrenamos z, un vector de solo d números, y lo llevamos al espacio grande con P, una matriz de D por d, al azar y fija. Por la regla de la cadena, el gradiente respecto de z es P transpuesta por el gradiente respecto de theta: el modelo solo se mueve en d direcciones. Y d noventa es el d más chico que alcanza el noventa por ciento del ajuste completo. En el dibujo, lo que medimos: con doscientos, ochenta y tres coma seis; con mil, noventa y cinco.

Ajustar moviendo solo d números

En código. Aplanamos todos los parámetros de la red base en un solo vector. P, la proyección al azar. z arranca en cero: al principio, el modelo es el base. Es lo único que se entrena. Cada forward arma theta a partir de z. functional call, de PyTorch, corre la red con esos parámetros en vez de los suyos. Y Adam recibe solo z. D, ochenta y cinco mil dos. Con diez números, cuarenta y seis coma nueve: casi nada. Con cincuenta, sesenta y cuatro. Con doscientos, ochenta y tres coma seis. Con mil, noventa y cinco coma dos.

red = cargar_base()
nombres = [n for n, _ in red.named_parameters()]
formas = [p.shape for p in red.parameters()]
tamanos = [p.numel() for p in red.parameters()]
theta0 = torch.cat([p.detach().flatten() for p in red.parameters()])
D = theta0.numel()                                     # todos los parámetros

def red_en_subespacio(d):
    torch.manual_seed(d)
    P = torch.randn(D, d) / D**0.5                     # proyección al azar, fija
    z = torch.zeros(d, requires_grad=True)             # lo único que se entrena
    def f(X):
        theta = theta0 + P @ z                         # θ = θ0 + P·z
        trozos = [t.view(s) for t, s in zip(theta.split(tamanos), formas)]
        return functional_call(red, dict(zip(nombres, trozos)), (X,))
    return f, z

print("D =", f"{D:,}")
for d in [10, 50, 200, 1000]:
    f, z = red_en_subespacio(d)
    entrenar(f, espejo(Xtr), ytr, pasos=300, lr=0.05, params=[z])
    print(f"d = {d:4}: {precision(f, espejo(Xte), yte):.1%}")

D = 85,002
d =   10: 46.9%
d =   50: 64.2%
d =  200: 83.6%
d = 1000: 95.2%

Precisión según la dimensión del subespacio

La verde, otra vez, el ajuste completo. La ámbar, el noventa por ciento de esa precisión: ochenta y cuatro coma dos. La violeta, entrenando solo z. La cruza entre doscientos y quinientos números, de ochenta y cinco mil. El serrucho del principio es ruido: cada d usa otra proyección al azar. La diferencia con LoRA: acá las direcciones son al azar; LoRA las aprende, matriz por matriz.

verde punteada: el ajuste completo, 93,5 %

Aghajanyan, Zettlemoyer y Gupta, 2020 — Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning · arXiv:2012.13255

200 parámetros · «by optimizing only 200 trainable parameters randomly projected back into the full space, we can tune a RoBERTa model to achieve 90% of the full parameter performance levels on MRPC.»

El experimento original. Con RoBERTa, un modelo de lenguaje de Facebook, y MRPC, una tarea que consiste en decidir si dos frases dicen lo mismo: entrenando solo doscientos parámetros, proyectados al azar al espacio completo, llegaron al noventa por ciento de la performance del ajuste completo. Que en nuestro juguete el número también ande por doscientos es casualidad. Y el paper encuentra que los modelos más grandes tienden a tener dimensión intrínseca más baja.

Lo que probamos y lo que no

• Probado: en este juguete, ΔW tiene rango intrínseco bajo
• Probado: rango 8 alcanza la precisión del ajuste completo
• No probado: que valga para toda tarea y todo modelo
• Se cree que en modelos grandes vale aún más

El alcance. Medimos que el cambio de un fine-tuning tiene casi toda su energía en ocho direcciones de doscientas cincuenta y seis. Y que recortarlo a rango ocho no pierde precisión. No que valga siempre: es una red chica y una tarea parecida a la original. El propio paper de LoRA avisa que no espera que un r chico funcione para cualquier tarea. Lo que la comunidad cree, apoyada en Aghajanyan, es que en modelos grandes y preentrenados la hipótesis funciona todavía mejor.

Resumiendo

• El cambio de un ajuste es una matriz: $\Delta W = W_1 - W_0$
• Rango intrínseco bajo: pocos $\sigma_i$ cargan casi toda la energía
• Juguete: 8 de 256 direcciones juntan el 92 %; al azar, el 11 %
• $W_0 + \Delta W_8$ acierta lo mismo que el ajuste completo (93,5 %)
• Dimensión intrínseca: ~200–500 números alcanzan el 90 %
• Próximo: la fórmula $W + \frac{\alpha}{r} BA$, por dentro

Papers del episodio

• Hu et al., 2021 — LoRA: Low-Rank Adaptation of Large Language Models
• Aghajanyan, Zettlemoyer y Gupta, 2020 — Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning