Lección 9 de 10 · MoE modernos · 16 min

Qué aprende cada experto

Abrir el router de Granite y mirar a quién manda cada token

Notebook del episodio

El código de esta lección, para correrlo y modificarlo en Google Colab, gratis y con tu cuenta. Para guardar tus cambios: Archivo → Guardar una copia en Drive.

Abrir en Colab Ver en GitHub

Mixture of Experts, episodio nueve: qué aprende cada experto. Hasta acá armamos un MoE, mixture of experts, de juguete y lo entrenamos. Hoy abrimos uno de verdad, Granite, y miramos a qué experto manda cada token. La pregunta es tentadora: ¿hay un experto de cocina y otro de astronomía? Spoiler: no exactamente.

El plan de hoy

Granite 3.1 1B-A400M: 24 capas, 32 expertos, top-8
Un hook en el router de cada capa
Medir carga, tema contra tipo de texto, tokens por experto
Comparar siempre contra el azar

El plan. Granite tres punto uno, de IBM, es un MoE abierto: veinticuatro capas, cada una con treinta y dos expertos, y cada token usa ocho. Unos mil trescientos millones de parámetros en total, y cuatrocientos millones activos por token. Vamos a colgar un hook, una función que PyTorch llama cada vez que un módulo termina su forward, en el router de cada capa. Con eso medimos tres cosas: cuánto trabaja cada experto, si el ruteo depende del tema o del tipo de texto, y qué tokens recibe cada uno. Y la regla del episodio: cada número se compara contra lo que daría un router que elige al azar.

Parte 1 · Escuchar al router

Primero hay que poder ver la decisión del router, token por token y capa por capa.

Un hook en el router de la capa 12

En Granite, cada capa tiene un módulo block sparse moe, y adentro un router con una matriz de treinta y dos por mil veinticuatro: una fila por experto. El router devuelve tres cosas: los logits, un número por experto; los pesos de los ocho elegidos; y sus índices. Nuestro hook solo los guarda. Corremos un pedacito de código Python y miramos el cuarto token, range. Los logits salen de ocho tokens por treinta y dos expertos, y los índices, de ocho por ocho. range eligió a los expertos veinticinco, diecinueve, trece y otros cinco, con pesos que van de cero coma veinte a cero coma cero ocho, y suman uno. El primer token, for, en cambio, le da cero coma ochenta y dos a un solo experto: el comienzo del texto es especial.

python

capa = model.model.layers[12].block_sparse_moe   # el MoE de la capa 12
print(type(capa.router).__name__, tuple(capa.router.weight.shape))

guardado = {}
def hook(modulo, entrada, salida):
    logits, pesos, idx = salida          # lo que devuelve el router
    guardado.update(logits=logits, pesos=pesos, idx=idx)
h = capa.router.register_forward_hook(hook)

ids = tok("for i in range(10):", return_tensors="pt")
model(**ids)
h.remove()
toks = tok.convert_ids_to_tokens(ids.input_ids[0])
print("logits", tuple(guardado["logits"].shape), "idx", tuple(guardado["idx"].shape))
t = 3                                            # el token "range"
print("expertos de", repr(toks[t]))
print(guardado["idx"][t].tolist())
print(" ".join(f"{p:.2f}"[1:] for p in guardado["pesos"][t]))
print("suman", f"{guardado['pesos'][t].sum():.2f}")
print("'for': peso máx", f"{guardado['pesos'][0].max():.2f}")

text

GraniteMoeTopKRouter (32, 1024)
logits (8, 32) idx (8, 8)
expertos de 'Ġrange'
[25, 19, 13, 30, 18, 8, 26, 24]
.20 .19 .15 .12 .09 .09 .08 .08
suman 1.00
'for': peso máx 0.82

Lo que calcula el router, con los números de «range»

\ell = W_r\, x \in \mathbb{R}^{32}

W_r: 32 × 1024; x: el token, 1024 números

S = \mathrm{top}_8(\ell)

los índices de los 8 logits más grandes

g_i = \frac{e^{\ell_i}}{\sum_{j \in S} e^{\ell_j}},\quad i \in S

softmax solo entre los 8 elegidos

Con números reales. x es el vector del token, mil veinticuatro números. W sub r es la matriz del router, treinta y dos por mil veinticuatro. En el dibujo: la fila i de W sub r por x da un solo número, ele sub i, el logit del experto i. S es el conjunto de los ocho índices con logit más alto. Y los pesos g salen de un softmax, pero solo entre esos ocho: exponencial de cada logit dividido la suma de las exponenciales de los ocho. Equivale a un softmax sobre los treinta y dos, renormalizado entre los ocho. En el dibujo, los ocho logits reales de range: uno coma doce para el experto veinticinco, uno coma cero cuatro para el diecinueve, y bajando. Pesos de cero coma veinte a cero coma cero ocho.

Parte 2 · ¿Cuánto trabaja cada experto?

Ahora pasamos mil cincuenta y cinco tokens de textos sintéticos: recetas, astronomía, código, números y puntuación.

La carga por experto, capa por capa

Los textos son nuestros, escritos para la serie: nada sale de un corpus real. Juntamos los índices de todos los tokens: mil cincuenta y cinco. Para cada capa, una matriz de mil cincuenta y cinco por ocho. bincount cuenta cuántas veces aparece cada experto, y dividimos por la cantidad de tokens: es la fracción de tokens que eligió a ese experto. Si fuera pareja, daría ocho de treinta y dos, veinticinco por ciento. Imprimimos el más y el menos usado. Nada parejo. En la capa cero, el experto treinta y uno lo eligen el sesenta y ocho por ciento de los tokens, y el veinte, el dos por ciento. En la última, el experto treinta lo elige el noventa y dos por ciento.

python

T = torch.cat(elegidos[0]).shape[0]          # tokens de todos los textos
print("tokens:", T, "· uniforme: 8/32 =", f"{8 / 32:.0%}")
for c in [0, 6, 12, 18, 23]:
    idx = torch.cat(elegidos[c])             # [T, 8]
    carga = torch.bincount(idx.reshape(-1), minlength=32) / T
    print(f"capa {c:2d}: máx {carga.max():.0%} (exp. {carga.argmax()})"
          f" · mín {carga.min():.0%} (exp. {carga.argmin()})")

text

tokens: 1055 · uniforme: 8/32 = 25%
capa  0: máx 68% (exp. 31) · mín 2% (exp. 20)
capa  6: máx 72% (exp. 21) · mín 2% (exp. 8)
capa 12: máx 61% (exp. 19) · mín 4% (exp. 27)
capa 18: máx 54% (exp. 9) · mín 3% (exp. 0)
capa 23: máx 92% (exp. 30) · mín 4% (exp. 23)

Carga por experto en cinco capas

La misma cuenta, completa, para cinco capas. Cada fila es una capa y cada columna, uno de los treinta y dos expertos; más intenso, más tokens. Capa cero. Capa seis. Capa doce. Capa dieciocho. Y la última, con el experto treinta encendido. ¿Es un problema? No necesariamente. El balanceo de carga que vimos en el episodio cuatro se mide sobre batches enormes y variados de entrenamiento. Acá son mil tokens de cinco tipos de texto: un corpus angosto produce cargas desparejas. Ojo también con comparar columnas entre filas: el experto cinco de la capa seis y el de la capa doce no tienen nada que ver. Son redes distintas, con el mismo número de casualidad.

Parte 3 · ¿Tema o tipo de texto?

La pregunta de fondo: ¿el router separa por tema, o por la forma del texto?

Cuánto se parecen dos repartos

p_i = \frac{\#\{\text{asignaciones al experto } i\}}{8\,T}

el reparto de un grupo de textos: suma 1

\mathrm{TV}(p, q) = \tfrac{1}{2} \textstyle\sum_{i=1}^{32} \lvert p_i - q_i \rvert

total variation: 0 iguales, 1 sin nada en común

\mathrm{TV}(\text{mitad}_1, \text{mitad}_2) \approx \text{ruido}

la vara: dos mitades del mismo tema

Necesitamos un número que diga cuánto se parecen dos repartos. Para un grupo de textos, p sub i es la fracción de todas las asignaciones que fueron al experto i: hay ocho por token, así que dividimos por ocho T. Es un vector de treinta y dos números que suma uno. En el dibujo, un ejemplo con seis expertos. La distancia de total variation es la mitad de la suma de las diferencias en valor absoluto. Da cero si los repartos son idénticos y uno si no comparten ningún experto. En el dibujo, las diferencias del ejemplo; sus tamaños suman cero coma tres, y la mitad es cero coma quince. Y como con mil tokens hay ruido, la vara es comparar dos mitades del mismo tema: eso es lo que da el azar del muestreo.

Cocina contra astronomía contra código

La función tv implementa la fórmula. Primero, la vara: las seis recetas partidas en dos mitades de tres. Después, recetas contra astronomía: mismo idioma y misma sintaxis, tema distinto. Y recetas contra código y contra números: otro tipo de texto. Mirá la tabla. Entre mitades de recetas, cero coma cero seis. Recetas contra astronomía: entre cero coma cero cinco y cero coma cero ocho, o sea, igual que el ruido. En cambio, contra código da alrededor de cero coma treinta y siete, y contra números, entre cero coma cincuenta y cero coma sesenta y cuatro. El tema casi no mueve al router; el tipo de texto, muchísimo.

python

def reparto(idx):                         # distribución de las asignaciones
    d = torch.bincount(idx.reshape(-1), minlength=32).float()
    return d / d.sum()
def tv(a, b):                             # total variation: 0 igual, 1 disjuntas
    return 0.5 * (reparto(a) - reparto(b)).abs().sum().item()

coc = TEXTOS["cocina"]
A, B = rutas(coc[:3]), rutas(coc[3:])      # dos mitades del mismo tema
C, D = rutas(coc), rutas(TEXTOS["astronomia"])
P, N = rutas(TEXTOS["codigo"]), rutas(TEXTOS["numeros"])
print("capa mitad astro código núm.")
for c in [0, 12, 23]:
    print(f"{c:4d} {tv(A[c], B[c]):5.2f} {tv(C[c], D[c]):5.2f}"
          f" {tv(C[c], P[c]):6.2f} {tv(C[c], N[c]):5.2f}")

text

capa mitad astro código núm.
   0  0.06  0.05   0.38  0.64
  12  0.06  0.08   0.36  0.50
  23  0.04  0.05   0.37  0.50

Distancia a las recetas (capa 12)

En barras, la capa doce. El piso del ruido: dos mitades de recetas. Astronomía queda pegada a ese piso. El código se despega. Y los números, todavía más. Para el router, una receta y un texto sobre cometas son casi lo mismo: palabras en castellano, artículos, comas. Lo que cambia el ruteo es la forma de los tokens.

otra mitad de recetas: 0,06
astronomía: 0,08
código Python: 0,36
números: 0,50

Jiang et al., 2024 — Mixtral of Experts · arXiv:2401.04088

sintaxis > tema · «We do not observe obvious patterns in the assignment of experts based on the topic»; la selección parece «more aligned with the syntax rather than the domain».

No es una rareza de Granite. En el paper de Mixtral, un MoE de ocho expertos por capa que elige dos, midieron el reparto sobre distintos dominios de un corpus público: papers de arXiv escritos en LaTeX, resúmenes de biología, filosofía. No encontraron patrones obvios por tema. Lo que sí vieron fue estructura sintáctica: la indentación del código siempre va a los mismos expertos, y una palabra como self, en Python, también.

Parte 4 · Qué tokens recibe cada experto

Si no es el tema, ¿qué es? Miremos los tokens que le llegan a algunos expertos de la capa doce.

Los tokens de tres expertos

Tomamos la capa doce. Para cada experto, contamos qué tokens le llegaron; el guión bajo marca el espacio antes de la palabra. Una función chica decide si un token es un número. Y miramos tres expertos. El experto diez recibió setenta y nueve tokens y el ochenta y siete por ciento son dígitos: uno, tres, dos, cero, ocho. El veintisiete casi no ve números: sus favoritos son the, el, la. Y el veintitrés recibe de, el, la, y paréntesis. Son especialistas, sí, pero en clases de tokens, no en temas: dígitos, artículos, preposiciones.

python

idx = torch.cat(elegidos[12])                 # capa 12: [T, 8]
por_experto = collections.defaultdict(collections.Counter)
for token, fila in zip(toks, idx.tolist()):
    for e in fila:
        por_experto[e][token.replace("Ġ", "_")] += 1

def es_digito(t): return t.strip("_").isdigit()
for e in [10, 27, 23]:
    c = por_experto[e]
    dig = sum(n for t, n in c.items() if es_digito(t)) / sum(c.values())
    top = " ".join(t for t, _ in c.most_common(5))
    print(f"e{e}: {sum(c.values())} tokens, {dig:.0%} dígitos")
    print("   ", top)

text

e10: 79 tokens, 87% dígitos
    1 3 2 0 8
e27: 46 tokens, 2% dígitos
    _the _el _la _f _su
e23: 83 tokens, 0% dígitos
    _de _el _la ( _the

Cómo leer esto sin engañarse

Son 1055 tokens: alcanza para ver patrones grandes, no finos
Un token pasa por 8 expertos: ninguno «es» el de los dígitos
El experto i de una capa no tiene relación con el i de otra
Cada número, contra su vara: el azar o dos mitades

Antes de sacar conclusiones, cuatro cuidados. Mil cincuenta y cinco tokens alcanzan para ver efectos grandes, como dígitos contra palabras, pero no para afirmar cosas finas. Cada token pasa por ocho expertos a la vez, con pesos distintos: decir este es el experto de los números es una simplificación. El índice de un experto no significa nada entre capas: comparar el experto diez de la capa doce con el diez de la capa trece no tiene sentido. Y siempre, una vara: el azar, o dos mitades del mismo grupo.

Parte 5 · Tokens vecinos, expertos vecinos

Una última medición, también del paper de Mixtral: ¿dos tokens seguidos van a los mismos expertos?

La vara del azar

P(\text{misma primera elección}) = \tfrac{1}{32} \approx 3{,}1\,\%

si el router eligiera al azar, parejo

\mathbb{E}[\,|S_t \cap S_{t+1}|\,] = \textstyle\sum_{i \in S_{t+1}} P(i \in S_t)

contar los comunes uno por uno

= 8 \cdot \tfrac{8}{32} = 2

dos expertos en común, por azar

Primero, qué daría el azar. Si el router eligiera al azar y parejo, la primera elección de un token coincidiría con la del vecino una vez cada treinta y dos: tres coma uno por ciento. Para el conjunto entero de ocho: S sub t son los ocho expertos del token t. El valor esperado de cuántos comparten, esa E de doble trazo, es la suma, sobre los ocho expertos del token siguiente, de la probabilidad de que cada uno esté también entre los ocho del anterior. En el dibujo, una barra por cada uno de los ocho, y cada una vale un cuarto, porque ocho de treinta y dos es un cuarto. Ocho por un cuarto: dos expertos en común, por pura casualidad.

Repetición entre tokens seguidos

Dentro de cada texto, alineamos cada token con el que sigue. Comparamos la primera elección. Y contamos cuántos de los ocho expertos comparten. Para tres capas. Contra el tres coma uno por ciento del azar, la primera elección se repite el trece coma cinco por ciento en la capa cero, el veinticinco en la doce y casi el treinta en la última. Y en lugar de dos expertos en común, entre tres y cuatro. Los tokens vecinos tienden a ir a los mismos expertos, y más en las capas altas. El paper de Mixtral encontró lo mismo, y señala que esa localidad sirve para guardar expertos en una caché, pero también puede sobrecargar a algunos cuando los expertos viven en máquinas distintas. De eso se trata el último episodio.

python

def repeticion(textos):
    misma, comunes = [], []
    for idx in textos:                       # [T, 8] de un texto
        a, b = idx[:-1], idx[1:]             # token i y token i+1
        misma += (a[:, 0] == b[:, 0]).tolist()   # ¿misma 1.ª elección?
        comunes += [len(set(x.tolist()) & set(y.tolist()))
                    for x, y in zip(a, b)]
    return sum(misma) / len(misma), sum(comunes) / len(comunes)

print(f"azar: misma 1.ª {1 / 32:.1%} · comunes {8 * 8 / 32:.1f}")
for c in [0, 12, 23]:
    r1, comun = repeticion(elegidos[c])
    print(f"capa {c:2d}: misma 1.ª {r1:.1%} · comunes {comun:.1f}")

text

azar: misma 1.ª 3.1% · comunes 2.0
capa  0: misma 1.ª 13.5% · comunes 3.1
capa 12: misma 1.ª 25.1% · comunes 3.9
capa 23: misma 1.ª 29.4% · comunes 4.1

Qué nos llevamos de Granite

El router especializa por tipo de token, no por tema
En un corpus angosto, la carga es muy despareja
Tokens vecinos repiten expertos: localidad
Un experto no es un módulo con un tema: es una pieza de una combinación de 8

Juntemos. En Granite, como en Mixtral, el router separa clases de tokens: dígitos, artículos, puntuación, código. El tema casi no se nota. La carga en un corpus angosto es muy despareja, aunque en entrenamiento esté balanceada. Los tokens vecinos repiten expertos bastante más que el azar. Y la imagen correcta de un experto no es la de un especialista en un área, sino la de una pieza que se combina con otras siete, distinta en cada token.

Resumiendo

Hook en el router: logits, pesos e índices por capa
Granite: softmax solo entre los 8 elegidos
Carga real muy despareja en textos angostos
TV: recetas ≈ astronomía; código y números, lejos
Vecinos: 1.ª elección repetida 13–29 % vs. 3,1 % al azar
Próximo: el costo de verdad (memoria, batch, all-to-all, upcycling)

Papers del episodio

Jiang et al., 2024 — Mixtral of Experts

¿Qué tal esta lección?

Sé el primero en calificarla

Entra para calificar

¿Terminaste la lección?

Márcala para sumar +50 XP y mantener tu racha.

Comentarios

Entra para comentar en el minuto exacto del video y responder dudas.Entrar