Combiner : Transformateur d'attention complète avec un coût de calcul peu élevé

Résumé :

Les transformateurs fournissent une classe d'architectures expressives qui sont extrêmement efficaces pour la modélisation de séquences. Cependant, la principale limitation des transformateurs est leur complexité quadratique en mémoire et en temps (L2) par rapport à la longueur de la séquence dans les couches d'attention, ce qui limite leur application à des séquences extrêmement longues. La plupart des approches existantes tirent parti de la sparsité ou des hypothèses de faible rang dans la matrice d'attention pour réduire le coût, mais sacrifient l'expressivité. Au lieu de cela, nous proposons Combiner, qui fournit une capacité d'attention complète dans chaque tête d'attention tout en maintenant une faible complexité de calcul et de mémoire. L'idée principale est de traiter le mécanisme d'auto-attention comme une espérance conditionnelle sur les incorporations à chaque emplacement, et d'approximer la distribution conditionnelle avec une factorisation structurée. Chaque emplacement peut prêter attention à tous les autres emplacements, soit par une attention directe, soit par une attention indirecte aux abstractions, qui sont à nouveau des attentes conditionnelles des incorporations des régions locales correspondantes. Nous montrons que la plupart des modèles d'attention clairsemée utilisés dans les transformateurs clairsemés existants sont capables d'inspirer la conception d'une telle factorisation pour l'attention complète, résultant en un même coût sub-quadratique ((Llog(L)) ou (LL‾‾√)). Combiner est un remplacement drop-in pour les couches d'attention dans les transformateurs existants et peut être facilement implémenté dans les frameworks communs. Une évaluation expérimentale sur des tâches de séquences autorégressives et bidirectionnelles démontre l'efficacité de cette approche, donnant des résultats de pointe sur plusieurs tâches de modélisation d'images et de textes.