Résumé technique : Une perspective d'optimisation contrainte des Transformers déroulés

Énoncé du problème

Bien que la littérature récente interprète les architectures Transformer comme des algorithmes de descente itératifs résolvant des problèmes d'optimisation (ex. : Yang et al., 2022 ; De Weerdt et al., 2023), l'entraînement standard via la minimisation du risque empirique (ERM) échoue souvent à imposer le comportement de descente monotone attendu de tels optimiseurs. Lorsqu'ils sont entraînés avec l'ERM, les représentations intermédiaires des Transformers déroulés ne diminuent pas nécessairement la fonction de perte de manière monotone à travers les couches. Cette incohérence provient du fait que l'entraînement standard ne garantit pas que les paramètres appris satisfont les conditions requises pour la minimisation de l'énergie, ni qu'il prend en compte la nature bilatérale de l'entraînement d'une architecture déroulée. Par conséquent, ces modèles peuvent présenter des motifs de perte non monotones et une robustesse réduite aux perturbations d'entrée et aux changements de distribution.

Méthodologie

Les auteurs proposent un cadre qui traite l'entraînement des Transformers comme un problème d'apprentissage contraint. Au lieu de minimiser la perte moyenne sur la sortie de l'ensemble du réseau, la méthode impose que la perte attendue diminue de manière monotone à chaque couche.

1. Formulation contrainte

La contribution centrale est la formulation d'un problème d'optimisation contraint où la sortie de chaque couche $l$ est tenue de réduire la perte attendue $f$ d'un facteur $(1 - \alpha_l)$ par rapport à la couche précédente $l-1$ :
$$\begin{aligned} T^* = \arg\min_{T} \quad & \mathbb{E}[f(X, \Phi(X; T))] \\ \text{sous contrainte de} \quad & \mathbb{E}[f(X, \Phi_l(X; T))] \leq (1 - \alpha_l) \mathbb{E}[f(X, \Phi_{l-1}(X; T))], \quad \forall l \end{aligned}$$
Ici, $\Phi_l$ représente la sortie de la couche $l$, et $\alpha_l$ est un calendrier de pas de taille (stepsize schedule). Cette formulation force le Transformer à se comporter comme un algorithme de descente, garantissant que les représentations intermédiaires améliorent progressivement l'objectif.

2. Algorithme d'entraînement primal-dual

Pour résoudre ce problème contraint non convexe, les auteurs emploient un schéma d'entraînement primal-dual dans le domaine dual. Ils construisent une fonction lagrangienne :
$$\mathcal{L}(T, \lambda) = \mathbb{E}[f(X, \Phi(X; T))] + \sum_{l=1}^L \lambda_l \mathbb{E}[f(X, \Phi_l(X; T)) - (1 - \alpha_l)f(X, \Phi_{l-1}(X; T))]$$
La procédure d'entraînement (Algorithme 1) alterne entre :

• Mise à jour primale : Minimisation du lagrangien empirique par rapport aux paramètres du Transformer $T$ en utilisant la descente de gradient stochastique.
• Mise à jour duale : Maximisation de la fonction duale par rapport aux multiplicateurs de Lagrange $\lambda$ en utilisant la montée de gradient projetée.

Les auteurs s'appuient sur des résultats récents montrant que les réseaux de neurones profonds présentent de faibles écarts de dualité (Chamon et al., 2023), permettant au problème dual de produire des solutions quasi-optimales et quasi-admissibles. Pour assurer une faisabilité stricte, une relaxation résiliente avec une variable d'écart (slack variable) et une pénalité quadratique est introduite.

Principales contributions

Le papier présente quatre contributions principales :

1. Formulation de l'apprentissage contraint : Les auteurs formulent un problème d'apprentissage où les couches du Transformer sont explicitement contraintes à réduire la perte attendue de manière monotone, divergeant ainsi de l'ERM standard.
2. Algorithme d'entraînement dual et garanties de convergence : Ils développent un algorithme d'entraînement primal-dual et fournissent des garanties théoriques (Théorème 2) selon lesquelles le Transformer déroulé contraint converge asymptotiquement vers la valeur optimale de la perte statistique. Ils établissent également des garanties de généralisation hors distribution (OOD) (Théorème 4), prouvant que les propriétés de descente sont maintenues sous les changements de distribution.
3. Robustesse aux perturbations des embeddings : Dans les tâches de classification de texte (IMDb, MNLI), les auteurs démontrent que les modèles contraints maintiennent une précision plus élevée sous des perturbations gaussiennes, uniformes et discrètes des embeddings de tokens par rapport aux bases non contraintes.
4. Robustesse dans le débruitage vidéo : Appliquée au débruitage vidéo (CUHK Avenue, UCSD, ShanghaiTech), la méthode produit des modèles plus robustes aux niveaux de bruit OOD, mesurés par l'erreur quadratique moyenne (RMSE), tout en préservant les performances en distribution.

Résultats expérimentaux

Les auteurs ont évalué leur méthode sur des Transformers pré-entraînés conventionnels (DistilBERT, RoBERTa, Llama 3.1) et des architectures déroulées (Deep Unfolded Sequential Transformer - DUST, Unrolled Transformer - UT).

• Descente monotone : Les modèles contraints ont présenté des diminutions de perte par couche lisses et monotones, tandis que les bases non contraintes montraient des motifs non monotones avec des pics (par exemple, un pic à la couche 19 pour RoBERTa).
• Robustesse OOD :
  • Classification de texte : À mesure que les niveaux de perturbation ($\gamma$) augmentaient, les modèles contraints se dégradaient plus graduellement. Par exemple, sur RoBERTa avec IMDb, le modèle contraint conservait une précision nettement plus élevée sous des niveaux de bruit élevés par rapport à la version non contrainte.
  • Débruitage vidéo : Les modèles contraints ont systématiquement obtenu un RMSE plus faible sous des niveaux de perturbation de test élevés ($\gamma \geq 1.0$) à travers divers jeux de données et architectures.
  • Fine-tuning de LLM : Sur Llama 3.1 8B, le modèle contraint a maintenu ses performances en distribution (taux de victoire de 50 % sur AlpacaEval sans bruit) tout en surpassant significativement le modèle non contraint sous bruit (taux de victoire de 69,93 % avec $\gamma=1,5$).
• Compromis : La méthode préserve généralement les performances en distribution avec une dégradation négligeable (souvent < 0,3 %), suggérant que l'imposition de contraintes de descente ne sacrifie pas la précision standard au profit de la robustesse.

Signification et affirmations

Le papier affirme qu'imposer des contraintes de descente aligne la dynamique interne des Transformers avec les propriétés de stabilité des algorithmes d'optimisation classiques. Cet alignement fournit une base théorique et empirique pour une meilleure robustesse face aux perturbations d'entrée et aux changements de distribution.

Les auteurs positionnent ce travail comme un pont entre le déroulement algorithmique (où les réseaux résolvent des problèmes d'optimisation) et l'apprentissage contraint. Ils soutiennent qu'en entraînant explicitement les Transformers pour descendre, on peut obtenir des modèles intrinsèquement plus robustes au bruit, une propriété particulièrement précieuse pour les applications impliquant la confidentialité différentielle locale, où les perturbations d'entrée sont nécessaires pour la vie privée des utilisateurs. Ce travail suggère que la perspective "déroulée" n'est pas seulement une interprétation théorique, mais un objectif d'entraînement pratique qui apporte des améliorations tangibles en termes de généralisation et de robustesse sans nécessiter de changements architecturaux.