FlexAttention est une nouvelle API rendue publique par l'équipe PyTorch en juillet 2024 qui fournit une interface flexible permettant d'implémenter de nombreuses variantes d'attention en quelques lignes de code PyTorch typique et en utilisant torch.compile Cela se réduit à un noyau FlashAttention fusionné, offrant une flexibilité sans sacrifier les performances. Les résultats pertinents de l'article sontFlexAttention pour des modèles vision-langage haute résolution efficaces", a été accepté par l'ECCV 2024.

FlexAttention est un mécanisme d'attention flexible conçu pour améliorer l'efficacité des modèles de vision-langage haute résolution. Le mécanisme proposé réduit considérablement le coût de calcul en codant les étiquettes d'images haute et basse résolution et en calculant la carte d'attention en utilisant uniquement les étiquettes basse résolution et quelques étiquettes haute résolution sélectionnées. La sélection de points de repère haute résolution est effectuée par un module de sélection haute résolution, qui peut récupérer les points de repère des régions pertinentes en fonction de la carte d'attention d'entrée. Les jetons haute résolution sélectionnés sont ensuite introduits dans une couche d'auto-attention hiérarchique avec les jetons basse résolution et les jetons de texte, et la carte d'attention générée par cette couche est utilisée pour la sélection des jetons haute résolution à l'étape suivante. Ce processus est répété à chaque couche d’attention. Les expériences montrent que FlexAttention surpasse les modèles de langage de vision haute résolution existants sur les benchmarks multimodaux tout en réduisant considérablement le coût de calcul de près de 40%.