Enhancing Layer Interaction Using Key-Correlated Layer Attention

Résumé Technique : Amélioration de l'interaction entre couches via l'Attention de Couche Corrélationnelle de Clé (KCLA)

Énoncé du Problème
Les avancées récentes dans la conception des réseaux de neurones profonds (DNN) ont utilisé l'Attention de Couche (Layer Attention) pour améliorer les interactions inter-couches en traitant les sorties des couches comme des jetons de séquence. Cependant, l'attention de couche standard souffre d'une complexité computationnelle quadratique ($O(D^2)$) par rapport à la profondeur du réseau $D$, ainsi qu'une croissance significative de la mémoire due à la nécessité de mettre en cache les sorties de toutes les couches précédentes.

Pour remédier à ces problèmes, des travaux antérieurs ont introduit l'Attention de Couche Récurrente (RLA) et des approximations linéaires comme le MRLA-L. Bien que la RLA réduise la charge de calcul en mettant en cache les représentations de Clé ($K$) et de Valeur ($V$), elle fonctionne comme un mécanisme statique où les caractéristiques des couches précoces restent fixes une fois générées, limitant ainsi les mises à jour dynamiques d'informations. De plus, les approximations linéaires telles que le MRLA-L tentent de réduire la complexité en approximant linéairement la sortie de l'attention. Les auteurs observent que cette approche provoque une atténuation exponentielle des contributions des couches antérieures (comme le montre leur analyse des coefficients de décroissance $\beta_l$), rendant difficile l'établissement de dépendances inter-couches à longue portée efficaces. En outre, les méthodes de linéarisation existantes échouent souvent à traiter le goulot d'étranglement de la mémoire lié au cache des sorties des couches précédentes ou les défis spécifiques de l'application de l'attention linéaire aux structures DNN hiérarchiques.

Méthodologie : Key-Correlated Layer Attention (KCLA)
La Key-Correlated Layer Attention (KCLA) proposée est une nouvelle méthode d'attention de couche linéarisée conçue pour atteindre une complexité computationnelle linéaire tout en préservant les mises à jour dynamiques et la modélisation des dépendances à longue portée. La méthodologie repose sur trois observations et mécanismes fondamentaux :

1. Observation de la Corrélation de Clé : Une analyse empirique révèle que les représentations de Clé à travers différentes couches dans l'attention de couche standard présentent une haute similitude cosinus. En exploitant cela, la KCLA approxime la sortie de l'attention de couche standard en factorisant la transformation linéaire de la Valeur ($V$) et en approximant les représentations de Clé comme étant directionnellement corrélées. Cela permet l'agrégation des sorties de couches dans un état récurrent ($U^t$) sans stocker l'historique complet des matrices $K$ et $V$.
2. Restauration Dynamique du Softmax : Un défi critique dans la linéarisation de l'attention est la perte de l'opération softmax, qui est essentielle pour normaliser les scores d'attention et permettre des ajustements de distribution dynamiques. La KCLA introduit un mécanisme d'attention multi-têtes modifié pour restaurer cette capacité :
   • Elle utilise plusieurs têtes avec des facteurs d'échelle différents (températures) dérivés d'une séquence logarithmiquement uniforme.
   • Elle utilise une transformation linéaire apprenable $\Phi$ pour mapper la norme de la Requête ($\|Q^t\|$) et l'angle de corrélation ($\cos \theta$) en poids spécifiques à chaque tête.
   • Cela permet au modèle d'ajuster dynamiquement la distribution d'attention à chaque couche, simulant efficacement le rôle du paramètre de température du softmax sans calculer la matrice complète $N \times N$.
3. Adaptation Trans-Stade : Pour s'adapter aux DNN ayant des tailles de cartes de caractéristiques et des dimensions de canaux variables selon les stades, la KCLA implémente une procédure d'adaptation spécifique (Algorithme 1). Cela implique de normaliser les caractéristiques accumulées, d'appliquer une convolution $1 \times 1$ avec un pas (stride) pour le sous-échantillonnage, et de re-calculer les facteurs de normalisation pour combler les disparités dimensionnelles entre les stades.

Contributions Clés

• Observation et Preuve : Les auteurs démontrent que les représentations de Clé à travers différentes couches dans l'attention de couche standard sont hautement corrélées en direction, fournissant la base théorique pour l'approximation linéaire.
• Linéarisation Novatrice : La KCLA est proposée comme une méthode d'attention de couche linéarisée qui évite la nature statique de la RLA et l'atténuation exponentielle du MRLA-L. Elle atteint une complexité computationnelle linéaire et une complexité spatiale fixe indépendante de la profondeur du réseau.
• Approximation de Softmax Dynamique : Une nouvelle méthode est introduite pour approximer l'opération softmax via le passage par des échelles multi-têtes et des transformations apprenables, permettant un ajustement dynamique et adaptatif des distributions d'attention à chaque couche.

Résultats Expérimentaux
Les auteurs ont évalué la KCLA sur trois tâches de vision fondamentales : la classification d'images, la détection d'objets et la segmentation d'images médicales.

• Classification d'Images : Sur CIFAR-100, la KCLA a obtenu les meilleures performances sur les profondeurs de ResNet 20, 56 et 110, surpassant le MRLA-L, le DLA-L et l'attention de couche standard. Sur ImageNet-1K, l'intégration de la KCLA dans ResNet-50 et ResNet-101 a amélioré la précision Top-1 de +1,7 % et +1,4 % respectivement par rapport à la ligne de base. Elle a maintenu un avantage léger mais constant sur le MRLA-L tout en utilisant nettement moins de paramètres que la méthode dynamique DLA-L.
• Détection d'Objets : Sur le jeu de données COCO2017 utilisant Faster R-CNN et Mask R-CNN, la KCLA a démontré des performances compétitives. Avec ResNet-101, elle a surpassé l'architecture plus complexe DLA-L sur plusieurs métriques (AP$_{75}$, AP$_S$, AP$_M$) tout en utilisant 2,2 millions de paramètres en moins.
• Segmentation d'Images Médicales : Sur les jeux de données ISIC2017, ISIC2018 et Kvasir-SEG, la KCLA a systématiquement surpassé le MRLA-L et la ligne de base ResNet-UNet, atteignant des scores mIoU et DSC supérieurs avec un surcoût de paramètres minimal.
• Analyse de Complexité : Contra part de la MRLA-B, dont la mémoire et les FLOPs croissent de manière quadratique avec la profondeur, la KCLA présente une croissance linéaire tant en charge de calcul qu'en coût de mémoire.

Signification et Revendications
L'article affirme que la KCLA comble avec succès l'écart entre l'efficacité de l'attention linéaire et la puissance de représentation de l'attention standard. En approximant directement le mécanisme d'attention de couche standard basé sur les corrélations de Clé, la KCLA évite les problèmes de représentation de caractéristiques statiques de la RLA et la dégradation des dépendances à longue portée du MRLA-L.

Les auteurs soulignent que leur méthode maintient la capacité de modéliser efficacement les dépendances inter-couches à longue portée, offrant un équilibre souhaitable entre efficacité et performance. Ils notent que la méthode proposée pour restaurer la fonction softmax est un facteur critique pour atteindre ces résultats. Bien que l'implémentation actuelle repose sur des paramètres de température pré-établis, les auteurs identifient la conception d'une attention de couche linéaire multi-tête à température auto-adaptative et dynamique comme une direction pour une optimisation future. Le code est rendu public pour faciliter la reproductibilité.