Résumé technique : Lorsque les pertes s'alignent : pondération basée sur le gradient des pertes composites pour un préentraînement efficace

Formulation du problème

Les systèmes d'apprentissage profond modernes utilisent fréquemment un paradigme en deux étapes : un préentraînement sur des données à grande échelle, souvent non étiquetées, suivi d'une adaptation à des tâches en aval. Dans de nombreuses chaînes de préentraînement, la fonction objectif n'est pas une perte unique, mais une combinaison pondérée de plusieurs termes de perte (par exemple, la prédiction du token suivant, des pertes de contraste ou des objectifs de régression). Les poids relatifs de ces termes sont des hyperparamètres critiques qui influencent considérablement la qualité des représentations apprises et les performances en aval.

Actuellement, la sélection de ces poids repose sur un réglage manuel, une recherche aléatoire ou une optimisation bayésienne. Ces méthodes sont prohibitives en termes de calcul pour le préentraînement à grande échelle car elles nécessitent de nombreuses exécutions de formation indépendantes. De plus, des valeurs par défaut naïves, telles qu'un poids uniforme, échouent souvent à prendre en compte les contributions variables et les conflits entre les différents termes de perte. L'article répond au besoin d'une méthode efficace pour régler les poids linéaires des pertes dans les objectifs de préentraînement composites afin de maximiser les performances des tâches en aval, sans engager le coût de formations complètes répétées.

Méthodologie : Préentraînement aligné sur le gradient (GraP)

Les auteurs proposent le Préentraînement aligné sur le gradient (GraP), une méthode d'optimisation hiérarchique (bilevel) basée sur le gradient qui apprend les poids de perte en ligne. L'idée centrale est d'aligner le gradient de préentraînement composite sur le gradient d'un objectif en aval.

1. Formulation de l'optimisation hiérarchique

Le problème est formulé comme une tâche d'optimisation hiérarchique où la boucle externe optimise les poids de perte $w$ pour minimiser la perte en aval $L_{down}$, et la boucle interne optimise les paramètres du modèle $\theta$ pour minimiser la perte de préentraînement composite $L_{pre}(w)$.
$$\min_{w} L_{down}(\theta^*(w)) \quad \text{sous contrainte} \quad \theta^*(w) = \arg\min_{\theta} \sum_{k=1}^K w_k L_k(\theta)$$

2. Règle de mise à jour basée sur le gradient

En utilisant une approche de rétropropagation tronquée (SGD à étape unique), les auteurs dérivent une règle de mise à jour pour les poids. Le gradient de la perte en aval par rapport aux poids est proportionnel au produit scalaire entre le gradient en aval et les gradients individuels de préentraînement :
$$\frac{\partial L_{down}}{\partial w_k} \propto -\eta \, g_{down}^\top g_k$$
où $g_{down}$ est le gradient en aval et $g_k$ est le gradient de la $k$-ième perte de préentraînement. L'objectif est de maximiser l'alignement entre le gradient composite $\sum w_k g_k$ et le gradient en aval.

3. Approximation dans l'espace d'embedding

Une implémentation naïve nécessiterait de calculer $K$ passages arrière complets à travers le modèle pour obtenir $g_k$. Pour réduire le coût de calcul, GraP calcule les gradients par rapport à l'embedding de base partagé $z$ plutôt que par rapport à l'espace complet des paramètres.
$$\tilde{g}_k = \nabla_z L_k(h_k(z))$$
Puisque les pertes sont appliquées via des têtes séparées au-dessus du même embedding, le calcul de ces gradients est considérablement moins coûteux que des passages arrière complets. L'objectif d'optimisation devient la maximisation de l'alignement dans cet espace de représentation partagé :
$$\arg\max_{w} \frac{w^\top \tilde{G} \tilde{g}_{down}}{\|w^\top \tilde{G}\|}$$

4. Stratégie de normalisation

L'article s'oppose à la normalisation des poids eux-mêmes (par exemple, $\sum w_k = 1$) ou de leur norme $L_2$, car cela peut déformer la direction optimale du gradient composite. Au lieu de cela, GraP impose une contrainte sur la norme $L_2$ du gradient composite lui-même ($\|\sum w_k g_k\| = 1$). Cela garantit que l'optimisation se concentre sur la direction du vecteur de mise à jour plutôt que sur sa magnitude, retrouvant ainsi la direction d'alignement optimale sous l'objectif local.

5. Efficacité de l'algorithme

L'algorithme résultant ne nécessite qu'un seul passage avant et un seul passage arrière à travers la base par itération. Les gradients calculés pour les têtes sont réutilisés à la fois pour mettre à jour les poids de perte et pour mettre à jour les paramètres de la base. Cela réduit la surcharge du réglage des hyperparamètres à environ 30 % au-dessus d'une seule formation, par rapport au coût de plusieurs ordres de grandeur de l'optimisation bayésienne.

Contributions clés

1. Formulation du problème : Les auteurs formulent la sélection des poids de perte de préentraînement comme un problème d'optimisation hiérarchique ciblant explicitement les performances en aval.
2. Insight sur la normalisation : Ils analysent les stratégies de normalisation et démontrent que contraindre la norme du vecteur de mise à jour (gradient composite) est supérieur à contraindre les poids eux-mêmes, préservant ainsi la direction d'alignement optimale.
3. Algorithme efficace : Ils proposent GraP, un algorithme de complexité $O(1)$ concernant le nombre d'étapes avant-arrière par rapport au nombre d'hyperparamètres. Il exploite l'alignement des gradients dans l'espace de représentation partagé pour éviter plusieurs passages arrière.
4. Efficacité empirique : La méthode atteint des performances comparables ou supérieures à l'optimisation standard des hyperparamètres (par exemple, recherche bayésienne) tout en étant près de 50 fois plus efficace, permettant un réglage dans des contextes où les méthodes traditionnelles sont computationnellement irréalisables.
5. Identification de la redondance : L'approche démontre la capacité d'identifier et de réduire le poids des termes de perte redondants (par exemple, attribuer un poids proche de zéro à des pertes de contraste spécifiques dans des tâches de vision), agissant efficacement comme un mécanisme automatique de sélection de pertes.

Résultats expérimentaux

La méthode a été évaluée sur deux domaines : la modélisation de séquences d'événements et la vision par ordinateur auto-supervisée.

Modélisation de séquences d'événements

• Ensembles de données : Churn, AgePred, AlfaBattle, MIMIC-III et Taobao.
• Configuration : Le préentraînement impliquait un objectif composite de prédiction du token suivant (NTP) sur plusieurs champs et une perte de contraste.
• Performance : GraP a égalé ou surpassé des bases de référence soigneusement réglées (incluant GradNorm, DWA, MGDA et PCGrad) et des poids réglés par Optuna. Il a systématiquement obtenu les meilleurs résultats sur les ensembles de données Churn et MIMIC-III à travers les bases RNN, Transformer et HT-Transformer.
• Efficacité : Par rapport à l'optimisation bayésienne (Optuna), qui nécessite des dizaines d'exécutions indépendantes, GraP a atteint des performances similaires avec une seule formation et une surcharge modérée (ralentissement moyen de 44 % du débit).

Vision par ordinateur auto-supervisée

• Configuration : Évalué sur CIFAR-10, CIFAR-100 et ImageNet-100 en utilisant le cadre All4One (combinant les objectifs Barlow Twins, BYOL et NNCLR).
• Performance : Dans ce contexte, où les objectifs sont souvent bien équilibrés, GraP a atteint des performances comparables aux paramètres par défaut d'All4One et à d'autres méthodes d'apprentissage multi-tâches.
• Sélection de pertes : Notamment, GraP a attribué un poids proche de zéro à la perte Attn-NNCLR. Le retrait manuel de cette perte a produit des performances similaires, validant la capacité de la méthode à identifier des objectifs à faible impact ou redondants.

Études d'ablation

• Fuite d'étiquettes : Les expériences variant la fraction de données étiquetées utilisées pour le réglage des poids ont montré une robustesse, suggérant que les étiquettes fournissent un signal de guidage faible plutôt qu'une supervision directe pour la base.
• Normalisation : La normalisation du gradient composite proposée a surpassé la normalisation standard de la somme des poids.
• Formation auxiliaire : Le préentraînement avec GraP suivi d'un affinage supervisé a surpassé la formation auxiliaire conjointe, probablement en raison d'un surapprentissage réduit sur les données étiquetées limitées.

Importance et affirmations

L'article affirme que l'optimisation des poids de perte basée sur le gradient est une alternative pratique et efficace à l'apprentissage multi-tâches classique et à l'optimisation des hyperparamètres basée sur la recherche. Sa principale importance réside dans le fait qu'elle rend le réglage des hyperparamètres réalisable pour des scénarios de préentraînement à grande échelle avec plusieurs termes de perte, là où les méthodes traditionnelles sont prohibitives en termes de calcul.

Les auteurs soulignent que GraP ne nécessite pas plusieurs exécutions de formation, ce qui le rend adapté aux contextes avec des budgets de calcul limités. Bien que la méthode suppose des têtes séparées pour chaque perte, elle gère efficacement les compromis entre objectifs conflictuels en alignant la direction de mise à jour du préentraînement sur l'objectif en aval. Le travail suggère que dans des chaînes de préentraînement complexes, un alignement automatique et en ligne des poids de perte peut produire des représentations qui transfèrent aussi bien, voire mieux, que celles obtenues par un réglage manuel coûteux ou une recherche en boîte noire.