Résumé Technique : Renforcement du Contrôle en Apprentissage par Renforcement (CRL)

1. Énoncé du Problème

Les Autoencodeurs Sparsifiés (SAE) sont devenus un outil puissant pour décomposer les activations des modèles de langage en caractéristiques interprétables et lisibles par l'humain. Cependant, les méthodes d'analyse existantes sont largement statiques : elles identifient quelles caractéristiques s'activent sur des entrées spécifiques, mais échouent à révéler quelles caractéristiques, lorsqu'elles sont amplifiées, modifient réellement les sorties du modèle. Les techniques actuelles d'ingénierie des activations reposent souvent sur des ensembles de données contrastifs pour la sélection de caractéristiques ou sur un stockage étendu des activations pour l'optimisation des coefficients, manquant de mécanismes de rétroaction adaptatifs pour ajuster les stratégies de pilotage en fonction de la qualité de la génération.

L'article comble le fossé entre l'identification statique des caractéristiques et l'intervention dynamique et interprétable. Le défi central consiste à développer un cadre qui non seulement sélectionne des caractéristiques SAE pour piloter un modèle, mais fournit également des journaux d'intervention au niveau des tokens expliquant pourquoi une caractéristique a été sélectionnée et comment elle a affecté la sortie, transformant ainsi le pilotage en un outil d'interprétabilité mécaniste.

2. Méthodologie : Renforcement du Contrôle en Apprentissage par Renforcement (CRL)

Les auteurs proposent le Renforcement du Contrôle en Apprentissage par Renforcement (CRL), un cadre qui traite le contrôle des représentations de transformateurs comme un Processus de Décision Markovien (MDP) sur les caractéristiques SAE.

2.1. Formulation MDP

• État ($S$) : L'observation est l'activation du flux résiduel ($x \in \mathbb{R}^d$) à une couche spécifique $\ell$ pour la position du token actuel.
• Action ($A$) : Un vecteur binaire sélectionnant une ou plusieurs caractéristiques SAE à amplifier. La politique sélectionne les $k$ meilleures caractéristiques (généralement $k=1$ pour l'interprétabilité).
• Transition ($P$) : Déterministe, régie par le passage avant du transformateur avec l'intervention de caractéristique sélectionnée appliquée.
• Récompense ($R$) : Récompenses spécifiques à la tâche basées sur la qualité de la sortie finale (par exemple, exactitude dans MMLU, refus dans les tâches de sécurité).

2.2. CRL-Token et Pilotage de Caractéristiques

À chaque étape de génération $t$, le cadre effectue des interventions au niveau des tokens :
$$\tilde{x}_t = x_t + c \cdot a_t W_{dec}$$
Où $a_t$ est le vecteur de sélection binaire sparse, $c$ est le coefficient de pilotage, et $W_{dec}$ est le décodeur SAE. Le réseau de politique $\pi_\theta$ mappe les activations du flux résiduel vers les logits de sélection de caractéristiques, implémenté comme un MLP à 2 couches.

2.3. Masquage Adaptatif des Caractéristiques (AFM)

Pour empêcher la politique de se contracter vers un ensemble restreint de caractéristiques fréquemment actives, le CRL introduit le Masquage Adaptatif des Caractéristiques.

• Un masque booléen par échantillon restreint les choix de sélection de caractéristiques.
• Le masque commence avec un petit sous-ensemble de caractéristiques et s'étend progressivement au fur et à mesure que la génération avance, intégrant les caractéristiques qui ont été naturellement activées lors des étapes précédentes.
• Ce mécanisme encourage la découverte diversifiée de caractéristiques tout en maintenant l'interprétabilité des interventions sur une seule caractéristique.

2.4. Optimisation

La politique et un réseau critique (estimant les valeurs d'état) sont entraînés en utilisant l'Optimisation de Politique Proximale (PPO).

• Coefficient de Pilotage ($c$) : Déterminé de manière adaptative en moyennant les magnitudes d'activation des échantillons d'entraînement correctement prédits, éliminant ainsi le besoin de réglage manuel à travers les couches.
• Réseau Critique : Un MLP à 2 couches qui estime la fonction de valeur $V_\phi(s)$, permettant de séparer les limitations de la politique des erreurs d'estimation de valeur lors de l'analyse.

3. Contributions Clés

L'article expose quatre contributions principales :

1. Journaux d'Intervention Interprétables au Niveau des Tokens : Un cadre produisant des journaux identifiant exactement quelle caractéristique SAE a été sélectionnée à chaque token et son effet sur la sortie, permettant une analyse comportementale fine.
2. Outils Diagnostiques :
   • Suivi des Points de Branchement : Identifie les points de décision critiques où des préfixes identiques produisent des résultats différents selon la sélection de caractéristiques.
   • Analyse de la Trajectoire du Critique : Utilise la fonction de valeur apprise pour distinguer les échecs de la politique (sélection de caractéristiques) des échecs du critique (estimation de récompense).
3. Masquage Adaptatif des Caractéristiques (AFM) : Un mécanisme qui équilibre exploration et exploitation pour assurer une découverte diversifiée de caractéristiques sans sacrifier l'interprétabilité des interventions sur une seule caractéristique.
4. Caractérisation Empirique des Sémantiques des Caractéristiques : Démontre que les premières couches encodent la structure syntaxique tandis que les couches ultérieures capturent des sémantiques abstraites, et que différentes tâches présentent des motifs de goulot d'étranglement distincts.

4. Résultats Expérimentaux

Le cadre a été évalué sur Gemma 2 2B (et LLaMA-3.1 8B) en utilisant les SAE Gemma Scope et LLaMA Scope sur cinq benchmarks : MMLU, BBQ, GSM8K, HarmBench et XSTest.

• Gains de Performance : Le CRL a obtenu des améliorations sur tous les benchmarks. Les gains notables incluent :
  • HarmBench : +7,66 points de pourcentage (41,46 % $\to$ 49,12 %).
  • BBQ (Ambigu) : +5,69 points de pourcentage.
  • MMLU : +3,31 points de pourcentage.
  • GSM8K : +1,03 points de pourcentage.
• Complémentarité : Le CRL apporte des gains additifs lorsqu'il est appliqué au-dessus de modèles affinés par apprentissage supervisé (SFT) (par exemple, la précision MMLU est passée de 55,2 % à 56,1 % avec SFT + CRL).
• Sémantiques Spécifiques aux Couches :
  • Couche 10 : Les caractéristiques capturent des propriétés concrètes et syntaxiques (par exemple, "notation mathématique", "références pronominales").
  • Couche 20 : Les caractéristiques capturent des sémantiques abstraites et contextuelles (par exemple, "structure de dérivation logique", "preuves mathématiques").
  • Les couches d'intervention optimales varient selon la tâche (par exemple, Couche 5 pour BBQ, Couche 24 pour MMLU).
• Analyse du Critique :
  • BBQ : Montre une coordination efficace politique-critique avec une discrimination claire des échantillons.
  • MMLU : Présente un "goulot d'étranglement du critique" où les échantillons corrigés et erronés sont presque indiscernables, suggérant que le critique peine à estimer la valeur pour les tâches de raisonnement complexe.
  • HarmBench vs XSTest : HarmBench montre des écarts clairs entre les échantillons corrects et incorrects, tandis que XSTest montre des motifs d'estimation inversés, indiquant des défis dépendants de la tâche dans l'estimation de la récompense.

5. Signification et Revendications

L'article positionne le CRL comme une avancée significative en interprétabilité mécaniste. Sa signification principale réside dans le changement de paradigme de l'analyse statique des caractéristiques (observer ce qui s'active) vers les sondes d'intervention dynamique (observer ce qui modifie le modèle).

• Interprétabilité : En imposant la sélection d'une seule caractéristique et en journalisant les interventions, le CRL fournit une vue "mécaniste" du comportement du modèle, révélant que des caractéristiques comme "opérations mathématiques" peuvent se généraliser lexicalement à des unités quantitatives comme "bébés" dans des problèmes de mots.
• Capacité Diagnostique : Le cadre sert d'outil diagnostique à la fois pour le modèle et pour le SAE. Si une politique se contracte vers des caractéristiques sémantiquement incohérentes malgré l'AFM, cela signale une déficience du SAE (par exemple, concepts manquants ou polysémie) plutôt qu'un échec de la méthode de contrôle.
• Praticité : Le CRL fonctionne avec une faible surcharge d'inférence (8–9 % de latence) et ne nécessite pas de réentraîner les poids du modèle de base, offrant un mécanisme d'améliération post-hoc léger.

Les auteurs concluent que le pilotage de caractéristiques appris, couplé à l'apprentissage par renforcement, établit un outil pratique pour comprendre et contrôler les modèles de langage, complétant l'analyse statique par des insights dynamiques et spécifiques à la tâche.