Résumé Technique : MASC (Margin Self-Correction) pour un désapprentissage rapide à grande échelle

Problématique

Les grands modèles de langage (LLM) sont sujets à la mémorisation de données d'entraînement sensibles, notamment des informations privées et du contenu protégé par le droit d'auteur, ce qui pose des risques importants pour la sécurité et la confidentialité. Bien que le « désapprentissage machine » (machine unlearning) vise à mettre à jour un modèle pour qu'il se comporte comme s'il n'avait jamais vu certaines séquences d'entraînement spécifiques, les méthodes actuelles de désapprentissage approximatif font face à deux limitations majeures de scalabilité :

1. Budgets d'entraînement inefficaces : Les méthodes actuelles (ex. : Gradient Ascent, NPO) s'exécutent généralement pour un nombre d'époques fixe et prédéfini. Cela conduit souvent à un entraînement au-delà du compromis souhaité entre oubli et rétention, ou à un arrêt prématuré, obligeant les praticiens à sélectionner rétrospectivement un modèle parmi les points de contrôle (checkpoints) sauvegardés.
2. Sélection de modèle coûteuse : La sélection du point de contrôle optimal nécessite des évaluations répétées sur des données de validation, ce qui engendre un surcoût de stockage et de calcul. De plus, ces méthodes manquent d'une règle d'arrêt intrinsèque et en ligne pour identifier l'équilibre optimal sans métriques externes.

De plus, l'impact de l'échelle du modèle sur l'efficacité du désapprentissage reste sous-exploré. Il n'est pas clair si les modèles plus grands, qui peuvent internaliser l'information plus fortement, répondent différemment aux procédures de désapprentissage par rapport aux modèles plus petits.

Méthodologie : MArgin Self-Correction (MASC)

Les auteurs introduisent MASC, une méthode de désapprentissage efficace qui remplace l'entraînement à budget fixe par une règle d'arrêt adaptative et en ligne. MASC repose sur le principe que la reproduction exacte d'une séquence oubliée nécessite que le modèle assigne une probabilité élevée au jeton suivant réel sur de nombreuses positions. Au lieu de minimiser globalement la vraisemblance, MASC impose une condition de dominance au niveau du jeton.

Mécanisme central

1. Comparaison de jetons restreinte : Pour une séquence à oublier $(x, y)$ à la position $t$, MASC identifie l'ensemble $S_{\theta,k}(c_t)$ des $k$ meilleurs jetons alternatifs non cibles du modèle.
2. Marge restreinte : Il calcule une marge restreinte $m^{(k,\beta)}_\theta(x, y, t)$, définie comme l'écart de logit entre le jeton cible $y_t$ et l'agrégat log-sum-exp des $k$ meilleures alternatives :
   $$m^{(k,\beta)}_\theta(x, y, t) = \beta z_\theta(y_t | c_t) - \log \sum_{v \in S_{\theta,k}(c_t)} \exp(\beta z_\theta(v | c_t))$$
   Cette marge mesure la force avec laquelle le modèle préfère le jeton oublié par rapport à un agrégat d'alternatives plausibles.
3. Perte d'auto-correction (Self-Correction Loss) : MASC minimise une perte de substitution basée sur une fonction de type "hinge" appliquée à cette marge. La perte n'est active que lorsque la marge dépasse un seuil $\tau_\rho$, indiquant que le jeton cible domine encore les alternatives. Si la marge est déjà faible (c'est-à-dire que le jeton est supprimé), le gradient s'annule.
4. Règle d'arrêt en ligne : L'algorithme surveille le taux de violation $\hat{V}_\rho$ sur un petit sous-ensemble de l'ensemble à oublier ($D_{val}$). L'entraînement s'arrête au premier pas où le taux de violation moyen descend en dessous d'une tolérance $\alpha$. Cela élimine la nécessité d'une évaluation des points de contrôle en aval.

Objectif d'optimisation

La méthode minimise un objectif lagrangien combinant la perte de désapprentissage MASC avec un régularisateur de rétention (divergence KL de l'original sur les données de rétention) :
$$\min_\theta \mathbb{E}_{(x,y)\sim D_{ret}} [\text{KL}(\pi_{\theta_0} \parallel \pi_\theta)] + \lambda \mathbb{E}_{(x,y)\sim D_{fg}} [\text{HingeLoss}(m^{(k,\beta)}_\theta)]$$
sous la contrainte que le taux de violation sur l'ensemble à oublier soit $\le \alpha$.

Principales contributions

1. Algorithme MASC : Un nouvel objectif de désapprentissage qui supprime les jetons cibles uniquement lorsqu'ils restent nettement plus probables qu'un agrégat de $k$ alternatives. Cela fournit une borne supérieure théoriquement fondée sur la probabilité de reproduction exacte (Proposition 1).
2. Arrêt adaptatif : MASC introduit une règle d'arrêt intrinsèque basée sur la même condition de marge optimisée pendant l'entraînement. Cela permet au modèle de stopper l'entraînement dès que la condition d'oubli est remplie, évitant ainsi des calculs inutiles et des évaluations en aval.
3. Analyse de mise à l'échelle (Scaling) : Le papier présente une étude systématique sur la famille Qwen2.5 (de 0,5B à 7B de paramètres) utilisant le jeu de données TOFU. Il distingue la mémorisation exacte (reproduction verbatim) de la mémorisation de connaissances (récupération sous forme de paraphrase).
   • Phase d'apprentissage : Les modèles plus grands amplifient la mémorisation exacte plus fortement que la récupération basée sur la paraphrase.
   • Phase de désapprentissage : Après le désapprentissage, les métriques de mémorisation résiduelle deviennent largement invariantes à l'échelle, tandis que l'utilité de rétention s'améliore avec la taille du modèle. Cela suggère que l'échelle bénéficie principalement au côté utilité du compromis plutôt que d'exacerber la mémorisation résiduelle.

Résultats expérimentaux

Les auteurs ont évalué MASC sur TOFU, MUSE News et MUSE Books, en le comparant à des bases telles que Gradient Ascent (GA), GradDiff, NPO, NPO+KLR, RMU et SimNPO.

• Efficacité : MASC atteint des compromis oubli-rétention compétitifs en une fraction du temps de calcul des bases existantes. Par exemple, sur TOFU, MASC s'est terminé en ~88 secondes contre plus de ~900 secondes pour GradDiff et NPO+KLR.
• Performance :
  • TOFU : MASC a obtenu les meilleurs scores 1-ROUGE-L, Truth Ratio et scores de rétention MU, tout en maintenant des scores 1-Prob compétitifs.
  • MUSE Books : MASC a préservé l'utilité de rétention proche des bases les plus fortes (non effondrées) avec un coût de calcul nettement inférieur.
  • MUSE News : MASC a obtenu une mémorisation de connaissance plus faible sur l'ensemble à oublier que la plupart des bases, bien que l'utilité de rétention ait été occasionnellement légèrement inférieure.
• Robustesse : Les expériences ont montré que MASC reste efficace après une quantification en 4-bits, maintenant une faible mémorisation de l'ensemble à oublier là où d'autres méthodes (comme NPO et SimNPO) souffrent d'une dégradation significative des performances.
• Mise à l'échelle (Scaling) : L'étude a confirmé que si les modèles plus grands mémorisent davantage pendant l'entraînement, le processus de désapprentissage normalise efficacement les métriques du côté de l'oubli à travers les échelles, les modèles plus grands conservant une utilité plus élevée sur l'ensemble de rétention.

Signification et affirmations

Le papier affirme que MASC déplace le paradigexe du désapprentissage d'un mode d'entraînement à longueur fixe vers un arrêt précoce ciblé. En alignant la fonction de perte, la contrainte et le critère d'arrêt, MASC offre une procédure auto-terminante qui est substantiellement plus rapide que ses homologues à budget fixe sans sacrifier le compromis oubli-rétention.

Les auteurs notent avec modestie que, bien que MASC améliore empiriquement les métriques de l'oubli au niveau de la paraphrase, la borne théorique contrôle actuellement la reproduction exacte. Ils suggèrent que la conception d'objectifs pour le désapprentissage sémantique (invariant aux reformulations de surface) reste une direction pour les travaux futurs. Le papier souligne que MASC fournit une solution pratique et scalable pour le désapprentissage dans les grands modèles, répondant à l'inefficacité computationnelle et à l'absence de mécanismes de sélection en ligne dans les approches actuelles.