El router que colapsa

Mixture of Experts. Episodio cuatro: el router que colapsa. Ya tenemos una capa MoE, mixture of experts, que funciona. Hoy la entrenamos en una tarea donde todos los tokens se parecen, y vemos cómo la mitad de los expertos se queda sin trabajo. Después, la pérdida auxiliar que lo evita, escrita y medida.

Qué es la carga de un experto

• Carga: cuántas asignaciones le tocan a cada experto
• Ideal: T·k/E para cada uno
• Un experto sin carga: parámetros que no sirven
• Un experto sobrecargado: cuello de botella

La carga de un experto es cuántas asignaciones le tocan en un batch: cuántos pares token y experto lo tienen a él. En inglés, load. Si el reparto fuera perfecto, cada uno recibiría T por k dividido E: con doscientos cincuenta y seis tokens, k igual a dos y ocho expertos, sesenta y cuatro cada uno. Un experto con carga cero es memoria ocupada que no aporta nada. Y el más cargado marca el tiempo: todos lo esperan.

Parte 1 · El colapso

Entrenamos la capa de juguete del episodio tres, ocho expertos y dos por token, en una tarea sin regiones: una sola función suave para todo el espacio.

Carga después de entrenar, sin balanceo

Esto es lo que queda después de mil quinientos pasos, medido sobre dos mil cuarenta y ocho tokens de test. Dos expertos, el cinco y el siete, se llevan casi el ochenta por ciento de las asignaciones. Y cuatro, el uno, el tres, el cuatro y el seis, prácticamente nada: menos del uno por ciento cada uno. La mitad del modelo está muerta. Las barras rojas son parámetros que guardamos en memoria y no usamos nunca.

• E0: 14 %
• E1: 0 %
• E2: 7 %
• E3: 0 %
• E4: 0 %
• E5: 43 %
• E6: 0 %
• E7: 36 %

Por qué pasa

• Tokens parecidos → el router al azar ya prefiere a algunos
• El no elegido no recibe gradiente
• Sin gradiente no mejora, y no lo eligen
• Nada en la pérdida pide repartir

Cuatro piezas, todas ya conocidas. Si los tokens se parecen entre sí, un router con pesos al azar ya los manda a casi todos a los mismos pocos expertos: en el primer paso de este entrenamiento, el más cargado tenía casi cuatro veces el promedio. Y del episodio dos: un experto que no es elegido no recibe gradiente, ni él ni su fila del router. Sin gradiente no mejora, y como no mejora, el router no tiene motivo para elegirlo. Y la pérdida solo mide el error de la tarea: no hay nada que pida repartir el trabajo.

El círculo vicioso

Es un círculo. Un experto recibe pocos tokens. Entonces recibe poco gradiente. Entonces no mejora. Entonces el router no lo elige, y vuelve a recibir pocos tokens. Esa es la explicación habitual. En nuestro juguete se ve así: en los primeros cien pasos el router reacomoda quién es favorito, y desde el paso doscientos, más o menos, cuatro expertos quedan en casi cero y ya no vuelven. Para romperlo hay que agregar a la pérdida algo que castigue el desbalance.

Parte 2 · La pérdida de balanceo

La receta que usan casi todos sale del paper de Switch Transformer. Son dos promedios y un producto punto.

f, P y la pérdida auxiliar

fracción de asignaciones al experto $i$

probabilidad media que el router le da a $i$

producto punto de $f$ y $P$, por $E$ y por $\alpha$

Tres definiciones. T es la cantidad de tokens del batch, k los elegidos por token, E los expertos. f sub i es la fracción de las asignaciones que fueron al experto i: contamos cuántas veces aparece i entre los elegidos y dividimos por T por k. El uno con corchetes vale uno si la condición se cumple y cero si no. En el dibujo, el modelo sin balanceo: ceros y dos barras gigantes. P sub i es la probabilidad promedio que el router le da al experto i, sobre todos los tokens. En el dibujo, se parece a f, pero nunca es cero. La pérdida auxiliar multiplica f y P experto por experto, suma, y escala por E y por alfa. En el dibujo: los productos suman cero coma veintiséis; por ocho, dos coma once.

Fedus, Zoph & Shazeer, 2021 — Switch Transformers, sección 2.2 · arXiv:2101.03961

α = 10⁻² · «sufficiently large to ensure load balancing while small enough to not to overwhelm the primary cross-entropy objective». Su f usa solo el experto top-1.

Usan alfa igual a diez a la menos dos, cero coma cero uno: según el paper, lo bastante grande para asegurar el balanceo y lo bastante chico para no tapar el objetivo principal. Una aclaración: Switch elige un solo experto por token, así que su f cuenta solo el elegido. Nosotros, con k igual a dos, contamos las asignaciones y dividimos por T por k.

Cuánto vale, parejo y colapsado

Dos casos armados a mano, con ocho tokens y cuatro expertos. Parejo: el router da un cuarto a cada experto, y las dieciséis asignaciones se reparten cuatro y cuatro. Colapsado: el router le da cero coma cuarenta y cinco a los expertos cero y uno, y todos los tokens van a esos dos. Parejo da uno exacto. Colapsado, uno coma ocho. Sin alfa, la pérdida auxiliar vale uno cuando todo está repartido, y crece cuando se concentra.

# parejo: cada experto recibe 4 de las 16 asignaciones, P = 1/4
probs = torch.full((T, E), 1 / E)
idx = torch.tensor([[0, 1], [2, 3]] * 4)
print("parejo:    ", perdida_balanceo(probs, idx, E).item())
# colapsado: todo a los expertos 0 y 1, con prob. 0,45 cada uno
probs = torch.tensor([[0.45, 0.45, 0.05, 0.05]] * T)
idx = torch.tensor([[0, 1]] * T)
print("colapsado: ", round(perdida_balanceo(probs, idx, E).item(), 2))

parejo:     1.0
colapsado:  1.8

Por qué vale 1 si todo es parejo

E términos iguales

el caso colapsado del snippet

los expertos vacíos no suman nada

Las dos cuentas, sin saltear pasos. Si todo es parejo, f y P valen un cuarto para cada uno. Cada producto es uno sobre E al cuadrado, hay E productos, y multiplicamos por E: E por E dividido E al cuadrado, uno. En el dibujo, cuatro barras de un dieciseisavo que suman un cuarto, y por cuatro, uno. Por eso Switch multiplica por E: así el valor parejo es uno, tenga los expertos que tenga. El caso colapsado: f es un medio para los dos primeros y cero para el resto. Los expertos vacíos multiplican por cero. Quedan dos productos de cero coma dos dos cinco; por cuatro, uno coma ocho.

En los modelos entrenados

Ahora en los modelos de verdad, entrenados una vez y guardados. Cargamos dos: uno sin balanceo y uno entrenado con alfa igual a cero coma uno; ya vamos a ver cómo. f: contamos con bincount y dividimos por el total de asignaciones. P: el promedio de las probabilidades del router. Y la fórmula, sin alfa. Sin balanceo, f tiene cuatro ceros y la pérdida da dos coma once. Con balanceo, todas las f están entre el once y el trece por ciento, y la pérdida da uno coma cero cero: el mínimo práctico.

pct = lambda v: " ".join(f"{round(100 * x)}" for x in v.tolist())
for alfa in ("0.0", "0.1"):              # entrenados en preparar_ep04
    red = cargar(alfa)
    with torch.no_grad():
        red(Xt)
    u = red.capa.ultimo
    f = torch.bincount(u["idx"].reshape(-1), minlength=8) / u["idx"].numel()
    P = u["probs"].mean(0)               # prob. media por experto
    L = 8 * (f * P).sum()
    print(f"α={alfa} L_aux={L:.2f} | f% {pct(f)} | P% {pct(P)}")

α=0.0 L_aux=2.11 | f% 14 0 7 0 0 43 0 36 | P% 17 4 11 2 2 32 5 26
α=0.1 L_aux=1.00 | f% 12 13 11 13 12 13 13 13 | P% 12 12 12 12 12 13 14 13

Parte 3 · Por qué f no tiene gradiente

f es lo que queremos emparejar, pero sale de un conteo. Contar no tiene derivada. ¿Cómo empuja entonces esta pérdida?

Autograd sobre la pérdida auxiliar

Quinientos doce tokens con logits al azar, y les sumamos un empujón a los expertos cero y uno, para tener un router desbalanceado. f sale de bincount: un conteo de enteros. P sale del softmax: una cuenta suave. Calculamos la pérdida y hacemos backward. Y miramos el gradiente respecto de cada P. f no tiene gradiente; P sí. El experto cero tiene el cuarenta y siete por ciento de las asignaciones, y el gradiente respecto de su P es tres coma setenta y ocho: exactamente ocho por cero coma cuarenta y siete. El más cargado es el más empujado para abajo.

logits = torch.randn(T, E, requires_grad=True)
logits_ = logits + torch.tensor([3.0, 2.0, 0, 0, 0, 0, 0, 0])  # 0 y 1, favoritos
probs = logits_.softmax(-1)
idx = logits_.topk(k).indices
f = torch.bincount(idx.reshape(-1), minlength=E) / (T * k)
P = probs.mean(0)
print("f tiene gradiente:", f.requires_grad, "| P:", P.requires_grad)
L = E * (f * P).sum()
P.retain_grad()
L.backward()
print("f       :", [round(v, 2) for v in f.tolist()])
print("dL/dP_i :", [round(v, 2) for v in P.grad.tolist()])   # = E · f_i

f tiene gradiente: False | P: True
f       : [0.47, 0.38, 0.02, 0.03, 0.03, 0.03, 0.03, 0.01]
dL/dP_i : [3.78, 3.03, 0.14, 0.25, 0.22, 0.23, 0.23, 0.12]

Cómo empuja

$f$ se trata como una constante

bajar más la probabilidad de los más cargados

lo que pierden unos, lo ganan otros: los vacíos suben

La lectura del gradiente. Como f no tiene derivada, para autograd es una constante. Entonces la derivada de la pérdida respecto de P sub i es alfa por E por f sub i. En el dibujo, la salida del snippet: las barras más altas son las de los expertos más cargados. Un paso de descenso por gradiente resta eta, el learning rate, por esa derivada: baja la probabilidad de cada experto, y la baja más cuanto más cargado está. Y como las probabilidades del softmax siempre suman uno, lo que pierden los cargados lo ganan los vacíos. Con el tiempo, sus logits suben hasta que entran en el top-k, empiezan a recibir tokens, y el círculo vicioso se rompe.

Lepikhin et al., 2020 — GShard, sección 2.2 · arXiv:2006.16668

(c_e/S)² → m_e · (c_e/S) · «because c_e is derived from top-2 operation and is not differentiable, we use the mean gates per expert m_e as a differentiable approximation»

¿Por qué justo f por P? GShard, también de Google, lo explica. Querríamos minimizar la suma de las fracciones al cuadrado, mínima cuando todo es parejo. Pero la fracción c sub e sobre S sale de un top-2 y no es diferenciable. Entonces reemplazan uno de los dos factores por m sub e, la probabilidad media del experto, que sí lo es. Esa es nuestra P: f por P es f al cuadrado con un factor cambiado por algo que se puede derivar.

Parte 4 · Entrenar con balanceo

Falta lo más importante: sumarla a la pérdida de la tarea y ver qué pasa con la carga y con el error.

El loop de entrenamiento

El loop completo, corto. Dos corridas: sin balanceo y con alfa igual a cero coma uno. La pérdida de la tarea es el error cuadrático medio, MSE, mean squared error. La capa ya dejó calculada su pérdida auxiliar en el forward. La pérdida total es la suma, con la auxiliar multiplicada por alfa. Es la única línea nueva. Al final, la carga del último batch: quinientas doce asignaciones. Sin balanceo, en solo cuatrocientos pasos: dos expertos con más de doscientas cada uno y otros con cero, uno o dos. Con alfa, todos entre cuarenta y nueve y setenta y cinco.

for alfa in (0.0, 0.1):
    torch.manual_seed(0)
    red = Red(16, 32, 8, MoE(32, 32, E=8, k=2))
    opt = torch.optim.Adam(red.parameters(), lr=3e-3)
    for paso in range(400):
        i = torch.randint(0, len(X), (256,))
        mse = F.mse_loss(red(X[i]), Y[i])
        aux = red.capa.ultimo["l_aux"]           # E · Σ f_i · P_i
        loss = mse + alfa * aux
        opt.zero_grad(); loss.backward(); opt.step()
    carga = red.capa.ultimo["carga"].tolist()     # de 512 asignaciones
    print(f"α={alfa}: carga {carga}")

α=0.0: carga [212, 0, 19, 2, 6, 230, 42, 1]
α=0.1: carga [74, 72, 49, 57, 53, 68, 64, 75]

Desbalance durante el entrenamiento

Las corridas largas, de mil quinientos pasos. El eje vertical es la carga del experto más cargado dividida por la carga promedio: uno es perfecto. Sin balanceo arranca en casi cuatro, porque el router al azar ya está desbalanceado. Baja a dos y medio en los primeros cien pasos, mientras el router reacomoda sus favoritos, y vuelve a subir: desde el paso ciento cincuenta queda entre tres y tres coma siete. Nunca se arregla. Con alfa igual a cero coma cero uno, el valor de Switch, baja rápido a alrededor de uno coma cinco. Con cero coma uno, a alrededor de uno coma dos.

Carga final por experto (test, semilla 0)

La carga final de cada experto, una fila por alfa. Sin balanceo, cuatro celdas en cero y dos que concentran casi todo. Con cero coma cero uno, todos trabajan, entre el seis y el dieciocho por ciento. Con cero coma uno, todos entre el once y el trece por ciento: el ideal es doce coma cinco.

Tres alfas, tres semillas

El resumen de las nueve corridas. El error de test, promedio de tres semillas. Expertos muertos: los que reciben menos del uno por ciento de las asignaciones. Y el cociente entre la carga máxima y la media. Sin balanceo: cuatro coma tres expertos muertos en promedio, de ocho. Con balanceo, ninguno. Pero mirá el error: sube de cero coma cero uno seis seis a cero coma cero uno nueve uno con alfa cero coma cero uno, y a cero coma cero dos uno seis con cero coma uno.

r = json.load(open(os.path.join(DATOS, "ep04_resultados.json")))
print(f"{'α':<6}{'MSE test':>9}{'muertos':>9}{'máx/medio':>11}")
for alfa, v in r.items():                         # 3 semillas cada uno
    mse = sum(v["test"]) / 3
    muertos = sum(sum(c < 0.01 * sum(cs) for c in cs)
                  for cs in v["carga_final"]) / 3          # < 1 % de carga
    ratio = sum(max(cs) / (sum(cs) / 8) for cs in v["carga_final"]) / 3
    print(f"{alfa:<6}{mse:>9.4f}{muertos:>9.1f}{ratio:>11.2f}")

α      MSE test  muertos  máx/medio
0.0      0.0166      4.3       3.44
0.01     0.0191      0.0       1.34
0.1      0.0216      0.0       1.09

El balanceo no es gratis

• En el juguete, balancear sube el error
• Tarea sin regiones: repartir no ayuda a aprender
• En escala, el balanceo paga en memoria y tiempo
• α grande compite con la tarea

Lo honesto. En este juguete, balancear empeora un poco el error. Tiene sentido: es una sola función suave, y no hay nada que ganar repartiéndola entre ocho expertos. En un modelo grande el argumento es otro: cada experto vive en memoria, y con varias placas el experto más cargado frena a todos. Y la tensión queda a la vista: cuanto más grande alfa, más pesa el balanceo frente a la tarea. Balancear sin esta pérdida tiene su propio episodio.

Resumiendo

• Carga = asignaciones por experto; ideal $Tk/E$
• Sin presión: 4 de 8 expertos muertos
• No elegido → sin gradiente → no mejora → no elegido
• $\mathcal{L}_{aux} = \alpha E \sum_i f_i P_i$, vale 1 si es parejo
• f no tiene gradiente; P sí: $\partial\mathcal{L}/\partial P_i = \alpha E f_i$
• Próximo: capacidad y tokens descartados

Papers del episodio

• Fedus, Zoph & Shazeer, 2021 — Switch Transformers
• Lepikhin et al., 2020 — GShard