Si bien los modelos Transformer son poderosos, su eficiencia durante la decodificación siempre ha sido un problema. Sin embargo, investigadores del Instituto de Ciencia y Tecnología de Corea, LG y DeepMind nos han sorprendido con una nueva arquitectura Transformer llamada Block Transformer, ¡que aumenta la velocidad de decodificación entre 10 y 20 veces!
¿Cómo lo han logrado? Han "dividido en bloques" el mecanismo de atención del Transformer. Esto elimina la ineficiencia del Transformer original, que accedía a la caché global KV cada vez que generaba un token.
Los investigadores analizaron las deficiencias del Transformer original: la utilización de la GPU era inferior al 1%, mientras que el 99% restante se dedicaba al acceso a la memoria. Esto es claramente ineficiente. Por lo tanto, propusieron Block Transformer, una nueva arquitectura que, mediante la descomposición de la atención a nivel de bloque y dentro del bloque, aumenta exponencialmente el rendimiento de inferencia.
En concreto, el flujo de trabajo de Block Transformer es el siguiente: primero, se divide la secuencia en bloques; luego, un Embedder convierte cada bloque en un vector de incrustación. El Block Decoder procesa los vectores de incrustación de los bloques, capturando las dependencias globales entre bloques; el Token Decoder procesa las dependencias locales entre tokens, generando la secuencia de tokens.
Este método no solo mejora la velocidad de inferencia, sino que también reduce significativamente el consumo de memoria. Algunos usuarios han comentado que habían tenido ideas similares, pero el rendimiento del modelo era insuficiente; este método parece reducir eficazmente la caché KV.
Además, la precisión de Block Transformer en varias tareas de cero-shot es comparable o incluso ligeramente superior a la del Transformer original del mismo tamaño, lo que demuestra que la mejora de la eficiencia no se ha producido a costa de la calidad.
La importancia de esta investigación va más allá. También reduce el coste de entrenamiento del modelo: el coste de acceso a la memoria del mecanismo de atención global se reduce en 16 veces, y la utilización de la GPU aumenta del 1% al 44%.
Enlace al artículo: https://arxiv.org/abs/2406.02657