Résumé technique : Routage résiduel multi-résolution pour les Transformers profonds de type Decoder-Only

Énoncé du problème

Les Transformers standards de type « decoder-only » reposent sur des connexions résiduelles PreNorm où les mises à jour des sous-couches sont agrégées avec des poids unitaires fixes. Bien que cette structure permette un entraînement stable, elle impose une règle d'accumulation rigide qui peut diluer la contribution des couches individuelles à mesure que la profondeur du réseau augmente. Des avancées récentes, telles que les Attention Residuals et les Block Attention Residuals (Block AttnRes), remédient à cela en remplaçant l'accumulation fixe par un routage appris et dépendant du contenu sur les sorties précédentes. Block AttnRes améliore spécifiquement l'efficacité en compressant plusieurs couches en résumés résiduels au niveau du bloc ($C_b$).

Cependant, les auteurs identifient une limitation critique dans Block AttnRes : un résumé de bloc unique ne représente que la composante basse fréquence (DC) de la trajectoire résiduelle au sein d'un bloc. En ne stockant que la somme des mises à jour ($\sum u_{b,i}$), le mécanisme néglige des structures directionnelles précieuses, telles que :

1. Le déséquilibre entre les sous-couches : Si le bloc est dominé par les mises à jour d'Attention ou de MLP.
2. La dynamique temporelle : Si les premières mises à jour au sein du bloc pointent dans une direction différente des mises à jour tardives.

Les auteurs émettent l'hypothèse qu'à mesure que la longueur de la trajectoire résiduelle augmente, cette perte d'information directionnelle devient cruciale pour un routage efficace selon la profondeur.

Méthodologie : WAV v1

Les auteurs proposent WAV v1 (Waveform Attention v1), une extension multi-résolution légère de Block AttnRes. Au lieu de représenter un bloc uniquement par sa somme totale, WAV v1 augmente chaque représentation de bloc avec deux bases de détails directionnels à somme nulle, tout en préservant les modules d'Attention et de MLP originaux.

1. Représentation multi-résolution

Pour un bloc $b$ contenant $m$ mises à jour de sous-couches $\{u_{b,i}\}_{i=1}^m$, WAV v1 calcule :

• Base d'état ($C_b$) : La somme standard du bloc, représentant l'état basse fréquence.
  $$C_b = \sum_{i=1}^m u_{b,i}$$
• Base de phase ($D^{phase}_b$) : Une source de détails contrastant les mises à jour d'Attention vs MLP. Elle assigne un signe de $+1$ aux mises à jour d'Attention et de $-1$ aux mises à jour de MLP.
  $$D^{phase}_b = \sum_{i=1}^m a_i u_{b,i}, \quad \text{où } a_i \in \{+1, -1\}$$
• Base de division ($D^{split}_b$) : Une source de détails contrastant la première moitié vs la seconde moitié du bloc. Elle assigne $+1$ à la première moitié et $-1$ à la seconde moitié.
  $$D^{split}_b = \sum_{i=1}^m r_i u_{b,i}, \quad \text{où } r_i \in \{+1, -1\}$$

Ces bases sont accumulées en ligne à mesure que les mises à jour sont produites, ce qui n'entraîne qu'un surcoût de calcul négligeable.

2. Mécanisme de routage

Le mélangeur softmax de profondeur (depth-wise softmax mixer) reçoit un pool de sources étendu. Pour Block AttnRes, le pool est $\{e, C_0, \dots, C_{b-1}, PC\}$. Pour WAV v1, il s'étend pour inclure les bases de détails pour chaque bloc terminé et le bloc partiel actuel :
$$S_{WAV} = \{e, C_0, \tilde{D}^{phase}_0, \tilde{D}^{split}_0, \dots, C_{b-1}, \tilde{D}^{phase}_{b-1}, \tilde{D}^{split}_{b-1}, PC, \tilde{P}^{phase}, \tilde{P}^{split}\}$$

3. Mécanismes de stabilité

Pour éviter la déstabilisation lors du début de l'entraînement, WAV v1 emploie deux stratégies conservatrices :

• Biais de détail négatif : Les sources de détails sont initialisées avec un fort biais négatif ($\beta_D = -2.0$), garantissant que le modèle se comporte comme le Block AttnRes standard à l'initialisation et n'utilise les détails que lorsqu'ils sont bénéfiques.
• Appariement RMS détaché (Detached RMS Matching) : Les tenseurs de détails sont mis à l'échelle par rapport à leur résumé de bloc correspondant ($C$) en utilisant un ratio RMS détaché, limité par un facteur d'échelle maximal $\rho$. Cela empêche les détails de dominer simplement en raison de magnitudes brutes plus grandes.
• Lecture conservatrice : Par défaut, le mélangeur de prédiction final ne lit que les sources d'embedding et $C$, et non les sources de détails, renforçant ainsi la stabilité.

Contributions clés

1. Formulation : Le papier cadre le routage résiduel par bloc comme un problème de représentation multi-résolution, distinguant l'information d'état basse fréquence ($C_b$) de l'information de trajectoire intra-bloc (détails directionnels).
2. Algorithme : Introduction de WAV v1, une extension prête à l'emploi qui augmente les résumés de blocs avec des bases de détails de phase et de division sans modifier les fonctions fondamentales d'Attention ou de MLP.
3. Preuves empiriques : Expériences préliminaires de mise à l'échelle démontrant que l'utilité du routage multi-résolution est dépendante de la profondeur, devenant significativement bénéfique uniquement dans les architectures plus profondes.

Résultats expérimentaux

Les auteurs ont évalué des modèles GPT decoder-only de 12, 24 et 48 couches sur les jeux de données de niveau caractère TinyStories et Text8.

• Tendance dépendante de la profondeur :
  • 12 Couches : WAV v1 est moins performant que Block AttnRes, suggérant que les détails directionnels sont redondants ou bruyants dans les réseaux peu profonds.
  • 24 Couches : WAV v1 devient compétitif, montrant de légères améliorations sur TinyStories et des performances comparables sur Text8.
  • 48 Couches : WAV v1 atteint la meilleure perte de validation sur les deux jeux de données, surpassant nettement Block AttnRes, ReZero et LayerScale.
• Gains quantitatifs (48 Couches) :
  • TinyStories : Réduction de la perte de validation de 0,4960 (Block AttnRes) à 0,4738 (WAV v1), soit une réduction de la perplexité (PPL) de 2,19 %.
  • Text8 : Réduction de la perte de validation de 0,9363 à 0,9305, soit une réduction de la PPL de 0,57 %.
• Dynamique d'entraînement : À 48 couches, WAV v1 se détache de Block AttnRes tôt dans l'entraînement et maintient un avantage persistant, indiquant que l'amélioration n'est pas un artefact de fin de parcours mais un changement fondamental de la trajectoire d'optimisation.

Signification et affirmations

Le papier affirme que le routage résiduel profond bénéficie non seulement de la sélection de quel bloc lire, mais aussi du routage de la structure directionnelle interne de chaque bloc. Les résultats suggèrent qu'à mesure que la trajectoire résiduelle s'allonge, la « forme » de la trajectoire de mise à jour (par exemple, la dominance de l'attention ou la dynamique début-fin) devient un signal informatif pour le mélangeur de profondeur.

Les auteurs positionnent WAV v1 comme une approche complémentaire aux méthodes existantes :

• Contrairement à ReZero ou LayerScale, qui modulent la magnitude des mises à jour, WAV v1 change la représentation disponible pour le routeur.
• Contrairement à Delta Attention Residuals, qui route sur des deltas individuels (coût élevé), WAV v1 offre une approximation au niveau du bloc et multi-résolution qui équilibre efficacité et richesse structurelle.

Les auteurs concluent que les évaluations futures des méthodes de routage résiduel devraient considérer la profondeur, car les bénéfices de l'information multi-résolution peuvent être liés à la longueur et à la structure de la trajectoire résiduelle. Ils reconnaissent les limites actuelles, notamment l'absence d'écart-type entre les graines (seeds), la restriction à de petits modèles de niveau caractère, et l'utilisation de bases de détails conçues à la main, suggérant que les travaux futurs devraient explorer des bases apprenables et des modèles de niveau token à plus grande échelle.