Résumé technique : Affinage supervisé préservant la rotation (RPSFT)

1. Énoncé du problème

L'affinage supervisé (SFT) est une étape standard post-entraînement pour les grands modèles de langage (LLM), utilisée pour adapter les modèles à des tâches spécifiques. Cependant, des travaux antérieurs indiquent que le SFT peut dégrader la généralisation hors domaine (OOD), un phénomène souvent attribué à l'« oubli catastrophique ». Des études récentes suggèrent que cette dégradation est liée à une dérive géométrique dans les matrices de poids préentraînées, spécifiquement la rotation excessive des sous-espaces singuliers dominants. Ces directions dominantes sont associées à des structures de haute variance et de haute courbure qui encodent des capacités à usage général.

Identifier et préserver directement les directions sensibles à la perte en utilisant l'information de Hessian ou de Fisher est prohibitif en termes de calcul à l'échelle des LLM modernes. Par conséquent, il existe un besoin d'une méthode de régularisation efficace qui limite la dérive géométrique inutile pendant le SFT tout en conservant la capacité d'adaptation à la tâche cible, améliorant ainsi le compromis entre les performances en domaine (ID) et la rétention hors domaine (OOD).

2. Méthodologie : Affinage supervisé préservant la rotation (RPSFT)

Les auteurs proposent l'Affinage supervisé préservant la rotation (RPSFT), une technique de régularisation qui ancre les sous-espaces singuliers dominants des matrices de poids préentraînées.

Mécanisme central

Le RPSFT repose sur l'hypothèse que les $k$ premiers vecteurs singuliers des matrices de poids préentraînées se chevauchent significativement avec les directions sensibles à Fisher (directions où les mises à jour provoquent de fortes augmentations de la perte OOD). Au lieu de figer les paramètres ou de projeter les gradients (ce qui peut être rigide ou coûteux), le RPSFT ajoute un terme de pénalité à la fonction de perte SFT standard.

Pour une matrice de poids sélectionnée $W$ avec une initialisation préentraînée $W_0$, la méthode précalcule les bases des vecteurs singuliers gauche et droit de rang $k$, $U^{(k)}_0$ et $V^{(k)}_0$, via la décomposition en valeurs singulières (SVD). Le terme de régularisation pénalise le changement du bloc projeté de la matrice de poids dans cette base :

$$\mathcal{L}_{RPSFT} = \mathcal{L}_{SFT} + \lambda \sum_{\ell \in \mathcal{M}} \| (U^{(k)}_{0, \ell})^\top (W_\ell - W_{0, \ell}) V^{(k)}_{0, \ell} \|_F^2$$

Où :

• $\mathcal{L}_{SFT}$ est la perte standard de vraisemblance négative.
• $\lambda$ est l'intensité de la régularisation.
• Le terme $\| (U^{(k)}_{0, \ell})^\top (W_\ell - W_{0, \ell}) V^{(k)}_{0, \ell} \|_F^2$ mesure la dérive de la matrice de poids dans le sous-espace singulier de rang $k$.
• $\mathcal{M}$ est l'ensemble des couches sélectionnées (généralement toutes les couches dans l'affinage à paramètres complets).

Justification de la conception

• Protection partielle : Contrairement à l'ancrage de rang complet (équivalent à une régularisation $L_2$ autour des poids initiaux), le RPSFT n'ancre que le bloc de rang $k$. Les directions complémentaires restent libres de s'adapter, permettant un apprentissage spécifique à la tâche.
• Efficacité computationnelle : La méthode nécessite de précalculer et de mettre en cache les bases singulières ($O((m+n)k + k^2)$ de mémoire par couche) et de calculer le bloc projeté pendant l'entraînement ($O(mnk)$ FLOPs par évaluation). Les auteurs notent que cela ajoute une surcharge minimale par rapport au SFT standard, en particulier lorsque la projection est calculée tous les $s$ pas.
• Justification théorique : Les auteurs analysent le développement local d'ordre deux de la perte. Ils soutiennent que tandis que le terme du premier ordre pilote l'adaptation, le terme du second ordre (courbure) pilote l'oubli. Puisque l'information de Fisher (un proxy de la courbure) est coûteuse à calculer, ils démontrent empiriquement que les blocs singuliers supérieurs des poids préentraînés capturent une partie significative (environ 20 %) de l'énergie du gradient projeté par Fisher, ce qui en fait un proxy structurel valide et peu coûteux pour les directions sensibles à la perte.

3. Configuration expérimentale

Les auteurs ont évalué le RPSFT sur trois familles et échelles de modèles : Llama-3.1-8B-Instruct, Qwen2.5-7B-Instruct et Qwen2.5-3B-Instruct.

• Données d'entraînement : L'affinage supervisé a été effectué sur OpenR1-Math (données de raisonnement mathématique).
• Lignes de base : Des comparaisons ont été faites avec le SFT standard à paramètres complets, l'Affinage Dynamique (DFT), le SFT pondéré par l'importance (IW) et l'ancrage d'initialisation $L_2$.
• Évaluation :
  • En domaine (ID) : Benchmarks mathématiques (AIME, AMC, MATH-500, etc.).
  • Hors domaine (OOD) : Benchmarks de raisonnement général, de sécurité et de connaissances (GPQA, MMLU-Pro, TruthfulQA, etc.).
  • RL en aval : Les points de contrôle SFT ont servi d'initialisations pour l'Affinage par Apprentissage par Renforcement (RLFT) utilisant DAPO (une variante de GRPO).

4. Résultats clés

Compromis en domaine et hors domaine

• Le RPSFT améliore systématiquement le compromis ID/OOD. Pour toutes les tailles de modèles, le RPSFT a obtenu les meilleures performances ID (Avg@k et Pass@k) ou des performances compétitives, tout en réduisant significativement l'oubli OOD par rapport au SFT standard.
• Comparaison avec les lignes de base : Bien que le DFT et l'IW aient montré une certaine rétention OOD, ils ont souvent souffert de sous-ajustement sur la tâche ID. Le RPSFT a maintenu une forte adaptation ID tout en préservant mieux les capacités OOD que le SFT standard. Par exemple, sur Llama-8B, le RPSFT a amélioré l'Avg@k OOD de +1,83 % par rapport au modèle de base, tandis que le SFT standard l'a amélioré de +1,73 % mais avec une moindre robustesse sur d'autres métriques.
• Sensibilité au rang : Des expériences variant le rang protégé $k$ ont montré que les rangs intermédiaires (par exemple, $k=768$) offrent le meilleur équilibre. Des rangs excessivement grands (approchant l'ancrage de rang complet) ont entravé l'adaptation à la tâche, tandis que $k=0$ revenait au SFT standard.

Performance RL en aval

• Initialisations plus fortes : Les points de contrôle RPSFT ont servi d'initialisations supérieures pour le RLFT en aval. Les modèles initialisés avec RPSFT ont atteint des performances RL finales plus élevées sur les benchmarks ID et OOD par rapport à ceux initialisés avec le SFT standard ou d'autres lignes de base.
• Préservation des capacités : Les améliorations de la rétention OOD au stade SFT n'ont pas été effacées par l'étape RL ultérieure ; elles se sont au contraire transférées, conduisant à des scores finaux constamment plus élevés.

Analyse géométrique et des représentations

• Réduction de la rotation : L'analyse de la dérive dans l'espace des poids a montré que le RPSFT a considérablement réduit la rotation des sous-espaces singuliers gauches dominants par rapport au SFT standard, au DFT et à l'IW.
• Stabilité des états cachés : Le RPSFT a mieux préservé la géométrie des représentations cachées. Les centroïdes des états cachés pour les modèles affinés avec RPSFT sont restés plus proches des centroïdes du modèle de base sur divers benchmarks, tandis que le SFT standard et le DFT ont présenté des dérives plus importantes. Cela suggère que contraindre la rotation de l'espace des poids se traduit directement par la préservation de l'espace de représentation interne du modèle.

5. Importance et affirmations

L'article affirme que le RPSFT fournit un régularisateur simple et efficace en termes de calcul qui atténue efficacement l'oubli des capacités à usage général pendant l'affinage supervisé.

• Mécanisme : Le travail fournit des preuves empiriques et théoriques montrant que la préservation des sous-espaces singuliers préentraînés dominants réduit la rotation géométrique, protège les représentations d'états cachés et améliore le compromis oubli–adaptation.
• Praticité : Contrairement aux méthodes nécessitant une estimation de Hessian ou des projections de gradients complexes, le RPSFT s'intègre directement dans les pipelines SFT standard à paramètres complets sans nécessiter de données supplémentaires ni de limites de tâches.
• Généralité : La méthode démontre des améliorations cohérentes sur différentes architectures de modèles (Llama et Qwen) et échelles (de 3 à 8 milliards de paramètres).

Les auteurs concluent que le contrôle de la dérive dans les sous-espaces singuliers préentraînés est une stratégie pratique et efficace pour maintenir la généralisation OOD tout en permettant une adaptation rapide aux tâches, offrant une base plus solide pour les étapes ultérieures d'apprentissage par renforcement.

6. Limites

Les auteurs reconnaissent que le RPSFT applique actuellement une règle de rang protégé uniforme à toutes les matrices de poids. Ils suggèrent que des travaux futurs pourraient investiguer une régularisation sélective par couche (par exemple, en se concentrant sur les premières couches d'attention automatique) pour réduire davantage la surcharge. De plus, les expériences actuelles se concentrent sur des ensembles de données d'entraînement mathématiques et des modèles dans la plage 3B–8B ; l'extension à des données non mathématiques et à des échelles de modèles plus grandes est identifiée comme une direction pour la recherche future.