Résumé Technique : Autoencodeur Creux Rationnel (RSAE)

Énoncé du Problème

Les autoencodeurs creux (SAEs) sont l'outil standard pour l'interprétabilité mécaniste des grands modèles de langage (LLMs), décomposant les activations internes en combinaisons linéaires creuses de dictionnaires de caractéristiques surcomplets. Cependant, les familles de SAE existantes reposent sur des non-linéarités d'encodage fixes et non lisses — spécifiquement ReLU (avec régularisation $\ell_1$), TopK et JumpReLU. Ces primitives d'activation fixes introduisent des pathologies spécifiques :

• SAEs ReLU : Souffrent d'une réduction de l'amplitude des caractéristiques actives et d'une population persistante de latents "morts".
• SAEs TopK : Remplacent les pénalités douces par des contraintes de cardinalité strictes, brisant le flux de gradient à travers les caractéristiques inactives et nécessitant des pertes de réveil auxiliaires.
• SAEs JumpReLU : Introduisent des seuils apprenables mais nécessitent des substituts de relaxation continue pour la rétropropagation à travers leurs portes indicatrices.

Ces mécanismes codés en dur contraignent la capacité du modèle à s'adapter à la géométrie des pré-activations observées, ce qui peut fausser l'équilibre entre la fidélité de reconstruction et la parcimonie (sparsity).

Méthodologie

Les auteurs proposent l'Autoencodeur Creux Rationnel (RSAE), qui remplace l'activation de l'encodeur fixe par une fonction rationnelle entraînable. Le RSAE conserve le squelette standard d'un SAE (une couche de pré-activation affine suivie d'un bloc d'activation creuse et d'un décodeur linéaire) mais modifie la fonction d'activation $\phi(\cdot)$ pour en faire une fonction rationnelle apprenable appliquée élément par élément.

Architecture

L'activation du RSAE est définie comme :
$$\phi(h) = C_{out} \cdot r_{(a,b)}\left(\frac{h}{C_{in}}\right)$$
où $h$ est le vecteur de pré-activation, $C_{in}$ et $C_{out}$ sont des paramètres de mise à l'échelle apprenables, et $r_{(a,b)}(t) = P(t)/Q(t)$ est une fonction rationnelle avec des coefficients entraînables $(a, b)$. Les paramètres de mise à l'échelle projettent la distribution de pré-activation dans l'intervalle de conception de la fonction rationnelle $[-1, 1]$ et projettent la sortie vers la magnitude de caractéristique attendue.

Fondement Théorique

L'article établit une asymétrie d'expressivité entre les fonctions rationnelles et les portes affines par morceaux utilisées dans les SAE actuels :

1. Borne Supérieure (Approximation Rationnelle) : En utilisant les fonctions de signe de Zolotarev, les auteurs prouvent que les activations rationnelles entraînables peuvent approximer les portes ReLU, JumpReLU et les portes TopK à seuil fourni avec une taille polylogarithmique par rapport à l'exactitude cible $\epsilon$ (spécifiquement $O(\log(1/\epsilon)\log(1/\delta)$ pour les portes discontinues avec une marge $\delta$).
2. Borne Inférieure (Limite Affine par Morceaux) : Inversement, ils démontrent qu'approximer certaines cartes cibles rationnelles à $O(1)$-paramètre en utilisant des encodeurs à sortie scalaire et à couches simples (ReLU/JumpReLU/TopK) nécessite $\Omega(\epsilon^{-1/2})$ coordonnées activées. Cela suggère que les activations rationnelles offrent une classe de représentation plus efficace pour l'encodeur du SAE.

Pipeline d'Entraînement

Le RSAE est construit via un pipeline en deux étapes conçu pour améliorer les SAE "enseignants" pré-entraînés :

1. Initialisation :
   • Échange de Remez : Une fonction rationnelle est ajustée à la primitive d'activation de l'enseignant (ReLU, JumpReLU ou TopK) sur une grille dense synthétique en utilisant l'algorithme d'échange de Remez relaxé afin de minimiser l'erreur $L_\infty$.
   • Calibration : Les coefficients rationnels et les paramètres de mise à l'échelle ($C_{in}, C_{out}$) sont adaptés à la distribution réelle de pré-activation de l'enseignant en minimant l'erreur quadratique entre la sortie rationnelle et la sortie de l'enseignant.
   • Héritage de Poids : Les poids de l'encodeur et du décodeur sont copiés textuellement de l'enseignant pré-entraîné.
2. Affinage (Fine-Tuning) : Tous les paramètres (poids, biais et coefficients rationnels) sont dégelés et optimisés conjointement sous l'objectif de reconstruction standard régularisé en $\ell_1$.

Contributions Clés

1. Architecture RSAE : Introduction d'un autoencodeur creux avec une activation rationnelle entraînable, fondée sur la théorie de l'approximation adaptée à l'encodeur peu profond du SAE.
2. Écart d'Expressivité Théorique : Preuve que les activations rationnelles fournissent des approximations compactes des portes de SAE fixes, tandis que les encodeurs affines par morceaux nécessitent beaucoup plus de coordonnées activées pour approximer des cibles rationnelles simples.
3. Algorithme d'Entraînement en Deux Étapes : Une procédure pratique d'initialisation et d'affinage qui permet aux RSAE d'être "injectés" pour remplacer les SAE pré-entraînés avec un surcoût négligeable.
4. Validation Empirique : Évaluation complète sur trois modèles de langage en poids ouverts (GPT-2 small, Pythia-160m, Gemma-2-2B) et sur les trois familles d'activation de référence.

Résultats

Les auteurs évaluent le RSAE par rapport aux lignes de base pré-entraînées sur les activations du flux résiduel. Les résultats indiquent des améliorations strictes après l'affinage :

• Métriques de Reconstruction : Le RSAE atteint systématiquement une erreur quadratique moyenne (MSE) plus faible, un $\ell_0$ (parcimonie) plus faible et une fraction plus élevée de caractéristiques "vivantes" par rapport aux modèles enseignants dans toutes les configurations testées.
• Métriques Aval (Downstream) : Le RSAE améliore le comportement en aval, montrant une réduction de la dégradation de l'entropie croisée ($\Delta CE$) et une récupération de perte ($LR$) plus élevée lors de l'interception du flux résiduel.
• Interprétabilité : Les expériences de sondage (probing) de parcimonie sur Pythia-160m montrent que le RSAE préserve l'interprétabilité au niveau des caractéristiques. La précision du sondage du dictionnaire complet sur huit tâches de classification binaire reste comparable à celle de l'enseignant, sans dégradation significative.
• Efficacité : La mise à niveau n'ajoute que quelques paramètres scalaires par autoencodeur. L'ensemble du pipeline (initialisation + 2 000 étapes d'affinage) s'exécute en quelques minutes sur un seul GPU grand public (NVIDIA RTX 5090).
• Robustesse : Les gains sont constants à travers différents modèles hôtes, familles d'activation de base et une large gamme de niveaux de parcimonie de base ($\ell_0 \in [51, 683]$).

Signification et Revendications

L'article affirme que le RSAE offre une modification "prête à l'emploi" (drop-in) pour les SAE pré-entraînés qui améliore strictement la fidélité du modèle à parcimonie égale sans sacrifier l'interprétabilité. En remplaçant les non-linéarités fixes par des fonctions rationnelles entraînables, le RSAE s'adapte à la géométrie spécifique des pré-activations du modèle hôte, surmontant ainsi les limites des primitives d'activation codées en dur.

Les auteurs soulignent que ces améliorations sont obtenues avec un coût de calcul négligeable (quelques minutes sur un seul GPU) et un surcoût de paramètres minimal. Ils positionnent le RSAE comme une stratégie de mise à jour universelle pour les SAE existants, capable d'améliorer à la fois la qualité de reconstruction et la performance des tâches en aval tout en maintenant les caractéristiques monosémantiques requises pour l'interprétabilité mécaniste. L'article note que bien que l'évaluation couvre trois modèles et trois familles d'activation, l'approche est générale et pourrait potentiellement être étendue à d'autres modèles de pointe ou variantes architecturales (par exemple, Gated, BatchTopK, Matryoshka SAEs).