Attention-Aware Routing: el router de un MoE mejora si mira la atención

Cambiar solo los parámetros del router de un modelo Mixture-of-Experts, sin tocar ni un peso del resto del transformer, basta para subir 3,37 puntos porcentuales en GSM8K. Es el resultado central de Attention-Aware Routing, un paper que ha aparecido hoy en arXiv y que se apoya en una observación sencilla: el router de un MoE decide con muy poca información sobre el contexto.

En un MoE convencional, cada token pasa por un router que elige a qué expertos enviarlo y con qué peso. Esa decisión se toma a partir del hidden state del token, un vector que resume lo que el modelo sabe de ese token en esa capa, pero que apenas codifica de forma explícita cómo se está distribuyendo la atención en la ventana reciente. Los autores parten de ahí: si la atención es el mecanismo por el que el modelo reparte su foco sobre el contexto, el router debería poder consultarla.

Qué propone AAR

Attention-Aware Routing, AAR, añade al router un conjunto de rasgos extraídos de una ventana deslizante de pesos de atención. Son rasgos de dos tipos:

Temporales: cómo evoluciona la atención a lo largo de las últimas posiciones.
Espectrales: la estructura de frecuencia de esos pesos, que resume patrones repetitivos o dispersos.

Ese resumen del estado contextual del modelo se mantiene separado del hidden state, de modo que el router recibe dos señales distintas en lugar de una. El diseño experimental es deliberadamente estrecho: el transformer base queda congelado por completo y solo se entrenan los parámetros de enrutamiento. Así, cualquier mejora se atribuye únicamente al router, sin ruido de un fine-tuning general.

Con esa configuración, AAR mejora GSM8K en 3,37 puntos porcentuales frente a una línea base que también hace SFT solo sobre el router, en OLMoE, el MoE abierto del Allen Institute for AI. La comparación es justa porque ambos entrenan exactamente los mismos parámetros. Lo único que cambia es qué información ve el router.

Enrutamiento y atención forman un circuito

La parte más interesante del trabajo no es la cifra, sino el análisis mecánico. Los autores muestran que un cambio de enrutamiento en la capa l se propaga por el residual stream y amplifica los attention sinks de la capa l+1. Los attention sinks son esas posiciones, típicamente el primer token, que absorben una fracción desproporcionada de la atención y actúan como válvula de escape del mecanismo.

Es decir, modificar solo el router reconfigura la atención de la capa siguiente sin que nadie haya tocado el mecanismo de atención. Routing y atención no son módulos independientes que se suman, sino un circuito acoplado en el que una decisión local sobre expertos altera cómo el modelo mira su contexto. Para quien se dedica a interpretabilidad, es un dato que complica la idea de estudiar cada componente por separado.

Respuestas erróneas más cortas

Hay un efecto secundario útil. AAR reduce lo que el paper llama generación divergente larga: las respuestas incorrectas se acortan, mientras que las correctas mantienen su longitud. En razonamiento matemático, un modelo que se equivoca tiende a alargarse en cadenas que no llevan a nada. Que el router ayude a cortarlas antes reduce tokens generados y, en un despliegue real, coste.

No vale en todas las capas

El último hallazgo es una advertencia. AAR es muy sensible a la profundidad: aplicarlo indiscriminadamente en todas las capas puede degradar la recuperación factual, mientras que el razonamiento matemático responde de forma distinta, según el propio abstract. El detalle por capa está en el cuerpo del paper, pero el mensaje es claro: elegir en qué capas se aplica forma parte del método, no es un ajuste opcional.

Para quién es relevante

Este trabajo no cambia nada para quien consume modelos por API. Sí interesa a:

Equipos que entrenan o afinan MoE abiertos como OLMoE, Mixtral o las variantes MoE de Qwen3, porque propone una mejora barata en parámetros que se puede probar sin reentrenar el modelo completo.
Investigadores de interpretabilidad, por la evidencia de que routing y attention sinks están conectados a través del residual stream.
* Quien optimiza inferencia, por la reducción de longitud en salidas erróneas, que se traduce en menos cómputo desperdiciado.

Quedan preguntas abiertas: el paper se evalúa en OLMoE, un modelo relativamente pequeño, y no sabemos si el efecto escala a MoE de decenas de expertos y cientos de miles de millones de parámetros. Tampoco está claro el coste de calcular los rasgos espectrales en cada paso de inferencia.

Desde ElephantPink nos parece más valioso el hallazgo del circuito acoplado que los tres puntos de GSM8K, porque explica algo que hasta ahora se trataba como dos problemas separados. Si se confirma en modelos más grandes, será una razón más para dejar de pensar en el router como un simple selector.