Bien que puissants, les modèles Transformer souffrent d'un problème d'efficacité lors du décodage. Cependant, des chercheurs du Korea Advanced Institute of Science and Technology (KAIST), de LG et de DeepMind ont apporté une solution surprenante : une nouvelle architecture Transformer, baptisée Block Transformer, qui accélère le décodage de 10 à 20 fois !

Comment cela est-il possible ? Ils ont « découpé » le mécanisme d'attention du Transformer. Cela révolutionne l'approche inefficace du Transformer original, qui accédait au cache KV global à chaque génération de jeton.

image.png

Les chercheurs ont analysé les faiblesses du Transformer original : l'utilisation du GPU était inférieure à 1 %, les 99 % restants étant consacrés à l'accès à la mémoire. Cela étant irrationnel, ils ont proposé Block Transformer. Cette nouvelle architecture, grâce à la décomposition de l'attention en blocs et à l'attention intra-bloc, améliore considérablement le débit d'inférence.

Concrètement, Block Transformer fonctionne ainsi : la séquence est d'abord découpée en blocs, puis un Embedder convertit chaque bloc en vecteur d'embedding. Le Block Decoder traite les vecteurs d'embedding des blocs, capturant les dépendances globales entre les blocs ; le Token Decoder traite les dépendances locales entre les jetons, générant la séquence de jetons.

image.png

Cette méthode améliore non seulement la vitesse d'inférence, mais réduit également considérablement la consommation de mémoire. Des internautes ont mentionné avoir eu des idées similaires, mais avec des performances insuffisantes. Cette méthode semble effectivement réduire efficacement le cache KV.

De plus, la précision de Block Transformer sur plusieurs tâches à zéro-shot est comparable, voire légèrement supérieure, à celle d'un Transformer original de taille équivalente, prouvant qu'il améliore l'efficacité sans sacrifier la qualité.

Cette recherche a une portée plus large. Elle réduit également le coût d'entraînement du modèle : la surcharge d'accès mémoire quadratique de l'attention globale est réduite de 16 fois, et l'utilisation du GPU passe de 1 % à 44 %.

Adresse de l'article : https://arxiv.org/abs/2406.02657