Résumé Technique : La Restauration Sélective de Couches Restaure la Diversité

Problématique

Les techniques de post-entraînement, telles que l'ajustement d'instructions (instruction tuning) et l'apprentissage par renforcement à partir du feedback humain (RLHF), améliorent considérablement les capacités de suivi d'instructions et la serviabilité des grands modèles de langage (LLM). Cependant, ces méthodes induisent souvent un effondrement de mode (mode collapse), un phénomène où le modèle assigne une masse de probabilité disproportionnellement élevée à une plage étroite de réponses, supprimant ainsi des alternatives diverses mais valides. Cela conduit à des sorties répétitives dans des contextes ouverts (ex. : écriture créative, QA ouverte, raisonnement), limitant l'utilité du modèle dans des applications qui valorisent des réponses distinctes.

Les stratégies de mitigation existantes se répartissent en trois catégories, chacune présentant des limites :

1. Les méthodes basées sur le décodage (ex. : mise à l'échelle de la température, échantillonnage min-p) augmentent la stochasticité mais ne peuvent pas récupérer la diversité perdue en raison des changements de paramètres lors du post-entraînement.
2. Les méthodes basées sur le prompt (ex. : demander plusieurs réponses) augmentent les coûts d'inférence et sont sensibles à la formulation du prompt.
3. Les méthodes basées sur l'entraînement (ex. : redéfinition de l'objectif) nécessitent un réentraînement et une collecte de données coûteux.

Les auteurs recherchent une méthode qui récupère efficacement la diversité, qui est efficace en termes de calcul lors de l'adaptation, et qui n'introduit aucun coût supplémentaire lors de l'inférence.

Méthodologie

L'approche proposée, la Restauration Sélective de Couches (SLR - Selective Layer Restoration), est une intervention dans l'espace des poids sans entraînement, basée sur deux observations empiriques :

1. Les couches des LLM présentent des rôles fonctionnels distincts (ex. : les associations factuelles sont souvent localisées dans les couches de profondeur intermédiaire).
2. Le post-entraînement oriente principalement quels comportements pré-entraînés sont exprimés, supprimant souvent les autres (oubli catastrophique), plutôt que d'acquérir de nouvelles connaissances entièrement nouvelles.

Les auteurs émettent l'hypothèse que l'effondrement de mode est localisé dans des couches spécifiques et que la restauration d'une plage de couches soigneusement choisie dans un modèle post-entraîné ($M_{post}$) à leurs poids pré-entraînés ($M_{pre}$) peut restaurer la diversité tout en maintenant la qualité.

Le cadre SLR

Étant donné un modèle avec $N$ couches, la SLR construit un modèle hybride $M_{i:j}$ en remplaçant un intervalle contigu de couches $[i, j]$ dans $M_{post}$ par les couches correspondantes de $M_{pre}$. L'objectif est de trouver l'intervalle $(i, j)$ qui maximise la diversité ($D$) sous réserve d'un seuil de qualité minimal ($Q \ge q_{min}$).

La tâche de substitution CRC

L'optimisation directe de l'intervalle de restauration sur des tâches complexes est coûteuse en calcul et difficile à évaluer en raison de la subjectivité de la « qualité ». Pour y remédier, les auteurs introduisent le Caractère Aléatoire Contraint (CRC - Constrained Random Character) comme tâche de substitution.

• Mécanisme : Le CRC consiste à générer des caractères aléatoires (chiffres ou lettres) au sein d'un ensemble valide strictement défini (ex. : « Générez un entier aléatoire dans [0, 5] »).
• Métriques :
  • Qualité ($Q_{CRC}$) : La masse de probabilité assignée aux sorties valides (mesurant l'adhérence aux instructions).
  • Diversité ($D_{CRC}$) : L'entropie de la distribution sur les sorties valides (mesurant la variété).
• Processus de sélection : Les auteurs effectuent une recherche par grille sur tous les intervalles de couches contigus possibles en utilisant le CRC. Ils sélectionnent l'intervalle qui maximise $D_{CRC}$ tout en satisfaisant $Q_{CRC} \ge 0,9 \times Q_{CRC}(M_{post})$. Cet intervalle sélectionné est ensuite appliqué aux tâches en aval sans ajustement supplémentaire.

Contributions Clés

1. Restauration Sélective de Couches (SLR) : Une méthode sans entraînement qui restaure un intervalle spécifique de couches contiguës d'un LLM post-entraîné vers ses paramètres pré-entraînés. Elle préserve l'architecture originale et le nombre de paramètres, n'incurant aucun coût d'inférence supplémentaire. Il est démontré qu'elle est complémentaire aux méthodes de diversité par décodage et par prompt.
2. Caractère Aléatoire Contraint (CRC) : Une tâche de substitution avec un ensemble de validité explicite et un objectif de diversité naturel. Le CRC permet une exploration efficace et à faible coût des intervalles de restauration et fournit une preuve empirique pour l'hypothèse de localisation des couches.
3. Validation Empirique : Évaluation complète à travers trois familles de modèles (Llama, Qwen, Gemma) et trois tâches en aval (écriture créative, QA ouverte, raisonnement multi-étapes).

Résultats Expérimentaux

Les auteurs ont évalué la SLR par rapport à la ligne de base post-entraînée et à une ligne de base « Model Soup » (moyenne pondérée des paramètres du modèle pré-entraîné et post-entraîné).

• Écriture Créative : La SLR améliore systématiquement la diversité (augmentation moyenne de +42,5 % de la dissimilarité basée sur les embeddings) avec une perte de qualité minimale (perte moyenne de -2,1 % dans les scores de l'LLM-judge). En revanche, le Model Soup montre des gains de diversité plus faibles (+27,3 %) et une dégradation de la qualité plus importante (-6,0 %).
• QA Ouverte : La SLR augmente substantiellement la diversité des réponses, montrant une augmentation de +112,7 % de l'entropie et une augmentation de +169,5 % de la couverture (fraction de réponses correctes uniques), tout en maintenant une précision élevée (baisse moyenne de -1,8 %).
• Raisonnement (GSM8K) : Sur des problèmes où les modèles post-entraînés échouent sous décodage glouton (greedy decoding), la SLR améliore le Pass@k pour tous les budgets d'échantillonnage ($k \in \{1, \dots, 64\}$), indiquant un déplacement de la masse de probabilité vers des chemins de raisonnement diversifiés mais corrects.
• Composabilité : La SLR reste efficace lorsqu'elle est combinée avec un décodage à température plus élevée ($T=1,5$) et des méthodes d'échantillonnage verbalisé (méthodes basées sur le prompt), suggérant que la SLR modifie la distribution sous-jacente du modèle tandis que les autres méthodes manipulent le processus d'échantillonnage ou l'entrée.
• Études d'Ablation : La restauration des couches guidée par la tâche de substitution CRC a produit des résultats supérieurs à la restauration d'intervalles arbitraires (ex. : les couches les plus précoces ou les plus tardives). La restauration de couches arbitraires a souvent échoué à améliorer la diversité ou a dégradé la qualité, confirmant que la sélection de l'intervalle spécifique est critique.

Signification et Revendications

L'article affirme que la SLR offre une voie simple mais efficace pour atténuer l'effondrement de mode sans la surcharge de calcul d'un réentraînement ou les coûts d'inférence de prompts complexes. En démontrant que la diversité peut être récupérée par des interventions ciblées dans l'espace des poids, ce travail soutient une vision modulaire des LLM post-entraînés :

1. Les fonctionnalités et les comportements (y compris l'effondrement de mode) sont localisés dans des couches spécifiques plutôt que distribués uniformément.
2. Le post-entraînement oriente principalement l'expression des comportements pré-entraînés plutôt que d'altérer fondamentalement la base de connaissances centrale du modèle.

Les auteurs concluent que les interventions fines dans l'espace des poids comme la SLR offrent une voie prometteuse pour contrôler le compromis qualité-diversité dans les LLM, permettant de récupérer des propriétés pré-entraînées désirables (comme la diversité) tout en conservant les bénéfices du post-entraînement (comme l'adhérence aux instructions). L'article reste modeste quant à la scalabilité, notant que les résultats sont actuellement limités aux modèles d'environ 7 à 9 milliards de paramètres et que des travaux futurs sont nécessaires pour déterminer si ces conclusions se maintiennent à plus grande échelle ou avec des interventions plus granulaires (ex. : restauration de sous-couches).