Résumé Technique : SoftSAE – Sélection Dynamique Top-K pour les Autoencodeurs Épars Adaptatifs

Énoncé du Problème

Les Autoencodeurs Épars (SAE) sont devenus un paradigme dominant en interprétabilité mécaniste, visant à décomposer les activations neuronales polysémantiques en caractéristiques monosémantiques compréhensibles par l'humain. Cependant, les architectures existantes de l'état de l'art, telles que les SAE TopK, les SAE BatchTopK et les SAE Matryoshka, reposent sur un niveau de parcimonie global fixe ($K$). Cette contrainte rigide impose le même nombre de caractéristiques actives pour toutes les entrées, indépendamment de la complexité inhérente de l'entrée.

Les auteurs soutiennent que les données naturelles résident sur des variétés présentant des Dimensions Intrinsèques Locales (LID) variables. Les entrées simples (par exemple, un fond uniforme ou une conjonction linguistique simple) nécessitent peu de concepts actifs, tandis que les entrées complexes (par exemple, des scènes encombrées ou une terminologie de domaine dense) exigent un ensemble plus riche de caractéristiques explicatives. Une valeur $K$ fixe force un compromis sous-optimal : elle soit hallucine des caractéristiques bruyantes pour les entrées simples, soit omet des nuances sémantiques critiques pour les entrées complexes. Les méthodes actuelles comme BatchTopK permettent une variation incidente via des seuils globaux, mais cela ne constitue pas une adaptation raisonnée à la complexité spécifique de l'entrée.

Méthodologie : SoftSAE

Pour remédier à la rigidité des contraintes de capacité fixe, les auteurs proposent SoftSAE, une architecture qui calibre dynamiquement son niveau de parcimonie ($\hat{k}$) en fonction de l'entrée spécifique traitée.

1. MLP de Parcimonie Dynamique

L'innovation centrale est un module neuronal dédié, le MLP de Parcimonie Dynamique, qui estime un niveau de parcimonie dépendant de l'entrée $\hat{k} \in (0, k_{max})$.

• Entrée : Les activations pré-encodeur $x \in \mathbb{R}^n$.
• Mécanisme : Le MLP prédit un scalaire continu $\hat{k}$ en utilisant la formule :
  $$\hat{k} := \sigma(\text{MLP}(x)) \cdot k_{max}$$
  où $\sigma$ est la fonction sigmoïde et $k_{max}$ est la borne supérieure de la taille du dictionnaire.
• Biais Inductif : La première couche du MLP est initialisée avec les poids de l'encodeur SAE ($W_{enc}$) pour aligner le prédicteur de capacité avec l'espace des caractéristiques, bien que les paramètres soient libres de diverger durant l'entraînement.

2. Sélection Top-K Douce Différentiable

Puisque les opérations Top-K standard sont non différentiables, SoftSAE emploie un opérateur Top-K Doux Différentiable durant l'entraînement pour permettre la propagation du gradient à travers le mécanisme de sélection.

• Phase d'Entraînement : L'opérateur génère des poids de sélection continus $p \in (0, 1)^d$ tels que $\sum p_i = \hat{k}$. Ces poids approximent un masque binaire appliqué à la représentation latente $z$.
• Phase d'Inférence : Pour assurer une parcimonie exacte et une interprétabilité, le modèle revient à une sélection Top-K dure, où $\hat{k}$ est discrétisé (arrondi).
• Stabilisation : Pour prévenir l'instabilité numérique et le « masquage » d'informations dans des poids proches de zéro lorsque le paramètre de température $\alpha \to 0$, le MLP de Parcimonie Dynamique est gelé lors de la phase finale d'entraînement, et la sélection douce est remplacée par un Top-K dur.

3. Objectif d'Entraînement et Budget de Parcimonie

La perte d'entraînement $L(x)$ équilibre la fidélité de reconstruction avec un budget de parcimonie dynamique :
$$L(x) := \|x - \hat{x}\|_2^2 + \lambda S(\hat{k}) + \gamma L_{aux}$$

• Perte de Reconstruction : Distance euclidienne quadratique standard.
• Pénalité de Budget de Parcimonie ($S(\hat{k})$) : Au lieu d'une pénalité ReLU dure qui provoque une instabilité du gradient près de la frontière de contrainte, les auteurs utilisent une approximation Softplus :
  $$S(\hat{k}) = \frac{1}{\beta} \ln(1 + e^{(E[\hat{k}] - k) \cdot \beta})$$
  Cela lisse la transition, encourageant le modèle à rester dans le budget de parcimonie attendu $k$ sans changements brusques de gradient. Cela agit également comme un régularisateur pour empêcher le modèle de dépasser la parcimonie cible lors de l'inférence.
• Perte Auxiliaire ($L_{aux}$) : Un composant standard pour prévenir les « neurones morts » en encourageant les caractéristiques sous-utilisées à capturer les erreurs résiduelles.

Contributions Clés

1. Architecture SoftSAE : Une approche novatrice utilisant une formulation Top-K Doux continue pour prédire et optimiser dynamiquement le niveau exact de parcimonie ($\hat{k}$) pour chaque entrée, plutôt que de dépendre d'un hyperparamètre statique.
2. Fondement Théorique : Ce travail relie la capacité d'explication directement à la dimension intrinsèque locale de l'espace d'activation neuronale, fournissant une base théorique pour la parcimonie adaptative en interprétabilité mécaniste.
3. Qualité Améliorée des Caractéristiques : En ajustant dynamiquement $\hat{k}$, le modèle atteint un équilibre supérieur entre fidélité de reconstruction et interprétabilité, découvrant des caractéristiques monosémantiques qui reflètent plus fidèlement la complexité multi-échelle des données naturelles que les modèles à capacité fixe.

Résultats Expérimentaux

Les auteurs ont évalué SoftSAE sur deux modalités :

• Vision : Embeddings CLIP (ViT-B/16) entraînés sur CC3M et évalués sur ImageNet-1k/100.
• Langage : Activations Gemma-2-2B (Couche 12) entraînées sur FineWeb.

Comparaisons : SoftSAE a été mis en comparaison avec TopK SAE, BatchTopK SAE et Matryoshka SAE.

• Reconstruction vs Parcimonie : SoftSAE correspond aux performances reconstruction-parcimonie de TopK et BatchTopK sur CLIP et se comporte de manière comparable à Matryoshka sur Gemma-2-2B. Il ne sacrifie qu'une quantité négligeable de performance de reconstruction pour gagner une structure de caractéristiques améliorée.
• Qualité des Caractéristiques (SAE Bench) :
  • Absorption/Séparation : SoftSAE démontre une résistance aux anomalies d'absorption et de séparation de caractéristiques, montrant un avantage clair dans la réduction de l'absorption de caractéristiques à des niveaux de parcimonie plus élevés par rapport aux bases de référence.
  • Désentanglement : Dans les tâches de Suppression de Corrélations Spurious (SCR) et de Perturbation de Sonde Ciblée (TPP), SoftSAE surpasse TopK et BatchTopK et reste compétitif ou dépasse Matryoshka à des niveaux de parcimonie plus élevés, indiquant un fort désentanglement et un contrôle précis des caractéristiques apprises.
• Analyse Qualitative : Les visualisations confirment que le MLP de Parcimonie Dynamique attribue correctement de faibles scores $\hat{k}$ aux entrées simples (par exemple, une entité unique, un fond uniforme) et de hauts scores aux entrées complexes (par exemple, plusieurs entités, contexte riche), alignant la longueur de l'explication avec la densité d'information de l'entrée.

Signification et Revendications

L'article revendique que SoftSAE représente un changement de la parcimonie en tant qu'hyperparamètre fixe vers un composant flexible et apprenable. Sa signification principale réside non seulement dans des améliorations marginales de la reconstruction ou de la qualité des caractéristiques, mais dans la capacité à allouer dynamiquement la capacité de représentation en fonction de la complexité de l'entrée.

Ce mécanisme permet au modèle de produire des explications plus courtes et plus claires pour les entrées simples et des représentations plus riches pour les entrées complexes, alignant efficacement la « longueur de l'explication » avec la dimension intrinsèque locale de la variété des données. Bien que les auteurs reconnaissent des limitations — spécifiquement un coût computationnel accru dû à l'opérateur Top-K Doux différentiable et la nécessité de gérer la transition douce-dure — ils postulent que cette approche ouvre une nouvelle direction pour l'allocation de capacité adaptative dans les systèmes neuronaux, améliorant potentiellement à la fois l'efficacité et l'interprétabilité dans les pipelines génératifs plus larges.