La fórmula

LoRA, low-rank adaptation. Episodio tres: la fórmula. En el episodio dos vimos que el cambio de un fine-tuning cabe en pocas direcciones. LoRA lo entrena directamente en esa forma: W más alfa sobre r por B por A. Hoy la leemos pieza por pieza: cómo se calcula, por qué B arranca en cero, y qué hace de verdad el alfa sobre r.

Parte 1 · La fórmula, con números chicos

Tres dimensiones y rango uno. Todo se puede hacer a mano.

LoRA con una entrada de tres números

x, la entrada, es un vector de tres números: uno, dos, tres. W sub cero, los pesos congelados: acá la identidad, que devuelve x igual. A es una matriz de una fila por tres columnas: uno, cero, uno. B, de tres filas por una columna: cero coma cinco, cero, menos cero coma cinco. Alfa y r valen uno, así que la escala s es uno. Primero A por x: un solo número. Después B por ese número, por s: tres números, el cambio. A por x da cuatro: uno más tres. El cambio es dos, cero, menos dos. Y h, la salida, es tres, dos, uno: el cambio dio vuelta el vector. Lo mismo que armar la matriz W sub cero más s B A y multiplicar.

x = torch.tensor([1.0, 2.0, 3.0])          # la entrada: 3 números
W0 = torch.eye(3)                          # pesos congelados: identidad
A = torch.tensor([[1.0, 0.0, 1.0]])        # r × k = 1 × 3
B = torch.tensor([[0.5], [0.0], [-0.5]])   # d × r = 3 × 1
alfa, r = 1.0, 1
s = alfa / r                               # la escala

Ax = A @ x                                 # r números: lo que lee
cambio = s * B @ Ax                        # d números: lo que escribe
h = W0 @ x + cambio
print("A x      =", Ax.tolist())
print("s B (Ax) =", cambio.tolist())
print("h        =", h.tolist())
print("igual con (W0 + s BA) x:", torch.equal(h, (W0 + s * B @ A) @ x))

A x      = [4.0]
s B (Ax) = [2.0, 0.0, -2.0]
h        = [3.0, 2.0, 1.0]
igual con (W0 + s BA) x: True

Lo que calculó el código

A lee: proyecta x en r direcciones

B escribe: reparte esos r números en las d salidas

$s = \alpha / r$

A por x. A tiene una fila: es un producto punto, uno por uno más cero por dos más uno por tres. Da cuatro. En el dibujo, la fila de A recorre la columna x. A lee: resume la entrada en r números. B por ese número. B tiene una columna: cada salida es un solo producto. En el dibujo, B escribe: reparte lo que leyó A en las d salidas. Y h, la salida, es lo de siempre, W sub cero por x, más el cambio. En el dibujo, las dos columnas sumadas. s es alfa sobre r: alfa, la letra griega, es un número que elegís vos.

LoRA es un cuello de botella al costado

Otra forma de verlo: como una red chiquita al costado de W sub cero. Las tres entradas. Una capa del medio con r neuronas, acá una sola, sin activación: es A por x, el cuatro. Todo lo que el cambio sabe de la entrada pasa por ese cuello de botella. Y las tres salidas, B por ese número: dos, cero, menos dos, que se suman a lo que da W sub cero. Los colores son esquemáticos.

Con la q_proj real de Gemma

Ahora con una matriz de verdad. La proyección de queries de la capa doce de Gemma, leída del archivo de pesos. A, de ocho por dos mil trescientos cuatro, al azar, dividida por la raíz de k para que A por x no sea enorme. B, de dos mil cuarenta y ocho por ocho, en cero. El forward. Las filas de x son tokens, por eso va x por W transpuesta, la convención de PyTorch. Y el paréntesis importa: primero x por A transpuesta, ocho números por token; después por B transpuesta. Contamos multiplicaciones. Las formas, las esperadas. La diferencia con el modelo original: cero exacto. Y por token, W sub cero cuesta cuatro millones setecientas mil multiplicaciones; el camino de LoRA, treinta y cuatro mil ochocientas.

W0 = cargar("model.layers.12.self_attn.q_proj.weight")   # real, congelada
dout, k = W0.shape                                       # 2048 × 2304
r, alfa = 8, 16
A = torch.randn(r, k) / k**0.5                           # r × k, al azar
B = torch.zeros(dout, r)                                 # d × r, en cero
x = torch.randn(4, k)                                    # 4 tokens

h0 = x @ W0.T                                            # sin LoRA
h = x @ W0.T + (alfa / r) * (x @ A.T) @ B.T              # con LoRA
print("W0", tuple(W0.shape), " A", tuple(A.shape), " B", tuple(B.shape))
print("diferencia máxima con B = 0:", (h - h0).abs().max().item())
print(f"multiplicaciones por token: W0 {dout * k:,}"
      f" | LoRA {r * (k + dout):,}")

W0 (2048, 2304)  A (8, 2304)  B (2048, 8)
diferencia máxima con B = 0: 0.0
multiplicaciones por token: W0 4,718,592 | LoRA 34,816

Las formas de q_proj, con r = 8

8 × 2304 = 18.432 multiplicaciones

2048 × 8 = 16.384 multiplicaciones

del tamaño de W: no hace falta armarla para entrenar

Las cuentas con las formas reales. A por x: ocho filas de dos mil trescientos cuatro. Dieciocho mil cuatrocientas treinta y dos multiplicaciones, y salen ocho números. B por eso: dos mil cuarenta y ocho filas de ocho. Dieciséis mil trescientas ochenta y cuatro. Sumadas, las treinta y cuatro mil ochocientas del código. Si en cambio armáramos B por A, sería una matriz del tamaño de W: en el dibujo, cada celda es la fila de B por la columna de A, apenas r productos, pero son cuatro millones setecientas mil celdas. Por eso, mientras se entrena, LoRA nunca la arma. Armarla tiene sentido una sola vez, al final, para fusionarla con W: episodio siete.

Parte 2 · Arrancar sin cambiar nada

B en cero hace que el modelo ajustado empiece idéntico al original. ¿Y por qué no A en cero también?

Los gradientes de A y de B

$\delta$: d números · cuánto mueve la pérdida cada salida

d × r: producto exterior

r × k: pasa por B

el primer paso mueve solo B

Para saber qué se mueve hay que mirar los gradientes. Delta, la letra griega minúscula, es el gradiente de la pérdida L respecto de la salida h: d números, cuánto subiría L si sube cada salida. En el dibujo, un esquema con tres. Por la regla de la cadena, de Álgebra lineal, episodio nueve: el gradiente respecto de B es s por delta, una columna, por A x transpuesta, una fila. Un producto exterior. No depende de B: si A no es cero, B recibe gradiente. El de A es s por B transpuesta por delta, por x transpuesta. Pasa por B: si B es cero, A no recibe nada en el primer paso. En el dibujo, el valor real que vamos a medir: cero para A, tres coma ochenta y nueve para B.

Tres formas de inicializar

Lo probamos con matrices chicas. Antes de entrenar, medimos si B por A ya es distinto de cero. En el primer paso, guardamos el tamaño de los dos gradientes, y entrenamos diez pasos. La de LoRA: A al azar, B en cero. Las dos en cero. Las dos al azar. La de LoRA arranca en cero, en el primer paso solo B tiene gradiente, y al final el cambio creció a cero coma noventa y ocho. Con las dos en cero, todo queda en cero para siempre: nadie recibe gradiente, porque cada uno depende del otro. Y con las dos al azar aprende, pero el modelo arranca ya cambiado: cuatro coma sesenta y ocho de ruido.

def probar(nombre, A, B, pasos=10):
    inicio = (B @ A).norm().item()                   # ¿arranca cambiado?
    A.requires_grad_(); B.requires_grad_()
    opt = torch.optim.Adam([A, B], lr=0.01)
    for paso in range(pasos):
        h = x @ W0.T + (x @ A.T) @ B.T               # forward con LoRA
        perdida = ((h - objetivo) ** 2).mean()
        opt.zero_grad(); perdida.backward()
        if paso == 0:
            gA, gB = A.grad.norm().item(), B.grad.norm().item()
        opt.step()
    fin = (B @ A).norm().item()
    print(f"{nombre}: |BA| {inicio:.2f} → {fin:.2f}"
          f" | paso 1: gA={gA:.2f} gB={gB:.2f}")

probar("A azar, B cero", torch.randn(r, k), torch.zeros(d, r))
probar("A cero, B cero", torch.zeros(r, k), torch.zeros(d, r))
probar("A azar, B azar", torch.randn(r, k), torch.randn(d, r))

A azar, B cero: |BA| 0.00 → 0.98 | paso 1: gA=0.00 gB=3.89
A cero, B cero: |BA| 0.00 → 0.00 | paso 1: gA=0.00 gB=0.00
A azar, B azar: |BA| 4.68 → 4.00 | paso 1: gA=2.16 gB=3.06

Qué pasa con cada inicialización

En tabla. La elección de LoRA: arranca idéntico al modelo original y aprende. En el primer paso se mueve solo B; desde el segundo, como B ya no es cero, también A. Todo en cero es una trampa: el gradiente de A pasa por B y el de B pasa por A. Todo al azar aprende, pero le suma ruido al modelo antes de empezar. Por eso una de las dos va en cero y la otra no. Se podría hacer al revés, A en cero y B al azar; el paper eligió B.

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

B = 0 · «We use a random Gaussian initialization for A and zero for B, so ∆W = BA is zero at the beginning of training.»

Es una línea del paper, sección cuatro punto uno. Inicialización gaussiana al azar para A, y cero para B, así que el cambio, B por A, es cero al empezar el entrenamiento. Gaussiana quiere decir sacada de una campana de Gauss, como torch punto randn.

Parte 3 · Por qué alfa sobre r

Alfa es un número fijo; r, el rango. Dividir por r cambia qué pasa cuando probás otro rango.

Qué controla la escala

un promedio de r capas de rango 1, no una suma

no depende de $s$: entonces $s$ multiplica el paso, como $\eta$

más capas que suman

dividir por r corrige de más

B por A es una suma de r capas de rango uno: columna j de B por fila j de A, como vimos en Álgebra lineal, episodio tres. Dividir por r la convierte en un promedio. En el dibujo, cuatro capas, cada una pesando un cuarto. Y alfa, con Adam, adaptive moment estimation, hace de learning rate. En el primer paso, A todavía no se mueve, y Adam mueve cada número de B lo mismo, eta, para un lado o para el otro, sin importar el tamaño del gradiente. Entonces s multiplica el cambio igual que eta. En el dibujo, medido: s dos con eta da lo mismo que s uno con el doble de eta. Sin escala, el cambio crece con r. Con alfa sobre r, baja.

Medir el primer paso según r

El experimento detrás de esos números. Una capa de doscientos cincuenta y seis por doscientos cincuenta y seis con LoRA de rango r y escala s, y B en cero, como en el paper. Un solo paso de Adam, con el mismo learning rate para todo r, y medimos cuánto cambió la salida. Primero, la equivalencia: s por eta es lo que importa. Después, sin escala y con alfa sobre r, alfa ocho. s dos y el doble de learning rate: cero coma cero nueve cero uno las dos. Sin escala, el cambio va de cero coma cero catorce con r uno a cero coma ciento cuarenta y seis con r sesenta y cuatro: diez veces más. Con alfa sobre r, de cero coma once a cero coma cero dieciocho: seis veces menos.

Cada línea, su término:

• def primer_paso(r, s, lr=1e-3) → $h = W_0 x + s\, B(Ax)$ — LoRA de rango r con escala s
• opt.step() → $\Delta B = -\eta\, \text{signo}(\nabla_B L)$ — un paso de Adam
• s=1, lr × 2 → $s \cdot \eta = 2 \cdot 10^{-3}$ — s y η se multiplican
• primer_paso(r, alfa / r) → $s = \alpha / r,\;\; \alpha = 8$ — sin escala contra α/r

def primer_paso(r, s, lr=1e-3):                     # s: la escala de BA
    torch.manual_seed(1)
    A = (torch.randn(r, k) / k**0.5).requires_grad_()
    B = torch.zeros(d, r, requires_grad=True)
    opt = torch.optim.Adam([A, B], lr=lr)
    h = X @ W0.T + s * (X @ A.T) @ B.T
    ((h - Y) ** 2).mean().backward()
    opt.step()                                      # un paso de Adam
    return (s * X @ (B @ A).T).norm(dim=1).mean().item()   # |ΔW x|

print(f"r=8, s=2, lr     : {primer_paso(8, 2.0):.4f}")
print(f"r=8, s=1, lr × 2 : {primer_paso(8, 1.0, lr=2e-3):.4f}")
alfa = 8
for r in [1, 4, 16, 64]:
    sin, con = primer_paso(r, 1.0), primer_paso(r, alfa / r)
    print(f"r={r:2}  sin escala {sin:.3f}   con α/r {con:.3f}")

r=8, s=2, lr     : 0.0901
r=8, s=1, lr × 2 : 0.0901
r= 1  sin escala 0.014   con α/r 0.110
r= 4  sin escala 0.032   con α/r 0.063
r=16  sin escala 0.065   con α/r 0.032
r=64  sin escala 0.146   con α/r 0.018

El tamaño del primer paso, según r

En curvas. Sin escala, el cambio crece con r, más o menos como la raíz de r: duplicar el rango es como subir el learning rate. Con alfa sobre r baja, también más o menos como uno sobre raíz de r: corrige de más. Las dos se cruzan en r igual a alfa, ocho, donde s vale uno. Lo que el paper buscaba es práctico: elegís alfa una vez, y al cambiar r no hace falta reajustar todo. Que la corrección justa sea dividir por la raíz de r, y no por r, es la idea de rsLoRA, rank-stabilized LoRA, una de las variantes del episodio nueve.

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

α / r · «This scaling helps to reduce the need to retune hyperparameters when we vary r.»

La justificación del paper, en la misma sección. Con Adam, ajustar alfa es más o menos lo mismo que ajustar el learning rate. Por eso fijan alfa en el primer r que prueban y no lo tocan más; esa escala ayuda a no tener que reajustar hiperparámetros cuando cambia r. Hiperparámetros son los números que elegís vos y no aprende el modelo: r, alfa, el learning rate. En peft, alfa se llama lora alpha: episodio seis.

Parte 4 · La fórmula, entera

Ya tenemos todas las piezas. Las juntamos.

El forward, completo

El recorrido entero de una capa con LoRA. Entra x. El camino de LoRA: A lee r números, B los lleva a d, y alfa sobre r los escala. Se le suma el camino original, W sub cero por x, como siempre, y sin gradiente para W. Y sale h. Al final del entrenamiento, el camino de LoRA se puede fundir con W sub cero y desaparece.

x → A x: r números → B (A x): d números → × α/r → + W₀ x (congelada) → h

Las cuatro piezas

• $h = W_0 x + \frac{\alpha}{r} B (A x)$: $A$ lee, $B$ escribe
• $W_0$ congelada; solo $A$ ($r \times k$) y $B$ ($d \times r$) entrenan
• $B = 0$: arranca igual; $A$ al azar: hay gradiente
• $\alpha / r$: $\alpha$ hace de learning rate; $1/r$ promedia

Juntemos todo. La salida es la de siempre más un camino lateral: A lee r números de la entrada, B los escribe en las d salidas, y alfa sobre r escala. W sub cero no se toca; se entrenan A, de r por k, y B, de d por r. B en cero hace que el primer paso parta del modelo original; A al azar hace que haya gradiente. Y alfa sobre r: alfa funciona como un learning rate extra, y el uno sobre r promedia las r capas.

Resumiendo

• Forward: $W_0 x + \frac{\alpha}{r} B(Ax)$, sin armar $BA$
• q_proj, r = 8: 34.816 multiplicaciones por token contra 4.718.592
• $B = 0$: salida idéntica (diferencia 0,0) y $\nabla_A L = 0$ en el paso 1
• Todo en cero no aprende nunca; todo al azar arranca con ruido
• $\alpha/r$: $\alpha$ hace de learning rate; dividir por $r$ corrige de más
• Próximo: dónde ponerlo y con qué rango

Papers del episodio

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