La lentille logit est une technique d'interprétabilité en apprentissage automatique utilisée pour inspecter les représentations internes des transformeurs basés sur les grands modèles de langage. Elle fonctionne en prenant l'état caché à une couche donnée, en appliquant la matrice de désintégration finale du modèle (qui mappe les états cachés aux logits sur le vocabulaire), puis en appliquant un softmax pour obtenir une distribution de probabilité. Cette distribution révèle ce que le modèle prédirait s'il s'arrêtait à cette couche, offrant une fenêtre sur la façon dont les prédictions évoluent couche par couche. La méthode a été introduite par nostalgebraist en 2020 et est depuis devenue un outil standard pour étudier les calculs internes de modèles comme GPT-2 et GPT-3.
Contrairement aux classificateurs de sondage qui nécessitent l'entraînement de modèles auxiliaires, la lentille logit utilise la tête de sortie propre du modèle, ce qui en fait une approche sans paramètre et sans entraînement. Elle est particulièrement efficace pour les modèles avec des flux résiduels, car les états cachés à chaque couche sont additifs et peuvent être directement projetés. La technique a été utilisée pour identifier quand les modèles s'engagent sur des prédictions spécifiques, détecter un décodage « tardif » par rapport à « précoce », et analyser des phénomènes tels que le rappel factuel et le raisonnement arithmétique. Cependant, son applicabilité est limitée aux modèles avec des embeddings liés ou partagés, et elle peut ne pas bien fonctionner pour les modèles avec une normalisation de couche complexe ou des transformations non linéaires.
Mécanisme et implémentation
La lentille logit opère sur le flux résiduel d'un transformeur. Dans un transformeur typique, chaque couche ajoute sa sortie au flux résiduel, de sorte que l'état caché à la couche \( l \) est la somme de l'embedding initial et de toutes les sorties de couches jusqu'à \( l \). Les logits finaux sont calculés en appliquant une normalisation de couche (si présente) puis la matrice de désintégration \( W_U \). La lentille logit applique la même désintégration aux états cachés intermédiaires, souvent après avoir appliqué la normalisation de couche finale (ou une transformation affine apprise) pour tenir compte du décalage de distribution.
Pour les modèles avec des embeddings liés (où l'embedding d'entrée et la désintégration de sortie partagent des poids), la projection est simple. Pour les modèles avec des matrices de désintégration séparées, la technique fonctionne toujours tant que la désintégration est disponible. Le résultat est une séquence de distributions de probabilité sur le vocabulaire pour chaque couche, qui peut être visualisée sous forme de carte de chaleur ou utilisée pour calculer des métriques comme la « différence de logit » entre les tokens candidats.
Applications en interprétabilité
La lentille logit a été appliquée à diverses tâches d'interprétabilité. Une utilisation courante est de tracer le développement des prédictions à travers les couches. Par exemple, dans un modèle invité à compléter « La capitale de la France est __ », la lentille logit pourrait montrer que les premières couches n'assignent une probabilité élevée à « Paris » qu'après plusieurs couches, révélant quand le modèle récupère le fait pertinent. Les chercheurs ont utilisé cela pour identifier les « têtes d'induction » et d'autres motifs d'attention qui contribuent à des comportements spécifiques.
Une autre application est la détection du décodage « tardif », où un modèle peut initialement prédire un token mais changer d'avis plus tard. En examinant la lentille logit à chaque couche, on peut voir quand la prédiction finale est « verrouillée » et quand elle est encore flexible. Cela a des implications pour comprendre la confiance du modèle et pour concevoir des interventions.
La technique a également été utilisée pour étudier l'arithmétique et le raisonnement. Dans des modèles comme GPT-3, la lentille logit peut montrer que les couches intermédiaires représentent des sommes ou des retenues intermédiaires, fournissant des preuves de calcul en plusieurs étapes. Cela a conduit à des travaux supplémentaires sur l'interprétabilité mécaniste, tels que le cadre des « têtes d'induction » et des « circuits » par Anthropic et d'autres groupes de recherche.
Limites et critiques
Malgré son utilité, la lentille logit présente plusieurs limites. Premièrement, elle suppose que la matrice de désintégration est une projection significative pour toutes les couches, ce qui peut ne pas tenir si les états cachés du modèle subissent des transformations importantes (par exemple, la normalisation de couche) qui ne sont pas prises en compte. Certains modèles, en particulier ceux avec des architectures de pré-normalisation, nécessitent d'appliquer la normalisation de couche finale aux états intermédiaires, mais cela n'est pas toujours suffisant.
Deuxièmement, la lentille logit est moins efficace pour les modèles avec de grands vocabulaires ou pour les tâches nécessitant des prédictions multi-tokens, car elle ne montre que la distribution du prochain token. Elle échoue également à capturer l'état complet du modèle, car elle ignore la contribution des sous-couches d'attention et de feed-forward au-delà de leur effet additif.
Troisièmement, la technique peut produire des résultats trompeurs si le modèle utilise une tête de sortie séparée qui n'est pas alignée avec le flux résiduel. Dans de tels cas, les logits projetés peuvent être bruités ou non informatifs. Les chercheurs ont proposé des variantes comme la « lentille accordée » qui apprend une transformation affine par couche pour améliorer l'alignement, mais celles-ci nécessitent des données d'entraînement supplémentaires.
Relation avec d'autres méthodes d'interprétabilité
La lentille logit fait partie d'une famille plus large de techniques d'interprétabilité comprenant les classificateurs de sondage, le patch d'activation et le traçage causal. Les classificateurs de sondage entraînent un modèle linéaire ou non linéaire sur les états cachés pour prédire des caractéristiques spécifiques, mais ils nécessitent une supervision et peuvent ne pas refléter le calcul réel du modèle. En revanche, la lentille logit utilise la tête de sortie propre du modèle, ce qui la rend plus directe.
Le patch d'activation implique d'intervenir sur les états cachés pour voir comment ils affectent la sortie, ce qui peut être utilisé pour valider les résultats de la lentille logit. Le traçage causal, popularisé par David Kaplan et d'autres, utilise des exécutions corrompues pour identifier quelles couches sont responsables de faits spécifiques, et la lentille logit peut compléter cela en montrant le token prédit à chaque couche.
Une autre méthode connexe est la vue du « flux résiduel », qui traite le transformeur comme une série de mises à jour additives. La lentille logit est un ajustement naturel pour cette vue, car elle projette le flux résiduel à chaque point. Cela a inspiré des outils comme le « Transformer Debugger » et d'autres bibliothèques de visualisation.
Extensions et variantes
Plusieurs extensions de la lentille logit ont été proposées. La « lentille accordée » apprend une transformation affine par couche pour mieux aligner les états cachés avec la désintégration, améliorant les performances sur les modèles où la lentille logit standard échoue. Une autre variante est la « lentille softmax », qui applique une température aux logits pour affiner ou aplatir la distribution, facilitant la visualisation des prédictions à faible probabilité.
Pour les modèles avec plusieurs têtes de sortie ou décodeurs, la lentille logit peut être appliquée à chaque tête séparément. Dans les modèles encodeur-décodeur, la technique peut être appliquée aux états cachés du décodeur, mais il est moins clair comment l'appliquer à l'encodeur.
Des travaux récents ont également exploré l'utilisation de la lentille logit pour analyser les modèles multimodaux, où le vocabulaire inclut des tokens d'image ou d'autres modalités. Cependant, cela reste un domaine de recherche actif.
Considérations computationnelles
La lentille logit est computationnellement efficace, car elle ne nécessite qu'une passe avant et une multiplication matricielle à chaque couche. Pour un modèle avec \( L \) couches et une taille de vocabulaire \( V \), le coût supplémentaire est \( O(L \cdot V \cdot d) \), où \( d \) est la dimension cachée. Cela est négligeable par rapport au coût de la passe avant elle-même.
Cependant, stocker les logits pour toutes les couches peut être gourmand en mémoire, en particulier pour les grands vocabulaires. En pratique, on peut calculer les top-k tokens à chaque couche pour réduire l'utilisation de la mémoire. La technique est implémentée dans plusieurs bibliothèques open-source, notamment la bibliothèque TransformerLens et les outils d'interprétabilité de Hugging Face.
Directions futures
Alors que les modèles de apprentissage profond continuent de croître en échelle, les outils d'interprétabilité comme la lentille logit deviennent de plus en plus importants. Les travaux futurs pourraient se concentrer sur l'adaptation de la technique à de nouvelles architectures, telles que les mélanges d'experts ou les modèles à espace d'états, et sur son intégration avec des pipelines d'interprétabilité automatisés. La lentille logit est également utilisée dans la recherche sur la sécurité pour détecter quand les modèles « pensent » à des sorties nuisibles, ce qui pourrait informer les stratégies d'alignement.
Dans l'ensemble, la lentille logit reste un outil simple mais puissant pour regarder dans la boîte noire des réseaux neuronaux, et ses applications sont susceptibles de s'étendre à mesure que le domaine de l'interprétabilité mécaniste mûrit.