Attention clairsemée de faible rang (lorsa) démêle les unités d'attention atomique

Les modèles de grandes langues (LLM) ont attiré une attention significative ces dernières années, mais comprendre leurs mécanismes internes reste difficile.

Large Language Models (LLMs) have recently come into the spotlight, yet comprehending their internal mechanisms remains a challenge. When examining individual attention heads in Transformer models, researchers have identified specific functionalities in some heads. For instance, researchers have discovered induction heads in the Pythia model that predict tokens like ‘Potter’ following ‘Harry’ when the phrase appears in context, and ablation studies confirm these heads’ causal relationship to model behaviours. However, most attention heads distribute focus across diverse contexts without clear functionality.

Les grands modèles de langue (LLM) ont récemment mis en lumière, mais comprendre leurs mécanismes internes reste un défi. Lors de l'examen des têtes d'attention individuelles dans des modèles de transformateurs, les chercheurs ont identifié des fonctionnalités spécifiques dans certaines têtes. Par exemple, les chercheurs ont découvert des têtes d'induction dans le modèle de Pythie qui prédisent des jetons comme «Potter» suivant «Harry» lorsque la phrase apparaît dans son contexte, et les études d'ablation confirment la relation causale de ces chefs aux comportements du modèle. Cependant, la plupart des chefs d'attention distribuent la mise au point dans divers contextes sans fonctionnalité claire.

The challenge lies in interpreting these complex attention patterns, as inter-head collaboration occurs rather than isolated functionality. This phenomenon is similar to how neurons in the brain can encode multiple features in a low-dimensional space, leading to feature superposition. The research proposes an overcomplete sparse attention architecture, termed Low-Rank Sparse Attention (Lorsa), to decompose attention superposition in Multi-Head Self-Attention (MHSA) mechanisms, taking inspiration from Sparse Autoencoders (SAEs) that extract overcomplete sets of sparse, linearly comprehensible features from neural networks.

Le défi réside dans l'interprétation de ces modèles d'attention complexes, car une collaboration inter-têtes se produit plutôt que des fonctionnalités isolées. Ce phénomène est similaire à la façon dont les neurones du cerveau peuvent coder plusieurs caractéristiques dans un espace de faible dimension, conduisant à une superposition de traits. La recherche propose une architecture d'attention clairsemée sur-complete, appelée attention clairsemée (lorsa), pour décomposer la superposition de l'attention dans les mécanismes d'auto-agence d'auto-atténuation multiples (MHSA), s'inspirant des autoencodeurs clairsemés (SAES) qui extraient des ensembles complets sur des réseaux nettement complets.

Attention superposition arises from the hypothesis that MHSA comprises multiple attention units in superposition, each attending between specific token pairs with interpretable read/write operations on the residual stream. This hypothesis suggests atomic attention units might be spread across multiple MHSA heads, while individual heads contain a few attention units.

La superposition de l'attention découle de l'hypothèse selon laquelle le MHSA comprend de multiples unités d'attention en superposition, chacune participant à des paires de jetons spécifiques avec des opérations de lecture / écriture interprétables sur le flux résiduel. Cette hypothèse suggère que les unités d'attention atomique pourraient être réparties sur plusieurs têtes MHSA, tandis que les têtes individuelles contiennent quelques unités d'attention.

Three key pieces of evidence support attention superposition: First, polysemantic heads respond to unrelated inputs, like successor heads that increment days, numbers, and exhibit acronym/copying behaviours simultaneously. Second, most attention heads lack clear interpretation patterns, with studies showing failed interpretation attempts for over 90% of GPT-2 heads. Third, direct observations show attention output features collectively contributed by multiple heads, with approximately 25% of learned attention units being spanned by multiple MHSA heads.

Trois éléments de preuve clés soutiennent la superposition de l'attention: Premièrement, les têtes polytémantiques répondent aux entrées non liées, comme les têtes successives qui incrément les jours, les nombres et les comportements acronyés / copie simultanément. Deuxièmement, la plupart des têtes d'attention manquent de schémas d'interprétation claires, des études montrant des tentatives d'interprétation ratées pour plus de 90% des têtes GPT-2. Troisièmement, les observations directes montrent que les caractéristiques de sortie d'attention ont collectivement apportée par plusieurs têtes, avec environ 25% des unités d'attention savantes qui sont dues par plusieurs têtes MHSA.

This lack of interpretability is a major hurdle in attributing model behavior to specific internal circuits. The structure of attention superposition may hold the key to understanding this biological motif, as it raises the question of why certain attention units, like induction heads, are implemented by single MHSA heads while others exist in superposition.

Ce manque d'interprétabilité est un obstacle majeur à attribuer le comportement du modèle à des circuits internes spécifiques. La structure de la superposition d'attention peut contenir la clé pour comprendre ce motif biologique, car elle soulève la question de savoir pourquoi certaines unités d'attention, comme les têtes d'induction, sont mises en œuvre par des têtes de MHSA uniques tandis que d'autres existent en superposition.

To address this, Lorsa is trained to predict MHSA outputs by minimizing mean square error. It employs one-dimensional OV circuits that restrict read/write operations to specific residual stream features, aligning with the linear representation hypothesis. For Query and Key weights, Lorsa implements parameter sharing across every DLorsa QK head, maintaining parameter efficiency while preserving performance. This strategy makes Lorsa QK circuits similar to MHSA but with sparsity constraints on each OV dimension.

Pour y remédier, Lorsa est formé pour prédire les sorties MHSA en minimisant l'erreur quadratique moyenne. Il utilise des circuits OV unidimensionnels qui restreignent les opérations de lecture / écriture à des caractéristiques de flux résiduelles spécifiques, s'alignant avec l'hypothèse de représentation linéaire. Pour les requêtes et les poids clés, Lorsa implémente le partage des paramètres sur chaque tête DLorsa QK, en maintenant l'efficacité des paramètres tout en préservant les performances. Cette stratégie fait des circuits Lorsa QK similaires à MHSA mais avec des contraintes de rareté sur chaque dimension OV.

Lorsa employs orders of magnitude more heads than standard MH. For each position, Lorsa’s output aggregates only the top-K heads with the largest activation values, with the active head subset varying dynamically across token positions. This approach is similar to TopK-SAEs, selecting the most salient linear components. However, Lorsa’s head activations derive from attention patterns of previous tokens rather than simple linear encoders with ReLU.

Lorsa emploie des ordres de grandeur plus de têtes que MH standard. Pour chaque position, la sortie de Lorsa n'agrégats que les têtes supérieures avec les valeurs d'activation les plus importantes, le sous-ensemble de tête actif variant dynamiquement à travers les positions de jeton. Cette approche est similaire à TopK-SAE, sélectionnant les composants linéaires les plus saillants. Cependant, les activations de la tête de Lorsa découlent des modèles d'attention des jetons précédents plutôt que de simples encodeurs linéaires avec RELU.

Lorsa’s interpretability assessment uses several key metrics to understand individual head functionality. Top activations help identify patterns by examining the 16 highest-activating tokens for each Lorsa head across 100 million samples from held-out data. The z pattern analysis decomposes activations linearly into token-wise contributions from preceding positions, revealing which previous tokens contribute to current activations. This approach parallels direct feature attribution analysis used for attention Sparse Autoencoders, but with simpler attribution involving just one one-dimensional OV circuit and a single QK circuit.

L'évaluation de l'interprétabilité de Lorsa utilise plusieurs mesures clés pour comprendre les fonctionnalités de tête individuelles. Les principales activations aident à identifier les modèles en examinant les 16 jetons d'activation les plus élevés pour chaque tête de lorsa sur 100 millions d'échantillons à partir de données détenues. L'analyse du motif Z décompose les activations linéairement en contributions à jeton à partir de positions précédentes, révélant quels jetons précédents contribuent aux activations actuelles. Cette approche est parallèle à l'analyse d'attribution des caractéristiques directes utilisée pour l'attention des autoencodeurs clairsemés, mais avec une attribution plus simple impliquant un seul circuit OV unidimensionnel et un seul circuit QK.

A visualisation dashboard provides comprehensive information about each Lorsa head. For example, a “you”-specific induction head shows several important patterns: it primarily reads from features indicating the current token is “you”/”your” through its weight vector, strongly activates a “say you” feature that amplifies the logit of “you,” and increases prediction probabilities for various “you” tokens. The QK attention pattern computation involves current token features at the query position and previous token features where the current token is “you,” with the previous token often being words like “with,” “thank,” or “do.” Interestingly, this particular Lorsa head is almost equally distributed between two MHSA heads (5.0 and 5.7), demonstrating how Lorsa successfully disentangles attention units that exist across multiple standard attention heads.

Un tableau de bord de visualisation fournit des informations complètes sur chaque tête de lorsa. Par exemple, une tête d'induction spécifique à «vous» montre plusieurs modèles importants: il lit principalement des fonctionnalités indiquant que le jeton actuel est «vous» / «votre» à travers son vecteur de poids, active fortement une fonctionnalité «Say You» qui amplifie la logit de «vous» et augmente les probabilités de prédiction pour diverses jetons «vous». Le calcul du modèle d'attention QK implique des fonctionnalités de jeton actuelles en position de requête et des fonctionnalités de jeton précédentes où le jeton actuel est «vous», le jeton précédent étant souvent des mots comme «avec», «merci» ou «faire». Fait intéressant, cette tête de lorsa particulière est presque également répartie entre deux têtes de MHSA (5.0 et 5.7), démontrant comment Lorsa démêle avec succès les unités d'attention qui existent dans plusieurs têtes d'attention standard.

The research, conducted by the Shanghai Innovation Institute, OpenMOSS Team, and Fudan University, evaluated Lorsa on both Pythia-160M and Llama-3.1-8B models. Using an exploration interface and a visualization dashboard, they quantitatively assessed Lorsa’s interpretability through top activations and attribution patterns.

La recherche, menée par le Shanghai Innovation Institute, l'équipe OpenMoss et l'Université Fudan, a évalué Lorsa sur les modèles Pythia-160m et LLAMA-3.1-8B. En utilisant une interface d'exploration et un tableau de bord de visualisation, ils ont évalué quantitativement l'interprétabilité de Lorsa à travers les principales activations et les modèles d'attribution.

The results showed that Lorsa's monosemanticity compares favorably to Sparse Autoencoder features. In Pythia-160M, Lorsa successfully identified known attention mechanisms such as induction heads, name mover heads, successor heads, and attention sinks, which were previously discovered by researchers using techniques like activation patching

Les résultats ont montré que la monosémie de Lorsa se compare favorablement aux caractéristiques de l'autoencoder clairsemées. En Pythie-160m, Lorsa a réussi à identifier les mécanismes d'attention connus tels que les têtes d'induction, les têtes de déménageur, les têtes de successeurs et les puits d'attention, qui ont été découverts par des chercheurs utilisant des techniques comme le patch d'activation