Después de AdamW

Optimizadores, con foco en AdamW. Episodio diez, el último: después de AdamW. Adam, adaptive moment estimation, guarda dos números por cada peso, y en un modelo grande eso es mucha memoria. Hoy, tres respuestas: Adafactor, que comprime v; Lion, que tira v y se queda con el signo; y Muon, que trata a cada matriz como matriz. Cada uno escrito a mano, en chico, y con su memoria medida.

Qué guarda AdamW por cada peso

• El peso y su gradiente
• m: el promedio del gradiente
• v: el promedio del gradiente al cuadrado
• Tres ideas: factorizar v, usar el signo, ortogonalizar

Durante el entrenamiento, cada peso ocupa su valor y su gradiente. AdamW suma m, del mismo tamaño que los pesos. Y v, otra vez del mismo tamaño. Una cuenta: un modelo de siete mil millones de pesos, con m y v en float32, de cuatro bytes, son siete mil millones por dos por cuatro: cincuenta y seis mil millones de bytes, solo de estado del optimizador. Las tres ideas de hoy atacan esto desde lugares distintos: guardar v en forma comprimida, no guardar v, o cambiar la forma del paso para las matrices.

Parte 1 · Adafactor: factorizar v

Para una matriz de n × m, en vez de n·m promedios, guardar n + m.

v de una matriz, con sumas de filas y columnas

Un gradiente de tres filas por cuatro columnas, con números redondos. V es lo que guardaría Adam: el cuadrado de cada elemento, doce números. Para simplificar, un solo paso, sin promedio móvil. R suma cada fila: tres números. C suma cada columna: cuatro. Y la reconstrucción de Adafactor: el producto exterior de R y C, dividido la suma total. Y contamos cuánto se guarda. R vale seis coma veinticinco, veinticinco y diez coma cinco. En la fila uno, el real es uno, cuatro, cero coma veinticinco, uno, y la reconstrucción, cero coma setenta y nueve, tres coma catorce, uno coma cincuenta y tres, cero coma setenta y nueve: el orden se respeta, los valores no del todo. La fila tres, que tiene otra forma, sale peor. De doce números, guardamos siete.

G = torch.tensor([[1.0, 2.0, 0.5, 1.0],
                  [2.0, 4.0, 1.0, 2.0],
                  [0.5, 1.0, 3.0, 0.5]])           # gradiente de una matriz 3 × 4
V = G ** 2                                         # lo que guardaría Adam: 12 números
R = V.sum(dim=1)                                   # suma de cada fila: 3 números
C = V.sum(dim=0)                                   # suma de cada columna: 4 números
V_hat = torch.outer(R, C) / R.sum()                # Adafactor: rango 1
print("R:", R.tolist(), " C:", C.tolist())
print("V real    fila 1:", V[0].tolist())
print("V factor. fila 1:", [round(x, 2) for x in V_hat[0].tolist()])
print("V real    fila 3:", V[2].tolist())
print("V factor. fila 3:", [round(x, 2) for x in V_hat[2].tolist()])
print(f"números guardados: {V.numel()} → {R.numel() + C.numel()}")

R: [6.25, 25.0, 10.5]  C: [5.25, 21.0, 10.25, 5.25]
V real    fila 1: [1.0, 4.0, 0.25, 1.0]
V factor. fila 1: [0.79, 3.14, 1.53, 0.79]
V real    fila 3: [0.25, 1.0, 9.0, 0.25]
V factor. fila 3: [1.32, 5.28, 2.58, 1.32]
números guardados: 12 → 7

Rango uno: una fila por una columna

$R$: un número por fila

$C_j$: la suma de la columna j

exacto solo si las filas son proporcionales (rango 1)

R sub i es la suma de la fila i de V. En el dibujo, la fila uno: uno más cuatro más cero coma veinticinco más uno, seis coma veinticinco. C sub j, lo mismo por columna. La reconstrucción: cada celda es R sub i por C sub j, dividido la suma de todas las R. En el dibujo, una columna de tres números por una fila de cuatro da una matriz de tres por cuatro: eso es un producto exterior, una matriz de rango uno. La celda uno dos da tres coma catorce. Para una matriz de n por m, Adam guarda n por m números; Adafactor, n más m. Mil por mil: un millón contra dos mil. El precio: si las filas de V no son proporcionales, la reconstrucción es aproximada.

Shazeer y Stern, 2018 — Adafactor: Adaptive Learning Rates with Sublinear Memory Cost · arXiv:1804.04235

n + m, no n·m · «we propose maintaining only the per-row and per-column sums of these moving averages» … «requiring space proportional to n + m rather than nm»

Es la propuesta de Shazeer y Stern, de dos mil dieciocho. Guardan solo las sumas por fila y por columna de los promedios móviles, y estiman cada v a partir de ellas: espacio proporcional a n más m en lugar de n por m. Además apagan el momentum, poniendo beta uno en cero, así que tampoco hay m. PyTorch lo trae como torch punto optim punto Adafactor.

Parte 2 · Lion: solo el signo

Un optimizador encontrado por una búsqueda automática de programas. Guarda un solo promedio.

Lion a mano

Lion, por EvoLved Sign Momentum, de Chen y coautores. Lo escribimos como un torch punto optim punto Optimizer, con los defaults del paper. Primero, c: una mezcla con beta uno, cero coma nueve, del momentum guardado y el gradiente nuevo. El weight decay, desacoplado, como en AdamW. El paso: eta por el signo de c. No hay división ni raíz: cada peso se mueve exactamente eta, para un lado o para el otro. Y el momentum se actualiza con otro beta, cero coma nueve nueve. Es el único estado: un número por peso, contra los dos de AdamW.

Cada línea, su término:

• c = b1 * m + (1 - b1) * p.grad → $c = \beta_1\, m + (1 - \beta_1)\, g$ — $\beta_1 = 0{,}9$: mezcla momentum y gradiente nuevo
• p.mul_(1 - grupo["lr"] * grupo["wd"]) → $\theta \leftarrow (1 - \eta\lambda)\,\theta$ — el decay de AdamW, igual
• p.sub_(grupo["lr"] * torch.sign(c)) → $\theta \leftarrow \theta - \eta\,\mathrm{signo}(c)$ — cada peso se mueve exactamente η
• m.mul_(b2).add_((1 - b2) * p.grad) → $m \leftarrow \beta_2\, m + (1 - \beta_2)\, g$ — $\beta_2 = 0{,}99$: el único estado

class Lion(torch.optim.Optimizer):
    def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0):
        super().__init__(params, dict(lr=lr, betas=betas, wd=weight_decay))

    @torch.no_grad()
    def step(self):
        for grupo in self.param_groups:
            b1, b2 = grupo["betas"]
            for p in grupo["params"]:
                m = self.state[p].setdefault("m", torch.zeros_like(p))
                c = b1 * m + (1 - b1) * p.grad          # mezclar g y m
                p.mul_(1 - grupo["lr"] * grupo["wd"])   # decay desacoplado
                p.sub_(grupo["lr"] * torch.sign(c))     # solo el signo
                m.mul_(b2).add_((1 - b2) * p.grad)      # momentum, con β₂

Dos pasos de Lion, con números

m arranca en 0; g = (0,5; −4); η = 0,1

$1 - \beta_2 = 0{,}01$: m se mueve despacio

g₁ = −0,2 cambió de signo: θ₁ sube 0,1

Theta vale uno y menos dos, el gradiente cero coma cinco y menos cuatro, y m arranca en cero. Entonces c es cero coma uno por g. Su signo es uno y menos uno, y theta se mueve cero coma uno en cada coordenada. En el dibujo, los c son ocho veces distintos y los pasos, iguales. m guarda solo el uno por ciento del gradiente: cero coma cero cero cinco y menos cero coma cero cuatro. En el segundo paso, el gradiente de theta uno cambió de signo: menos cero coma dos. c sub uno da menos cero coma cero uno cinco cinco. En el dibujo, el gradiente nuevo gana, el signo se da vuelta y theta uno vuelve a uno.

Lion contra AdamW, en el valle

El valle del episodio uno, muy plano en x y empinado en y, desde el mismo punto de partida. Cien pasos: AdamW con learning rate cero coma uno, con su weight decay por defecto, y Lion con la mitad, porque su paso es más largo. Medimos el largo de cada paso en x. AdamW termina en cero coma cero cuatro, cero coma cero uno, casi en el mínimo; su paso en x empieza en cero coma uno cero cuatro y termina en dos milésimas. Lion termina en cero coma cinco: se pasó. Su paso en x vale cero coma cero cinco al principio y al final: el signo no sabe que está cerca.

def f(p): return 0.1 * p[0] ** 2 + 3 * p[1] ** 2          # el valle del episodio uno

def camino(Clase, lr, pasos=100):
    p = torch.tensor([-4.5, 1.5], requires_grad=True)
    opt = Clase([p], lr=lr)
    puntos = [p.tolist()]
    for _ in range(pasos):
        opt.zero_grad(); f(p).backward(); opt.step()
        puntos.append(p.tolist())
    return np.array(puntos)
cam = {"AdamW": camino(torch.optim.AdamW, 0.1), "Lion": camino(Lion, 0.05)}
for k, c in cam.items():
    dx = np.abs(np.diff(c, axis=0))                        # largo de cada paso
    print(f"{k:5}  final ({c[-1, 0]:5.2f}, {c[-1, 1]:5.2f})"
          f"  pasos en x: {dx[0, 0]:.3f} → {dx[-1, 0]:.3f}")

AdamW  final ( 0.04,  0.01)  pasos en x: 0.104 → 0.002
Lion   final ( 0.50,  0.00)  pasos en x: 0.050 → 0.050

Cien pasos en el valle

AdamW: baja en diagonal, cruza el valle, y frena al acercarse a la estrella, porque m sombrero se achica antes que la raíz de v. Lion: avanza en escalones de cero coma cero cinco en cada eje, a ritmo constante, y cruza el mínimo sin frenar. En y también se pasa de largo al principio: el momentum tarda en cambiar de signo. Por eso Lion se usa con un schedule que baje el learning rate al final, como el coseno del episodio ocho: el que frena es el schedule, no el optimizador.

Chen et al., 2023 — Symbolic Discovery of Optimization Algorithms · arXiv:2302.06675

solo momentum · «It is more memory-efficient than Adam as it only keeps track of the momentum.» … «It also requires a smaller learning rate than Adam due to the larger norm of the update produced by the sign function.»

El paper de Lion es de Chen y coautores, en dos mil veintitrés. Dicen que es más eficiente en memoria que Adam porque solo guarda el momentum, y que requiere un learning rate más chico que Adam, por la norma más grande del paso que produce la función signo. Es lo que vimos: un paso de largo eta en todas las coordenadas.

Parte 3 · Muon: ortogonalizar

Para los pesos que son matrices: igualar el tamaño del paso en todas sus direcciones.

Newton–Schulz, en cinco pasos

Muon, por MomentUm Orthogonalized by Newton-Schulz, es un método reciente, publicado en el blog de Keller Jordan en diciembre de dos mil veinticuatro, y ya está en torch punto optim. Toma el momentum de una matriz y, antes de usarlo como paso, lo ortogonaliza. Los tres coeficientes del blog. Dividimos por la norma para que los valores singulares queden por debajo de uno. Cinco veces, un polinomio de grado cinco en X. Medimos los valores singulares antes y después. Y comparamos con U por V transpuesta, de la SVD, singular value decomposition. G tenía valores singulares veintiuno coma uno, cuatro coma siete cuatro, uno coma siete uno y cero coma uno tres. Después: uno, cero coma nueve nueve, cero coma siete tres, cero coma siete. El coseno con U por V transpuesta es cero coma nueve ocho siete.

def newton_schulz(G, pasos=5):                       # del blog de Keller Jordan
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G / G.norm()                                  # valores singulares ≤ 1
    for _ in range(pasos):
        A = X @ X.T
        X = a * X + (b * A + c * A @ A) @ X           # polinomio de grado 5
    return X

G = torch.randn(4, 6) * torch.tensor([[10.0], [3.0], [1.0], [0.1]])
sv = lambda M: [round(x, 2) for x in torch.linalg.svdvals(M).tolist()]
print("valores singulares de G:  ", sv(G))
print("después de Newton–Schulz: ", sv(newton_schulz(G)))
U, S, Vt = torch.linalg.svd(G, full_matrices=False)
cos = torch.nn.functional.cosine_similarity(
    newton_schulz(G).flatten(), (U @ Vt).flatten(), dim=0)
print(f"coseno con U·Vᵀ: {cos:.3f}")

valores singulares de G:   [21.1, 4.74, 1.71, 0.13]
después de Newton–Schulz:  [1.0, 0.99, 0.73, 0.7]
coseno con U·Vᵀ: 0.987

Qué hace la ortogonalización

Álgebra, ep. 6: rotar, estirar, rotar

mismas direcciones, todos los estiramientos en 1

actúa sobre cada valor singular por separado

La SVD, de Álgebra lineal para LLMs, episodio seis: toda matriz es una rotación U, un estiramiento S y otra rotación V transpuesta. S es diagonal, con los valores singulares. En el dibujo, los de G: el primero es ciento sesenta veces el último. Ortogonalizar es reemplazar S por unos: U por V transpuesta. Mismas direcciones, todas del mismo tamaño. En el dibujo, lo que dio Newton-Schulz: no exactamente uno, entre cero coma siete y uno. A es X por X transpuesta, y cada término del polinomio actúa sobre cada valor singular por separado: sigma pasa a a sigma más b sigma al cubo más c sigma a la quinta. En el dibujo, un paso: cero coma uno sube a cero coma tres cuatro, cero coma cinco a uno coma diecinueve, y uno baja a cero coma setenta.

Dónde va Muon y dónde no

• Solo matrices 2D de las hidden layers
• Embeddings y la capa de salida: AdamW
• Sesgos y escalas: AdamW
• Un número de estado por peso: el momentum

Las reglas de uso, del mismo blog. Muon es para los parámetros de dos dimensiones de las hidden layers: las matrices de attention y de las redes densas. Para los embeddings y la capa final, que convierte en logits, el blog recomienda AdamW. Lo mismo para los escalares y vectores, como sesgos y escalas de normalización. Muon y AdamW conviven, cada uno con su grupo de parámetros, como en el episodio siete. Y en memoria, Muon guarda solo el momentum: uno por peso.

Parte 4 · La memoria, medida

Un millón de pesos, seis optimizadores, y el estado que cada uno guarda después de un paso.

Contar el estado de cada optimizador

Le preguntamos a cada uno qué guardó. Seis optimizadores: SGD, stochastic gradient descent, sin y con momentum, AdamW, nuestro Lion, y los Muon y Adafactor de PyTorch. Una matriz de mil por mil: un millón de pesos. Un paso con un gradiente al azar, porque el estado se crea en el primer step. Y contamos los números de todos los tensores del estado. SGD sin momentum, nada. Con momentum, un millón: uno por peso. AdamW, dos millones. Lion y Muon, un millón cada uno. Y Adafactor, dos mil: dos milésimas por peso, una fila y una columna.

opciones = [("SGD", torch.optim.SGD, dict(lr=0.1)),
            ("SGD + momentum", torch.optim.SGD, dict(lr=0.1, momentum=0.9)),
            ("AdamW", torch.optim.AdamW, {}),
            ("Lion", Lion, {}),
            ("Muon", torch.optim.Muon, {}),
            ("Adafactor", torch.optim.Adafactor, {})]
res = {}
for nombre, Clase, kw in opciones:
    W = torch.nn.Parameter(torch.randn(1000, 1000))   # un millón de pesos
    opt = Clase([W], **kw)
    W.grad = torch.randn(1000, 1000); opt.step()      # el estado nace aquí
    n = sum(v.numel() for v in opt.state[W].values()
            if torch.is_tensor(v) and v.dim() > 0)    # tensores guardados
    res[nombre] = n / W.numel()
    print(f"{nombre:15} {n:9,} números  →  {n / W.numel():.3f} por peso")

SGD                     0 números  →  0.000 por peso
SGD + momentum  1,000,000 números  →  1.000 por peso
AdamW           2,000,000 números  →  2.000 por peso
Lion            1,000,000 números  →  1.000 por peso
Muon            1,000,000 números  →  1.000 por peso
Adafactor           2,000 números  →  0.002 por peso

Números de estado por peso (matriz de 1000 × 1000)

AdamW, dos números por peso. SGD con momentum, Lion y Muon, uno. Adafactor, casi nada, porque guarda una fila y una columna por matriz, sin momentum. Y SGD a secas, cero. Ahorrar estado tiene precio: Lion pide otro learning rate, Adafactor aproxima v, y Muon solo sirve para matrices. Por eso AdamW sigue siendo el punto de partida.

• AdamW: 2
• SGD + momentum: 1
• Lion: 1
• Muon: 1
• Adafactor: 0,002
• SGD: 0

Qué cambia cada uno

AdamW guarda m y v, normaliza el paso de cada peso y aplica el decay aparte. Adafactor guarda v comprimida en una fila y una columna. Lion guarda solo m y usa el signo. Muon guarda m de cada matriz y la ortogonaliza antes de dar el paso.

Resumiendo

• AdamW: 2 números de estado por peso
• Adafactor: $\hat V = R\,C^\top/\sum R$, de $n \cdot m$ a $n + m$
• Lion: $\theta \leftarrow \theta - \eta\,\mathrm{signo}(c)$, solo m
• Muon: el momentum de cada matriz, ortogonalizado
• Serie: del paso de descenso a AdamW, todo a mano
• Seguí con: «LoRA», entrenar con pocos parámetros

Papers del episodio

• Shazeer y Stern, 2018 — Adafactor: Adaptive Learning Rates with Sublinear Memory Cost
• Chen et al., 2023 — Symbolic Discovery of Optimization Algorithms