[Traducción] MoD (Mixture-of-Depths): un enfoque para optimizar el cómputo en modelos de lenguaje basados en Transformer, y MoDE (MoD+MoE)

(discuss.pytorch.kr)

7 puntos por ninebow 2024-04-13 | Aún no hay comentarios. | Compartir por WhatsApp

PyTorchKR

Recientemente se están usando mucho los métodos MoE (Mixture-of-Experts), que procesan los tokens de entrada seleccionando uno o más entre varios bloques de “expertos”, como Jamba, Qwen1.5-MoE o DBRX. Este enfoque MoE muestra mejor desempeño al asignar de forma eficiente los recursos de cómputo mediante el control dinámico de una parte de los enormes LLM.
Si MoE ajusta dinámicamente el ancho (width) del modelo seleccionando algunos expertos entre varios, el MoD (Mixture-of-Depths) que presentamos esta vez adopta un enfoque para ajustar dinámicamente la profundidad (depth) del modelo. Es decir, en lugar de ejecutar el cómputo sobre todas las capas del modelo para ciertos tokens, realiza el cómputo de forma dinámica saltándose algunas capas. Veámoslo juntos.

MoD (Mixture-of-Depths): un enfoque para optimizar el cómputo en modelos de lenguaje basados en Transformer (Mixture-of-Depths: Dynamically allocating compute in transformer-based language models)

Introducción

Tradicionalmente, los modelos Transformer asignan los mismos recursos de cómputo a todas las partes de la secuencia de entrada. Sin embargo, el paper de MoD (Mixture-of-Depth) Transformer que presentamos aquí intenta asignar recursos de cómputo de forma dinámica a posiciones específicas de la secuencia en un Transformer, y variar esa asignación a través de distintas capas. Para usar esta estrategia de Mixture-of-Depths (MoD), los autores emplean un mecanismo de enrutamiento top-k para decidir qué tokens procesar y así controlar el volumen total de cómputo. Esta técnica iguala o supera el rendimiento de modelos existentes, requiere muchos menos FLOPs y ofrece hasta un 50% de mejora de velocidad durante el muestreo después del entrenamiento.

La arquitectura Transformer presentada en el paper Attention is All You Need

Los modelos Transformer recientes y los LLM (Large Language Model) siguen aumentando en escala y complejidad. Estos modelos utilizan miles de millones de parámetros para acercarse a capacidades de comprensión y generación de lenguaje a nivel humano, y como resultado muestran un desempeño impresionante. Sin embargo, estos avances vienen acompañados de costos computacionales y consumo energético significativos, y el aumento del tamaño del modelo incrementa drásticamente el tiempo de entrenamiento e inferencia.

Esto se debe a que el Transformer realiza la misma cantidad de cómputo para todos los tokens de la secuencia de entrada. Pero la idea de partida es que no todos los tokens aportan la misma cantidad de información: algunos cumplen un papel clave para entender el contexto, mientras que otros son relativamente menos importantes.

El MoD Transformer permite al modelo concentrar los recursos donde realmente se necesitan al asignar cómputo dinámicamente a cada token. Esto reduce el costo computacional total del modelo y, al mismo tiempo, aumenta la eficiencia del procesamiento de información al prestar más atención a los tokens importantes. En modelos grandes como los LLM, el enfoque de MoD puede convertirse en una forma importante de ahorrar tiempo y recursos tanto en entrenamiento como en inferencia.

Introducción a la técnica MoD (Mixture-of-Depths)

La técnica MoD integra un mecanismo de enrutamiento top-k que selecciona dinámicamente qué tokens procesar en cada capa del Transformer. Gracias a esto, solo una parte de los tokens seleccionados pasa por todo el proceso de cómputo, mientras que los demás omiten una o más capas.

Figura 1 / Izquierda: vista general del Transformer MoD (Mixture-of-Depths): al igual que en una estructura MoE (Mixture-of-Experts), la ruta de cómputo se selecciona mediante un router. A diferencia de MoE, aquí se decide si ejecutar o no el bloque estándar (Self-Attention y MLP). Si el token de entrada ($X_i$) no activa el cómputo completo según la decisión del router, no pasa por el cómputo de esa capa.
Figura 1 / Derecha: resultado de las decisiones de enrutamiento para una secuencia corta reducida a 64 tokens. (Eje X: secuencia, eje Y: capa). El color morado indica que se ejecutó todo el cómputo de esa capa, y el naranja indica que se omitió el cómputo de esa capa.

Como se explica en la figura anterior, MoD implementa un router predictivo dentro de la arquitectura Transformer para decidir, para cada token, si se ejecuta o no el cómputo de una capa. Esto permite concentrar los recursos computacionales en los tokens relativamente “importantes”.

Entrenamiento del MoD Transformer

El proceso de entrenamiento del MoD Transformer es, en esencia, similar al de un Transformer estándar. La secuencia de entrada se suministra al modelo y, al pasar por cada capa, se realiza el cómputo necesario. Sin embargo, el MoD Transformer incluye un paso adicional en el que el router decide qué ruta computacional seguirá cada token.

Cálculo de los pesos del router: el router calcula, en cada capa del modelo, un peso para cada token. Este peso representa la cantidad de cómputo que debe recibir ese token y se basa en su importancia.
Selección de los tokens principales: con base en los pesos calculados, el router selecciona los tokens principales para cada capa. Estos tokens son guiados por la ruta donde se realiza el cómputo completo.
Asignación de la ruta de cómputo: los tokens seleccionados se asignan a la ruta que ejecuta el cómputo estándar (Self-Attention y MLP). Los tokens restantes pasan por la conexión residual y avanzan a la siguiente capa sin cómputo adicional.
Cálculo de la pérdida y backpropagation: la salida del modelo se evalúa con la función de pérdida final, y los gradientes se propagan hacia atrás por el modelo para actualizar los pesos. En este proceso, el router también aprende, de modo que con el tiempo puede tomar decisiones de enrutamiento más eficientes.

Métodos de enrutamiento del MoD Transformer (Routing Schemes)

La implementación del enrutamiento es una parte central del MoD Transformer. Los autores compararon tres métodos de enrutamiento y confirmaron que el enrutamiento de tokens del enfoque Expert-Choice MoD ofrece el mejor rendimiento. Veamos cómo se compara con los otros métodos presentados:

Enrutamiento de selección por token (Token-choice)

En el enrutamiento de selección por token, cada token elige por sí mismo la ruta por la que pasará. El modelo genera, para cada token, una distribución de probabilidad sobre varias rutas de cómputo posibles (por ejemplo, distintas capas o bloques de procesamiento) y, con base en esa distribución, envía el token por la ruta correspondiente.

El enrutamiento de selección por token permite una gran flexibilidad, ya que cada token puede elegir la ruta computacional más adecuada. Además, el modelo puede considerar con mayor precisión la importancia o el contexto de cada token, lo que permite optimización a nivel individual.

Enrutamiento de selección por experto (Expert-choice)

En el enrutamiento de selección por experto, el modelo elige directamente qué tokens procesará cada ruta (o “experto”). Esto se realiza mediante un mecanismo top-k que selecciona los tokens principales con base en los pesos del router de cada token.

Con este enrutamiento de selección por experto, los tokens pueden distribuirse de manera uniforme entre las rutas, por lo que la carga computacional se reparte de forma equilibrada y se aprovechan mejor los recursos del modelo. Además, como los tokens principales quedan definidos con claridad, mejora la previsibilidad y estabilidad del modelo.

Expert-choice MoD

Expert-choice MoD aplica el concepto de enrutamiento de selección por experto a la estructura del MoD Transformer. En este método, el modelo selecciona los tokens principales para cada capa o etapa de procesamiento, y solo esos tokens son procesados por la ruta que ejecuta el cómputo. Los tokens restantes siguen una ruta de menor costo a través de la conexión residual.

Como solo los tokens seleccionados realizan el cómputo, el uso total de FLOPs del modelo se reduce de manera importante. Además, este enfoque puede mantener o incluso mejorar el desempeño del modelo al garantizar el cómputo de todas las capas para los tokens importantes, mientras minimiza el cómputo innecesario. Finalmente, gracias al enrutamiento dinámico mediante el router, el modelo puede asignar recursos computacionales de forma flexible según las características de la secuencia de entrada.

Implementación del enrutamiento

La implementación del enrutamiento puede realizarse a través de las siguientes tres etapas principales:

Embeddings de tokens y determinación de los pesos del router: para cada capa del modelo, existe un conjunto de embeddings de tokens de la secuencia dada, $X^i$. Para cada embedding de token $x^l_i$, el router calcula, mediante una transformación lineal, un peso escalar $r^l_i$ que representa la probabilidad de que ese token participe en el cómputo.
Determinación de los pesos superiores: en el conjunto de pesos del router de los tokens $R^l$, se busca el percentil correspondiente al peso número $\beta$, denotado como $P_\beta(R^l)$. Aquí, $\beta$ determina la proporción de tokens que el modelo procesará en cada capa. Por tanto, los tokens cuyo peso del router sea mayor que $P_\beta(R^l)$ participarán en el cómputo.
Selección de la ruta de cómputo: cada token compara su peso del router $r^l_i$ con $P_\beta(R^l)$; si es mayor o igual, se procesa en esa capa, y si es menor, se envía directamente a la siguiente capa. De este modo, solo los tokens seleccionados realizan cómputo, mientras que los demás pueden seguir una ruta que no genera costo computacional.

Al implementar este enrutamiento, deben considerarse como puntos clave los siguientes aspectos:

Decisión dinámica basada en pesos: el proceso de enrutamiento se decide dinámicamente según los pesos de cada token, lo que refleja su importancia. Los tokens importantes reciben más cómputo para optimizar el desempeño del modelo, mientras que los menos importantes reducen costos al pasar por la ruta residual.
Eficiencia computacional y optimización del rendimiento: con este mecanismo de enrutamiento, el MoD Transformer puede usar eficientemente los recursos de cómputo mientras mantiene o mejora el desempeño del modelo. En particular, al procesar solo los tokens importantes en las capas de mayor costo computacional, es posible reducir el uso total de FLOPs.
Simplicidad y eficiencia de la implementación: como el enrutamiento se implementa mediante transformaciones lineales y cálculo de percentiles, es relativamente simple y eficiente. Esto permite una asignación dinámica de recursos computacionales sin añadir una carga significativa al entrenamiento o la inferencia del modelo.

Ejemplo de funcionamiento del MoD Transformer

Con base en lo visto hasta ahora, expliquemos el funcionamiento del MoD Transformer con un ejemplo.

Por ejemplo, supongamos que la secuencia de entrada tiene 100 tokens y que el router ha configurado el valor en 10. En ese caso, en cada capa el router selecciona los 10 tokens más importantes, y solo esos tokens se procesan siguiendo la ruta de cómputo estándar.

Los otros 90 tokens omiten el cómputo y se pasan directamente a la siguiente capa a través de la conexión residual. Mediante este proceso, el MoD Transformer reduce el uso total de FLOPs y al mismo tiempo garantiza el cómputo necesario para los tokens importantes.

Comparación de rendimiento del MoD Transformer

Rendimiento de entrenamiento

Optimización de hiperparámetros: mediante experimentos que entrenaron una variante del MoD Transformer con bloques de enrutamiento de capacidad del 12.5% colocados de forma alternada con bloques completos de attention, se encontró la mejor configuración de hiperparámetros. Esta configuración ofreció un desempeño superior a la mejor línea base en el análisis isoFLOP.

Eficiencia computacional: una variante específica de MoD mostró una velocidad de procesamiento 66% mayor, al lograr el mismo rendimiento que el modelo existente usando menos FLOPs.

Análisis del enrutamiento

Esparsidad de las decisiones de enrutamiento: al usar enrutamiento con capacidad del 12.5%, la mayoría de los tokens evitó el bloque y solo una pequeña cantidad de tokens importantes lo atravesó. Esto también pudo observarse en la distribución de pesos del router, y la red incluso tendió a preferir y enrutar ciertos tokens según la profundidad.

Rendimiento de inferencia

Inferencia autorregresiva: al cambiar de un esquema de enrutamiento top- no causal (non-causal) durante el entrenamiento a un enfoque causal (causal) basado en predicción durante la inferencia, casi no hubo pérdida de rendimiento. Esto parece deberse a que el aprendizaje de ese problema de predicción es sencillo y logró una precisión superior al 97% desde etapas tempranas del entrenamiento.

Mixture-of-Depths-and-Experts (MoDE) = MoD + MoE

MoDE (Mixture-of-Depths-and-Experts) se refiere a la integración de la técnica MoD (Mixture-of-Depths) y la técnica MoE (Mixture of Experts). Este enfoque busca combinar las ventajas de ambos métodos para optimizar al mismo tiempo la eficiencia computacional y el rendimiento de los modelos Transformer. Mientras que MoD selecciona si realizar un cómputo más profundo o menos profundo para cada token, MoE procesa el token a través del experto más adecuado entre varios “expertos”. MoDE integra ambos enfoques para permitir que un token salte ciertos bloques o realice selectivamente un cómputo más profundo y, al mismo tiempo, pueda elegir el mejor método de procesamiento entre distintos expertos.

Método de implementación de MoDE

MoDE tiene dos métodos principales de implementación: MoDE por etapas (Staged) y MoDE integrado (Integrated).

Staged MoDE (MoDE por etapas): en el MoDE por etapas, primero se aplica la técnica MoD para determinar la ruta de cómputo de cada token, y luego se aplica la técnica MoE para enrutar los tokens que serán procesados en la ruta seleccionada hacia sus respectivos bloques de expertos. De este modo, el modelo puede primero ajustar la profundidad del cómputo y después seleccionar el experto óptimo en cada etapa.
Integrated MoDE (MoDE integrado): el MoDE integrado combina MoD y MoE en una sola etapa de enrutamiento. Es decir, al enrutar un token, una sola decisión selecciona simultáneamente la profundidad de la ruta computacional y el experto que lo procesará. Este método puede ser más simple y eficiente de implementar, y ofrece la ventaja de decidir al mismo tiempo la ruta de cómputo y el experto más adecuados para cada token.

Ventajas de MoDE

Eficiencia computacional: MoDE gestiona de forma eficaz el costo computacional total del modelo al decidir dinámicamente, para cada token, cuánta computación necesita y qué experto lo procesará.
Optimización del rendimiento: al asignar cada token a la ruta computacional y al experto más adecuados, MoDE maximiza el rendimiento del modelo. Esto puede ser especialmente útil al resolver problemas complejos o procesar distintos tipos de datos.
Flexibilidad y escalabilidad: MoDE ofrece flexibilidad para aplicarse a diversas arquitecturas y tareas, y proporciona un framework escalable capaz de mejorar simultáneamente la eficiencia computacional y el rendimiento del modelo.

MoDE sugiere una dirección futura para el desarrollo de modelos Transformer y hace una contribución importante a la investigación que explora cómo maximizar el rendimiento del modelo minimizando al mismo tiempo el costo computacional.

Leer más

https://arxiv.org/abs/2404.02258

Este texto se basa en un resumen elaborado con un modelo GPT, por lo que podría contener partes organizadas de manera distinta al contenido o a la intención del texto original. Si el tema te interesa, te recomendamos consultar también el original. Si al leer encuentras algo extraño o incorrecto, te agradeceremos que lo compartas en los comentarios.

⚠️Publicidad⚠️: ¿Te resultó útil este artículo recopilado por la comunidad de usuarios de PyTorch Korea? Si te registras como miembro, te enviaremos los artículos principales por correo electrónico. (Por defecto es Weekly, pero también puedes cambiarlo a Daily).