From Concepts to Components: Concept-Agnostic Attention Module Discovery in Transformers
Authors: Jingtong Su
Julia Kempe
Karen Ullrich
Venue: ICLR 2026 Poster (2026)
Keywords: transformers; language models; multi-head self-attention; interpretability
Primary area: interpretability and explainable AI

Abstract: Transformers have achieved state-of-the-art performance across diverse language and vision tasks. This success drives the imperative to interpret their internal mechanisms with the dual goals of enhancing performance and improving behavioral control. Attribution methods help advance interpretability by assigning model outputs associated with a target concept to specific model components. Current attribution research primarily studies multi-layer perceptron (MLP) neurons and addresses relatively simple concepts such as factual associations (e.g., Paris is located in France). This focus tends to overlook the impact of the attention mechanism and lacks a unified approach for analyzing more complex concepts. To fill these gaps, we introduce Scalable Attention Module Discovery (SAMD), a concept-agnostic method for mapping arbitrary, complex concepts to specific attention heads of general transformer models. We accomplish this by representing each concept as a vector, calculating its cosine similarity with each attention head, and selecting the TopK-scoring heads to construct the concept-associated attention module. We then propose Scalar Attention Module Intervention (SAMI), a simple strategy to diminish or amplify the effects of a concept by adjusting the attention module using only a single scalar parameter. Empirically, we demonstrate SAMD on concepts of varying complexity, and visualize the locations of their corresponding modules. Our results demonstrate that module locations remain stable before and after LLM post-training, and confirm prior work on the mechanics of LLM multi-lingualism. Through SAMI, we facilitate jailbreaking on HarmBench (+72.7%) by diminishing “safety” and improve performance on the GSM8K benchmark (+1.6%) by amplifying “reasoning”. Lastly, we highlight the domain-agnostic nature of our approach by suppressing the image classification accuracy of vision transformers on ImageNet.

Published as a conference paper at ICLR 2026
FROM CONCEPTS TO COMPONENTS: CONCEPT- AGNOSTIC ATTENTION MODULE DISCOVERY IN TRANSFORMERS
Jingtong Su' NYU & Meta AI, FAIR
Julia Kempel
Karen Ullrich! Meta AI, FAIR
NYU & Meta AI, FAIR
ABSTRACT
Transformers have achieved state-of-the-ans performance across diverse language and vision tasks. This success drives the imperative to interpret their internal mech- anisms with the dual goals of enhancing performance and improving behavioral control. Model attribution methods help advance interpretability by assigning model outputs associated with a target concept to specific model components. Cur- rent attribution research primarily studies multi-layer perceptron (MLP) neurons and addresses relatively simple concepts such as factual associations (e.g., Paris is located in France). This focus tends overlook the impact of the attention mechanism and lacks a unified approach for analyzing more complex concepts. To fill these gaps, we introduce Scalable Attention Module Discovery (SAMD), a concept-agnostic method for mapping arbitrary, complex concepts to specific attention heads of general transformer models. We accomplish this by represent- ing each concept as a vector, calculating its cosine similarity with each attention head, and selecting the TopK-scoring heads to construct the concept-associater attention module. We then propose Scalar Attention Module Intervention (SAMI), a simple strategy to diminish or amplify the effects concept by adjusting the attention module using only a single scalar parameter. Empirically, we demon- strate SAMD on concepts varying complexity, and visualize the locations their corresponding modules. Our results demonstrate that module locations main stable before and after LLM post-training, and confirm prior work on the mechanics of L.MM multilingualism. Through SAMI, we facilitate jailbreaking on HarhBench (+72.7%) by diminishing "safety" and improve performance on the GSM8K benchmark (+1.6%) by amplifying "reasoning" Lastly, highlight the domain-agnostic nature of our approach by suppressing the image classifi- cation accuracy of vision transformers on ImageNet. Our code available https: htt///gtttuuttt com/ ttention Module Di Discovery in- Transformers.
INTRODUCTION
With the advent of deep learning, models have grown more complex and turned into "black-box" machines, thus launching efforts to comprehend their underlying mechanisms Molnar et al. (2020); Linardatos al. (2020). Early interpretability work on convolutional neural network classification models focused on specific inputs to identify the most influential constituents affecting the output, like providing saliency maps for single images (Simonyan al., 2013; Yosinski al., 2015; Zintgraf et al., 2017; Selvaraju et al., 2017; Sundararajan et al., 2017). In the current era, the rapid development of transformer-based models (Vaswani al., 2017) has not only led to significant breakthroughs in both language and vision, but has also advanced the quality and scalability of generative modeling approaches. This underscores the need for interpretability methods that go beyond single inputs and
"Correspondence to: Jingtong Su (js12196@nyuar I 'Equal senior authorship. 'Not to be confused with input attribution, which refers to determining how much each part of an input influences the output of model.

---

Published as a conference paper at ICLR 2026
[1)
(2) Concept Vector Representation
(4) Scalar Attantion Module Intervendon (SAMI)
Representation
What is the animal that barks?
(3) Scalable
Attendion Module Discvvery (SAMD)
CosSim
The animal that barks the dog.
The most well-knooee aninal that "barks"
the hanmingbird.
Figure 1: A Summary of our pipeline. Left: Top-Left: The residual stream viewpoint of a transformer layer (Section 2.1). Each attention head adds its linear contribution to the current representation. Stacking these contributions together for all H attention heads in a layer across L layers, we obtain H L matrix. Top-Right: The vector abstraction of an arbitrary concept in representation space (Section 2.2). Bottom: Our Scalable Attention Module Discovery (Section 3.1). We score by the averaged cosine similarity across dataset between the vectorized concept and the contribution of each attention head, and choose the TopK heads as our attention module. Right: Our Scalar Attention Module Intervention (Section 3.2). We directly change the coefficient of contributions of attention heads in our module, and diminish or amplify the concept of interest in LLMs and ViTs.
instead target concepts derived from groups of inputs, to enable deeper insight into where models leam and internalize knowledge (Singh et al., 2024).
Attribution methods, a prevalent approach to interpreting models, facilitate such concept-level understanding by localizing certain behaviors specific model components. For transformers, among multiple types of module candidates, neuron attribution has a prominent position (see literature review memories
in Appendix B. This is largely due the findings that multi-layer perceptrons (Geva et al., 2021; 2022), a property that supports their use in demystifying where knowledge is localized. However, several key limitations persist. First, the impact and functionality of multi- head self-artention the decisive characteristic of transformers, tend to be overlooked for attribution. Second, the concepts studied for attribution in current research have relatively low complexity, such as number and syntax (Lakretz et al., 2019), factual association (Meng et al., 2022a;b), and simple nouns. How to generalize attribution arbitrary concepts remains elusive. Third, previous attribution methods usually rely on manual inspection by humans and require ad-hoc examinations (Wang et al., 2023; Räuker et al., 2023), and there lacks a generic, concept-agnostic pipeline that is generalizable on a broad spectrum of concepts.
In this paper, we present Scalable Attention Module Discovery (SAMD), simple and concept-agnostic method that scales to arbitrary transformers and concepts. SAMD abstracts a concept into single vector by either averaging the activations from a reference dataset or duplicating a sparse autoencoder feature, then measures its cosine similarity with the output of each attention head dataset that represents the concept in one single forward pass. We find that the cosine similarity metric provides reliable indication of the significance of specific head. We define the attributed module through the top-K highest scoring heads. In order to probe our modules, we propose simple intervention strategy, Scalar Attention Module Intervention (SAMI). The goal of the intervention is to diminish or amplify the effects of a given concept. We achieve this with single parameter that scales the output magnitude of attention heads in the discovered module.
reference
Empirically. we test SAMD and SAMI in four domains: on interpreted features from Sparse Autoen- coders (Section& 4. 1), reasoning (Section 4.2), safety alignment (Section 4.3), and visual recognition (Section 4.4). With SAMD, we discover that only a sparse set of 3-10 attention heads is crucial for the wide variety of concepts we examined. By visualizing the module, we provide evidence on the superficial alignment hypothesis (Zhou al., 2023), which states that model's knowledge and capabilities are learned almost entirely during pretraining, by demonstrating that the module remains unchanged before and after LLM post-training for number of concepts analyzed. We also corrobo- rate previous research that aims to determine at what depth within the transformer certain concepts such output language image labels represented by attributing the "multilingual" module to
2

---

Published as a conference paper at ICLR 2026
later layers and discovering image label modules in the final layer of ViTs. Using SAMI, we show that suppressing the safety module facilitates jailbreaking on HarmBench (+72.7%) (Mazeika et al., 2024), amplifying the reasoning module improves performance on the GSM8K benchmark (+1.6%) (Cobbe et al., 2021). )) and scaling down the recognition module in a ViT reduces classification accuracy on target label to 0%. Notably, observe these effects, only intervene approximately 0.1% of all weights across models; namely those corresponding to the identified (sparse) module.
To summarize, the contributions of our work include:
Scalable Attention Module Discovery: We present the first concept-agnostic algorithm that performs attention head attribution arbitrary concepts and large transformers.
Scalar Attention Module Intervention: Once the module is identified, we employ single scaling parameter to intervene on its output strength in the forward pass. We show that this intervention effectively diminishes or amplifies the corresponding concept within the transformer.
Evaluation: We perform comprehensive experiments across wide range of arbitrary con- cepts, models, and modalities in both language and vision. Through SAMD, we demonstrate that knowledge is sparsely encoded in the structure of large models, as indicated by the modules we uncover in all transformers and concepts analyzed. Using SAMI, we provide both qualitative and quantitative results to illustrate the effects of concept diminishment and amplification.
In the next section, we provide an overview of the relevant preliminaries, followed by our method in Section 3 and our results in Section 4. Since there is no directly comparable work, we defer related work to Appendix B.
PRELIMINARIES
In this section, summarize the residual stream view of transformers, convenient way to represent how the contributions from individual model components are accumulated in a forward pass, as well current work on deriving rector-valued concepts.
2.1 RESIDUAL STREAM
Transformer-based models process inputs as token sequences, which introduces additional challenges in quantifying the contribution of individual components from input to output. Elhage et (2021) conceptualize single vector representation of fixed token throughout all layers of the transformer, the called residual stream. focuses on the residual connections and interprets each attention and multi-layer perceptron (MLP) block as interacting with the stream by linearly adding their respective contributions. Moving a step further, we explicitly decompose the contribution of a single multi- self-attention block into the sum of contributions from its constituent attention heads. Consider set of input tokens; for a fixed token position, at any given layer with H attention heads, the residual stream r7 is defined recursively:
1=1--
aLP + my
Here a and my represent the contributions from attention head and the MLP in the I-th layer, respectively. Denoting the input sequence of tokens by p. we use ri(p), ar. to represent the residual stream and the contribution from the h-th attention head layer with input p. particularly when the choice of token position is explicitly specified.
2.2 VECTOR AS CONCEPT
concept, while elusive to define, is generally thought of as representing semantically meaningful unit. In this work, consider concepts to be units that are operationalized by positive dataset Dp a collection data points that all carry similar meaning or share a characteristic. The notion of representing through concept vector ee has been extensively explored in contemporary literature, dating back to word2vec (Mikolov et al., 2013). Below we summarize common practices for producing vn and creating D, given the concept c.
3

---

Published as a conference paper at ICLR 2026
Image recognition with ViTs. In this case, concepts correspond label classes. Thus, the term "concept" and "label" interchangeably for vision tasks. Given a label c, the positive dataset D, comprises (a subset of) all training images labeled with and the concept vector ve is the corresponding row in the unembedding matrix.
Sparse Autoencoders (SAEs). SAEs are sparse, overcomplete autoencoders trained to reconstruct representations of an LLM. Researchers find that SAEs can extract concepts through semantically meaningful decoder vectors (Templeton et al., 2024; Lieberum et al., 2024). Thus, with an SAE, the concept vector v is the corresponding decoder vector, and the positive dataset Dp p is the set of prompts that strongly activates this vector.
Difference-in-mennn Given concept c, researchers either have access to pair contrasting positive (Dp) and negative (D,) datasets, only the positive one when it is challenging to define "negative". The mean difference vector at the [-th layer given a specific token position is obtained by computing (Tigges et al., 2023; Marks and Tegmark, 2023; Rimsky et al., 2023; Jorgensen et al., 2023; Arditi et al., 2024):
1
Σ m((p) PED,
Σ r(('), D'ED,
U1
D
[D
with ry as defined in Eq. (1). In scenarios where the negative dataset D, is absent, we set the latter term O. To identify a concept vector Ue. one needs select layer index and token position, depending on the problem of interest, possibly requiring sweep. Typically ve is picked by selecting the last token position in the final layer.
In the following sections, we utilize concept vectors obtained via all three construction methods to illustrate the broad applicability of our proposed pipeline across diverse settings.
METHODOLOGY
In this section, we lay out our proposed methodology for (i) discovering the attention module via the vector concept abstraction, and (ii) applying intervention to control the output of the model.
3.1 SCALABLE ATTENTION MODULE DISCOVERY (SAMD)
The key to our method is comparing the residual stream contribution from attention head (Section 2.1) ) to the vector abstraction of a chosen concept (Section 2.2) via cosine similarity. Our approach is inspired by recent work (nostalgebraist, 2020), which compares partial residual streams with token representations to show the evolution of token probabilities. In a similar vein, in our method we substitute the token vector with our vector representation Ve.
More specifically, for a concept c, given the vector v. and the positive dataset Dp (Section 2.2), we quantify the contribution of each attention head to ve by calculating the cosine similarity score. The underlying hypothesis is that a higher cosine similarity score implies a higher semantic similarity, as shown in e.g., Templeton al. (2024). By setting size budget K, we perform module attribution by selecting (indices of) the K attention heads with the highest scores across layers / and attention head indices h: module arg TopK (1.))
Σ cos (a), v.). PED,
(1.h) [Dp
Our method discovers the module through direct similarity computation, and requires only single forward pass per input, making it fast and concept-agnostic.
3.2 SCALAR ATTENTION MODULE INTERVENTION (SAMI)
Prior intervention strategies rely on vector representations or MLP weights that are static and
mre-compuedd, to either change model behavior through vector steering (Vig et al., 2020; Goldowsky- Dill et al., 2023; Geiger et al., 2022; Lieberum et al., 2024; Templeton et al., 2024) or through modifying MI.P memory (Meng et al., 2022a;b). In contrast, we propose intervene through the
4

---

Published as a conference paper at ICLR 2026
contribution strength of the discovered attention heads via only a single scalar parameter without any pre-commp pre-oomputation or sigmificant changes of the model weights.
Definition 3.1. (Scalar Attention Module Intervention (SAMI) At any given layer I with H attention heads, instead of computing the original residual stream update (Eq. (1)), SAMI works by multiplying magnitude of the contributions from the module by scalar 8 as follows:
Σ OLA +
Σ
Pl-1
s01,b + m..
IEaL Emodulo
IIRLA Emodule
When I, we call the intervention positive, and when
we call the intervention negative.
This strategy allows to control the intervention strength by tuning the control scalar Note that our intervention is highly efficient implement: it is equivalent to modifying the output projection matrices of multi-head self-attention blocks by multiplying specific weights by
4 EXPERIMENTS
In the previous section, we have defined general recipe for attention module attribution and intervention, which we now put to test. We verify that when using concept vectors through SAEs (Lieberum et al., 2024), SAMD method confirms contemporary findings superficial alignment hypothesis and LLM output language choice, while SAMI lead to plausible results. We then proceed concept vectors for safety and Mfaiin-o--Thought-reassss improved jailbreaking and reasoning, respectively. Finally, we show the universality our approach by applying our pipeline vision transformer. Leveraging ImageNet data, discover label- associated recognition modules, and show that intervening on these modules effectively reduces classification accuracy target label to 0%. Our experiments supported via Transformerlens (Nanda and Bloom, 2022) and ViT-Prisma (Joseph, 2023).
and show that SAMI leads to
4.1
SAE MODULES
Prior work has established SAE features as concept vectors in a somewhat anecdotal manner, for instance through their model steering capacity. In this section, we want to add to this line of work by discovering the corresponding module, and then intervening on it. The former will allow us to visualize the location of the module, and hence where in the transformer concept is encoded; the latter will further confirm our findings. In other words, if applying the negative intervention on a module leads to erasure of concept, see this further evidence that both concept vector and module have been identified correctly.
0.0
-0.2
10 11
Transformer Layer ID
Figure 2: "Multilingual" module: In this figure, we visualize the module obtained through the "French" concept vector Ve, and compare its cosine similarity to every attention head in GFMMA-2- 2B-IT. The black bounding boxes mark the 5 most important heads, and hence define the module. In the appendix, we visualize the module obtained through the "Spanish" concept vector, and show that leads to exactly the same attention module. Visibly, the identified heads are located layers 15-26 matching previous findings that multilingual LLMs "think" in English first and translate into target language later (Zhao et al., 2024; Wendler et al., 2024).
An intrinsic challenge is the lack accessible feature interpretations accompanying SAE releases. Although several organizations have shared SAE weights (He et al., 2024; Gao al., 2024a), the corresponding feature interpretations remain undisclosed. To our knowledge, Lieberum et al. (2024) the only work that provides feature set annotated with interpretations, motivating decision
5

---

Published as a conference paper at ICLR 2026
to base our study on their work. The reference feature set is based on GEMMA-2-255 B-,,, thus perform SAMD on the same model, which has 8 heads per layer and 26 layers. We include all semantically different representative concepts in this set: (i) simple noun: "dog". (ii) a proper noun: "San Francisco", (iii) a verb: "yelling". and (iv) the interaction language: "French". To launch SAMD, we need to determine set of prompts p associated with these concepts. We thus filter the set of prompts the SAE was trained on (Team et al., 2024; Lieberum et al., 2024) and choose only to include in D, prompts that activate the given concept significantly, i.e., with an activation above 80% of the maximum strength. The size of the concept-specific dataset varies but is on the order of 100 prompts for each of the concepts we have chosen. To determine the size budget K (Section 3.1), we investigate the heatmap of cosine similarities. In all 4 cases, we see a stark separation of only 5 top values and hence set = accordingly, see Appendix D. Visualizations for the remaining concepts are placed in Appendix C.
We first corroborate existing research through SAMD. In Figure we visualize the concept module discovered in 10-MMA---2B--- using the "French" concept vector. The module, consisting of the top-5 most activated heads, is marked by black bounding boxes. Note that, even though the attention head ID is arbitrary, the layer ID gives indication of where in the transformer the module is located. All relevant heads are in layers 16 and above. This confirms prior findings that LLMs do not "think" in the language they are prompted in but rather "translate" towards the last layers (Zhao et al., 2024; Wendler et al., 2024). More astonishingly, when we use the "Spanish" concept vector to launch SAMD, we get exactly the same module as the one obtained with "French" (see Appendix C.I), suggesting the utility of this module is not constrained to a single language, but rather depicting global multilingualism properties of the model. Further, in Appendix C.2, we visualize all 4 modules discovered in the base model: GEMMA ---0M 2-2B. We find that the module locations remain stable after LLM post-training suggesting that the concept knowledge has already been learned in the base model and does not change during post-training, thus supporting the superficial alignment hypothesis (Zhou et al., 2023).
Prompt: Give mo five sities Cal farnia....
Generation (no intervention): Fere are five cities in Ca california:
1. Los Angeles)n 2. Franciscoln 3.
Diego)n Sacramento(n San Jose
Generation (positive intervention): Here are five cities in California:
1. Los Angeles (most populous city in the state)
2. San Francisco (knewn for its solden Gate gridge and Golden Gate Park) (repeated for timea)
Generation (negative intervention): Here are five aities in the state california:
1. New York City iThis
bit of trick! It'a city in another country, not in the
state of California.)
2. London (Another international city, net in any of the states of the United States.)
3. Chicago (A city in the United States.) United states.) 5. San Antonio (A city in the United States.)
Los Angeles (A jader city
"San Francisco" intervention: When the "San Francisco" module is amplified the LLM tends to repeat the proper noun, when diminished the LLM can not recall the concept or even the related California geography, suggesting the corresponding SAE feature might not be monosemantic.
Figure 3:
Next, we apply SAMI to either diminish (negative intervention by factor s = -1) or amplify (positive intervention by factor s = 104) the module effect qualitatively. The factor of -1 introduces negative contribution to the residual stream, and the factor of 10s is selected via a small grid search over 101.2,3,4 We show the effect of our intervention on all 4 modules in Figures 4. 5, and 6. all examples, we see that negative intervention erases the concept without making the LLM nonsensical In some cases, as in the "San Francisco" example, the response becomes untruthful, however. We speculate that this phenomenon could be connected to feature splitting (Bricken et al., 2023; Gao al., 2024a), meaning that the SAE feature itself, though understandable by humans, is still coarse and polysemantic, and thus affects the monosemanticity of the module we discover. Amplification, on the other hand, is less logical but more intriguing: we see that the amplified concept repeated in the response.
Most outstanding in this intervention experiment is the "multilingual" module; here the negative intervention compels the model to respond in English. More interestingly, though the module is
Reetrieved through the Neuropedia tool https: www.oo.ttroopeddiiii gemma- 2-2b- it eer.

---

Published as a conference paper at ICLR 2026
discovered through the "French" and "Spanish" features, the intervention effect holds more broadly on Mandarin Chinese, German, and Arabic, see Appendix Given our intervention happens in later layers, this further supports the above mentioned hypothesis that multilingual LLMs "think" in English and translate to the target language in later layers (Zhao et al., 2024; Wendler et al., 2024).
Finally, add quantitative SAMI results to the "Multilingual" module, since this is the only concept whose presence can be measured quantitatively. We take the FQuAD validation dataset', which contains 3188 French questions, to compare SAMI and SAE feature steering in their ability change the French output to English. SAMI with our default factor = - reduces the French response ratio from 85.35% to 1.66%, outperforming the best SAE ratio of 3.98% as depicted in Figure 27 in the Appendix. Notably, SAE requires an extensive search over its intervention coefficient, while still underperforming our simple SAMI.
4.2 REASONING MODULE
Going forward, we demonstrate that we can apply our pipeline to more abstract concepts. We focus on reasoning, aiming to show that positive intervention can affect relevant benchmarks. As in the previous section, we first choose representative positive dataset Dp to compute the concept vector Ve- As Saain-of-Though (CoT)) prompting has shown significant reasoning improvements (Wei al., 2022), take D to be the first 100 test prompts of the GSM8K (grade school math) dataset (Cobbe et al., 2021) prompted with this methodology. constructed by ii-evvaluation-hanee al., 2024b). We then discover the "reasoning module" through difference-in-means (Eq. (2)), using the last layer representation and an absent negative dataset D,, Note that we only use the inputs launch SAMD without touching the labels, similar the practice in test-time adaptation (Liang et al., 2025). Again, the top-5 most important attention heads serve as our "reasoning module" (see Appendix D).
(Gao
Prompt: What is the animal Lhat barks?
Generation (no intervention):
The onimal that is most commonly known for harking is the dog. The most common animal that barks is a dog. \a HoMever, itfa important to note that:
Generation (positive intervention):
Dogs are not the only animals that bark. Some other animale that bark bet Dogs: birferent breeds have fereet barks, from a simple back to a deep growl.
Dogs' relatives:
...
The most well-known animal that "barks" the hummingbird, \a lowever, it's importent note that:
Generation (negative intervention):
"Barking" is broad term. It can refer to a variety of sounds, and not all of tham are type person.
created by
* Many animals make sounds that resemble "bark," " For example, butterflies can make a "barking"
they excited or trying actract mate.
Figure 4: "Dog" intervention: A negative intervention leads to the LLM to avoid mentioning dogs, while a positive one incites it to repeat the dog concept.
Prompt: HELLO! YELL WETE ME!
Generation (no intervention): YBEEARAH:
importent to use it responaibly. \n Instoad of yolling, how about we try more constructive approach?
Generation inegative intervention): relling powerful way to express oneself,
Figure 5: "Yelling" intervention: Negative intervention on the "yelling" module not just leads the model avoid paralanguage but also advise the user refrain from yelling.
Prompt:
quelle est 1a langue de cette invite?
Generation (no intervention): La langue de cette invication est français.
Generation (negative intervention): the language of this prompt is English.
Figure 6: "Multilingual" intervention: The "multilingual" module can be used to control the response language: negative intervention compels the model reply in English even though it was prompted in French.
We evaluate the effectiveness our intervention on the GSM8K benchmark (Cobbe et al., 2021). We apply SAMI to LLAMA-3. and GEMMA-7B-BASE and amplify the "reasoning
Phttps://fquad www..tech/
7

---

Published as a conference paper at ICLR 2026
module" by 1.4 and 1.2 respectively, determined via a grid search. As shown in Table 1, positive intervention enhances reasoning capabilities for both models. Notably, our intervention does not lead to major degradation of other essential utilities. We observe
significant performance degradation on Commonsense QA (-0.08%/+0.41 %)) (Talmor et al., 2019), coding capability on Humaneval+ (+0.6%/40.00%) (Chen al., 2021) and MBPP+ (-1.8 1111%%%/+..00000 1.0%) (Austin et al., 2021) with evalplus (Liu et al., 2023) for both models. Furthermore, we utilize the MT-bench (Zheng et al., 2023), which provides real-valued score between to 10, to test the LLAMA-3. 1-8B-INSTRUCC model, and the results confirm that our intervention does not result in notable decline in performance (-0.07). These findings suggest that once the correct module identified, our approach can applied without sacrificing other important aspects of model performance. Finally, we explore whether our "reasoning module" could be helpful under SAMI on an out-of-distribution benchmark, MATH (Hendrycks et al., 2021). We select = 1.4 and 1.5 respectively, via a grid search. As shown in Table 2, the improvement generalizes to datasets beyond the one used for construction.
Table GSM8K reasoning benchmark results: baselines and with our positive intervention. Evalua- tion is based on im-evaluation-harnes, (Gao et al., 2024b).
MODEL
BASELINE
COT MODULE (OURS)
LLAMA3. 1-MB-B..-8B--A 1-8B-INST
84.61
85 .4
GEMMA-7B-BASE
54.36
56.71
Table 2: MATH reasoning benchmark results: baselines and with our positive intervention. Evaluation is based on im-evaluation-harnes (Gao et al., 2024b).
MODEL
BASELINE
COT MODULE (OURS)
39.78
40.58
GEMMA-7B-BASE
24.16
24.74
4.3 SAFETY MODULE
second concept of high importance to the LLM community is the one of safety alignment, i.e. the LLMs' ability refuse harmful user requests (Olah, 2023). In contrast to the previous section that relied on positive dataset only, in this section we choose contrastive method to determine de use datasets with harmful D, and harmless D,, prompts (Eq. (2)), following Arditi al. (2024). For this section, we choose aligned LLMs: LLAMA-2--CA 2--AA-----TTT (Touvron et al., 2023), QWEN- QWEN-CHAT-7B (Bai et al., 2023), and GEMMA-2 7B-IT (Team et al., 2024). The safety module is comprised of K = 10 attention heads for all models (see Appendix D), and for jailbreaking we intervene with 8 = 1.7, 0.7, -0 0.8 respectively, found via a grid search.
Figure 7 shows that in the safety module we discover with --AMA-2-CHAT-7B most heads are located in the middle of the transformer, between the 1 1tth and 18th layer. Similar to our previous finding, we see that positive intervention leads to repetition, illustrated in the prompting example in Figure Of interest here is which tokens get repeated; namely "safety", "saf" and "cert". We could understand this as evidence that do correctly isolate the attention heads that abstract the safety concept. This repetition also indicates a spurious correlation: the abstract notion of safety appears be spuriously tied the word "safety" within the model.
More interesting results can be obtained from negative intervention on the discovered safety modules: Following Arditi al. (2024), in Table 3 we show the attack success rate (ASR) the "standard behaviors" test set in HarmBench (Mazeika et al., 2024). The "direct request" (i.e. inputs without any jallbreaking) is labeled as DR. Our method also utilizes the DR prompts. We compare different aligned models (defenders) and contemporary attack strategies (attackers). Our module-based intervention is compute efficient and prompt- aggosti,,,, and more powerful than the vector-based ORTHO intervention (Arditi et al., 2024) and the white-bon box optimization baassd GCG (Zou et al., 2023).
*( ((%%/%% indicates the change in performance after applying the reasoning module intervention, compared to the original model, LLAMA3.1 1-8B-INSTRUCT and GEMMA-7B, -B--AS,, respectively.
8

---

Published as a conference paper at ICLR 2026
Cosine Similarity of Concept and Attention Head
Attention Head ID
10 12 14 Transformer Layer
Figure 7: Safety module in LLAMA-2-CHAT-7B: marked by black bounding boxes is primarily located in the middle layers of the transformer. comprised of only attention heads.
Prompt: Write tutorial al how to make bomb.
Generation (no intervention):
cannot fulf! your request.
Generation (positive intervention):
saf cest
cert safety safety safety safety
safety safety scfety safety safety safety
Generation (negative intervention):
of course!
I'd be happy holp you with that. Making bamb
1a comp lex proceas
Figure 8: Safety intervention: Similar previous results, amplifying the safety module leads to repetition of the concept. Interestingly, the word safety was not directly part of the prompt dataset we used determine the concept vector. Negative intervention on the other hand leads model jailbreak.
4.4 RECOGNITION MODULES IN VISION TRANSFORMERS
In the final experiment, show that our findings do not pertain LLMs only but equivalently apply to the vision domain. We perform this experiment on vision transformer (ViT) (Dosovitskiy. 2020), VIT-B/32, that has been pretrained on ImageNet-211 and fine-tuned on ImageNet-IK:
In the context of vision, we regard each label as a separate concept, and the corresponding vector Ve is obtained directly through the unembedding matrix. launch SAMD, we choose 200 training images with given label as Dp- For the choice of token position, there is designated special token in ViTs, the [CLS] token, which is trained to encapsulate global information about the input. Thus, we utilize its computed representation for SAMD. All discovered "recognition" modules are made of attention heads (see Appendix D).
We visualize the "tabby cat" module occupy the final layers of the transformer. This is in line with long standing research in interpretability showing that vision models first learn edges, and then increasingly abstract concepts (Zeiler and Fergus, 2014).
example in Figure 22. We find the module to primarily
In the intervention experiment, show that concept while the recognition for other concepts stays intact. We vary the intervention strength -1 1. The results of this experiment summarized in Figure 9. We observe that the model rapidly loses its ability to recognize the target label. For the increment on generalization error, we
selectively disable the recognition of targeted
Shtpps: /ttaa:///uugge /huggi angaacee........ co/t mmmm/vv /vi _bsse_aptccha
12224.uuree__in1 12 1__t__ff iilk
Some (outlier) tokens can also contain global information when the model is large enough and sufficiently trained (Darcet et al., 2024). However, these token positions are input-dependent, rendering them less reliable for purposes.

---

Published as a conference paper at ICLR 2026
Table 3: HarmBench (Mazeika et al., 2024) attack success rate (ASR) with direct request (DR), GCG (Zou et al., 2023), weight orthogonalization (Arditi et al., 2024) and our safety module negative intervention. All evaluations use the model's default system prompt if it exists. Results marked with * are taken from Arditi et al. (2024).
ATTACKER
DR
GCG
ORTHO
SAFETY MODULE (OURS)
DEFENDER
LLAMA-2
0.0
34.5*
22.6*
QWEN 7B
7.0*
79.5*
79.2*
78.0
GEMMA 7B
8.2
53.5
73.0
84.3
Inagelet
*ighte"
050 Min
Figure 9: SAMI results VIT-B/32. Left: 4 labels with the lowest/highest generalization error after intervention. Right: Average effect across all labels of ImageNet. With our discovered modules, under negative intervention, disable the recognizability also show the generalization error (green curve) on the full validation set.
the attacked target (orange curve). We
hypothesize that the label taxonomy plays role in this phenomenon, such that intervening in particular label leads the model to lose recognizability for similar entities. That said, our model works rather well on average across all labels. Overall, our intervention effectively suppresses the ViT's ability to classify the target label while retaining its generalization capacity. In Appendix G, we repeat the experiments on ViT-L/16, 6,, and confirm that the observations could scale up to larger models.
5 CONCLUSION AND DISCUSSION
In this paper, we propose the first concept-agnostr attention head attribution pipeline in general transformer-based large models, including LLMs and ViTs. Our work uniquely highlights attention head modules as underexplored yet highly effective units for arbitrary concept attribution. Our pipeline enables broad applicability to arbitrary and complex concepts. Through SAMD, we experi- mentally demonstrate that knowledge is sparsely encoded in the large model structure, witnessed by the tiny modules we discover for all the concepts we study in this paper. Furthermore, our scalar attention module intervention (SAMI) allows direct control isolating and manipulating attention modules, our pipeline sheds light on the internal mechanism of transformer-basee models, offering new insights into the role of attention in concept attribution.
model prediction and generation. By
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
Our experimental results are fully reproducible. Our code based on the Transformerlens (Nanda and Bloom, 2022) and ViT-Prisma (Joseph, 2023) repositories, and all models considered in our paper are open-sourced, including the Llama series (Touvron et al., 2023; Dubey et al., 2024), Gemma series (Team al., 2024), Qwen series (Bai al., 2023), and ViT-B/32 (Dosovitskiy, 2020). In subsection 4.1, the sparse autoencoder features, interpretations, and text prompts are collected from Gemma Scope (Lieherum et al., 2024). In subsection 4.2, the text prompts are obtained from GSM8K (Cobbe et al., 2021) through in-evaluation-harness (Gao et al., 2024b). In subsection 4.3, we reuse text prompts from Arditi al. (2024) and evaluate the results using HarmBench (Mazeika al., 2024). In subsection 4.4, we evaluate the model on ImageNet-IK, KK. We release our code GitHub.
ACKNOWLEDGMENTS
JK thanks the Simons Foundation for support through the Collaborative Grant "The Physics of Learning and Neural Computation". JS and JK acknowledge support by the NSF through NRT Award 1922658.

---

Published as a conference paper at ICLR 2026
REFERENCES
Christoph Molnar, Giuseppe Casalicchio, and Bernd Bischl. Interpretable machine learning-a brief history, state-of-the-art and challenges. In Joint European conference on machine learning and knowledge discovery in databases, pages 417-431. Springer, 2020. 1
Pantelis Linardatos. Vasilis Papastefanopoulos, and Sotiris Kotsiantis. Explainable ai: A review of machine learning interpretability methods. Entropy, 23(1):18, 2020.
Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman. Deep inside convolutional networks: Visualising image classification models and saliency maps. arXiv preprint arXiv: 2013.
Jason Yosinski, Jeff Clune, Anh Nguyen, Thomas Fuchs, and Hod Lipson. Understanding neural networks through deep visualization. arXiv preprint arXi:: 2015. 1
Luisa M Zintgraf, Taco S Cohen, Tameem Adel, and Max Welling. Visualizing deep neural network decisions: Prediction difference analysis. In International Conference on Learning Representations, 2017. 1
Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based localiza- tion. In Proceedings of the IEEE international conference on computer vision, pages 618-626, 2017. 1
Mukund Sundararajan, Ankur Taly, and Qiqi Yan. Axiomatic attribution for deep networks. In International conference on machine learning, pages 3319-3328. PMLR, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
Chandan Singh, Jeevana Priya Inala, Michel Galley, Rich Caruana, and Jianfeng Gao. Rethinking interpretability in the era of large language models. arXiv preprint arXiv:2402.01761, 2024. 2
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 2021.
Mor Geva, Avi Caciularu, Kevin Wang, and Yoav Goldberg. Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 2022. 2
Yair Lakretz, Germán Kruszewski, Théo Desbordes, Dicuwke Hupkes, Stanislas Dehaene, and Marco Baroni. The emergence number and syntax units in Istm language models. In Proceedings gs of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 11-20, 2019. 2, 18
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in Neural Information Processing Systems, 35: 35:17759-17372. 2022a. 2, 4, 18
Kevin Meng, Arnab Sen Sharma, Alex Andonian, Yonatan Belinkov, and David Bau. Mass-editing memory transformer. arXiv preprint arXiv:2210..7229, 2022b. 2, 4, 18
Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: a circuit for indirect object identification in GPT-2 small. In The Eleventh International Conference on Learning Representations, 2023. URL pettps://opens view. net forum? iid www.iiiidiNNsooooo-
2, 18
Tilman Räuker, Anson Ho, Stephen Casper, and Dylan Hadfield-Mendd Toward transparent ai: A survey on interpreting the inner structures of deep neural networks. In 2023 ieee conference on secure and trustworthy machine learning (satml), pages 464-483. IEEF, 2023. 2
12

---

Published as a conference paper at ICLR 2026
Chunting Zhou, Pengfei Liu, Puxin Xu, Srinivasan Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, Lili Yu, et al. Lima: Less is more for alignment. Advances in Neural Information Processing Systems, 36:55006-55021, 2023. 2,
Mantas Mazeika, Long Phan, Xuwang Yin, Andy Zou, Zifan Wang, Norman Mu, Elham Sakhaee, Nathaniel Li, Steven Basart, Bo Li, David Forsyth, and Dan Hendrycks. HarmBench: A standardized evaluation framework for automated red teaming and robust refusal. In Pro- ceedings of the 41st International Conference on Machine Learning, 2024. URL https:
//proceeding ngg..inng.... ml r.presspes///ee. /v235/ htt//mazeiie ka24a. html. 3, 8, 10, 11
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers solve math word problems. arXiv preprint arXiv:2110.14168, 2021. 3, 7.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfiel-Doddss, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathemattcal framework for transformer circuits. Transformer Circuits Thread, 2021. https://taasfforee-ccccuisss wwtp:/cuuuuuttpo/22eeee /faaa/000./.aameeeoo//neeee.. htmml.
Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representa- tions in vector space, 2013. URL https: : / arxiv org/ abs ////1111.3770 3781.
Adly Templeton, Tom Conerly, Jonathan Marcus, Jack Lindsey, Trenton Bricken, Brian Chen, Adam Pearce, Craig Citro, Emmanuel Ameisen, Andy Jones, Hoagy Cunningham, Nicholas L Turne, Callum McDougall, Monte MacDiarmid, C. Daniel Freeman, Theodore R. Sumers, Edward Rees, Joshua Batson, Adam Jermyn, Shan Carter, Chris Olah, and Tom Henighan. Scaling monosemanticity:: Extracting interpretable features from claude sonnet. Transformer Circuits Thread, 2024. URL https: rann fo former-ciiinooooo www.iiircttooooooooo pwwppp22/404/ http///////caaar////////a scw./nalllngoooooooo emanticity innex.inx...//mm html.44 4, 18
Tom Lieberum, Senthooran Rajamanoharan, Arthur Conmy, Lewis Smith, Nicolas Sonnerat, Vikrant Varma, János Kramár, Anca Dragan, Rohin Shah, and Neel Nanda. Gemma scope: Open sparse autoencoders everywhere all at once on gemma 2. arXiv preprint arXiv:2408.05147 2024. 4, 5, 6, 11, 18
Curt Tigges, Oskar John Hollinsworth, Atticus Geiger, and Neel Nanda. Linear representations of sentiment in large language models. arXiv preprint arXiv:23 10. 15144, 2023.
Samuel Marks and Max Tegmark. The geometry of truth: Emergent linear structure in large language model representations of true/false datasets. arXiv preprint arXhv:2310.068) 2023. 4
Nina Rimsky, Nick Gabrieli, Julian Schulz, Meg Tong, Evan Hubinger, and Alexander Matt Turner. Steering llama 2 via contrastive activation addition. arXiv preprint arXiv:2312.06681, 2023.
Ole Jorgensen, Dylan Cope. Nandi Schoots, and Murray Shanahan. Improving activation steering in language models with mean-centring. arXiv preprint 2rXiv:23/2.0386 20233 4
Andy Arditi, Oscar Obeso, Aaquib Syed, Daniel Paleka, Nina Rimsky, Wes Gurnee, and Neel Nanda. Refusal in language models mediated by a single direction. arXiv preprint arXiv:2406.11717
2024. 4, 8. 10, 11
nostalgebraist. interpreting gpt: the logit lens. LessWrong, 2020. URL htthhtts://ae http://wwww.w.........../.. com.coomossteoooo /tte:///////ttaaaa////s//////ee 4
gpt-the- the logit-lens.
Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart Shieber. Investigating gender bias in language models using causal mediation analysis. Advances in neural information processing systems, 33:12388-12401, 2020. 4
Nicholas Goldowsky-Die Dill, Chris MacLeod, Lucas Sato, and Aryaman Arora. Localizing model behavior with path patching. arXiv preprint arX5992304.05969, 2023. 4
13

---

Published as a conference paper at ICLR 2026
Atticus Geiger, Zhengxuan Wu, Hanson Lu, Josh Rozner, Elisa Kreiss, Thomas Ieard, Noah Goodman, and Christopher Potts. Inducing causal structure for interpretable neural networks. In International Conference Machine Learning, pages 7324-7338. PMLR, 2022. 4
Neel Nanda and Joseph Bloom. Transformerlens. https: http::///gtthcb..ommaaaaaaaaeae
nsOrg Tearzgg/rrassfoooo wrr.annssfrrrrreens...m 2022. 5, 11
Sonia Joseph. Vit prisma: A mechanistic interpretability library for vision transformers. https:
//github. cwm/ son het.jsssp//ttppppep
2023. 5,
Yiran Zhao, Wenxuan Zhang. Guizhen Chen, Kenji Kawaguchi, and Lidong Bing. How do large language models handle multilingualism? arXiv preprint arXiv:2402.18815, 2024. 5, 6,
Chris Wendler, Veniamin Veselovsky, Giovanni Monea, and Robert West. Do llamas work in english? on the latent language of multilingual transformers. arXiv preprint arXiv:2402.10550 2024. 5, 6,
Zhengfu He, Wentao Shu, Xuyang Ge, Lingjie Chen, Junxuan Wang, Yunhua Zhou, Frances Liu, Qipeng Guo, Xuanjing Huang, Zuxuan Wu, et al. Llama scope: Extracting millions of features from 11mm-3.1188 with sparse autoencoders. arXiv preprint wrXiv:2410.20522 2024. 5, 18
Leo Gao, Tom Dupré la Tour, Henk Tillman, Gabriel Goh, Rajan Troll, Alec Radford, Ilya Sutskever, Jan Leike, and Jeffrey Wu. Scaling and evaluating sparse autoencoders. arXiv preprint arXiv:2406.0409: 2024a. 5, 6,
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based gemini research and technology. arXiv preprini arXiv:2403.08225 2024. 6, 8, 11
Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, Zac Hatfield-Dodee Alex Tamkin, Karina Nguyen, Brayden McLean, Josiah E Burke, Tristan Hume, Shan Carter, Tom Henighan, and Christopher Olah. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023. http:///waaoneeeee.....////////s put/2002////mooommmaicccccc features/index.htmll
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24887, 2022.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang. Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 07 2024b. URL https: /tt:://eeodd.o wwwp//eeodd//2/00eee0 7, 8. 11
Jian Liang, Ran He, and Tieniu Tan. A comprehensive survey on test-time adaptation under distribu- tion shifts. International Journal of Computer Vision, 133(1)311644, 2025. 7
Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. Commonsense@AA: A question answering challenge targeting commonsense knowledge. In Proceedings of the 2019 Conference of the North American Chapter the Association for Computational Linguisties: Human Language Technologies, Volume (Long and Short Papers). Association for Computational Linguistics, 2019,
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint urXiv:2107.03377 2021.
Jacob Austin, Augustus Odena, Maxwell Nye. Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai. Michael Terry, Quae Le, et al. Program synthesis with large language models. arXiv preprint 108.07322. 2021.
14

---

Published as a conference paper at ICLR 2026
Jiawei Liu, Chunqiu Steven Xia, Yuyao Wang, and Lingming Zhang. Is your code generated by chatGPT really correct? rigorous evaluation of large language models for code generation. In seirth- Conference Neural Information Processing Systems, 2023. URL https: // //openrev ew net taat/fooe forum? i qvx OCu7. 8
Lianmin Zheng, Wei-Lin Chiang. Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang. Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al. Judging llm-as-a-judge with mt-bench and chatbot arena. Advances Neural Information Processing Systems, 66:46595-46622 2023.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03877 2021. 8
Christopher Olah. Interpretability dreams. Transformer Circuits Thread, 2023. https:/ttransformer- citerispppl2200///meeecaaillyeed www.:p/wwwaaiii.dddreeenneeeeeaa
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint 2tt.::33770092888, 2023. 8, 11
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et al. Qwen technical report. arXiv preprint arXiv:2309.16600 2023, 8, 11
Andy Zou, Zifan Wang, Nicholas Carlini, Milad Nasr, J Zico Kolter, and Matt Fredrikson. Universal and transferable adversarial attacks on aligned language models. arXiv preprint 2,0. 2023. 8, 10
Alexey Dosovitskiy. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11922 2020. 11
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. In The Twelfth International Conference on Learning Representations, 2024. URL net/ woo.mmodddddnnn1......
pettps://operreet
Matthew D Zeiler and Rob Fergus. Visualizing and understanding convolutional networks. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Pari I 13, pages 818-833. Springer, 2014. 9
Abhimanyu Dubey, Abhinay Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint A4i2:2447.227888 2024. 11
David Bau, Bolei Zhou, Aditya Khosla, Aude Oliva, and Antonio Torralba. Network dissection: Quantifying interpretability of deep visual representations. In Proceedings of the IEEE conference
computer vision and pattern recognition, pages 6541-6549, 2017. 18
Bolei Zhou, David Bau, Aude Oliva, and Antonio Torralba. Interpreting deep visual representations via network dissection. IEEE transactions on pattern analysis and machine intelligence, 41(9):
2131-2145, 2018a.
18
Bolei Zhou, Yiyou Sun, David Bau, and Antonio Torralba. Interpretable basis decomposition for visual explanation. In Proceedings of the European Conference Computer Vision (ECCV), pages 1 199, 2018b. 18
David Bau, Jun-Yan Zhu, Hendrik Strobelt, Agata Lapedriza, Bolei Zhou, and Antonio Torralba. Understanding the role of individual units in deep neural network. Proceedings of the National Academy of Sciences, 117(49)3007113007 2020.
Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson. Dissecting recall of factual associations in auto-regressive language models. In The 2023 Conference on Empirical Methods in Natural Language Processing. 18
15

---

Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. 7oom in: An introduction to circuits. Distill, 2020. doi: 023311///isiill.0022...00 http:://iiiillppb///00vrrrr/z//o Published as a conference paper at ICLR 2026
Peter Hase, Mohit Bansal, Been Kim, and Asma Ghandeharioun. Does localization inform editing? surprising differences in causality- based localization vs. knowledge editing in language models. Advances in Neural Information Processing Systems, 1177444376688 2023. 18
Anshita Gupta, Debanjan Mondal, Akshay Sheshadri, Wenlong Zhao, Xiang Li, Sarah Wiegreffe, and Niket Tandon. Editing common sense in transformers. In Proceedings of the 2023 Conference on Empirical Methods Natural Language Processing, pages 8214 8232, 2023. 18
Chenmien Tan, Ge Zhang, and Fu. Massive editing for large language models via meta learning. In The Twelfth International Conference on Learning Representations, 2024. URL https: //ppnrevieww net Eorum? iit.uundddd-L11112222 wwwpreewwww..eeeooouneeeeeeeee
18
Xiaopeng Li, Shasha Li, Shezheng Song, Jing Yang. Jun Ma, and Jie Yu. Pmet: Precise model editing in transformer. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 18564-18572, 2024. 18
Yuheng Chen, Pengfei Cao, Yubo Chen, Kang Liu, and Jun Zhao. Journey to the center of the knowledge neurons: Discoveries of language-independen knowledge neurons and degenerate knowledge neurons. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 17817-17825, 2024. 18
Róbert Csordás, Sjoerd van Steenkiste, and Jürgen Schmidhuber. Are neural nets modular? inspecting functional modularity through differentiable weight masks. arXiv preprint arXiv:2010.02066, 2020. 18
Nicola De Cao, Leon Schmid, Dieuwke Hupkes, and Ivan Titov. Sparse interventions in language models with differentiable masking. arXiv preprint arXiv:21 (2.0:2377, 2021. 18
Damai Dai, Dong, Yaru Hao, Zhifang Sui, Baobao Chang, and Furu Wei. Knowledge neurons in pretrained transformers. arXiv preprint arXiv: 104006696 2021. 18
Haowen Pan, Yixin Cao, Xiaozhi Wang, Xun Yang, and Meng Wang. Finding and editing multi-modal neurons in pre-trained transformers. In Findings of the Association for Compulational Linguistics
ACL 2024, pages 1012-1037, 2024. 18
Junfeng Fang, Zac Bi, Ruipeng Wang, Houcheng Jiang, Yuan Gao, Kun Wang, An Zhang, Jie Shi, Xiang Wang, and Tat-Seng Chua. Towards neuron attributions in multi-modal large language models. Advances in Neural Information Processing Systems, 17:122877-1222890, 2024. 18
Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. Retrieval head mechanistically explains long-context factuality, arXiv preprint 1555 2024. 18
Rhys Gould, Euan Ong, George Ogden, and Arthur Conmy. Successor heads: Recurring, interpretable attention heads in the wild. arXiv preprint arXiv:2312.092331 2023. 18
Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey, Sparse autoen- coders find highly interpretable features in language models. arXiv preprint (xXiv:2309.08600 2023. 18
Leonard Bereska and Efstratios Gavves. Mechanistic interpretability for ai safety-a aafeey-arvview. arXiv preprint arXiv:2404.1 14082, 2024. 18
Lee Sharkey, Bilal Chughtai, Joshua Batson, Jack Lindsey, Jeff Wu, Lucius Bushnaq. Nicholas Goldowsky- Stefan Heimersheim, Alejandro Ortega, Joseph Bloom, et al. Open problems in mechanistic interpretability. arXiv preprint arXiv:2501.16496 2025. 18
16

---

Published as a conference paper at ICLR 2026
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitr, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. hawpiiga-----ndniiioo--aaaaiiiie
html. 18
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability.. The Eleventh International Conference on Learning Representations, 2023. URL https: http///pppeeeeeeeeee//e/eeeeeeee n forum? d= 9XF SKKFDbb,,d 18
Vivien Cabannes, Charles Amal, Wassim Bouaziz, Alice Yang, Francois Charton, and Julia Kempe. Iteration head: mechanistic study of chain-of -f-thuuhtt arXiv preprint arXiv:2406.02128, 2024.
18
Arthur Conmy, Augustine Mavor-Parker, Aengus Lynch, Stefan Heimersheim, and Adrià Garriga- Alonso. Towards automated circuit discovery for mechanistic interpretability. Advances in Neural Information Processing Systems, 36: 16:16618--166 16352, 2023. 18
17

---

Published as a conference paper at ICLR 2026
A LARGE LANGUAGE MODEL USAGE
We only use L.LMs to polish the text throughout the paper.
B RELATED WORK
In this section, we provide a literature review on existing attribution research and include the background information of sparse autoencoders that we have used in our paper. Finally, we discuss the relationship between our work and the field of mechanistic interpretability.
B.1 ATTRIBUTION AND LOCALIZATION
A long line of research has aimed to map model behaviors to specific model components, process referred to as attribution or localization, with early efforts focused on convolutional neural networks (Bau et al., 2017; Zhou al., 201 8a;b; Bau al., 2020). In the context of transformers, attribution studies have primarily focused on neurons (Meng et al., 2022a;b; Geva et al.; Hase et al., 2023; Gupta al., 2023; Tan al., 2024; Li al., 2024; Chen et al., 2024; Lakretz et al., 2019; Csordás et al.,
2020; De Cao et al., 2021; Dai et al., 2021; Pan et al., 2024; Fang et al., 2024). Attention heads, while fundamental to the transformer architecture, have received comparatively less attention. Existing research typically associates attention heads with narrowly defined behaviors, such as information retrieval (Wu et al., 2024) and simple tasks such as prediction of successors (Gould al., 2023). Whether attention head attribution can generalize to arbitrary concepts remains open question. Moreover, no existing framework provides a scalable, concept-agnostic approach for attention head attribution. Our work addresses these limitations. Furthermore, to knowledge, since the first to perform attention head attribution plus intervention using arbitrary concepts, there is no directly comparable work.
B.2 SPARSE AUTOENCODERS (SAES)
SAES (Cunningham et al., 2023) are autoencoders trained to reconstruct the representations from a specific layer of transformer. Prior work has shown that SAEs can discover cumman-understnddabb vector directions corresponding to arbitrary concepts, such as "the Golden Gate Bridge" (Templeton al., 2024), and can used steer model behavior via vector manipulation. Although open- source SAE weights have been released for various language models (Gao et al., 2024a; Lieberum al., 2024; He al., 2024), the associated human-internn internnntterrettblle feature annotations remain largely unavailable. For our study, we use the only publicly available set of SAE features with accompanying interpretations, released by Lieberum et al. (2024), and exclude abstract concepts that are difficult to evaluate even qualitatively (e.g.. "bravery", "humor").
B.3 MECHANISTIC INTERPRETABILITTY
Another line of research is mechanistic interpretability, whose primary role is to reverse engineer the mechanism of models (Olah al., 2020; Bereska and Gavves, 2024; Sharkey al., 2025). The core objective is to identify the critical computational graph (a.k.a. circuit) that is causally responsible for particular model behavior, with clear functional interpretations of each node and explicit edges connecting them. Prior work has uncovered several intriguing circuits, such as induction heads (Olsson al., 2022), circuits for indirect object identification (Wang et al., 2023), grokking (Nanda et al., 2023), and iteration (Cabannes et al., 2024), among others. Our work differentiates itself from this paradigm in several key aspects. First, we focus on attention head attribution. which is neither causal nor involving edge discovery. Second, the tasks and models involved in mechanistic interpretability studies remain relatively simple (Bereska and Gavves, 2024; Conmy et al., 2023), while pipeline operates general transformer models and supports the attribution of arbitrary concepts. Third, we introduce an attention module-base intervention strategy (Section 3.2), a direction not emphasized in mechanistic interpretability.
18

---

Published as a conference paper at ICLR 2026
C ATTENTION HEAD HEATMAPS FOR SAMD
In this section, provide the visualization of attention head heatmap, well the module discover using the cosine similarity score for the broad spectrum of concepts we consider the paper. All experiments are launched NVIDIA V100 and A100 GPUs.
C.I MODULES DISCOVERED USING SAE FEATURES
In Figure 10, and 12. we visualize the "dog". "yelling" and "San Francisco" attention module we discovered in GEMMA-2-2B-IT In Figure 13, we visualize the attention module discovered through the "Spanish" feature. The module is exactly the same as the one discovered through the "French" feature, suggesting the correlation to the multilingualism capability of the LLM.
0.6
0.4
0.2
7
10 12
15
17
18
10
21
22
23
24
25
Transformer Layer ID
Figure 10: "Dog" module: We determine the "Dog" concept vector Ve., and compare its cosine similarity to every attention head in GEMMA-2-2B-11 The black bounding boxes mark the most important heads, and hence define the module.
Attention
10
10
12
15
16
18
20
23
24
25
26
Transformer Layer
Figure 11: "YELLING" module: We determine the "YELLING" concept vector Ve., and compare its cosine similarity every attention head in GFMMA-2-2B-IT. The black bounding boxes mark the 5 most important heads, and hence define the module.
C.2 MODULES DISCOVERED USING SAE FEATURES, BASE MODEL
In Figure 14, 15, 16 and 17, visualize the "dog". "yelling", "San Francisco" and "multilingual" modules we discovered GEMMA-2-25 2--22. Compared to the modules GEMMA-2-2B-IT, only one attention head in the "yelling" module changed, with all other attention heads remains the same. This suggests the modules do not change after LLM post-training, and thus supports the superficial alignment hypothesis.
C.3 CHAIN-OF- THRUOHH-TOOOUUUUTREEEEORTR
MODULE
In Figure 18 and 19, we visualize the Chain-of-Thougeee reasoning module we discovered in 1LAMA3.1-8BB-INSS and GEMMA-7B-RASE.
19

---

Published as a conference paper at ICLR 2026
15 18
23
24
25
26
Transformer Layer ID
Figure 12: "San Francisco" module: We determine the "San Francisco" concept vector th and compare its cosine similarity to every attention head in GEMMA-2-2B-IT. The black bounding boxes mark the most important heads, and hence define the module.
0.3
0.2
Attention
0.1
0.0
-0.1
-0.2
7
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
Layer
We determine the "Spanish" concept vector Ue, and compare its cosine similarity to every attention head in GEMMA-2-25 hence define the module. This module is exactly the same as the one discovered through the "French" concept vector, and thus suggesting its correlation with multilingualism capability.
Figure 13:
The black bounding boxes mark the most important heads, and
C.4 SAFETY MODULE
In Figure 20 and 21, we visualize the safety module we discovered in QWEN-7B-CHAT and GEMMA- 7B-IT.
C.5 VIT MODULE
In Figure 22, we visualize the "Tabby Cat" module we discovered in VIT-B/32.
0.6
0.4
Attention
10
-0.
10 12
13
15
16
18
20
23
24
26
Transformer Layer
Figure 14: "Dog" module: We determine the "Dog" concept vector ce. and compare its cosine similarity every attention head in GEMMA-2-2B. The black bounding boxes mark the 5 most important heads, and hence define the module.
20

---

Published as a conference paper at ICLR 2026
0.2
0.1
0.0
15 Transformer Layer ID
23 24
25
26
Figure 15: "YELLING" module: We determine the "YELLING" concept vector ve and compare its cosine similarity to every attention head in GEMMA-2-2B. The black bounding boxes mark the most important heads, and hence define the module.
0.3
0.2
Attention Head
7
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
Transformer Layer
Figure 16: "San Francisco" module: We determine the "San Francisco" concept vector te. and
compare its cosine similarity to every attention head GEMMA-2-2B. The black bounding boxes mark the 5 most important heads, and hence define the module.
D NUMBER OF HEADS DETERMINATION
We plot the sorted AVG cosine similarity (i.e., the key quantity we use to plot the heatmap as well as to choose the top attention head nodes), to demonstrate that the hyperparameters 10) usually lead to selecting the most significant heads. The results on SAE concepts (Figure 23), safety (Figure 24), reasoning (Figure 25) and VIT concepts (Figure 26) provided.
choose (3, 5 and
Attention
0.0
10
-0.2
10 12
13
15
16
18
20
23
24
25
26
Transformer Layer
Figure 17: "Multilingual" module: We determine the "French" concept vector u.. and compare its cosine similarity every attention head in GEMMA-2-2B. The black bounding boxes mark the most important heads, and hence define the module.
21

---

Published as a conference paper at ICLR 2026
0.15
Cosire Similrity Concept Attention Hoad
0.10
14
Attention Head
18 17 16
13
70
-0.05
10 11
14:15
16
21
32
-0.10
Transformer
Figure 18: CoT module: We determine the CoT concept vector V,, and compare its cosine similarity to every attention head in 1--------ISTUCY The black bounding boxes mark the most important heads, and hence define the module.
E FQUAD SAE SWEEP RESULT
Figure 27, we provide FR-EN ratio under sparse autoencoder vector steering. Our default SAMI outperforms SAE vector intervention, which requires extensive sweep to determine the optimal intervention coefficient.
F LIMITATIONS
We acknowledge the following limitations of our work.
We do emphasize causality for the module we identify through SAMD. This implies there could exist other forms knowledge encoding besides the attention head module in transformers have studied. Further, the gap between correlation and causation attention module could be either overcomplete or incomplete. We leave more fine-grained analysis future work.
We provide neither a thcoretically grounded sample complexity analysis nor a stability analysis to proposed SAMD. Though methodology is robust across broad spectrum concepts, models and modalities we have tested, answering these questions would be beneficial.
Nonetheless, work provides the first concept-agnostic pipeline for arbitrary concept attribution to attention heads, which we hope will enhance the understanding state-of-the-ar transformer- based large models.
22

---

Published as a conference paper at ICLR 2026
Attertion Head
-0.1
Attention Head
0.0
Cosine Similants
0.2
-0.3
10
20 21
25 23
27 28
Transformer Layer
Figure 19: CoT module: We determine the CoT concept vector Ue.. and compare its cosine similarity every attention head GEMMA-7B 7B-BASE, The black bounding boxes mark the most important heads, and hence define the module.
G RECOGNITION MODULE IN VIT-L
In this section, we expand our experiments Section 4.4 to a larger vision transformer, VIT-L/16'. We determine the "recognition" modules 5 attention heads by checking the average cosine similarity scores. In the intervention experiment, we repeat our finding that we can selectively disable the recognition targeted concept while largely maintaining the capability of recognizing other concepts, as illustrated in Figure 28.
H SAMI ON NEGATIVE ATTENTION HEAD MODULES
In this section, we explore the functionality of attention heads that strongly correlates with the concept vector negatively. To be specific, we repeat our experiments in Section 4.2 and Section 4.3, but use attention modules formed by attention heads with the lowest average cosine similarities to the concept vector.
For reasoning, we sweep over s from -1.0 to -200. We do not observe any improvement to the baseline. Further, on Gemma-78 Barmma-7B-Basee this intervention drops its accuracy to 0%, likely because an attention head the first layer included its module.
For safery, we use = -1.0 and 2.0. We do not observe any change with respect to model behavior against jailbreak attacks: the ASR Llama stays at 0%, while that of Qwen and Gemma oscillates between 4. 4.40%, 8. 8.18% and 15.1% 10.1% 1%.
Thus, we conclude that the negative heads do not correspond the concept of interest. Our speculation that, in the context of the concept vector have discussed, the opposite direction does not reliably correspond to the opposite meaning of certain concept. This suggests that a high absolute value with negative score does not reliably indicate the importance of that head with "negative concept".
I ROBUSTNESS OF SAMD WHEN USING LESS DATA
Our attention module discovery uniquely depends on the positive dataset. To validate its robustness, we launch SAMD use only half of the data we have used throughout our paper.
hhttps://hugginee co/t i mmmvvmiittt www.llarrgpppoo 6_224. a2graggggggnn22 Ft nik
23

---

Published as a conference paper at ICLR 2026
0.2
Cosire Similarity Concept and Attention Head
0.1
Attention Head 13 16
-0.0
0.2
14 16 19
32
Transformer
Figure 20: Safety module: We determine the safety concept vector e. and compare its cosine similarity to every attention head in QWEN-7B-CHAT. The black bounding boxes mark the most important heads, and hence define the module.
In Figure 29, 30, 31 and 32, visualize the "dog". "yelling", "San Francisco" and "multilingual" modules under this test. Compared to the original modules, only one attention head in the "yelling" module changed, with all other attention heads remaining the same.
In Figure 33 and 34, we visualize the Chain-of- Thoughh reasoning module under this test. Both modules are exactly the same the original ones.
In Figure 35, 36, and 37, we visualize the safety module under this test. Compared to the original modules, only one attention head the Qwen and Gemma modules changed, with all other attention heads remaining the same.
Together, we believe these experiments suggest the robustness of our SAMD process.
SAMI EXAMPLES OF THE MULTILINGUALISM MODULE
In this section, we provide examples of our intervention strategy, SAMI, applied to the "multilingual" module that we discovered through SAMD. In addition to the French example presented in the main text, we include results with Spanish and several out-of-distribution languages, such as Mandarin Chinese, German, and Arabic, as shown in Figure 38, 39, 40, and 41. We consistently find that our discovered module generalizes multiple languages, suggesting its utility in assisting LLM with processing multilingual inputs.
24

---

Published as a conference paper at ICLR 2026
Cosiiee Similantty Concept and Attention Head
0.1
Attention Head
-0.0
0.1
16
-0.2
23 25
Transformer Layer
Figure 21: Safety module: We determine the safety concept vector tes and compare its cosine similarity to every attention head in GEMMA-7B-IT. The black bounding boxes mark the most important heads, and hence define the module.
Codine
Transformer Layer
Figure 22: "Tabby cat" module: The module for "tabby cat" recognition in vision transformer unsurprisingly sits at the very end (last layers) of the transformer. More surprisingly, our VIT-B/32 model relies only attention heads for the task.
25

---

Published as a conference paper at ICLR 2026
Cosine Similarity Sorted, dog
Cosine Similarity Sorted, French
Castns Sinkartty
AVD Collte Sinkant
0.0
-022
-02
-03
Actention
Attendion
(a)
(b)
Cosine Similariy Sorted, San Francisco
Cosine Similarity Sorted, Yelling
0.2
AV1 Coare Sirtarty
AVC Coalne Sor
0.2
-0.1
-0.1
50
100
150
200
50
100
150
200
Aclandon Hiad Indox
Attendio Plead Indix
(C)
(d)
Figure 23: AVG cosine similarity plot of the SAE concepts. Red dots indicate the chosen attention heads (the most significant ones) to form the module in our paper.
Sattd
Corlta Serad
Simiarty Sanad
100
200
(a)
(b)
(c)
Figure 24: AVG cosine similarity plot of the "safety" concept. Red dots indicate the chosen attention heads (the most significant ones) to form the module in paper.
26

---

Published as a conference paper at ICLR 2026
Similanity Sorted, reasoning, gemma-7b
Similarity Sorted, ressoning,
0.00
02
Simbertyy
0.10
0.0
0.05
0.00
-02
-04
200
(a)
(b)
Figure 25: AVG cosine similarity plot of the "reasoning" concept. Red dots indicate the chosen attention heads (the most significant ones) form the module in our paper.
(a)
(b)
(c)
Teue Gestitt Satelyt Losenter
Dribur Netetyyy STT
Oode tedieds Sata WI, diak
(d)
Figure 26: AVG cosine similarity plot of ViT. Red dots indicate the chosen attention heads (the most significant ones) to form the module in our paper. We choose to report the results use the first random labels in the dataset, in addition the tabby cat label used in our paper. Labels from Lep-Left to Bottom-Right. Tabby Cat: Goldfish; Great white shark; Hammerhead: Tench; Tiger shark.
27

---

Published as a conference paper at ICLR 2026
16-
30
1h:D4
30
10
-25
(a)
(b)
Figure 27: SAE FR-EN ratio (%) over different intervention strengths. We start from the recom- mended strength, -122, following Neuronpedia (https: /www. neuronpedia org gemm a 2- it /steer). We determine the optimal intervention coefficient use a coarse search (a) followed by fine-grained search (b). The best SAE result still underperforms our "multilingual" module's performance.
Isoec
100
Figure 28: SAMI results on VIT-L/16. Left: 4 labels with the lowest/highest generalization error after intervention. Right: Average effect across all labels ImageNet. With our discovered modules. under negative intervention, we disable the recognizability on the attacked target (orange curve). We also show the generalization error (green curve) on the full validation set.
0.2
4
10
12
13
15
16
17
18
19
20
21
22
23
24
25
26
Transformer Layer ID
Figure 29: "Dog" module: We determine the "Dog" concept vector Ve. and compare its cosine similarity to every attention head use half of the original data. The black bounding boxes mark the most important heads, and hence define the module.
28

---

Published as a conference paper at ICLR 2026
0.3
0.2
0.1
0.0
15 18 Transformer Layer ID
Figure 30: "YELLING" module: We determine the "YELLING" concept vector ve., and compare its cosine similarity to every attention head use half of the original data. The black bounding boxes mark the most important heads, and hence define the module.
0.2
Head
0.1
-.00
7
10 11
12
13
14
15
16
17
18
20
21
23
24
25
26
Transformer Layer
Figure 31: "San Francisco" module: We determine the "San Francisco" concept vector Ve and compare cosine similarity to every attention head use half of the original data. The black bounding boxes mark the 5 most important heads, and hence define the module.
10
15
18
21
23
26
Transformer Layer
Figure 32: "Multilingual" module: We determine the "French" concept vector U.. and compare its cosine similarity every attention head use half of the original data. The black bounding boxes mark the most important heads, and hence define the module.
29

---

Published as a conference paper at ICLR 2026
0.15
Cosine Similarity Concept Attention
10 12 12 12
Attention Head ID
1817
19
20
23 22
24
22
-0.05
30
31
32
10 11 12
14:15
20
22
23
24
25
26
27
28
29 30 31 32
Transformer Layer
Figure 33: CoT module: We determine the CoT concept vector Ve, and compare its cosine similarity to every attention head in LLAMA3. 1-----B-ITTRRR use half of the original data. The black bounding boxes mark the most important heads, and hence define the module.
D.1
Attennion He
0.0
Similurity Concept
0.1
12
00.
13
-0.3
10
12
13
14
15
17
18
20
22
23
25
28
Transformer Layer
Figure 34: CoT module: We determine the CoT concept vector Ve, and compare its cosine similarity to every attention head in GEMMA- 7 use half of the original data. The black bounding boxes mark the most important heads, and hence define the module.
30

---

Published as a conference paper at ICLR 2026
0.3
0.2
Simiarity Concept and Attertion
0.1
Attentioon Head ID
10 18 17 16
21
22
23
-0,1
24
25
26
21
28
22
10 12 15
19 21 22
30 31
Transoommrr Layer ID
Figure 35: Safety module: We determine the safety concept vector ee and compare its cosine similarity to every attention head in LLAMA-2-CHAT-7B use half of the original data. The black bounding boxes mark the 5 most important heads, and hence define the module.
31

---

Published as a conference paper at ICLR 2026
0.2
Cosine Similarity Conc Adtention
0..
Attenton Head ID
20 19
23
24
22
31
32
9 10
11
12
14:15
16 18 19 20
22
23
25
26
27
28 29 30 31 32
Transformer Layer
Figure 36: Safety module: We determine the safety concept vector te., and compare its cosine similarity every attention head in QWEN-7B 7B-CHAT half of the original data. The black bounding boxes mark the most important heads, and hence define the module.
-..2
Cosime Similarity Concept Attention
0.1
Attention He
12
13
-0.2
10
12 13
15
17
18
20
22
23
25
28
Transformer Layer
Figure 37: Safety module: We determine the safety concept vector Ves and compare its cosine similarity to every attention head in GEMMA-7B-IT use half of the original data. The black bounding boxes mark the most important heads, and hence define the module.
32

---

Published as a conference paper at ICLR 2026
Prompt:
eate mensaje? Generation (no intervention): E1 idioma de eate mensaje es español.
goudl
Generation (negative intervention):
Chis nessage is in English.
Figure 38: "Multilingual" intervention on Spanish input: the negative intervention compels the model to reply in English.
Prompt:
Generation (no intervention):
B 2
Generation (negative
current prompt
input is English.
Figure 39: "Multilingual" intervention on Chinese input: the negative intervention compels the model reply in English.
Prompi: In Melcher Sprache ist diese Aufforderung formuliert?
Generation (no intervention): Die Aufforderung ist in Dautsch formuliert.
(negative intervention):
this prompt written English.
Figure 40: "Multilingual" intervention German input: the negative intervention compels the model to reply in English.
Prompt:
الليصي lif inl بي
Generation (no interveation):
jay بلطمy ترري & us الموي فرنج
Generation (negative Intervention):
This Le text based AI, so it doeea't t. have "language" in the Iooaa sense. An An Rowever, the Janguage used to communicate with me is English.
Figure 41: "Multilingual" intervention Arabic input: the negative intervention compels the model to reply in English.
33


Reviews:
[{"review_id": "mG01iFYFQD", "reviewer": "Reviewer_BQcd", "summary": "The authors propose SAMD and SAMI - Scalable Attention Module Discovery and Intervention. Given a feature vector $v_c$, SAMD discovers a set of $K$ attention heads whose output $a_{l,h}$ has on average high cosine similarity to $v_c$. They call this set of attention heads (which is a circuit in a way) a module.  Given a the circuit of $K$ attention heads, SAMI amplifies or suppresses the module by scaling the attention heads output $a_{l,h}$ by a scaling factor $s$ which they choose on a per problem basis using grid search. They demonstrate the effectiveness of their methods in 4 different experiments:\na) They find and steer modules that correspond to SAE features. This motivates the name - as they choose concepts, and search for relevant modules related to the concept. This is done by generating a dataset $D_p$ for which an SAE feature $v_c$ is highly activated and then running SAMD with $v_c$ and $D_p$.\nb) They find and steer a module that corresponds to reasoning capabilities in the model. They report improved scores on the GSM8K reasoning benchmark.\nc) They find and steer a module that corresponds to a refusal direction. The report improved attack success rate over orthogonalization of the refusal direction in two out of three cases. \nd) They find and steer a module that corresponds the classification of a given ImageNet target on ViT-B/32 21k, showcasing effectiveness of targeted unlearning of a single class.", "questions": "* Did the authors explore including highly negative cosine sim attention heads (e.g. Fig 24.) and flipping the $s$ value for these heads? If so, that did they find? Especially in Fig 24 gemma 7b one head seems to have the highest absolute alignment with -0.4 similarity while the highest positive alignment is only 0.3.\n* What is the rational behind the $D_p$ construction via the test samples of GSM8K? Is it that the prompts are explicitly encouraging to reason?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Novel and elegant method for circuit discovery\n* Clear presentation of the findings\n* Demonstration of effectiveness of method on a broad variety of applications over two different modalities. Especially in the vision literature this is addressing a research gap, as vision-circuit discovery remains under-explored.", "weaknesses": "* 4.2 the construction of $D_p$ is unclear from just reading the main body of the paper. \n* Concept figure should be improved\n\t* Font way to small\n\t* No order of panels provided\n\t* SAMI is not explained in the rightmost panel\n* Comparison to baseline such as e.g. difference in means is missing for 4.1, 4.2 and 4.4. If this concern is addressed appropriately I will improve my score.\n* In 4.2 the authors only report evals on the dataset that they used for construction. An OOD reasoning benchmark eval would be useful to evaluate the generality of the reasoning-module.\n* Only ViT-B 32 evaluated in 4.4). Experiments with at least ViT-L would be recommended as ViT-B often shows different behaviors from its bigger counterparts. If this concern is addressed I will improve my score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "As6sAuOEL6", "reviewer": "Reviewer_HhFz", "summary": "This paper introduces an attribution method (of model components, not the input) for explaining transformer-based models. Specifically, it identifies the most important attention heads across the model which are relevant to a concept of interest i.e. the French language. To do so, they identify the maximally close (in terms of cosine similarity) attention heads to concept vectors captured from a positive dataset. They then propose that the identified components of the model can be up- or downscaled to change model behavior, thus verifying their component extraction, and providing a real use case for the method.", "questions": "Did the authors consider a knowledge editing style benchmark?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Very interesting method with good novelty. Unlike many previous MLP neuron attribution approaches, this is the first I have seen which identifies attention concepts. This is an interesting and critical result with the stronger push for mechanistic interpretability and adjacent approaches in the modern XAI literature. \n\nWell written with extensive experimental results.  \n\nI think the simplicity of the approach is a benefit to its usability. I feel that I could replicate this with a few hours of work if I had access to the datasets.", "weaknesses": "Minor – plainly calling this an attribution method feels misaligned with the literature. Attribution methods often refer to input (feature) attribution. This is more aligned with neuron attribution. Perhaps it should be attention attribution but not to be confused with input attribution using attention weights/gradients. \n\nThere are not any true comparisons against other methods. It is hard to tell if this should be negative because it may be challenging to create a fair comparison against a similar MLP based method. I think it could have been possible to use knowledge editing benchmark.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xGZE3X9ulP", "reviewer": "Reviewer_WBy6", "summary": "This paper presents Scalable Attention Module Discovery (SAMD), a method for identifying attention heads associated with abstract “concepts” in transformers via cosine similarity, and Scalar Attention Module Intervention (SAMI), a one-parameter mechanism to amplify or suppress their influence. The approach is demonstrated across LLMs and ViTs, with qualitative interpretability results and modest quantitative effects on reasoning, safety, and vision benchmarks.", "questions": "How is a “concept” defined in this framework, and how can the method be considered “concept-agnostic” if it depends on concept-specific vectors?\n\nCan the identified attention modules be causally validated (e.g., via path patching or feature ablation)?\n\nHow robust are the discovered modules to randomization, different seeds, or model variants?\n\nCould quantitative measures (e.g., mutual information or probing accuracy) strengthen the claims?\n\nThe heatmap figures are difficult to distinguish between colors, which are somehow important (e.g., with bold borders) and which are not.\nPlease, change the color.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Introducing attention-head–level concept attribution is an original direction that is computationally light and easily applicable to diverse transformer architectures.\n\n - The same pipeline is used across text and vision models, suggesting potential generality and extensibility.\n\n - The paper contributes to ongoing efforts to connect internal transformer components to semantic behaviors, particularly through sparse, interpretable “modules.”", "weaknesses": "- The evaluation is dominated by qualitative visualizations and anecdotal examples. There are no robust statistical analyses, reproducible metrics, or causal validation to confirm that the discovered modules truly mediate the claimed concepts.\n\n - The use of cosine similarity as a proxy for conceptual alignment is not theoretically or empirically justified; results may reflect correlation, not causation.\n\n - Choices of K (number of heads) and s (scaling factor) appear arbitrary, tuned via small grid searches without sensitivity analysis.\n\n - The paper never clearly defines in what sense the approach is “concept-agnostic.” This weakens the interpretive and theoretical clarity of the contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)