El router

Mixture of Experts. Episodio dos: el router. En el episodio uno vimos que una capa MoE, mixture of experts, tiene muchos expertos y usa pocos por token. Hoy abrimos la pieza que decide cuáles: el router. Lo hacemos a mano con números chicos, después lo miramos con autograd, y al final corremos el router de un modelo real.

Qué tiene que hacer un router

• Puntuar los E expertos para cada token
• Quedarse con los k mejores
• Dar un peso a cada uno de esos k
• Ser entrenable con gradientes

Antes del código, el pedido. Para cada token, el router tiene que darle un puntaje a cada uno de los E expertos. Quedarse con los k de mayor puntaje: eso es top-k. Decidir cuánto pesa cada uno de esos k en la suma final. Y todo eso tiene que poder aprenderse con gradientes, junto con el resto del modelo. La cuarta condición es la delicada, porque elegir los k mejores no es una operación suave.

Parte 1 · Del token a los pesos

Un router es una sola matriz y tres operaciones. Las hacemos a mano con cuatro expertos.

El camino de un token

El recorrido completo. El token entra como un vector x de d números, se multiplica por la matriz del router, W sub r, y sale un número por experto: los logits. Un softmax los convierte en probabilidades. Nos quedamos con las k más altas y las renormalizamos para que sumen uno: esos son los pesos g. Y la salida de la capa es la suma de los k expertos elegidos, cada uno multiplicado por su peso.

x: d números → x · W_r → logits: E → softmax → top-k → renormalizar → g: k pesos → Σ g_i · E_i(x)

Router, paso a paso

un logit por experto

softmax: positivos y suman 1

$k = 2$: $\mathcal{T} = \{2, 1\}$

x es un token de tres números: uno, cero coma cinco y menos uno. W sub r es una matriz de tres por cuatro: una columna por experto. El logit del experto i es x por la columna i. En el dibujo, la columna dos da cero coma cinco más cero más uno: uno coma cinco. El softmax eleva e a cada logit y divide por la suma de todas esas exponenciales. Así todos quedan positivos y suman uno. En el dibujo, el experto dos se lleva el cuarenta y siete por ciento. T es el conjunto de los k elegidos, acá el dos y el uno. Dividimos sus probabilidades por la suma de las dos: en el dibujo, los no elegidos caen a cero y quedan sesenta y dos contra treinta y ocho.

Lo mismo en torch

Las mismas cuentas, en seis líneas. El token. Los logits son un producto de vector por matriz. softmax sobre la última dimensión. topk devuelve los dos valores más altos y en qué posición estaban. Y renormalizamos dividiendo por su suma. Logits cero coma cinco, uno, uno coma cinco y menos cero coma cinco. Probabilidades de diecisiete, veintinueve, cuarenta y siete y seis por ciento. Elegidos: el dos y el uno. Y los pesos finales: cero coma sesenta y dos y cero coma treinta y ocho.

x = torch.tensor([1.0, 0.5, -1.0])            # un token, d = 3
W_r = torch.tensor([[1.0, 0.0, 0.5, -1.0],    # d × E = 3 × 4
                    [0.0, 2.0, 0.0, 1.0],
                    [0.5, 0.0, -1.0, 0.0]])
logits = x @ W_r                              # un número por experto
probs = logits.softmax(-1)                    # suman 1
top = probs.topk(2)                           # k = 2
g = top.values / top.values.sum()             # renormalizar: suman 1
print("logits:", logits)
print("probs: ", probs)
print("top-2: ", top.indices.tolist(), top.values)
print("g:     ", g)

logits: tensor([ 0.50,  1.00,  1.50, -0.50])
probs:  tensor([0.17, 0.29, 0.47, 0.06])
top-2:  [2, 1] tensor([0.47, 0.29])
g:      tensor([0.62, 0.38])

Renormalizar = softmax de los elegidos

reemplazamos cada $p$ por su definición

$Z$ arriba y abajo: se cancela

la cuenta del snippet

Un detalle que conviene ver, porque los modelos reales lo escriben de la otra forma. Cada p es e a la ele dividido Z, donde Z es la suma de las exponenciales de todos los expertos. Reemplazamos arriba y abajo. Z aparece dividiendo arriba y abajo, así que se cancela. Queda un softmax calculado solo sobre los logits elegidos. En el dibujo, softmax de uno coma cinco y uno: sesenta y dos y treinta y ocho, lo mismo que antes. Y con los números del snippet: cero coma cuarenta y siete dividido la suma, cero coma setenta y seis, da cero coma sesenta y dos. Por eso hay código que primero hace top-k sobre los logits y después el softmax: es la misma cuenta.

Combinar: solo se calculan los elegidos

Ahora la salida. Cuatro expertos, cada uno un MLP, un multilayer perceptron, de tres a cuatro a tres. La suma recorre solo los dos elegidos: el experto dos con peso cero coma sesenta y dos, y el uno con cero coma treinta y ocho. Los otros dos no se ejecutan. Para comparar, armamos el vector de pesos completo: scatter pone cada g en su posición y deja ceros en el resto. Y sumamos los cuatro expertos, con esos ceros. Los dos resultados son idénticos. El MoE es la fórmula de la suma completa, pero sabiendo de antemano qué términos valen cero y salteándolos.

Cada línea, su término:

• expertos = [Experto(d=3, h=4) → $E_i : \mathbb{R}^3 \to \mathbb{R}^3$ — cuatro MLP chicos
• y = sum(gi * expertos[i](x) → $y = \sum_{i \in \mathcal{T}} g_i\, E_i(x)$ — solo 2 expertos calculados
• g_todos = torch.zeros(4).scatter → $g = (0;\ 0{,}38;\ 0{,}62;\ 0)$ — el vector completo, con ceros
• y_denso = sum( → $y = \sum_{i=1}^{4} g_i\, E_i(x)$ — los 4: mismo resultado

expertos = [Experto(d=3, h=4) for _ in range(4)]
probs = (x @ W_r).softmax(-1)
top = probs.topk(2)
g = top.values / top.values.sum()
# solo los elegidos: 2 expertos calculados
y = sum(gi * expertos[i](x) for gi, i in zip(g, top.indices.tolist()))
# los 4, con g = 0 en los no elegidos
g_todos = torch.zeros(4).scatter(0, top.indices, g)
y_denso = sum(g_todos[i] * expertos[i](x) for i in range(4))
print("g completo:", [round(v, 2) for v in g_todos.tolist()])
print("y:      ", [round(v, 3) for v in y.tolist()])
print("y_denso:", [round(v, 3) for v in y_denso.tolist()])

g completo: [0.0, 0.38, 0.62, 0.0]
y:       [-0.344, 0.246, 0.404]
y_denso: [-0.344, 0.246, 0.404]

Parte 2 · Por dónde pasa el gradiente

Elegir los k mejores es un corte: no tiene derivada. ¿Cómo aprende entonces el router?

Autograd nos dice quién aprende

Preguntémosle a autograd. Una capa de juguete con cuatro expertos y dos por token. Un token, una pérdida cualquiera, y backward: el paso que calcula el gradiente de la pérdida respecto de cada peso. Miramos el tamaño del gradiente de los pesos de cada experto. Y el de la fila del router que corresponde a cada experto. Eligió el dos y el cero. Esos dos expertos reciben gradiente; el uno y el tres, cero exacto: para este token, no existieron. En el router pasa lo mismo: solo las filas del dos y del cero se mueven, y las dos con el mismo tamaño, uno coma tres cuatro seis. Ya vamos a ver por qué.

moe = MoE(d=8, h=8, E=4, k=2)
x = torch.randn(1, 8)                            # un token
moe(x).pow(2).sum().backward()                   # una pérdida cualquiera
print("elegidos:", moe.ultimo["idx"][0].tolist())
for i, e in enumerate(moe.expertos):
    g_exp = e.arriba.weight.grad
    n_exp = 0.0 if g_exp is None else g_exp.norm().item()
    n_rou = moe.router.weight.grad[i].norm().item()  # fila i de W_rᵀ
    print(f"E{i}: pesos {n_exp:.3f} | router {n_rou:.3f}")

elegidos: [2, 0]
E0: pesos 0.536 | router 1.346
E1: pesos 0.000 | router 0.000
E2: pesos 1.873 | router 1.346
E3: pesos 0.000 | router 0.000

Por qué solo los elegidos

$\theta_i$: los pesos del experto $i$

dividimos arriba y abajo por $e^{\ell_a}$

iguales y opuestos; cero fuera de los elegidos

Theta sub i son todos los pesos del experto i. El experto i entra en la salida multiplicado por g sub i. Si g sub i es cero, la derivada respecto de sus pesos también es cero. En el dibujo, los gradientes del snippet: solo dos barras. Con k igual a dos, el peso de un elegido depende de un solo número. Dividimos arriba y abajo por e a la ele a, y queda la sigmoide de la diferencia entre los dos logits elegidos. La sigmoide es la curva en S que lleva cualquier número a algo entre cero y uno. Por eso las derivadas respecto de los dos logits son iguales y de signo opuesto, y respecto de los no elegidos, cero. Las dos filas del router se mueven lo mismo.

La trampa de top-1

De esa misma fórmula sale una trampa. Probamos dos variantes, ahora con k igual a uno: un solo experto por token, con y sin renormalizar. Miramos los pesos de los primeros tokens, y el gradiente del router. Renormalizado, el único peso vale siempre uno, y el gradiente del router es cero exacto: el router nunca aprende. Sin renormalizar, el peso es la probabilidad cruda, alrededor de cero coma treinta y siete, y el router sí recibe gradiente. Con un solo elegido no hay nada que repartir: hay que usar la probabilidad directa.

for renorm in (True, False):
    torch.manual_seed(0)
    moe = MoE(d=8, h=8, E=4, k=1, renormalizar=renorm)
    moe(x).pow(2).sum().backward()
    pesos = moe.ultimo["pesos"][:3, 0].tolist()
    g = moe.router.weight.grad.norm().item()
    print(f"renormalizar={renorm!s:5} pesos {[round(p, 2) for p in pesos]}"
          f" grad router {g:.3f}")

renormalizar=True  pesos [1.0, 1.0, 1.0] grad router 0.000
renormalizar=False pesos [0.37, 0.37, 0.37] grad router 1.656

Top-1: renormalizar mata el gradiente

constante: el router no aprende

depende de todos los logits

En fórmulas. Con k igual a uno, el softmax de un solo logit es e a la ele dividido e a la ele: uno, siempre. En el dibujo, una barra que siempre llega a cien por ciento. Una constante tiene derivada cero. Sin renormalizar, el peso es la probabilidad del softmax sobre todos los expertos, que depende de todos los logits. En el ejemplo a mano, cero coma cuarenta y siete. Si ese experto ayuda, el gradiente empuja su probabilidad para arriba, y la de los demás para abajo.

Shazeer et al., 2017 — Outrageously Large Neural Networks, sección 2.1 · arXiv:1701.06538

k > 1 · «If we choose k > 1, the gate values for the top k experts have nonzero derivatives with respect to the weights of the gating network.»

El paper que lanzó los MoE modernos lo dice en una frase. Si elegimos k mayor que uno, los pesos de los k mejores expertos tienen derivadas distintas de cero respecto de los pesos del router. Es lo mismo que acabamos de ver con autograd. Gating network es el nombre que le dan al router.

Parte 3 · Explorar con ruido

El gradiente solo llega a los elegidos. Entonces, un experto que nunca es elegido nunca tiene la chance de mostrar que servía.

Noisy top-k gating

Una idea del mismo paper de Shazeer, que la presenta como ayuda al balanceo de carga: sumar ruido antes de elegir. El logit de siempre. Más un número al azar, épsilon, sacado de una normal de media cero y desvío uno, multiplicado por softplus de otra matriz entrenable. softplus es una curva suave que siempre da positivo: así el modelo aprende cuánto ruido quiere para cada experto. Después, los que no quedaron en el top-k pasan a menos infinito, y el softmax los manda a cero, porque e a la menos infinito es cero. Sin ruido, este token elige siempre el dos y el uno. Con ruido, en mil sorteos, el experto cero fue elegido el treinta y nueve por ciento de las veces, y el tres, el seis por ciento.

Cada línea, su término:

• limpio = x @ W_g → $(x \cdot W_g)_i$ — el logit de siempre
• ruido = torch.randn(4) * F.softplus → $\varepsilon \cdot \mathrm{softplus}\big((x \cdot W_{\text{noise}})_i\big)$ — $\varepsilon \sim \mathcal{N}(0, 1)$
• masked_fill(H < corte, float("-inf")) → $\mathrm{KeepTopK}(H, k)_i = -\infty \text{ si no está en el top-}k$ — los demás, a menos infinito
• return H.softmax(-1) → $G(x) = \mathrm{softmax}\big(\mathrm{KeepTopK}(H(x), k)\big)$ — $e^{-\infty} = 0$

def noisy_top_k(x, k=2):
    limpio = x @ W_g                                      # (x · W_g)_i
    ruido = torch.randn(4) * F.softplus(x @ W_noise)      # ε · softplus
    H = limpio + ruido
    corte = H.topk(k).values[-1]
    H = H.masked_fill(H < corte, float("-inf"))           # KeepTopK
    return H.softmax(-1)                                  # G(x)

veces = torch.zeros(4)
for _ in range(1000):
    veces += noisy_top_k(x) > 0
print("sin ruido, top-2:", (x @ W_g).topk(2).indices.tolist())
pct = (veces / 10).round().int().tolist()            # de 1000 a %
print("con ruido, % elegido:", pct)

sin ruido, top-2: [2, 1]
con ruido, % elegido: [39, 68, 87, 6]

Mil sorteos del mismo token (salida real)

En barras. Sin ruido, los expertos dos y uno ganarían el cien por ciento de las veces. Con ruido, el experto cero entra en el treinta y nueve por ciento de los sorteos. El uno, en el sesenta y ocho. El dos, en el ochenta y siete. Y el tres, que tiene el logit más bajo, en el seis por ciento. Cada vez que un experto entra, recibe gradiente y puede demostrar si servía. El paper presenta el ruido como una ayuda al balanceo de carga, un problema que tiene su propio episodio.

• experto 0: 39 %
• experto 1: 68 %
• experto 2: 87 %
• experto 3: 6 %

Parte 4 · El router de un modelo real

Todo esto en Granite, el MoE abierto de IBM: treinta y dos expertos, ocho por token.

El router de Granite, capa 0

Leemos la matriz del router de la capa cero: treinta y dos filas de mil veinticuatro. La cargamos en la clase de router que usa la librería transformers para Granite. Le pasamos un vector al azar: no es un token real, solo queremos ver la mecánica. El router devuelve los logits, los pesos y los índices. Y comparamos con nuestra receta: softmax de los treinta y dos y renormalizar. Ocho elegidos, con pesos entre cero coma doce y cero coma catorce que suman uno. Y la comparación da verdadero: Granite hace top-k sobre los logits y después el softmax de esos ocho, que es exactamente renormalizar.

nombre = "model.layers.0.block_sparse_moe.router.layer.weight"
for archivo in glob.glob(f"{carpeta}/*.safetensors"):
    with safe_open(archivo, "pt") as f:
        if nombre in f.keys():
            W_r = f.get_tensor(nombre).float()       # E × d = 32 × 1024
router = GraniteMoeTopKRouter(AutoConfig.from_pretrained(REPO))
router.weight.data = W_r                             # el router real, capa 0
x = torch.randn(1, 1024)                             # vector al azar
with torch.no_grad():
    logits, g, idx = router(x)
p = logits.softmax(-1)[0, idx[0]]                    # softmax de los 32
print("W_r:", tuple(W_r.shape), "| elegidos:", idx[0].tolist())
print("g:", [round(v, 2) for v in g[0].tolist()])
print("suma:", round(g.sum().item(), 3),
      "| igual a renormalizar:", torch.allclose(g[0], p / p.sum()))

W_r: (32, 1024) | elegidos: [25, 27, 0, 17, 28, 13, 15, 4]
g: [0.14, 0.13, 0.13, 0.13, 0.12, 0.12, 0.12, 0.12]
suma: 1.0 | igual a renormalizar: True

Las decisiones de diseño del router

• ¿Renormalizar los k pesos o usar p crudo?
• ¿Ruido para explorar, o no?
• ¿Cuántos expertos, cuántos por token?
• ¿Cómo evitar que elija siempre los mismos?

Con esto quedan a la vista las perillas de un router. Renormalizar los k pesos, como Granite, o usar la probabilidad cruda, que es obligatorio con k igual a uno. Sumar ruido para explorar, como en el paper de dos mil diecisiete, o no; el código de Granite que corrimos no suma ruido. Cuántos expertos y cuántos por token: Granite usa treinta y dos y ocho. Y la gran pregunta que sigue abierta: qué impide que el router se enamore de unos pocos expertos y deje a los demás sin trabajo.

Resumiendo

• Router: $\ell = xW_r$, una matriz $d \times E$
• softmax → top-k → renormalizar = softmax de los k logits
• Solo se calculan, y solo aprenden, los k elegidos
• Con k = 2, cada peso es σ(ℓ_a − ℓ_b)
• Top-1 renormalizado: peso 1, router sin gradiente
• Próximo: una capa MoE completa, desde cero

Papers del episodio

• Shazeer et al., 2017 — Outrageously Large Neural Networks