Résumé technique : CascadeFormer : Des Transformers à décroissance de profondeur motivés par l'asymétrie du fan-in de gradient

1. Énoncé du problème

Les Transformers profonds reposent généralement sur le principe de mise à l'échelle uniforme, où des blocs résiduels identiques sont empilés pour augmenter la profondeur du modèle. Cependant, des observations empiriques indiquent que les couches plus profondes contribuent souvent avec une valeur fonctionnelle décroissante, présentant une similitude et une redondance de représentation élevées. Bien que l'explication conventionnelle attribue cela à l'atténuation de la magnitude des gradients dans les couches profondes, les auteurs soutiennent que cela traite un symptôme plutôt que la cause profonde. Le papier pose que le problème fondamental n'est pas la magnitude du gradient, mais la diversité compositionnelle du signal de gradient, qui est structurellement limitée dans les couches profondes.

2. Hypothèse centrale : Asymétrie du Fan-in de Gradient (GFA)

Le papier introduit l'Asymétrie du Fan-in de Gradient (GFA) comme un compte structurel de la redondance des couches. Dans les piles résiduelles Pre-LayerNorm, le gradient à une couche $l$ donnée est la somme d'un chemin d'identité et de tous les chemins computationnels en aval.

• Déséquilibre structurel : Les couches peu profondes agrègent les gradients provenant d'un grand nombre de chemins de calcul en aval, ce qui résulte en un signal riche et compositionnellement diversifié. Inversement, les couches profondes agrègent à partir d'un ensemble de plus en plus restreint de chemins, ce qui produit un gradient structurellement simple et pauvre en informations.
• Formulation mathématique : Le fan-in de gradient ($\phi_l$) diminue linéairement avec la profondeur dans les architectures standards. Dans les architectures employant une supervision profonde (par exemple, LayerSkip), l'introduction de têtes de perte auxiliaires crée de nouvelles hiérarchies de gradients, transformant la disparité de fan-in de linéaire à quadratique, concentrant ainsi l'information encore plus lourdement dans les couches les plus superficielles.
• Distinction clé : Les auteurs soutiennent que la magnitude du gradient est un indicateur de ce déséquilibre structurel de l'information. Un optimiseur peut rescaler les magnitudes de gradient, mais il ne peut pas inventer l'information compositionnelle manquante qui est structurellement absente dans les couches profondes.

3. Méthodologie et interventions

Pour valider la GFA et l'exploiter à des fins d'efficacité, les auteurs ont mené des expériences sur des modèles entraînés de zéro jusqu'à 1,2 milliard de paramètres (incluant des Transformers de type Llama et ResNet-50).

3.1 Validation empirique

• Preuve corrélationnelle : Les auteurs ont mesuré la Part de gradient accumulée ($\bar{g}_i$) pendant l'entraînement et l'ont corrélée à l'Importance fonctionnelle ($\Delta \mathcal{M}_i$), définie par la dégradation des performances lors de l'ablation de couches. De fortes corrélations positives ont été trouvées à travers les Transformers classiques, les modèles LayerSkip et ResNet-50, indiquant que le flux de gradient d'entraînement prédit l'importance finale des couches.
• Tests interventionnels :
  1. Égalisation de magnitude (Ablatif) : Les auteurs ont artificiellement rescalé les gradients par couche pour qu'ils aient des normes L2 égales pendant l'entraînement. Cette intervention a échoué à restaurer l'importance fonctionnelle des couches profondes, suggérant que la magnitude seule n'est pas le facteur causal.
  2. Augmentation du nombre de chemins (Constructif) : Les auteurs ont augmenté le fan-in de gradient en aval pour les couches profondes en répétant les dernières couches avec des paramètres partagés (augmentant la "profondeur virtuelle"). Ce changement structurel a réussi à restaurer et élever l'importance fonctionnelle des couches profondes, confirmant que la diversité structurelle des chemins est le goulot d'étranglement.

3.2 Applications proposées

Motivés par la GFA, le papier propose deux méthodes d'efficacité :

1. CascadeFlow Pruning (CFP) :

   • Mécanisme : Une méthode d'élagage efficace qui utilise les normes de gradient d'entraînement accumulées ($\bar{g}_i$) comme un indicateur de l'importance structurelle.
   • Avantage : Contra irement aux heuristiques standards (ex: critères de Taylor, élagage de magnitude ou similitude des états cachés) qui nécessitent des analyses post-hoc coûteuses ou des passes avant, le CFP exploite les gradients collectés directement pendant l'entraînement.
   • Implémentation : Les couches élaguées sont implémentées via une simple fonction de passage (fonction identité) lors de l'inférence, ne nécessitant aucune modification architecturale.

2. CascadeFormer :

   • Mécanisme : Une architecture qui intériorise le principe de la GFA en réduant la largeur du modèle (têtes d'attention et/ou dimensions FFN) à mesure que la profondeur augmente. Cela aligne la capacité du modèle avec la décroissance naturelle de la diversité de la composition du gradient.
   • Conception : La réduction est régie par des étapes discrètes pour les dimensions d'attention et une réduction linéaire pour les dimensions FFN.
   • Objectif : Améliorer l'efficacité de l'inférence (latence et débit) à un budget de FLOP d'entraînement fixe en réduisant la capacité là où elle est structurellement moins nécessaire.

4. Résultats clés

• Performance de l'élagage : Sur un modèle de 1,2 milliard de paramètres, le CFP a surpassé les heuristiques d'élagage standards (Similitude, Taylor, Magnitude) en termes de perplexité (PPL) et de stabilité. Le CFP a maintenu des performances compétitives sur les tâches en aval (HellaSwag) même sous un élagage agressif, là où les autres méthodes montraient une dégradation brutale des performances et des classements instables selon les graines aléatoires.
• Efficacité de l'architecture : La variante CascadeFormer-A2, qui réduit les dimensions d'attention, a atteint une perplexité comparable à celle d'une ligne de base uniforme de 15 couches entraînée avec un budget de calcul équivalent.
  • Gains d'efficacité : Elle a réduit la latence d'inférence de 8,6 % et augmenté le débit de 9,4 % sur un GPU A100.
  • Comparaison : Les variantes réduisant uniquement les dimensions FFN ou les deux étaient également compétitives, mais la réduction de l'attention a donné les meilleurs résultats, probablement parce que le mécanisme d'auto-attention est la cible la plus efficace pour la réduction de capacité.

5. Signification et affirmations

Le papier affirme que la redondance des couches n'est pas un échec d'optimisation mais une conséquence du flux d'information structurel inhérent aux réseaux résiduels.

• Contribution théorique : Il fournit un compte structurel (GFA) expliquant pourquoi les couches plus profondes contribuent moins, soutenu par des preuves corrélationnelles et interventionnelles. Il conteste la vue selon laquelle la magnitude du gradient est le principal moteur de la redondance.
• Contribution pratique : Il offre deux méthodes immédiates et efficaces pour améliorer l'efficacité des modèles : une stratégie d'élagage robuste (CFP) et une nouvelle classe d'architecture (CascadeFormer) qui fait correspondre la capacité au flux d'information.
• Modestie et limites : Les auteurs déclarent explicitement que leurs preuves sont limitées aux modèles entraînés de zéro jusqu'à 1,2 milliard de paramètres. Ils reconnaissent que savoir si la GFA se maintient à l'échelle de 100B+ reste une question empirique ouverte. De plus, ils notent que la magnitude du gradient n'est qu'un indicateur du fan-in dans les régimes de haut rang ; mesurer le rang effectif et l'orthogonalité du gradient est une étape de mesure nécessaire. Enfin, le CFP nécessite l'accès aux gradients lors de l'entraînement, ce qui le rend inapplicable pour l'élagage post-hoc de modèles pré-entraînés à code fermé.

Les auteurs concluent que, bien que leurs méthodes soient motivées par le compte de la GFA, les méthodes elles-mêmes sont précieuses sur la base de leurs mérites mesurés, même si un lecteur reste sceptique quant au compte théorique sous-jacent.