Mixture of Depths (MoD)

Mixture of Depths (MoD) es una técnica de Transformer en la que un enrutador aprendido decide, bloque a bloque, qué tokens se procesan y cuáles se saltan el bloque mediante su conexión residual, de modo que distintos tokens reciben distintas cantidades de cómputo. David Raposo, Adam Santoro y sus colegas de Google DeepMind la presentaron en abril de 2024 en Mixture-of-Depths: asignación dinámica de cómputo en modelos de lenguaje basados en Transformer. Una capacidad fija k determina cuántos tokens procesa cada bloque enrutado, por lo que el presupuesto de cómputo y las formas de los tensores se conocen de antemano, aunque la elección de los tokens dependa de la entrada. Los modelos de lenguaje grandes (LLM) MoD igualaron la calidad de referencia con los mismos FLOPs de entrenamiento y tiempo de ejecución, aunque utilizaron una fracción de los FLOPs por pasada hacia delante y generaron muestras más de un 50 % rápido. Con un presupuesto de entrenamiento fijo, el mejor modelo MoD fue más grande que el mejor modelo de referencia y alcanzó una pérdida menor.

Cómo funciona el enrutamiento de MoD#

Antes de un bloque enrutado, un enrutador lineal asigna a cada token una puntuación escalar. Los k tokens con mayor puntuación pasan por la autoatención y el MLP del bloque; sus salidas se multiplican por las puntuaciones del enrutador y se vuelven a añadir al flujo residual. Multiplicar por la puntuación sitúa el enrutador en la ruta del gradiente, por lo que la retropropagación puede enseñarle qué tokens resultan beneficiosos. Los demás tokens pasan sin cambios y se excluyen de la atención de ese bloque; por tanto, los tokens posteriores no pueden atenderlos allí y el coste de atención se reduce con el cuadrado de la capacidad: la mitad de los tokens supone una cuarta parte de los productos consulta-clave. La mejor configuración del artículo enrutaba un bloque de cada dos con una capacidad del 12,5 % de la secuencia. Esta es una versión mínima de un bloque enrutado, solo con el MLP, en PyTorch:

import torch
from torch import nn

torch.manual_seed(0)
batch, seq_len, dim = 2, 64, 128
capacity = int(0.125 * seq_len)  # 12.5% of tokens, the paper's best setting
x = torch.randn(batch, seq_len, dim)
router = nn.Linear(dim, 1)
block = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

# The router scores every token, and only the top-k tokens enter the block
weights, idx = router(x).squeeze(-1).topk(capacity, dim=1)
idx = idx[..., None].expand(-1, -1, dim)
update = block(x.gather(1, idx)) * weights[..., None]  # router weight keeps the router trainable

# Selected tokens get the residual update; the rest skip the block unchanged
out = x.scatter_add(1, idx, update)
unchanged = (out == x).all(dim=-1).sum().item()
print(f"Block ran on {batch * capacity} of {batch * seq_len} tokens; {unchanged} skipped it unchanged")

Imprime Block ran on 16 of 128 tokens; 112 skipped it unchanged: el MLP del bloque realizó una octava parte de su trabajo habitual y el tensor de salida conserva la forma de la secuencia completa para el siguiente bloque.

Enrutamiento durante la generación#

La selección top-k no es causal: que un token quede entre los k primeros de su secuencia depende de las puntuaciones de los tokens que aparecen después, que aún no existen cuando un modelo genera texto token a token. Los autores hicieron causal el enrutamiento de dos maneras: mediante una función de pérdida auxiliar de entropía cruzada binaria que enseña a la salida del enrutador a predecir la pertenencia al top-k y que solo modificó la pérdida del modelado del lenguaje en un 0,2-0,3 %, o mediante un pequeño predictor MLP independiente que recibe entradas sin gradiente y alcanzó rápidamente una precisión del 97-99 %. En ambos casos, el muestreo autorregresivo apenas perdió calidad. Los enrutadores entrenados también muestran un patrón: algunos tokens activan todos los bloques enrutados a lo largo de la profundidad del modelo, mientras que otros los rodean siempre que pueden.

MoD frente a técnicas relacionadas#

• Mezcla de expertos (MoE): Enruta cada token a uno o varios MLP expertos dentro de una capa, lo que añade parámetros y mantiene fijo el cómputo por token. El artículo también combinó ambos enfoques en MoDE, donde un token puede dirigirse a un experto o a una operación nula, y las mejoras se acumularon.
• Salida anticipada: Un token se detiene tras una capa y se salta todas las demás. MoD puede saltarse un bloque y seguir utilizando los posteriores.
• Fusión de tokens (ToMe): Fusiona tokens similares en un Transformer de visión para que las capas posteriores procesen menos, sin un enrutador aprendido.
• Poda de modelos: Elimina pesos de forma permanente, de modo que todas las entradas reciben la misma red reducida.
• Enrutamiento de modelos: Envía una solicitud completa a uno de varios modelos; es la misma idea aplicada al nivel de los modelos completos.

La mayoría de las redes fijan su profundidad al construirse. Ultralytics YOLO26 la establece según la escala del modelo: el multiplicador de profundidad de su YAML del modelo es 0.50 para las escalas n, s y m, y 1.00 para l y x; así, una etapa C3k2 indicada con 2 repeticiones ejecuta 1 bloque en YOLO26n y 2 en YOLO26x, y todas las imágenes pasan por ellos.