Mixture-of-Depths: técnica para asignar dinámicamente recursos de cómputo en transformers

• Mixture-of-Depths (MoD) de Google DeepMind es un enfoque en el que los modelos de lenguaje basados en transformers no usan los mismos FLOPs para todos los tokens, sino que, en cada capa, solo los tokens necesarios participan en los cálculos de attention y MLP
• Un router por capa genera un peso escalar para cada token, y solo los tokens dentro de una capacidad top-k predefinida pasan por el bloque; el resto lo evita mediante una residual connection
• Al fijar k de antemano, el grafo de cómputo y los tamaños de tensores se mantienen estáticos, mientras que la distribución del cómputo por token cambia dinámicamente según el contexto
• En los experimentos, bajo el criterio isoFLOP, MoD logró una loss más baja que el transformer base con el mismo tiempo de entrenamiento, o bien redujo los FLOPs por forward pass para el mismo rendimiento, acelerando los steps de entrenamiento y muestreo
• La mejor configuración fue aplicar bloques con 12.5% de capacidad en bloques alternos (every other block), y en el muestreo auto-regressive redujo el problema de no causalidad del top-k mediante routing basado en un predictor

---

El desperdicio de cómputo al que apunta MoD

• Un modelo de lenguaje transformer típico usa la misma cantidad de cómputo para todos los tokens en el forward pass
• MoD parte de la premisa de que no todos los tokens y secuencias requieren el mismo tiempo y esfuerzo para la predicción, y distribuye el cómputo a los tokens necesarios mientras reduce el presupuesto total de cómputo
• La conditional computation existente puede crear grafos de cómputo dinámicos, lo que puede no encajar con los grafos de cómputo estáticos que prefiere el hardware actual
• Si el usuario define el presupuesto total de cómputo antes del entrenamiento, el modelo aprende dentro de ese límite dónde usar el cómputo por token y por capa
  • La cantidad de cómputo, el uso de memoria y los FLOPs por forward pass pueden predecirse de antemano
  • Los tokens objetivo del routing cambian según el contexto de entrada

Estructura de routing y forma de reducir cómputo

• MoD usa un router como un MoE transformer, pero en lugar de elegir uno entre varios experts, envía los tokens por una de dos rutas
  • Cómputo de bloque estándar, que incluye attention y MLP
  • Residual connection, que pasa el valor tal cual
• La residual connection tiene bajo costo computacional, y la salida del bloque queda determinada por la entrada
• Si la capacity del bloque se configura por debajo del número total de tokens T, los FLOPs por forward pass se reducen frente a un transformer base
  • Por ejemplo, si la capacity se reduce a T/2, la multiplicación de matrices query-key de attention pasa de T² a (T/2)², por lo que los FLOPs quedan en torno al 25%
• Si la capacity se fija en 0, es muy rápido, pero el rendimiento empeora porque no usa la mayoría de los parámetros del transformer; si se fija en T, es igual que un transformer base
• El objetivo es encontrar, entre esos dos extremos, un punto que iguale o supere el rendimiento del transformer base y a la vez aumente la velocidad de los steps

Routing top-k expert-choice

• En cada capa l, el router genera un peso escalar para cada embedding de token
• Según la capacity C definida por el usuario, solo los tokens correspondientes al top-k dentro de la secuencia participan en los cálculos de attention y MLP
• Los tokens no seleccionados saltan el cómputo del bloque y se entregan a la siguiente capa mediante la residual connection
• Este enfoque se basa en expert-choice routing
  • En token-choice routing, el token elige la ruta, por lo que pueden surgir problemas de load balancing
  • En expert-choice routing, cada ruta elige sus tokens top-k, por lo que la capacity se llena exactamente
• En MoD, como en la práctica las rutas son dos —“cómputo de bloque” y “bypass”—, una sola operación top-k puede dividir los tokens en dos conjuntos
• Los pesos del router se multiplican por la salida del bloque, quedan sobre el gradient path y reciben la influencia del descenso por gradiente durante el entrenamiento de language modeling

Problema de no causalidad en el muestreo

• La operación top-k de expert-choice routing es no causal (non-causal)
  • Para decidir qué token entra en el top-k, hay que conocer también los pesos del router de tokens posteriores
  • En el muestreo auto-regressive no se puede usar información de tokens futuros
• Para evitarlo, se probaron dos métodos
  • Una pérdida auxiliar de binary cross-entropy para que la sigmoid de la salida del router prediga si será seleccionada por top-k
  • Un pequeño predictor MLP auxiliar que recibe la misma entrada que el router, pero con stop-gradient aplicado
• El primer método afectó el objective principal de language modeling en aproximadamente 0.2~0.3%
• El segundo método basado en predictor no afectó el objective de language modeling y tampoco tuvo un impacto significativo en la step speed
• La tarea auxiliar de predicción alcanzó experimentalmente 99% de exactitud con rapidez

Resultados experimentales y observaciones

• Los experimentos iniciales se realizaron buscando hiperparámetros de MoD con un presupuesto total de 6e18 FLOPs
• La configuración óptima de MoD fue aplicar routing en bloques alternos y fijar top-k en 256
  • Con una longitud de secuencia de 2048, pasan por attention y MLP 256 tokens, es decir, solo 12.5%
  • Los 1792 tokens restantes, es decir, 87.5%, evitan el bloque
• Una variante MoD de 220M parámetros tuvo un rendimiento ligeramente mejor que el baseline óptimo isoFLOP, mientras que durante el entrenamiento la velocidad por step fue más de 60% mayor; en el mismo hardware, el wall-clock time total de entrenamiento fue casi igual
• También en el análisis isoFLOP de 6e18, 2e19 y 1e20 FLOPs, los modelos FLOP-optimal de MoD tendieron a tener más parámetros que el baseline
• Existen variantes de MoD que logran el mismo rendimiento o una loss más baja con menos FLOPs por forward pass
• La mejora en step speed proviene de dos factores
  • Como algunos tokens evitan el bloque, se reducen los FLOPs per parameter para el mismo tamaño de modelo
  • Se puede elegir una variante MoD más pequeña que alcance el mismo rendimiento que el baseline
• MoD con stochastic routing mostró un rendimiento mucho menor que el transformer base y que MoD normal, lo que mostró la importancia del routing aprendido
• En la evaluación auto-regressive se usaron 256,000 secuencias y 500M tokens de held-out data; incluso al cambiar a routing basado en predictor, la caída de rendimiento fue pequeña
• MoD también puede combinarse con MoE para formar Mixture-of-Depths-and-Experts (MoDE)
  • Staged MoDE decide antes de attention si envía el token al bloque o lo hace evitarlo
  • Integrated MoDE integra el routing de MoD insertando un expert “no-op” entre los MLP experts existentes
• En el análisis de routing se observó un patrón en el que algunos tokens pasan con frecuencia por varios bloques, mientras que otros los evitan siempre que pueden
• Un análisis preliminar mostró que los tokens que pasan con frecuencia por los bloques se correlacionan con casos de mayor entropy en la predicción de salida, lo que se relaciona con la posibilidad de que sean tokens más difíciles de predecir

[GN⁺]

Opiniones de Hacker News

• Una ruteo más complejo sin duda parece que se volverá más común.
  En particular, creo que algún día iremos hacia un ruteo recursivo que vuelva a pasar por el conjunto de expertos. En el futuro, el “proceso de pensamiento (chain-of-thought)” probablemente ocurrirá de forma recursiva dentro del modelo.

  • A estos objetos hipotéticos se les podría llamar Recursive Neural Networks.
  • Lo que se describe aquí se parece un poco a una línea de investigación de Universal Transformers. La idea es pasar los embeddings de entrada varias veces por un solo bloque transformer, y que un módulo separado decida si el embedding ya está lo suficientemente “cocido” como para sacarlo.
    Más cercano a la idea de “expertos” está el paper Sparse Universal Transformers del año pasado, que combina Universal Transformer con mezcla dispersa de expertos, donde un mecanismo de gating decide qué bloques transformer usar y en qué orden.
    No es mi área de especialidad, pero entiendo que es difícil entrenarlo bien y que, para obtener resultados similares a un transformer normal, requiere más cómputo total en inferencia. Aun así, es una dirección interesante, y creo que el hecho de que haya un límite superior al número de pasos de cómputo por token es una de las grandes desventajas de la arquitectura clásica de transformers.
  • Creo que la razón por la que esto todavía no se ha logrado es que, durante el entrenamiento, no hay forma de decidir cuántas veces hacer la recursión.
    Si eliges un número aleatorio de veces o pruebas varias profundidades recursivas, la salida se “difumina”. Es decir, ya no queda claro si la salida de cierta capa debe aportar información importante para el resultado final o producir la mejor entrada para la siguiente ronda de recursión.
  • La atención es básicamente ruteo, y estas otras formas de ruteo pueden darle al modelo opciones más gruesas, lo que potencialmente facilita el entrenamiento.
  • La tendencia sin duda va hacia más ruteo dinámico, pero creo que MoE/MoD/MoDE no tanto habilitan un razonamiento más profundo, sino que más bien permiten almacenar hechos adicionales reduciendo solapamientos dentro de los pesos.
    El razonamiento más profundo probablemente vendrá de la dinamicidad a nivel de token, más que de la dinamicidad a nivel de capa. Por ejemplo, está el paper reciente Quiet-STaR, que genera tokens de razonamiento que el modelo luego descarta: https://arxiv.org/abs/2403.09629

• Podría ser el paper más importante de 2024.
  La idea de querer modelos que no usen la misma cantidad de cómputo para todos los tokens existe desde hace tiempo, pero es la primera vez que veo un mecanismo convincente para lograrlo.

  Equipped with these new methods, we can sample autoregressively by choosing to route tokens to or around a block based on the router’s output, which does not depend on any information from future tokens. We provide empirical evidence that this is a relatively easy auxiliary task that quickly achieves 99% accuracy.
  ¿No es algo sorprendente?

  • Sparse Universal Transformer es más antiguo y ya hacía salida temprana basada en ruteo.
  • ¿Hasta el punto de ser el más importante? Que no todos los tokens necesiten toda la ventana de contexto debería ser una optimización obvia.

• Versión introductoria simple:
  Imagina que tienes un asistente inteligente que puede entender y procesar el lenguaje. Normalmente, este asistente presta la misma atención a todas las palabras, sin importar qué tan importante sea cada una para el significado completo.
  Ahora supón que encontramos una forma de enseñarle al asistente a usar sus “recursos cerebrales” de manera más inteligente. En vez de prestar la misma atención a todas las palabras, se concentra más en las palabras más importantes para entender el significado, y ajusta ese foco sobre la marcha según el contexto.
  También ponemos un límite a la cantidad total de “recursos cerebrales” que puede usar en un momento dado para que el asistente no se sobrecargue. Es como darle un presupuesto y decirle: “solo puedes gastar recursos en cierta cantidad de palabras a la vez”. Entonces el asistente tiene que decidir qué palabras son las más importantes.
  Incluso con ese límite, el asistente ajusta de forma flexible su uso de recursos. Usa más en algunas palabras y menos en otras, priorizando según la situación.
  Al entrenarse de esta manera, el asistente funciona de forma eficiente sin dejar de prestar atención de manera inteligente. Entiende tan bien como un asistente que presta la misma atención a todas las palabras, pero usa menos recursos cerebrales en total, por lo que responde y procesa nueva información mucho más rápido.

  • Entiendo que es ELI5, pero ¿no es eso ya lo que hace la atención? Se enfoca específicamente en las palabras contextualmente importantes de la secuencia anterior.

• Basándome en lo que entendí, hice un pequeño resumen aquí:
  https://lifeinthesingularity.com/p/googles-breakthroughs-in-...

  • Buen artículo. Me recuerda al estilo de New Scientist. Tiene esa sensación de “explicar sobre la marcha, pero manteniéndolo breve”, así que sirve para hacerse una idea del campo.

• Es muy parecido a una mezcla de expertos. Solo que, en lugar de rutear tokens a varios expertos, es como “asignarlos a un único experto que se puede saltar dinámicamente”.

  • Mezclar ambos sería bastante genial. Podría reducir aún más el cómputo de MoE manteniendo el rendimiento.

• “This is more computationally efficient than performing a full content-based lookup across an entire memory buffer for each step in the future, and could be one step towards drastically increasing the context-length available for making a prediction.”
  ¿Así es como se hace posible una ventana de contexto de 10 millones de tokens? ¿O se refiere a ventanas de contexto más largas en el futuro?

• Después de intentar entender e implementar algunos algoritmos RASP, la conclusión a la que llegué fue que ciertas funciones necesitan una cantidad determinada de capas transformer para funcionar.
  Siguiendo esta lógica, se ve que las funciones aprendidas por un transformer pueden distribuirse entre varias cabezas. Repetir esas funciones podría ser muy valioso para entender y resolver problemas, pero el método de inferencia actual no puede ejecutar repetidamente bloques consecutivos de cabezas. Este paper realmente parece apuntar en una dirección prometedora.
  [1] https://arxiv.org/pdf/2106.06981.pdf
  [2] https://www.youtube.com/watch?v=t5LjgczaS80

• Tal vez la única desventaja de que los LLM avancen tan rápido es que los papers salen más rápido de lo que cualquiera, al menos fuera de Google, puede aprender y probar las mejoras.
  Cuando empecé con deep learning, ReLU y dropout estaban en pleno auge, y podías cambiar una o dos líneas de código en una 1080 de consumo y probar en unas horas si había una mejora. Ahora parece que hay que esperar semanas hasta que alguien como Mistral lo pruebe.

  • Bienvenido a la pobreza de GPU.
    Yo me estoy enfocando en enfoques de cuantización y pruebo con GPUs viejas de generaciones anteriores.

• hu-po hizo una reseña en livestream profundo de papers de IA.
  Muy recomendable; aquí cubre el paper sobre mixture-of-depths que se está discutiendo: https://www.youtube.com/watch?v=Teru_qIdB8Y