Résumé technique : Entraînement aligné pour les autoencodeurs parcimonieux

Énoncé du problème

Les autoencodeurs parcimonieux (SAE) constituent un outil principal en interprétabilité mécaniste, conçus pour décomposer les activations des réseaux de neurones profonds (DNN) en caractéristiques parcimonieuses et interprétables. Malgré leur utilité, l'entraînement standard des SAE souffre de trois lacunes critiques et persistantes :

1. Caractéristiques mortes : Une fraction significative des caractéristiques apprises ne s'active jamais sur aucune entrée, représentant un échec à exploiter la capacité du dictionnaire.
2. Instabilité : Les caractéristiques apprises lors d'entraînements indépendants (différentes graines aléatoires) sont souvent incohérentes, convergeant vers des dictionnaires qualitativement différents malgré des données et des hyperparamètres identiques.
3. Reconstruction sous-optimale : L'entraînement standard produit souvent une qualité de reconstruction inférieure au potentiel théorique de l'architecture.

L'article identifie la cause racine de ces problèmes comme une dégénérescence négligée dans le paysage d'optimisation. Plus précisément, le produit scalaire entre la ligne d'encodeur d'une caractéristique ($W^{enc}_{i,\cdot}$) et sa colonne de décodeur correspondante ($W^{dec}_{\cdot,i}$), appelé score d'alignement ($a_i$), suit une distribution bimodale. Alors que les caractéristiques bien alignées ($a_i \approx 1$) sont corrélées à une haute interprétabilité, une fraction substantielle de caractéristiques sont des artefacts presque orthogonaux avec des scores d'alignement proches de zéro. L'entraînement standard laisse ce degré de liberté non contraint, permettant à la descente de gradient d'exploiter cette dégénérescence de manière préjudiciable à l'interprétabilité.

Méthodologie : Entraînement aligné

Les auteurs proposent l'Entraînement aligné, une reparamétrisation sans paramètre des SAE qui impose une contrainte géométrique pour éliminer cette dégénérescence.

Concept central

La méthode est motivée par l'observation qu'une reconstruction parfaite dans un cadre minimal force le score d'alignement à être exactement égal à un. Les auteurs émettent l'hypothèse que l'imposition de $a_i = 1$ pour toutes les caractéristiques $i$ élimine une source d'instabilité et améliore la qualité des caractéristiques.

Reparamétrisation

Au lieu d'entraîner directement la matrice d'encodeur $W^{enc}$, la méthode introduit une matrice entraînable non contrainte $A \in \mathbb{R}^{m \times n}$. Les lignes de l'encodeur sont définies comme la projection des lignes de $A$ sur un hyperplan affine défini par les colonnes du décodeur.

Pour chaque caractéristique $i$, la ligne d'encodeur est construite comme suit :
$$W^{enc}_{i,\cdot} := A_{i,\cdot} + \alpha_i W^{dec}_{\cdot,i}$$
où le scalaire $\alpha_i$ est calculé pour satisfaire la contrainte $W^{enc}_{i,\cdot} \cdot W^{dec}_{\cdot,i} = 1$ :
$$\alpha_i = \frac{1 - A_{i,\cdot} \cdot W^{dec}_{\cdot,i}}{\|W^{dec}_{\cdot,i}\|^2}$$

Propriétés clés

• Sans paramètre : Cette approche n'introduit aucun nouvel hyperparamètre. Les biais ($b^{enc}, b^{dec}$) et le décodeur ($W^{dec}$) restent des paramètres entraînables standards.
• Contrainte géométrique : La contrainte $W^{enc}_{i,\cdot} \cdot W^{dec}_{\cdot,i} = 1$ couple chaque ligne d'encodeur à sa colonne de décodeur, réduisant efficacement la variété des poids accessibles et biaisant l'optimisation vers un bassin plus cohérent.
• Compression : Puisque la contrainte réduit la dimensionnalité de l'espace des lignes d'encodeur d'une unité, le dernier élément de chaque ligne dans $A$ peut être fixé à zéro, résultant en légèrement moins de paramètres qu'un SAE standard tout en exprimant le même espace de poids accessibles.
• Indépendance de l'architecture : La méthode est applicable à diverses architectures de SAE, y compris les variantes ReLU, Top-K et BatchTop-K.

Résultats clés

Les auteurs ont évalué l'Entraînement aligné sur les modèles Pythia 160M et Gemma 2 2B à travers des tailles de dictionnaire de 4K, 16K et 65K, et divers niveaux de parcimonie, en utilisant le protocole SAEBench.

1. Améliorations de Pareto en reconstruction : L'entraînement aligné réalise des améliorations de Pareto dans la qualité de reconstruction (mesurée par la variance expliquée et la perte d'entropie croisée récupérée) sur tous les modèles testés, les tailles de dictionnaire et les niveaux de parcimonie.
2. Élimination des caractéristiques mortes : En empêchant structurellement les caractéristiques d'entrer dans le régime d'activation nulle, la méthode réduit la fraction de caractéristiques mortes à près de zéro. Cela est obtenu sans pertes auxiliaires, rééchantillonnage ou gradients fantômes.
3. Stabilité améliorée : La méthode améliore considérablement la cohérence des caractéristiques entre différentes graines aléatoires. Mesurée par la Similarité Cosine Minimale Moyenne (MMCS) entre des dictionnaires entraînés indépendamment, l'Entraînement aligné montre un écart constant et croissant en termes de stabilité par rapport à l'entraînement standard, en particulier pour les architectures Top-K.
4. Compatibilité : La méthode s'intègre parfaitement avec des techniques existantes telles que le p-Annealing (qui traite le rétrécissement des caractéristiques). La combinaison de l'Entraînement aligné et du p-Annealing produit des résultats supérieurs tant en réduction des caractéristiques mortes qu'en métriques de reconstruction par rapport à chaque méthode seule.
5. Désentrelacement des caractéristiques : Sur le benchmark de suppression des corrélations spuriaires (SCR), l'Entraînement aligné surpasse généralement les bases standards, les avantages devenant plus prononcés à mesure que les contraintes de parcimonie augmentent.

Importance et affirmations

L'article affirme que l'Entraînement aligné résout les problèmes fondamentaux d'instabilité et de caractéristiques mortes des SAE à leur source en contraignant directement la variété d'optimisation, plutôt que de corriger les symptômes a posteriori.

• Amélioration structurelle : Les auteurs soutiennent que l'amélioration de la stabilité est structurelle et non un artefact d'initialisation, car l'écart de performance s'ouvre tôt dans l'entraînement et s'élargit avec le temps.
• Aucun coût computationnel : La méthode améliore la qualité et la stabilité des caractéristiques sans ajouter de complexité computationnelle, de coût d'entraînement ou d'exigences de données.
• Réponse aux problèmes de reproductibilité : En traitant directement les préoccupations de cohérence des caractéristiques soulevées par la littérature récente (par exemple, Paulo et Belrose [23], Song et al. [30]), la méthode soutient le développement d'une interprétabilité mécaniste fiable.
• Généralisabilité : L'approche est présentée comme une solution générale applicable aux architectures modernes de SAE, offrant une alternative robuste à l'entraînement standard ne nécessitant aucun réglage supplémentaire d'hyperparamètres.

Les auteurs maintiennent une modestie concernant leurs affirmations, notant que leur évaluation empirique est actuellement limitée à des familles de modèles spécifiques (Pythia 160M, Gemma 2 2B) et que la persistance de ces bénéfices à des échelles de modèles plus grandes ou sous des régimes d'entraînement substantiellement différents reste une question ouverte pour les travaux futurs.