From Weights to Features: SAE-Guided Activation Regularization for LLM Continual Learning

Résumé technique : Des poids aux caractéristiques : Régularisation de l'activation guidée par SAE pour l'apprentissage continu des LLM

Énoncé du problème
L'apprentissage continu (CL) dans les grands modèles de langage (LLM) est confronté au défi de l'oubli catastrophique, où le réglage fin séquentiel écrase les représentations nécessaires aux tâches antérieures. Bien que les méthodes de régularisation basées sur les poids, comme l'Elastic Weight Consolisation (EWC), soient standards pour les modèles plus petits, elles sont nettement moins performantes pour les LLM. Les auteurs soutiennent que cet échec est structurel, provenant de la nature « polysémantique » des LLM. En raison de l'hypothèse de la superposition, les neurones individuels et leurs poids associés encodent simultanément plusieurs concepts sans rapport. Par conséquent, les estimations d'importance dans l'espace des poids (par exemple, l'information de Fisher) sont trop grossières ; protéger un poids pour préserver le Concept A contraint par inadvertance le Concept B, créant une pénalité non sélective qui étouffe la plasticité. De plus, les méthodes existantes échouent souvent à imposer explicitement la plasticité requise pour apprendre de nouvelles tâches, menant à une « intransigeance » où le modèle conserve les connaissances anciennes mais ne peut plus s'adapter.

Méthodologie
Le papier propose la régularisation de l'activation guidée par SAE, une méthode qui déplace la cible de la régularisation de l'espace des paramètres entrelacés vers un espace d'activation plus sélectif en utilisant des autoencodeurs creux (SAE) pré-entraînés.

1. Changement de système de coordonnées : Au lieu de pénaliser la dérive des poids, la méthode utilise un SAE pré-entraîné et gelé comme un dictionnaire de caractéristiques monosémantiques fixes. Les SAE décomposent les activations denses du modèle en directions de caractéristiques creuses qui sont souvent plus interprétables et moins entrelacées que les neurones individuels.
2. Formulation de l'optimisation contrainte : L'objectif d'entraînement est dérivé d'un problème d'optimisation contrainte plutôt que d'une conception de pénalité ad hoc. Le modèle doit minimiser la perte de la tâche actuelle sous deux contraintes :
   • Stabilité : Les caractéristiques non pertinentes pour la tâche (faible pertinence) ne doivent pas dériver au-delà d'un budget spécifique ($\varepsilon_s$).
   • Plasticité : Les caractéristiques pertinentes pour la tâche (haute pertinence) doivent dériver suffisamment ($\ge \varepsilon_p$) pour éviter des solutions dégénérées où le modèle n'apprend rien.
3. Fonction de perte : En utilisant la relaxation lagrangienne et une pénalité de type "squared-hinge", les auteurs dérivent une perte d'entraînement comprenant une perte de protection (pénalisant la dérive excessive des caractéristiques protégées) et une perte de guidage (pénalisant l'insuffisance de dérive des caractéristiques pertinentes). L'EWC est montré comme étant un cas particulier de ce cadre où l'espace de représentation est celui des poids et où la contrainte de plasticité est absente.
4. Construction du masque : Pour chaque nouvelle tâche, un masque de pertinence continu ($M_t$) est construit en passant les données de la tâche actuelle à travers le modèle de base gelé et le SAE. Ce masque identifie les caractéristiques pertinentes via la magnitude d'activation et propage la pertinence aux caractéristiques voisines en utilisant le k-NN sur les vecteurs du décodeur. Crucialement, une fois le masque construit, toutes les données de la tâche sont supprimées ; seul le masque compact est conservé.
5. Entraînement : Pendant l'entraînement, le modèle actuel (avec LoRA) et le modèle de base gelé sont tous deux encodés par le SAE. La dérive des caractéristiques est pondérée par le masque pour calculer les pertes de protection et de guidage.

Contributions clés

• Régularisation guidée par SAE : Une méthode de CL pratique qui régularise dans l'espace d'activation en utilisant des caractéristiques de SAE. Elle ne nécessite aucune donnée de tâche précédente, aucune activation stockée, aucun paramètre spécifique à la tâche, et aucun routage au moment de l'inférence après la construction du masque.
• Dérivation fondée sur des principes : Le compromis stabilité-plasticité est formulé comme un problème d'optimisation contrainte dans l'espace des caractéristiques SAE, produisant une perte d'entraînement "squared-hinge" via la relaxation lagrangienne.
• Preuve empirique de la polysémanticité : Le papier fournit une preuve mécaniste que la protection de l'espace des poids est non sélective. Les expériences montrent que les représentations pertinentes pour la tâche sont linéairement séparables dans l'espace des caractéristiques SAE (AUC ~0,88) mais impossibles à distinguer du hasard dans l'espace des poids/neurones (AUC ~0,50). De plus, protéger les poids pour une tâche contraint les caractéristiques pour une tâche suivante à un taux de 91 à 96 % de la tâche originale, tandis que la protection dans l'espace des caractéristiques réduit cette contrainte collatérale à 43–61 %.
• Scalabilité : En raison de la dimensionnalité plus faible de l'espace des caractéristiques par rapport à l'espace des paramètres, le stockage par tâche est réduit de gigaoctets (ancres de poids) à des sous-mégaoctets (masques de caractéristiques).

Résultats
La méthode a été évaluée sur TRACE-5000 (8 tâches trans-domaines) et MedCL (10 tâches biomédicales) en utilisant Gemma-2 9B-it avec LoRA.

• Performance : Sur TRACE, la régularisation guidée par SAE a obtenu la meilleure performance parmi les approches non architecturales (OP = 0,545), surpassant les régulariseurs d'espace de poids (EWC : 0,447, SI : 0,513), les méthodes de projection de gradient et les bases de replay (même avec 10 % de données). Sur MedCL, elle arrive en tête de tous les modèles qui ne conservent aucun exemple de tâche précédente (OP = 0,510).
• Compromis Stabilité-Plasticité : Contraên l'EWC, qui a atteint une haute rétention (ARA = 0,987) en sacrifiant sévèrement la plasticité (0,453), la méthode proposée maintient une haute plasticité (0,567, comparable au fine-tuning non protégé) tout en conservant 95 % de la performance apprise.
• Efficacité : La méthode ajoute un surcoût de calcul minimal (~1,88s/étape contre 1,70s pour la ligne de base) et réduit considérablement les besoins de stockage (412 Ko par tâche contre 6,5 Go pour les ancres EWC).

Signification et affirmations
Le papier affirme que la sous-performance de la régularisation de l'espace des poids dans les LLM n'est pas un problème de réglage d'hyperparamètres mais une conséquence structurelle de la polysémanticité. En déplaçant la contrainte de régularisation vers une base de caractéristiques monosémantiques, la méthode permet une protection plus sélective des connaissances. Les auteurs soutiennent que l'efficacité de la régularisation dépend de manière critique de l'espace de représentation dans lequel les contraintes sont imposées. La méthode démontre qu'il est possible d'obtenir de solides performances d'apprentissage continu sans stocker de données passées ou modifier l'architecture du modèle lors de l'inférence, à condition que la régularisation opère sur un système de coordonnées aligné avec la structure conceptuelle interne du modèle. Le papier conclut que la régularisation basée sur les SAE offre une alternative plus sélective et plus évolutive que les méthodes traditionnelles de l'espace des poids pour les LLM.