Investigadores de Google desarrollan nueva tecnología de inteligencia artificial que no desperdicia capacidad intelectual en palabras inútiles

Tabla de enlaces

1. Introducción
2. Fondo
3. Implementación de transformadores de mezcla de profundidades

   • 3.1. Definición de un presupuesto de cómputo

   • 3.2. Enrutamiento alrededor de los bloques del transformador

   • 3.3 Esquemas de enrutamiento

   • 3.4 Implementación de enrutamiento

   • 3.5. Métodos de muestreo y 3.6. Métodos de entrenamiento

4. Resultados
   • 4.1. Entrenamiento, comparaciones de isoFLOP
   • 4.2. Evaluación autorregresiva y 4.3. Combinación de profundidades y expertos (MoDE)
5. Discusión y referencias

Los modelos de lenguaje basados en transformadores distribuyen los FLOP de manera uniforme en las secuencias de entrada. En este trabajo demostramos que los transformadores pueden, en cambio, aprender a asignar dinámicamente los FLOP (o calcularlos) a posiciones específicas en una secuencia, optimizando la asignación a lo largo de la secuencia para diferentes capas en toda la profundidad del modelo. Nuestro método impone un presupuesto computacional total al limitar la cantidad de tokens (𝑘) que pueden participar en los cálculos de autoatención y MLP en una capa dada. Los tokens que se procesarán son determinados por la red utilizando un mecanismo de enrutamiento top-𝑘. Dado que 𝑘 se define a priori, este procedimiento simple utiliza un gráfico de cálculo estático con tamaños de tensor conocidos, a diferencia de otras técnicas de cálculo condicional. Sin embargo, dado que las identidades de los tokens 𝑘 son fluidas, este método puede gastar FLOP de manera no uniforme en las dimensiones de tiempo y profundidad del modelo. Por lo tanto, el gasto computacional es completamente predecible en total, pero dinámico y sensible al contexto a nivel de token. Los modelos entrenados de esta manera no solo aprenden a asignar cómputo de forma dinámica, sino que lo hacen de manera eficiente. Estos modelos igualan el rendimiento de referencia para FLOPS equivalentes y tiempos de reloj de pared para entrenar, pero requieren una fracción de los FLOP por pasada hacia adelante y pueden ser hasta un 50 % más rápidos para avanzar durante el muestreo posterior al entrenamiento.

1. Introducción

No todos los problemas requieren la misma cantidad de tiempo o esfuerzo para resolverse. Análogamente, en el modelado de lenguajes no todos los tokens y secuencias requieren el mismo tiempo o esfuerzo para hacer una predicción precisa. Y, sin embargo, los modelos de transformadores gastan la misma cantidad de cómputo por token en un pase hacia adelante. Idealmente, los transformadores utilizarían presupuestos computacionales totales más pequeños al no gastar cómputo innecesariamente.

El cálculo condicional es una técnica que intenta reducir el cálculo total gastándolo solo cuando es necesario (Bengio et al., 2016; Bengio, 2013; Bengio et al., 2013). Varios algoritmos ofrecen soluciones a cuándo y cuánto cálculo se debe utilizar (Ainslie et al., 2023; Bapna et al., 2020; Fedus et al., 2022). Sin embargo, las formulaciones generales de este desafiante problema pueden no funcionar bien con las restricciones de hardware existentes, ya que tienden a introducir gráficos de cálculo dinámicos (Dehghani et al., 2018; Graves, 2016). Los métodos de cálculo condicional más prometedores pueden ser, en cambio, aquellos que sean armoniosos con nuestra pila de hardware actual, que prioriza los gráficos de cálculo estático y los tamaños de tensor conocidos que se seleccionan para maximizar la utilización del hardware.

Aquí consideramos el problema del modelado de lenguajes utilizando un presupuesto de cómputo estático que puede ser menor que el utilizado por un transformador tradicional. La red debe aprender a asignar dinámicamente el cómputo disponible tomando decisiones por token, en cada capa, sobre dónde gastar el cómputo del presupuesto disponible. En nuestra implementación, el cómputo total está definido por el usuario y no cambia antes del entrenamiento, en lugar de ser una función de las decisiones sobre la marcha de la red. Por lo tanto, las ganancias de eficiencia del hardware (como la reducción de la huella de memoria o la reducción de FLOP por pase hacia adelante) se pueden anticipar y aprovechar con anticipación. Como demostraremos, estas ganancias se pueden obtener sin sacrificar el rendimiento general.

Utilizamos un enfoque similar a los transformadores Mixture of Experts (MoE), en el que se toman decisiones de enrutamiento dinámico a nivel de token en toda la profundidad de la red. A diferencia de MoE, elegimos aplicar un cálculo a un token (como sería el caso de un transformador estándar) o pasarlo a través de una conexión residual (permaneciendo sin cambios y ahorrando cálculo). También a diferencia de MoE, aplicamos este enrutamiento tanto a MLP de avance como a atención multi-head. Dado que esto también afecta las claves y consultas que procesamos, el enrutamiento toma decisiones no solo sobre qué tokens actualizar, sino también qué tokens están disponibles para atender. Nos referimos a esta estrategia como Mixture-of-Depths (MoD) para enfatizar cómo los tokens individuales pasan a través de diferentes cantidades de capas o bloques a través de la profundidad del transformador (ver figura 1).

La técnica MoD también permite equilibrar el rendimiento con la velocidad. Por un lado, se puede entrenar un transformador MoD que mejora a los transformadores tradicionales hasta en un 1,5 % en el objetivo de entrenamiento de probabilidad logarítmica final para FLOP de entrenamiento equivalentes (isoFLOP), y al mismo tiempo que requiere una cantidad equivalente de tiempo de reloj para entrenarse. Por otro lado, se puede entrenar un transformador MoD que logra la paridad de pérdida de entrenamiento con un transformador tradicional óptimo isoFLOP, pero que utiliza una fracción de los FLOP (más del 50 %) por paso hacia adelante y, por lo tanto, es más rápido de ejecutar. En conjunto, estos resultados implican que los transformadores MoD aprenden a enrutar de manera inteligente (es decir, omitiendo cálculos que son innecesarios) ya que pueden lograr probabilidades logarítmicas iguales o mejores por secuencia a pesar de una huella de FLOP más pequeña por paso hacia adelante.