Expertos finos y compartidos

Mixture of Experts, episodio siete: expertos finos y compartidos. Hasta ahora, un MoE, mixture of experts, tenía pocos expertos grandes y cada token elegía uno o dos. Hoy vemos las dos ideas de DeepSeekMoE: muchos expertos chicos, y unos pocos expertos que trabajan siempre, para todos los tokens.

Dos ideas

• Fine-grained experts: partir cada experto en m más chicos y elegir m·k
• Mismo cómputo, muchísimas más combinaciones posibles
• Shared experts: siempre activos, para lo que todos necesitan
• Medirlo en el juguete, y ubicar a Granite

El plan. Fine-grained experts: tomar cada experto, partirlo en m pedazos más chicos, y en lugar de elegir k, elegir m por k. El cómputo por token no cambia, pero las combinaciones posibles de expertos explotan. Shared experts: uno o más expertos que procesan todos los tokens, sin pasar por el router, para que lo común no se aprenda repetido en cada experto. Y después, lo medimos en nuestro MoE de juguete, y vemos dónde queda Granite, el MoE abierto de IBM que usamos en la serie.

Parte 1 · Partir un experto

Antes de partir expertos, hay que ver que partir un MLP, multilayer perceptron, no cambia lo que calcula.

Un MLP es una suma de neuronas

W₁: h × d, W₂: d × h, σ: la activación

cada neurona j suma una columna de W₂

agrupar las h neuronas en m bloques de h/m

El MLP es la red densa que hace de experto. W sub uno es una matriz de h por d: lleva el vector x, de d números, a h números, las neuronas de la hidden layer. Pasa por la activación sigma y W sub dos, de d por h, vuelve a d. En el dibujo, la fila j de W sub uno por x da la neurona j. Escrito de otra forma: la salida es la suma, sobre las h neuronas, de cuánto se activó cada una por su columna de W sub dos. En el dibujo, seis columnas pesadas por su neurona. Y una suma se puede agrupar como quieras: partimos las h neuronas en m bloques, y cada bloque es un MLP chico. En el dibujo, tres bloques de dos neuronas dan exactamente la misma salida.

Un experto grande = cuatro chicos sumados

Comprobémoslo. Un experto de dieciséis a sesenta y cuatro neuronas y de vuelta a dieciséis. Cuatro expertos chicos, de dieciséis neuronas cada uno. A cada chico le copiamos un cuarto de las filas de la matriz de subida, y las columnas correspondientes de la de bajada. El sesgo de salida se suma una sola vez, así que va solo en el primero. La suma de los cuatro chicos coincide con el grande hasta nueve por diez a la menos ocho. Los parámetros casi coinciden: dos mil ciento veintiocho contra dos mil ciento setenta y seis, por los sesgos de salida que repiten los chicos. Partir un experto no le saca ni le agrega capacidad: cambia cómo se puede elegir.

grande = Experto(d, h)                        # un experto: d → 64 → d
chicos = [Experto(d, h // m) for _ in range(m)]   # 4 expertos: d → 16 → d
with torch.no_grad():
    for j, c in enumerate(chicos):
        filas = slice(j * h // m, (j + 1) * h // m)   # un cuarto de la capa
        c.arriba.weight.copy_(grande.arriba.weight[filas])
        c.arriba.bias.copy_(grande.arriba.bias[filas])
        c.abajo.weight.copy_(grande.abajo.weight[:, filas])
        c.abajo.bias.copy_(grande.abajo.bias if j == 0 else 0 * c.abajo.bias)
    suma = sum(c(x) for c in chicos)
print(f"máx dif = {(grande(x) - suma).abs().max():.1e}")
print("parámetros:", parametros(grande), "vs",
      sum(parametros(c) for c in chicos))

máx dif = 8.9e-08
parámetros: 2128 vs 2176

Contar combinaciones

de cuántas formas se eligen K entre N

16 expertos, top-2

cada uno partido en 4: 64 expertos, top-8

¿Por qué partir? Por las combinaciones. N sobre K, el número combinatorio, cuenta de cuántas formas se pueden elegir K cosas entre N, sin importar el orden: N factorial dividido K factorial por N menos K factorial. Con dieciséis expertos y top-2: dieciséis por quince, dividido dos, ciento veinte combinaciones. Si partimos cada experto en cuatro, quedan sesenta y cuatro chicos, y para el mismo cómputo elegimos ocho: más de cuatro mil cuatrocientos millones de combinaciones. En el dibujo, en escala logarítmica, de dos a casi diez. Cada token puede armarse una mezcla mucho más a medida.

Las combinaciones, contadas

Con math punto comb, de la biblioteca estándar de Python. El caso de dieciséis expertos y top-2. El mismo cómputo, partido en cuatro. Y dos modelos reales: Granite, y DeepSeek-V3, que tiene doscientos cincuenta y seis expertos ruteados y elige ocho. comb hace la cuenta exacta. Ciento veinte; cuatro coma cuatro por diez a la nueve; Granite, uno coma cero cinco por diez a la siete, unos diez millones; y DeepSeek-V3, cuatro coma uno por diez a la catorce.

casos = [("16 expertos, top-2", 16, 2),
         ("partidos en 4: top-8", 64, 8),
         ("Granite", 32, 8),
         ("DeepSeek-V3", 256, 8)]
for nombre, N, K in casos:
    print(f"{nombre:21} C({N},{K}) = {comb(N, K):.3g}")

16 expertos, top-2    C(16,2) = 120
partidos en 4: top-8  C(64,8) = 4.43e+09
Granite               C(32,8) = 1.05e+07
DeepSeek-V3           C(256,8) = 4.1e+14

Dai et al., 2024 — DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models · arXiv:2401.06066

120 → 4.426.165.368 · Con N = 16, top-2 da 120 combinaciones; partiendo cada experto en 4, el ruteo fino «can yield 4,426,165,368 potential combinations».

Es exactamente el ejemplo del paper. DeepSeekMoE propone dos estrategias: partir los expertos en m N más chicos y activar m K, y aislar algunos como compartidos. Su argumento para la primera es este conteo: de ciento veinte combinaciones a más de cuatro mil millones, con el mismo cómputo. La hipótesis es que así cada experto se puede especializar más.

Parte 2 · Expertos compartidos

Si todos los expertos necesitan saber lo mismo, cada uno lo aprende por separado. Un experto compartido lo aprende una vez.

La capa de DeepSeek-V3

compartidos sin peso + ruteados con peso g

los no elegidos no se calculan

los valores de DeepSeek-V3

La fórmula del paper de DeepSeek-V3. u sub t es el vector del token t que entra a la capa, y h sub t, el que sale. Se le suma la salida de los N sub s expertos compartidos, sin ningún peso, y la de los N sub r expertos ruteados, cada una multiplicada por su peso g. FFN es feed-forward network, otro nombre del MLP. En el dibujo: la primera columna, el compartido, entra siempre con peso uno; de las otras cinco, solo tres tienen peso. g vale cero para los expertos que no quedaron en el top K sub r: esos ni se calculan. En DeepSeek-V3: un compartido, doscientos cincuenta y seis ruteados, y ocho elegidos por token.

Fino + compartido

El esquema, en dos caminos que corren en paralelo sobre el mismo token. El camino ruteado: el router puntúa muchos expertos chicos, elige k, y cada salida entra pesada por g. Esos expertos compiten por tokens, así que necesitan balanceo de carga. El camino compartido: unos pocos expertos que procesan todos los tokens, sin router, y su salida se suma tal cual. No compiten, los reciben todos. Las dos salidas se suman y vuelven al residual stream, la corriente principal del modelo. Para no gastar más cómputo, si agregás un compartido, elegís un ruteado menos.

Tres capas con el mismo presupuesto

Armamos tres capas con el MoE de juguete de la serie, que ya acepta expertos compartidos. Una función cuenta los parámetros activos: el router, los k ruteados y los compartidos. Grueso: ocho expertos de sesenta y cuatro neuronas, top-2. Fino: cada uno partido en cuatro, treinta y dos de dieciséis neuronas, top-8. Y fino con un compartido: treinta y un ruteados, top-7, más uno fijo. Los totales rondan los treinta y cuatro mil y los activos, entre ocho mil seiscientos y nueve mil seiscientos. La diferencia es el router, que con treinta y dos expertos tiene cuatro veces más columnas, y los sesgos. Y las dos últimas líneas son el camino fijo, leídas del código del MoE de juguete: para cada compartido, se suma su salida, sin peso.

def activos(moe):                       # router + k ruteados + compartidos
    return (moe.router.weight.numel() + moe.k * parametros(moe.expertos[0])
            + sum(parametros(c) for c in moe.compartidos))

configs = {"grueso": MoE(32, 64, E=8, k=2),        # 8 de h=64, top-2
           "fino": MoE(32, 16, E=32, k=8),          # 32 de h=16, top-8
           "fino+comp": MoE(32, 16, E=31, k=7, n_compartidos=1)}
print("           total  activos")
for nombre, moe in configs.items():
    print(f"{nombre:9} {parametros(moe):6d}  {activos(moe):7d}")
fuente = inspect.getsource(MoE.forward).splitlines()   # el camino fijo
i = next(j for j, l in enumerate(fuente) if "self.compartidos" in l)
print("\n".join(l.strip() for l in fuente[i:i + 2]))

           total  activos
grueso     33792     8640
fino       35328     9600
fino+comp  35296     9568
for c in self.compartidos:
y = y + c(x)

Parte 3 · ¿Ayuda? Medirlo

Entrenamos las tres en la tarea sintética de la serie y miramos dos cosas: el error y cuánto se parecen los expertos entre sí.

Error y parecido entre expertos

Para medir redundancia, pasamos los mismos quinientos doce vectores por cada experto. Normalizamos las salidas y calculamos el coseno entre cada par de expertos: uno si calculan lo mismo, cero si no tienen nada que ver. Promediamos sin la diagonal, que es cada experto consigo mismo. Las tres capas entrenan ochocientos pasos, con la pérdida de balanceo. El grueso llega a cero coma doscientos ocho de error cuadrático medio y sus expertos se parecen cero coma catorce. El fino, a cero coma doscientos veintiséis, pero sus expertos se parecen diez veces menos: cero coma cero uno tres. Con compartido, cero coma doscientos cincuenta, y parecido de cero coma cero dos. Con dos semillas más, probadas aparte, el orden del error se repite.

def parecido(moe, h):              # coseno medio entre expertos
    s = torch.stack([e(h).flatten() for e in moe.expertos])
    s = F.normalize(s, dim=-1) @ F.normalize(s, dim=-1).T
    E = len(moe.expertos)
    return ((s.sum() - E) / (E * (E - 1))).item()   # sin la diagonal

for nombre, c in configs.items():
    torch.manual_seed(0)
    moe = MoE(32, **c); red = Red(16, 32, 8, moe)
    entrenar(red, X, Y, pasos=800, alfa=0.01)
    with torch.no_grad():
        mse = F.mse_loss(red(Xt), Yt).item()
        h = red.norma(red.entrada(Xt[:512]))   # lo que entra a la capa MoE
        print(f"{nombre:9} MSE {mse:.3f} · parecido {parecido(moe, h):.3f}")

grueso    MSE 0.208 · parecido 0.139
fino      MSE 0.226 · parecido 0.013
fino+comp MSE 0.250 · parecido 0.020

Resultado en el juguete

En tabla. El grueso tiene el menor error. El fino tiene expertos mucho más distintos entre sí, que es lo que el paper busca, pero acá el error empeora un poco. Y el compartido no mejora nada. No hay que forzar la conclusión: este juguete tiene ocho grupos de datos, y ocho expertos grandes ya alcanzan para separarlos. No tiene conocimiento común que valga la pena compartir, y entrena ochocientos pasos.

Cómo leer un resultado negativo

• El juguete no tiene lo que la idea aprovecha: mucho conocimiento común
• Sí se ve el efecto mecánico: expertos finos más distintos
• Las ganancias del paper son a escala: 16B parámetros, 2T tokens
• Probar en tu problema antes de copiar una arquitectura

¿Qué hacemos con esto? La idea del compartido tiene sentido cuando hay mucho que todos los tokens necesitan: gramática, formato, lo frecuente. Nuestra tarea no lo tiene. Lo que sí se ve es el efecto mecánico: con expertos finos, cada uno calcula algo bien distinto de los otros. Las ganancias del paper se midieron en modelos de miles de millones de parámetros, entrenados con billones de tokens. Y la lección general: una arquitectura que gana a escala no garantiza nada en tu problema; se prueba.

Dai et al., 2024 — DeepSeekMoE · arXiv:2401.06066

≈ 40 % del cómputo · DeepSeekMoE 16B «achieves comparable performance with LLaMA2 7B, with only about 40% of computations».

Lo que reporta el paper a escala. Su modelo de dieciséis mil millones de parámetros, con expertos finos y compartidos, rinde parecido a LLaMA dos de siete mil millones, un modelo denso, usando más o menos el cuarenta por ciento del cómputo. También dicen que su versión de dos mil millones iguala a un MoE convencional estilo GShard de dos coma nueve mil millones.

Parte 4 · El mapa

Con estas dos ideas, cada MoE conocido se puede ubicar en un mapa: cuántos expertos, cuántos por token, y si hay compartidos.

Granite, mirado desde su config

Leemos la configuración de Granite, sin cargar los pesos. Número de expertos y cuántos por token. Buscamos alguna clave que hable de compartidos. El tamaño de un experto con compuerta: tres matrices de d por h. Y cuál sería su equivalente grueso. Treinta y dos expertos, ocho por token, y cada experto tiene una hidden layer de quinientas doce neuronas: la mitad del ancho del modelo, mil veinticuatro. Expertos finos. No hay ninguna clave de compartidos. Por capa activa doce coma seis millones de parámetros, lo mismo que ocho expertos cuatro veces más anchos con top-2, pero con diez millones de combinaciones en lugar de veintiocho.

cfg = AutoConfig.from_pretrained("ibm-granite/granite-3.1-1b-a400m-base")
E, k = cfg.num_local_experts, cfg.num_experts_per_tok
d, h = cfg.hidden_size, cfg.intermediate_size
print(f"E={E}, k={k}, d={d}, h={h}")
print("claves 'shared':", [c for c in cfg.to_dict() if "shared" in c])
por_experto = 3 * d * h                       # gate, up y down
print(f"activos/capa: {k} × {por_experto / 1e6:.2f} M = "
      f"{k * por_experto / 1e6:.1f} M")
print(f"grueso igual: {E // 4} de h={4 * h}, top-{k // 4}")
print(f"C({E},{k}) = {comb(E, k):,} · C(8,2) = {comb(8, 2)}")

E=32, k=8, d=1024, h=512
claves 'shared': []
activos/capa: 8 × 1.57 M = 12.6 M
grueso igual: 8 de h=2048, top-2
C(32,8) = 10,518,300 · C(8,2) = 28

El mapa de algunos MoE

El mapa, con datos de sus papers y del config. Switch: muchos expertos, uno por token. Mixtral: ocho expertos grandes, dos por token. Granite: treinta y dos expertos finos, ocho por token, sin compartidos. DeepSeek-V3: doscientos cincuenta y seis ruteados, ocho por token, y uno compartido. La tendencia de los últimos modelos va hacia la última fila: más expertos, más chicos, más elegidos. Y fijate que ninguno pasa de ocho por token: cada experto elegido es un vector que viaja, y en el último episodio vamos a ver cuánto cuesta ese viaje.

Resumiendo

• Un MLP es una suma de neuronas: partirlo no cambia la cuenta
• Fine-grained: m·N expertos chicos, top m·K, mismo cómputo
• C(16,2) = 120 → C(64,8) ≈ 4,4·10⁹ combinaciones
• Shared experts: siempre activos, sin peso del router
• Juguete: expertos finos 10× más distintos, pero no mejor error
• Próximo: balancear sin pérdida auxiliar (DeepSeek-V3)

Papers del episodio

• Dai et al., 2024 — DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models