Résumé technique : Les têtes Laplaciennes améliorent les Transformers en lissant les représentations de tokens

Énoncé du problème

L'architecture Transformer standard met à jour les représentations de tokens via l'attention multi-têtes et des connexions résiduelles. Dans cette formulation, chaque tête d'attention $i$ calcule une matrice d'attention $P^{(i)}$ de type softmax et met à jour l'entrée $X$ par une somme pondérée des projections de valeurs : $X \leftarrow X + \sum_i P^{(i)}XW_{Vi}W_{oi}$.

Bien que efficace, ce mécanisme manque de contrôle explicite sur la variance intra-séquence des représentations de tokens. La vision prédominante dans la littérature, notamment concernant les réseaux de neurones graphiques (GNN) et les Transformers profonds, considère le « lissage excessif » (oversmoothing) — où les représentations de tokens deviennent trop similaires, perdant leur spécificité — comme un mode de défaillance néfaste. Par conséquent, la plupart des travaux antérieurs se concentrent sur la prévention du lissage par la normalisation, l'élagage des arêtes ou des pertes de diversité. Cet article remet en question l'hypothèse selon laquelle le lissage excessif des tokens est intrinsèquement nuisible, proposant plutôt qu'un lissage contrôlé peut être bénéfique pour l'apprentissage de représentations.

Méthodologie

Les auteurs proposent une modification architecturale sans paramètres : le remplacement d'un sous-ensemble de têtes d'attention standard par des têtes Laplaciennes.

L'opérateur Laplacien

Dans la mise à jour standard, une tête calcule $P^{(i)}X$, qui est une moyenne pondérée par l'attention des tokens. La tête Laplacienne proposée remplace la matrice d'attention $P^{(i)}$ par l'opérateur Laplacien de graphe $I - P^{(i)}$, où $I$ est la matrice identité. L'équation de mise à jour devient :
$$X \leftarrow X + \sum_{i \in A} P^{(i)}XW_{Vi}W_{oi} + \sum_{i \in L} (I - P^{(i)})XW_{Vi}W_{oi}$$
où $A$ est l'ensemble des têtes d'attention standard et $L$ est l'ensemble des têtes Laplaciennes.

Motivations

L'article fournit deux motivations théoriques pour cette modification :

1. Contrôle direct de la moyenne et de la variance : Les têtes d'attention standard mettent à jour la moyenne des représentations de tokens. En revanche, le terme $(I - P^{(i)})X$ calcule la différence entre chaque token et sa moyenne pondérée par l'attention. L'ajout de cette différence permet au modèle de contrôler directement la variance intra-séquence des représentations de tokens, une capacité absente des Transformers standards.
2. Diffusion sur graphe : Si les tokens sont considérés comme des nœuds dans un graphe dont les poids d'arêtes sont définis par $P^{(i)}$, alors $I - P^{(i)}$ est le Laplacien de graphe normalisé par la marche aléatoire. L'application de cet opérateur correspond à une étape unique de diffusion thermique sur le graphe de tokens, lissant progressivement le signal (les représentations de tokens) à travers la séquence.

Contributions clés

1. Proposition architecturale : Une modification simple et sans paramètres de l'architecture Transformer qui intègre des têtes Laplaciennes aux côtés des têtes d'attention standard.
2. Validation empirique : Démonstration que cette modification améliore les performances dans trois paradigmes d'apprentissage distincts : classification supervisée, modélisation du langage et apprentissage auto-supervisé.
3. Caractérisation géométrique : Une analyse détaillée de la manière dont les têtes Laplaciennes modifient la géométrie des représentations de tokens apprises, montrant qu'elles induisent des formes bénéfiques de « collapse » et de lissage plutôt qu'un lissage excessif néfaste.

Résultats expérimentaux

1. Classification d'images (Apprentissage supervisé)

Utilisation de la famille DeiT-3 (ViT-B) sur CIFAR-10, CIFAR-100 et ImageNet-1k :

• Performance : L'ajout de têtes Laplaciennes a systématiquement amélioré la précision Top-1. L'amélioration était presque monotone avec le nombre de têtes Laplaciennes.
  • Sur ImageNet, le modèle avec 12 têtes Laplaciennes (ViT-B-12L) a atteint une précision de 82,32 % contre 81,45 % pour la référence.
  • Sur CIFAR-100, le gain a été d'environ 2 % (77,04 % contre 75,12 %).
• Géométrie :
  • ANOVA : Les têtes Laplaciennes ont considérablement réduit la variance intra-séquence tout en augmentant la variance inter-classe. Cela indique que les tokens au sein d'une même séquence convergent vers une moyenne commune, tandis que les moyennes des différentes séquences se séparent plus efficacement.
  • Collapse neuronal : Les modèles ont montré une convergence plus forte vers la géométrie du Collapse neuronal. Les représentations de tokens ont formé des clusters bien séparés, de type simplexe (Simplex ETF) pour chaque classe, les moyennes de séquence s'alignant sur la géométrie globale des classes.

2. Modélisation du langage (Autoregressive)

Utilisation de modèles de type GPT-2 (561 millions de paramètres) pré-entraînés sur FineWeb-Edu et affinés sur des benchmarks de raisonnement/code :

• Performance : Les modèles avec têtes Laplaciennes ont surpassé la référence en moyenne sur ARC-Easy, ARC-Challenge, MMLU, GSM8K et HumanEval.
  • Le modèle avec 5 têtes Laplaciennes a atteint la meilleure performance moyenne (30,06 %), avec des gains significatifs sur ARC-Easy (+16 %) et ARC-Challenge (+7 %).
• Géométrie :
  • Décroissance spectrale : Les têtes Laplaciennes ont produit des représentations de tokens avec des spectres à décroissance plus rapide. Le nombre de composantes principales nécessaires pour expliquer 99 % de la variance a diminué de 21 % (passant de 547,7 à 430,4) lorsque toutes les têtes étaient Laplaciennes.
  • Collapse linguistique : Les modèles ont montré une séparabilité accrue parmi les représentations de tokens partageant la même prédiction de token suivant, suggérant une forme plus forte de « collapse linguistique » où les tokens ayant des rôles sémantiques similaires se regroupent.

3. Apprentissage auto-supervisé (DINO)

Utilisation de modèles ViT-S sur ImageNet-1k pour le pré-entraînement DINO :

• Performance : Les têtes Laplaciennes ont amélioré la précision du sondage linéaire sur ImageNet et les performances de segmentation sur ADE20K.
  • ViT-S-5L a atteint la meilleure mIoU sur ADE20K (32,27 %) comparé à la référence (29,81 %).
• Géométrie :
  • Segmentation : Les visualisations par PCA ont montré que les têtes Laplaciennes produisaient des contours d'objets plus nets et préservaient des détails plus fins (par exemple, formes de casques, numéros de maillots) par rapport à la référence.
  • Décroissance spectrale : Similairement à la modélisation du langage, ces modèles ont présenté une décroissance spectrale plus rapide, concentrant les représentations de tokens intra-image dans un sous-espace de dimension inférieure.

Signification et affirmations

L'article soutient que la vision prédominante — selon laquelle le lissage excessif des tokens est universellement nuisible — est incomplète. À l'inverse, les auteurs affirment que le lissage est un mécanisme d'apprentissage de représentations qui peut être bénéfique lorsqu'il est contrôlé.

• Avantages spécifiques à la tâche : Le « lissage » induit par les têtes Laplaciennes n'est pas un mode de défaillance mais une fonctionnalité qui s'aligne sur les objectifs de la tâche :
  • Dans la classification supervisée, il facilite le Collapse neuronal en faisant converger les tokens au sein des séquences et en alignant les moyennes de séquence sur la géométrie des classes.
  • Dans la modélisation du langage, il améliore le Collapse linguistique, augmentant la séparabilité des tokens ayant des cibles de token suivant identiques.
  • Dans l'apprentissage auto-supervisé, il concentre les représentations dans des sous-espaces de dimension inférieure, aidant les tâches en aval comme la segmentation.
• Mécanisme : En incorporant explicitement la diffusion sur graphe (via le Laplacien) dans l'architecture, le modèle acquiert la capacité de contrôler directement la variance intra-séquence, conduisant à des représentations plus robustes et géométriquement structurées.

Les auteurs concluent que, bien que leur implémentation à stratégie fixe soit probablement sous-optimale, les résultats suggèrent que le lissage devrait être davantage investigué en tant que composant fondamental de l'apprentissage de représentations, plutôt que uniquement comme un problème à atténuer.