SafeMERGE: Preserving Safety Alignment in Fine-Tuned Large Language Models via Selective Layer-Wise Model Merging

1. Problématique

Le fine-tuning (ajustement fin) des grands modèles de langage (LLM) est une pratique courante pour les adapter à des domaines spécialisés (mathématiques, code, biomédecine). Cependant, ce processus présente un risque majeur : il peut éroder l'alignement de sécurité initial du modèle.

• Fragilité de l'alignement : Des études récentes montrent que même un fine-tuning avec des données bénignes peut involontairement dégrader la sécurité, rendant le modèle vulnérable aux demandes nuisibles ou aux "jailbreaks".
• Limites des solutions existantes : Les méthodes actuelles pour rétablir la sécurité souffrent de deux défauts principaux :
  1. Elles nécessitent souvent des algorithmes d'alignement personnalisés et complexes, difficiles à intégrer dans les bibliothèques open-source standards.
  2. Les méthodes plus simples tendent à sacrifier l'utilité du modèle (ses performances sur la tâche visée) pour retrouver la sécurité, ou inversement.

2. Méthodologie : SafeMERGE

Les auteurs proposent SafeMERGE, un cadre léger et post-fine-tuning qui restaure la sécurité sans réentraînement, en utilisant une fusion sélective au niveau des couches.

Principes clés :

1. Deux modèles complémentaires :
   • Un modèle fine-tuné sur la tâche spécifique (mais potentiellement dangereux).
   • Un modèle sécurisé (safe model), entraîné sur un petit ensemble de données d'alignement de sécurité (task-agnostic, réutilisable).
2. Détection des couches dangereuses :
   • SafeMERGE ne fusionne pas tout le modèle. Il identifie les couches spécifiques qui ont dévié du comportement sûr.
   • Critère de similarité cosinus : Pour chaque couche LoRA (Low-Rank Adaptation), le modèle calcule la similarité cosinus entre les mises à jour du modèle fine-tuné ($\Delta W_f$) et leur projection sur un sous-espace d'alignement de sécurité ($V_i$).
   • Le sous-espace $V_i$ est défini comme la différence de poids entre un modèle de base et un modèle aligné (instruct/chat).
   • Si la similarité cosinus $\rho_i$ est inférieure à un seuil $\tau$, la couche est considérée comme "dangereuse" (elle s'est éloignée de l'alignement de sécurité).
3. Fusion sélective :
   • Seules les couches identifiées comme dangereuses sont fusionnées avec les couches correspondantes du modèle sécurisé.
   • La fusion utilise une stratégie de pondération linéaire : $\Delta W_{merge} = \alpha \Delta W_f + (1-\alpha) \Delta W_s$.
   • Les couches sûres restent inchangées, préservant ainsi l'utilité de la tâche.

3. Contributions Clés

• Approche Post-Fine-Tuning sans Réentraînement : SafeMERGE opère après la phase d'entraînement, évitant la nécessité de modifier les pipelines d'entraînement existants ou d'utiliser des algorithmes complexes.
• Sélectivité au niveau des couches : Contrairement aux méthodes qui projettent globalement les vecteurs de mise à jour (comme SafeLoRA) ou fusionnent aveuglément tout le modèle (comme RESTA), SafeMERGE cible uniquement les couches dégradées. Cela permet un compromis optimal entre sécurité et performance.
• Indépendance vis-à-vis des bibliothèques : La méthode est simple à implémenter avec des bibliothèques standards (ex: HuggingFace PEFT) et peut être exécutée entièrement sur CPU.
• Robustesse et Généralisation : Le modèle sécurisé de référence n'a besoin d'être entraîné qu'une seule fois et peut être réutilisé pour diverses tâches de fine-tuning.

4. Résultats Expérimentaux

Les auteurs ont évalué SafeMERGE sur quatre LLMs populaires (Llama-2, Llama-3.1, Qwen-2, Qwen-2.5) et plusieurs tâches (GSM8K, PubMedQA, et des tâches télécoms hors distribution).

• Performance Sécurité vs Utilité :
  • SafeMERGE réduit significativement les sorties nuisibles (mesurées sur les benchmarks DirectHarm et HexPhi) par rapport aux modèles fine-tunés bruts.
  • Il maintient, voire améliore, les performances sur les tâches utilitaires (précision sur GSM8K, PubMedQA). Par exemple, sur Llama-3.1 pour GSM8K, SafeMERGE atteint une précision de 78,50% tout en réduisant le taux de nuisance à 8,80% (contre 28,30% pour le modèle fine-tuné).
• Comparaison avec les SOTA :
  • SafeMERGE surpasse systématiquement les baselines : SafeInstruct (ajout de données de sécurité), RESTA (soustraction de vecteurs de tâche), et SafeLoRA (projection globale).
  • Il offre le meilleur compromis sécurité-utilité, évitant la dégradation des performances observée avec RESTA et la sécurité insuffisante de SafeLoRA.
• Analyse des Ablations :
  • Similarité Cosinus vs Distance Euclidienne : L'utilisation de la similarité cosinus est cruciale. La distance euclidienne (basée sur la magnitude) pénalise les mises à jour utiles, entraînant une chute drastique de l'utilité (ex: 42% sur GSM8K pour Llama-3.1), tandis que la similarité cosinus préserve l'utilité tout en corrigeant la sécurité.
  • Seuil $\tau$ : Un seuil optimal se situe généralement entre 0,6 et 0,8. Un seuil trop bas laisse passer des couches dangereuses ; un seuil trop haut fusionne trop de couches sûres, réduisant l'utilité.
  • Stratégie de fusion : La fusion linéaire s'avère plus robuste et cohérente que les stratégies TIES ou DARE dans ce contexte.

5. Signification et Impact

SafeMERGE représente une avancée significative pour la sécurité des LLMs en production :

• Praticité : Il résout le problème de la "fragilité de la sécurité" sans imposer de coûts computationnels élevés ni de complexité algorithmique.
• Adoptabilité : En s'intégrant facilement dans les flux de travail standards et en ne nécessitant pas de réentraînement, il est immédiatement applicable par les praticiens.
• Principe de précaution : Il démontre que la sécurité peut être restaurée de manière ciblée et géométriquement interprétable, validant l'hypothèse que la dégradation de sécurité est souvent localisée à certaines couches spécifiques du modèle.

En conclusion, SafeMERGE établit une nouvelle référence pour les défenses post-fine-tuning, prouvant qu'il est possible de maintenir des modèles à la fois performants et sûrs grâce à une fusion intelligente et sélective des poids.