1 Introduction↩︎

Modern deep neural networks (DNNs) exhibit a notable phenomenon: DNNs trained at different starting initializations frequently yield solutions with similar prediction performance [1]–[3]. Yet, this begs a neglected corollary question today: do current transformers learn radically different sets of intrinsic computational schemes that happen to perform similarly, or do transformers learn internal representations that are largely the same? [4]

Further, the application scenarios explored in current large language model explainability research (LLM XAI) tend to be limited to narrow tasks and lack generalizability. The extracted “LLM circuits" are rarely validated in unseen settings or tasks [5]. Some authors argue [6] that current AI systems are too complex and that a top-down interpretability approach that captures emergent properties is preferable to attempts to reverse-engineer models”neuron by neuron” or in small functional blocks of neurons called “circuits”. In other words, by focusing on representations as the primary units of analysis—rather than neurons or small circuits—this approach seeks meaning in patterns of activity across many neurons at a time. Furthermore, existing mechanistic research into transformers often relies on analyzing single attention heads in isolation, effectively cherry-picking learned modules without establishing their representativeness.

These observations can be seen through the lens of the “universality hypothesis” from the mechanistic interpretability literature [4], [7]. This hypothesis asserts that models learn the same or similar representational features and circuits across different models when trained on similar tasks. However, the strength of the similarities between these learned model components has yet to be systematically quantified, leaving the critical distinction between strong and weak universality unresolved. [4]. Hence, in this work we make the notion of universality experimentally accessible by investigating whether the model converges to the requisite representational space—alignment of the attention head’s induced representation, that is, the subspace it attends to and writes into—and therefore to the same functional behavior.

Over the last months and years, we have seen a flurry of advances in LLM circuit discovery methods. These methods involve both automated [8] and manually discovered circuits [9], [10]. These methods make sweeping statements about the functional roles of particular heads in an LLM based on the patterns emerging in their attention score matrices [11]–[14]. However, we know that multiple theoretical solutions can exist for the same problem. For example, transformers have been shown to model “backup" circuitry that is utilized only when the primary circuit is compromised. And proxy models for interpretability such as transcoders and sparse autoencoders do not learn canonical concepts, even if they have high reconstruction fidelity and are monosemantic [15].

We face a similar problem in the context of LLM circuits and internal representations in general, where we might have different attention head representations for a retrained model but with different parameter initialization (i.e. a different seed). This hinders the application of LLMs in safety-critical areas like finance, nuclear, or healthcare, where the cost of a single point of failure might be huge [16]. We cannot rely on features or circuits extracted from a domain-specialized LLM as explanations of the domain’s intrinsic mechanisms, because these structures may be unstable across refits. For example, circuits identified in a model trained on disease-specific data should not be treated as a definitive account of the disease without demonstrated consistency across multiple random initializations of the same model. Without these replications, a seemingly “explainable" feature impacting diagnosis may not have any bearing on the true underlying biological mechanisms, and may in fact be entirely spurious [17].

It is therefore crucial to study the “seed stability” of LLM attention heads, the components most critical in circuits. We ask: if two different teams use the same LLM architecture and the same data, do they end up with the same attention heads? While this may be a straightforward question to articulate, comparing two refits of the same model is not trivial. There exist permutation symmetries of hidden units [18]; that is, one can swap any two units of a hidden layer in a network and, assuming weights are adjusted accordingly, network functionality will not change. Despite the same or similar output prediction performance, these permutations of parameter spaces can cause various transformations of the internal activation space of complex LLMs, which makes the task of comparing among refits even more challenging [19]. Cross-seed comparisons made using methods like CKA, CCA, and SVCCA [20]–[22] are not flawless and can be misleading [20], [23].

Given this, in the present paper, we compare attention head functions using their attention score matrices, which encode token-token relationships on a common basis and are therefore directly comparable across seeds [24]. This choice yields a simple, linear, and easily attributable metric for assessing the seed stability of attention heads.

Our main contributions and findings based on the experiments detailed in Methods (3) and Results (4) sections are the following:

• Middle-layer instability: Attention heads in middle transformer layers are the least stable across refits and the most representationally distinct.

• Depth dependence: Instability increases with the depth of the transformer model; alternative instances of deeper models exhibit stronger mid-depth divergence across refits.

• Functional implication of instability: With increasing depth of the layer, unstable heads become increasingly influential within their respective layer.

• Effect of weight decay: Using AdamW (decoupled weight decay) substantially improves attention-head seed stability across model refits.

• Residual-stream robustness: Across model refits, the residual stream is relatively much more stable than the attention heads.

From our comprehensive experiments, we discover that the seed stability of a GPT-like architecture is clearly tied to the various choices we make while deciding upon the training configuration of the model (e.g., optimiser, weight decay). These choices as such might have no appreciable impact upon the performance of the final trained model but can have a large impact on the universality of representational structures across model instantiations. This in turn has a cascading effect on the practical usefulness of these representational structures in the broader context of explainability.

2 Related Work↩︎

Prior studies exploring variation across refits of Transformer models have largely relied on handpicked task subsets or output-level metrics to assess stability. Similar to our work, these works train multiple seeds of transformer-based architectures, but their emphasis is typically on downstream behavior rather than internal alignment. For example, [25] trains multiple Pythia variants across nine seeds and primarily assess downstream performance stability and the representational stability of specific linguistic features via probing. [26] uses MultiBERTs (five BERT refits)[27] to conduct longitudinal analyses across training checkpoints, focusing on the evolution of an internal task subspace. Similarly, [28] performs a checkpoints-wise longitudinal study over six seeds for three Parc families (Parc-Pythia, Parc-Mamba, Parc-RWKV). Notably, most of these efforts are contingent on architectures other than the decoder-only transformer defined in [29]. While [30] reports training five GPT-family seeds at multiple scales, the paper does not analyze internal representations across refits.

By contrast, our methods performs inter-seed comparison in GPT-like (decoder-only) architectures spanning scales from 2-layer (7M parameters) to 12-layer GPT-2-small (124M parameters; [31]). Unlike prior LLM-XAI work, we compare internal representations directly and in a task-agnostic manner by operating on attention score matrices. To the best of our knowledge, no previous work has compared attention heads with their representations across seeds in decoder-only transformers.

3 Methods↩︎

3.1 Architectures and training details↩︎

We study attention head similarity in variants of GPT-like (decoder-only) language models. To isolate the intrinsic seed-stability behavior, we hold most training choices fixed at standard defaults, varying only fundamental architectural factors—depth (2, 4, 8) and heads per layer (8 or 16). Following [11], we also trained identical attention-only variants of the above architectures. In addition to this, we also trained a GPT2-small architecture with 12 layers and 12 heads configuration.

Our priority was attribution of seed stability in the context of each model’s overall architecture. Many hyperparameters such as learning rate schedules, warm up lengths, etc, were kept constant across refits and architectures. This allowed us to isolate and study the seed-stability as a direct outcome of the architecture itself rather than a combination of minor hyperparameter choices.

The only exception to this was the mirror set of identical variants trained with AdamW [32] optimizer instead of Adam [33], which is motivated by the “norm growth" rationale discussed in Results section 4.4.1. To probe the effect of over- and under-training regimes on seed stability [34], we also varied the pretraining corpus and size—the 12 layered GPT2-small models were trained on a bigger 9 billion tokens of OpenWebText [35] dataset, whereas the 2, 4, and 8 layered models were trained on a smaller 2 billion token subset of C4 [36], [37] dataset.

In total, we focused on 26 architectures. We trained 50 refits each for 2, 4, and 8 layered variants and 5 refits each for 12 layered (GPT2-small) architectures. Associated code and complete training/architecture configuration templates are provided in App. 7.2,  7.3 and the Supplementary Material. Consistent results across these architectural & dataset variations (4) strengthen our analysis & interpretation.

3.2 Attention-head stability across refits↩︎

We quantify a concrete notion of stability of an attention head by measuring how consistently it can be matched to a head in the same layer of independently trained refits of the same architecture. Let there be \(N\) refits (same architecture, hyperparameter set and data, yet different random seeds).

Setup and notation↩︎

• Let \(\mathcal{P}\) be a fixed set of prompts.

• Choose an anchor refit \(m\) and a pair refit \(n\) (\(n \neq m\)).

• \(h_{i}\) is a head belonging to layer \(l\) in anchor refit \(m\).

• \(h_j\) can be any head belonging to layer \(l\) in pair refit \(n\).

• All possible \(h_j\) heads are considered as candidate heads to be compared & matched with \(h_{i}\) head.

3.2.1 Step 1: Prompt-wise head similarity↩︎

For a head \(h_i\) of layer \(l\) in anchor refit \(m\) and one of its candidate heads \(h_{j}\) in pair refit \(n\), we define the prompt-wise similarity between \(h_i\) and \(h_j\) as the cosine similarity between their vectorized (flattened) post-softmax attention score matrices for prompt \(p \in \mathcal{P}\), that is (\(\!\big(A^{m}_{h_{i}}(p)\big), \!\big(A^{n}_{h_{j}}(p)\big)\)), from heads \(h_i\) and \(h_j\) respectively: \[s^{(m,n)}_{(h_{i},h_{j})}(p) = CosineSim\left( \operatorname{vec}\!\big(A^{m}_{h_{i}}(p)\big), \operatorname{vec}\!\big(A^{n}_{h_{j}}(p)\big)\right) \label{eq:eq1}\tag{1}\]

3.2.2 Step 2: Average similarity score across prompts↩︎

We aggregate over prompts to obtain an average similarity scores between the two heads (\(h_i\) and \(h_j\)): \[\bar{s}^{(m,n)}_{(h_{i},h_{j})} = \frac{1}{|\mathcal{P}|}\, \sum_{p \in \mathcal{P}} s^{(m,n)}_{(h_{i},h_{j})}(p) \label{eq:eq2}\tag{2}\]

3.2.3 Step 3: Best-match head in the pair refit↩︎

By repeating Step 2 (Eq. 2 ), We compare head \(h_{i}\) against all possible candidate heads \(h_{j}\) in the pair refit (\(\{1,\dots,H^n\}\)). Then, we choose the best matching head based on their average similarity scores: \[s^{(m,n)}_{h_{i}} = \max_{ h_{j} \in \{1,\dots,H^n\}} \;\bar{s}^{(m,n)}_{(h_{i},h{j})} \label{eq:eq3}\tag{3}\]

We interpret \(s^{(m,n)}_{h_{i}}\) as the stability of head \(h_{i}\) with respect to the single refit \(n\). For instance, a representative plot for an 8-layer architecture Fig. 7.5 visualizes the stability of all heads from an anchor refit with with respect to other individual pair refits.

3.2.4 Step 4: Stability over all refits↩︎

Finally, we average the single-refit stability (Eq. 3 ) across all pair refits to obtain the overall stability of head \(h_i\) across all \(N\) refits: \[S^{(m)}_{h_{i}} = \frac{1}{N-1}\, \sum_{n \ne m} s^{(m,n)}_{h_{i}} \label{eq:eq4}\tag{4}\] Interpretation: \(S^{(m)}_{h_{i}}\) indicates the propensity of anchor head \(h_i\) to be matchable to a head in the same layer with a similar attention-head representation across independently initialized refits. In this way, we obtain a principled, permutation-invariant notion of attention head stability that obviates the need for Hungarian matching [18].

Next, we can repeat this procedure for every head in the anchor refit \(m\).

3.3 Cross-layer best-match variant↩︎

We also consider a relaxation of the 3.2 procedure that drops the layer constraint when identifying a best match. For an anchor head (\(l,h_i\)) in refit \(m\), the candidate set includes all heads from all layers (rather than only layer \(l\)) in the pair refit \(n\). That is, \(h_j\) can be any head belonging to any layer in pair refit \(n\). The rest of the method remains identical.

4 Results↩︎

In this section we present the experimental results that underpin our main contributions. We discuss ten key results that probe cross-seed stability from complementary perspectives. For brevity, we report results only for representative architectures in the main body; the complete set of similar assessments for other model architectures is provided in App. 8.

4.1 Head-wise and layer-wise stability↩︎

Head-wise stability: Using the metric defined in 3.2 and a 100-prompts set, we compute the stability of every head in a single anchor instance drawn from different refits of a given architecture. The results allow for a bird’s eye view of cross-seed stability: many refits are compared against one anchor within the same architecture. For an 8-layer, 8-head MLP architecture (Fig. 2), we consistently observe a higher cross-seed stability in the earlier and later layers, with a pronounced dip in stability measures corresponding to attention heads in the middle layers of the architecture. At the head level, profiles vary: some layers are uniformly stable (or uniformly unstable) across seeds, whereas others exhibit patterns with variance.

Layer-wise stability: From the ensuing head-wise estimates, we compute a layer-wise mean by averaging the stability of all heads within a layer. Here, \(h_i \in H\), where \(H\) denotes the set of all heads in layer \(l\), and \(m\) denotes the anchor refit. \[S^{(m)}_{l} = \frac{1}{|H|}\,\sum_{h_i \in H}S^{(m)}_{h_{i}} \label{eq:layer95stab}\tag{5}\] Fig. 2 highlights the mid-depth instability. For example, \(S^{(m)}_{5}\) (layer 5) drops to \(\approx 0.70\), down from 1.0 in layer 1. This finding highlights the inherent challenge in identifying truly universal circuits in the intermediate layers of transformer models. Comparable figures for other architectures are provided in the 8.1.

4.2 Cross-layer best-match stability↩︎

Relaxing the layer constraint: As noted in 3.3, we also compute a cross-layer best-match stability by allowing each anchor head to match with any head of any layer in the pair refit. As expected, the resulting stability is slightly higher that the stability calculated with layer constraints ([results:head95layer95stab]), simply due to the larger candidate pool. However, overall trend was similar. Representative plots for a few of the architectures are provided in the App. 8.2.

Alignment Map & Middle-layer dispersion: To visualize “Which layer do the best matches come from?", we construct a row-normalized alignment matrix (H) (shown in Fig. 3) with entries (\(H_{i,j}\)). The y-axis denotes the anchor head’s layer index (\(i\)); the x-axis denotes the layer index (\(j\)) of the best-matched heads. Rows sum to 1. The heatmap exhibits strong diagonal dominance, that is, most heads match to a similar layer in paired refits, confirming a strong layer wise correspondence across seeds.

Compared to early and late layers, middle layers show a broader, flatter row distribution (lower kurtosis), with substantial off-diagonal mass into neighboring layers. This cross-layer spillover is also consistent with mid-depth instability and indicates a less sharply defined layer-to-layer correspondence for those heads.

4.3 Within-layer uniqueness of attention heads↩︎

Method: For anchor refit \(m\), head \(h_i\) in layer \(l\) is compared to all other heads in the same layer using attention score matrices over the prompt set \(\mathcal{P}\) (as in 3.2). Additional details are provided in App. 7.6. The mean similarity of an attention head to its peers quantifies the commonness (non-uniqueness) of the head.

Result & Interpretation: As shown in Fig. 4 and 8.3, mid-layer heads are most unique; early/late layers are much more prototypical—mirroring the head stability pattern described in [results:head95layer95stab].

4.4 Factors promoting attention head instability↩︎

Building on the layer-wise stability analysis in [results:head95layer95stab], we sought simple, model-internal correlates of head stability. We focused on parameters that directly shape the attention score matrices, namely query, key, value, and output projection matrices, as well as the pre-attention LayerNorm. Thus, we conducted an exploratory analyses (CKA-based summaries, diagnostic visualizations, and layer-wise statistics) on these objects. Two of the consistent signals that emerged were related to (i) \(\ell_{2}\) norm of the query weights and (ii) prompt length.

4.4.1 Correlation between Query-weight norm and layer-wise stability↩︎

We compared layer-wise stability from [results:head95layer95stab] to the mean (\(\ell_2\)) Frobenius norm of the attention heads’ query-weight matrices of the same layer. We observe a clear negative correlation—layers with larger mean query-weight norm tend to exhibit lower stability. As shown in fig. 8.5, this trend holds across architectures. This observation is especially intriguing in light of prior work linking parameter (\(\ell_2\)) norm growth to training instabilities and related failure modes [38], [39]. We reproduce these findings in the context of circuit identification and interpretability.

4.4.2 Effect of prompt length on stability↩︎

We next recomputed the layer-wise stability using multiple prompt sets whose length spanned 5-50 tokens per prompt. We find that stability decreases sharply with the prompt length, indicating that longer prompts induce greater token-token interaction variability and reduce cross-seed alignment (Fig. 5).

4.5 Stability comparison: Adam vs AdamW↩︎

Motivation: If cross-seed instability partly arises from uncontrolled norm growth, as explored in 4.4.1, then explicitly regularizing norms during pretraining may provide a means to improve head stability. Prior work shows that AdamW (decoupled weight decay) aids convergence, generalization, and norm control [32], [40], [41]. Given these findings, it is possible that AdamW also improves the universality of circuit identification across model refits. We therefore trained mirror sets of refits that are architecturally identical to the Adam runs but use AdamW as the optimizer.

Results: Studying Figures 6 and App. 8.6, we see that AdamW yields a clear increase in head seed-stability for deeper MLP models (8- and 12-layer) and for all attention-only architectures. In contrast, head seed-stability gains for shallow MLP models (2- and 4-layer) are small, possibly due to limited norm growth at low depth.

Performance Parity: Despite the stability gains, replacing Adam with AdamW typically leaves validation perplexity nearly unchanged, indicating that there are no performance tradeoffs when using AdamW. See  8.7 and [42].

Mechanistic check: To verify that AdamW controls norm growth [32], we compared the mean (\(\ell_2\)) norm of attention-head output activations between matched Adam and AdamW refits. Using AdamW consistently suppresses activation norms, supporting the hypothesis that controlling norm growth improves the cross-seed stability. (See Fig. 8.6.1).

4.6 Most- and least-stable layers↩︎

4.6.1 Stability gap between the most- and least-stable layers↩︎

Definition: For each of the 26 architectures (3.1), we compare layer-wise stability \(S_{l}\) ([results:head95layer95stab]) and identify the most-stable layer \(l_{\max}\) and least-stable layer \(l_{\min}\). We visualize these two layer stability extremities (\(S_{l_{\max}}, S_{l_{\min}}\)) for an anchor refit per model architecture.

To quantify the difference in attention head stability across layers in a single model architecture, we define the stability gap as:\[\Delta S = S_{l_{\max}} - S_{l_{\min}}\label{eq:stab95gap}\tag{6}\] Results: (Fig. 7) Across architectures, \(\Delta S\) widens with depth, most prominently in 8- and 12-layer models, and is modestly larger for 8-head than for otherwise comparable 16-head models. Using AdamW consistently narrows \(\Delta S\) relative to Adam, indicating that decoupled weight decay reduces cross-seed variability within the network itself.

4.6.2 Position of the most- and least-stable layer relative to architecture depth↩︎

Definition: For each of the 26 architectures (3.1), we locate the most-stable layer (\(l_{\max}\)) and the least-stable layer (\(l_{\min}\)) from the layer-wise stability curve for an anchor refit ([results:head95layer95stab]). We calculate their relative depth (\(r_l\)). Here, \(l\) is the layer index in \((1,.,l,..,L)\) and \(L\) is the total number of layers in the model. \[r_l = l/L \label{eq:rel95pos}\tag{7}\] Results: (Fig. 8) Across nearly all configurations—aside from the inherently shallow 2-layer cases—the least-stable layer occurs after the most-stable layer and, in 8- and 12-layer models, occurs around mid-depth (\(r \approx\) in range \([0.4,0.8])\), providing evidence that instability peaks in the middle of the network. For example, in the 8-layer, 8-head MLP (Adam) anchor refit (Fig. 8), the most-stable layer is (\(\ell=1\) with \(r=0.125\)) while the least stable layer is (\(\ell=5\) with \(r=0.625\)). Optimizer choice (Adam vs. AdamW) has no significant effect on these positions.

4.7 Stability of residual stream↩︎

Motivation: Attention heads exhibit notable cross-seed variability ([results:head95layer95stab]), especially at mid-depth, yet refits still produce similar functional behavior. This suggest that downstream representations in the model’s residual stream may be more stable than the individual head representations that feed into it, providing more stable ground for possible universal circuits discovery.

Method: Using a set of 100 prompts \(\mathcal{P}\), we compute layer-wise residual-stream stability between an anchor refit and a pair refit. For each layer, we use a set of residual stream activations after the attention sub layer is added to the residual stream (i.e. post-attention residual addition). We use these activations to compare anchor and pair refits by utilizing Centered Kernel Alignment (CKA). For CKA, we form centered Gram matrices for each representation set and compute their normalized similarity [20], [43]. CKA is invariant to isotropic rescaling and orthogonal transforms, but not to arbitrary affine or nonlinear transforms; thus it can underestimate similarity when refits differ by transformations outside the scope of CKA. Further details about the methodology used for this result are provided in App. 7.7.

Results: Across architectures and optimizers, the residual stream is consistently more stable than the corresponding attention heads (Fig. 9;  8.8; [20]). Optimizer choice has no major effect on residual-stream stability.

Interpretation: Despite head-level divergence, the objective and shared residual pathway (via skip connections) might nudge the residual stream toward seed-robust representations, helping explain “why independent refits yield similar outputs despite differing attention patterns"; consistent with [44]–[46].

4.8 Layer-wise correlation between the stability and post-ablation change in perplexity↩︎

Method: For each layer \(l\), we compute the Pearson correlation between head stability (\(S_{h_i}\)) and functional importance, quantified by post-ablation change in perplexity (\(\Delta \mathrm{PPL_{h_i}}\)); additional details are provided in App. 7.8.

Results: Although per-layer correlations are small and somewhat noisy, the aggregate trend is negative across layers (Figure 10 and  8.9).

Interpretation: These results suggest unstable heads are much more important in later layers for model performance than unstable heads in the earlier layers.

4.9 Geometric observation of attention head activations using meta-SNE↩︎

Objective: We sought a geometry-aware view of attention heads across architectures and refits by embedding their activation patterns into a shared space using meta-SNE [47]. Each point denotes one head from one refit, across architectures; see App. 7.9.

Analysis: (Fig. 11) When colored by relative depth (See 4.6.2), we observe a clear pattern where embedded heads are clustered by relative depth. This ordering reveals that the functional role of an attention head is determined more by its relative position in the computational stack than by specific architectural configurations (e.g. absolute number of layers).

5 Conclusion↩︎

Our comprehensive evaluation of the “seed stability" of transformer architectures draws a new boundary for mechanistic interpretability: the internal computational schemes (circuits) of LLMs are not as universal as their output prediction performance may suggest. We here demonstrate that structural universality of attention heads is a matter of continuous degree that varies by computational depth, optimizer choice during pre-training, and layer-wise uniqueness.

In particular, our results reveal a “stability dip" in middle layers of transformer architectures, where attention heads are the least stable across model refits, and also most unique. Paradoxically, we find that these unstable heads in deeper layers become more functionally influential, suggesting that the most critical components for model performance may be the hardest to replicate across instances. Still, there is hope for the isolation of more universal circuits in transformer models. First, the residual stream demonstrates higher cross-seed stability than individual attention heads, pointing towards the shared residual pathway as an important consolidator of universal interpretable functional units. Second, we find that the use of AdamW solvers substantially improves attention head stability with no model performance penalty. Simple adjustments to model training can hence have outsized impact on improving the universality extracted interpretable components from transformers.

Without establishing robustness across model instances, reported LLM “circuits" risk being stochastic artifacts rather than robust accounts of a model’s intrinsic mechanisms. Our findings open the door for more principled monitoring techniques, urging the community to move beyond single-instance analysis of transformer models and toward a more rigorous, stability-tested framework for AI safety.

Limitations: Our focus was to study and understand the mechanisms underlying trends in seed stability in LLMs. At the outset, we did not attempt an exhaustive hyperparameter sweep. The only exception was the inclusion of optimizer variants (Adam vs AdamW) for the reasons noted above. However, a study exploring and analyzing the relation between seed stability and a more comprehensive set of hyperparameters is left for future work. Budget and resource considerations also guided our choice of model sizes and variants.

6 Impact section↩︎

Our work has direct bearing for isolating and characterizing universal circuits in transformers and LLMs. These results may yield positive societal outcomes through an enhanced ability to build safe and human understandable AI. Conversely, as for any fundamental mechanistic interpretability research, it is possible that this research may be used to design more sophisticated jailbreak or hijacking attacks on LLMs, or may be used to evade LLM alignment. On net, we expect this interpretability work to have a positive impact for society.

Appendix Index↩︎

• A. Methods

• B. Results

7 Methods↩︎

7.1 Datasets↩︎

We train on two corpora:

• C4 (2B token subset): All models with 2, 4, or 8 layers are pretrained on a 2-billion token subset of C4. The dataset is available at [36]

• OpenWebText (9B tokens): All 12-layer models are pretrained on 9 billion tokens on OpenWebText, an open-source replication of OpenAI’s WebText used for GPT-2. The dataset is available at [35]

7.2 Architecture details↩︎

All architectures are decoder-only Transformers. Each architecture is trained across multiple independent seeds (refits). We consider both with-MLP and attention-only variants. We used TransformerLens API [48] to instantiate the models. Architecture configuration card templates for variants with a given number of layer (NUM_LAYERS) and number of heads per layer (NUM_HEADS) are provided below (Listing [lst:mlp] and  [lst:attn_only]). GPT2-small template has mostly similar configuration with a different tokenizer (Figure [lst:gpt_2]).

In total we trained 26 architectures. The architectural details are provide in the 1. We trained 50 independently initialized refits (random seeds) for every architectures, except GPT2-small types, for which we trained 5.

The supplementary material, related code and model weights for all trained refits, is provided in the repositories listed below.

• Code: https://github.com/karanbali/attention_head_seed_stability

• Weights: https://huggingface.co/karanbali/attention_head_seed_stability

{
'act_fn': 'gelu',
 'attention_dir': 'causal',
 'attn_only': False,
 'attn_scale': np.float64(8.0),
 'attn_scores_soft_cap': -1.0,
 'attn_types': None,
 'checkpoint_index': None,
 'checkpoint_label_type': None,
 'checkpoint_value': None,
 'd_head': 64,
 'd_mlp': 2048,
 'd_model': 512,
 'd_vocab': 48262,
 'd_vocab_out': 48262,
 'decoder_start_token_id': None,
 'default_prepend_bos': True,
 'device': device(type='cuda'),
 'dtype': torch.float32,
 'eps': 1e-05,
 'experts_per_token': None,
 'final_rms': False,
 'from_checkpoint': False,
 'gated_mlp': False,
 'init_mode': 'gpt2',
 'init_weights': True,
 'initializer_range': np.float64(0.035355339059327376),
 'load_in_4bit': False,
 'model_name': 'GELU_8L512W_C4_Code_8H',
 'n_ctx': 1024,
 'n_devices': 1,
 'n_heads': [NUM\_HEADS],
 'n_key_value_heads': None,
 'n_layers': [NUM\_LAYERS],
 'n_params': 25165824,
 'normalization_type': 'LN',
 'num_experts': None,
 'original_architecture': None,
 'output_logits_soft_cap': -1.0,
 'parallel_attn_mlp': False,
 'positional_embedding_type': 'standard',
 'post_embedding_ln': False,
 'relative_attention_max_distance': None,
 'relative_attention_num_buckets': None,
 'rotary_adjacent_pairs': False,
 'rotary_base': 10000,
 'rotary_dim': None,
 'scale_attn_by_inverse_layer_idx': False,
 'seed': 2,
 'tie_word_embeddings': False,
 'tokenizer_name': 'NeelNanda/gpt-neox-tokenizer-digits',
 'tokenizer_prepends_bos': False,
 'trust_remote_code': False,
 'use_attn_in': False,
 'use_attn_result': False,
 'use_attn_scale': True,
 'use_hook_mlp_in': False,
 'use_hook_tokens': False,
 'use_local_attn': False,
 'use_normalization_before_and_after': False,
 'use_split_qkv_input': False,
 'window_size': None
 }

{
'act_fn': 'gelu',
 'attention_dir': 'causal',
 'attn_only': True,
 'attn_scale': np.float64(8.0),
 'attn_scores_soft_cap': -1.0,
 'attn_types': None,
 'checkpoint_index': None,
 'checkpoint_label_type': None,
 'checkpoint_value': None,
 'd_head': 64,
 'd_mlp': 2048,
 'd_model': 512,
 'd_vocab': 48262,
 'd_vocab_out': 48262,
 'decoder_start_token_id': None,
 'default_prepend_bos': True,
 'device': device(type='cuda'),
 'dtype': torch.float32,
 'eps': 1e-05,
 'experts_per_token': None,
 'final_rms': False,
 'from_checkpoint': False,
 'gated_mlp': False,
 'init_mode': 'gpt2',
 'init_weights': True,
 'initializer_range': np.float64(0.035355339059327376),
 'load_in_4bit': False,
 'model_name': 'GELU_8L512W_C4_Code_8H_Attn_only',
 'n_ctx': 1024,
 'n_devices': 1,
 'n_heads': [NUM\_HEADS],
 'n_key_value_heads': None,
 'n_layers': [NUM\_LAYERS],
 'n_params': 8388608,
 'normalization_type': 'LN',
 'num_experts': None,
 'original_architecture': None,
 'output_logits_soft_cap': -1.0,
 'parallel_attn_mlp': False,
 'positional_embedding_type': 'standard',
 'post_embedding_ln': False,
 'relative_attention_max_distance': None,
 'relative_attention_num_buckets': None,
 'rotary_adjacent_pairs': False,
 'rotary_base': 10000,
 'rotary_dim': None,
 'scale_attn_by_inverse_layer_idx': False,
 'seed': 4,
 'tie_word_embeddings': False,
 'tokenizer_name': 'NeelNanda/gpt-neox-tokenizer-digits',
 'tokenizer_prepends_bos': False,
 'trust_remote_code': False,
 'use_attn_in': False,
 'use_attn_result': False,
 'use_attn_scale': True,
 'use_hook_mlp_in': False,
 'use_hook_tokens': False,
 'use_local_attn': False,
 'use_normalization_before_and_after': False,
 'use_split_qkv_input': False,
 'window_size': None
 }

{
'act_fn': 'gelu',
 'attention_dir': 'causal',
 'attn_only': False,
 'attn_scale': np.float64(8.0),
 'attn_scores_soft_cap': -1.0,
 'attn_types': None,
 'checkpoint_index': None,
 'checkpoint_label_type': None,
 'checkpoint_value': None,
 'd_head': 64,
 'd_mlp': 3072,
 'd_model': 768,
 'd_vocab': 50257,
 'd_vocab_out': 50257,
 'decoder_start_token_id': None,
 'default_prepend_bos': True,
 'device': device(type='cuda'),
 'dtype': torch.float32,
 'eps': 1e-05,
 'experts_per_token': None,
 'final_rms': False,
 'from_checkpoint': False,
 'gated_mlp': False,
 'init_mode': 'gpt2',
 'init_weights': True,
 'initializer_range': np.float64(0.02886751345948129),
 'load_in_4bit': False,
 'model_name': 'gpt2',
 'n_ctx': 1024,
 'n_devices': 1,
 'n_heads': 12,
 'n_key_value_heads': None,
 'n_layers': 12,
 'n_params': 84934656,
 'normalization_type': 'LN',
 'num_experts': None,
 'original_architecture': None,
 'output_logits_soft_cap': -1.0,
 'parallel_attn_mlp': False,
 'positional_embedding_type': 'standard',
 'post_embedding_ln': False,
 'relative_attention_max_distance': None,
 'relative_attention_num_buckets': None,
 'rotary_adjacent_pairs': False,
 'rotary_base': 10000,
 'rotary_dim': None,
 'scale_attn_by_inverse_layer_idx': False,
 'seed': 4,
 'tie_word_embeddings': False,
 'tokenizer_name': 'gpt2',
 'tokenizer_prepends_bos': False,
 'trust_remote_code': False,
 'use_attn_in': False,
 'use_attn_result': False,
 'use_attn_scale': True,
 'use_hook_mlp_in': False,
 'use_hook_tokens': False,
 'use_local_attn': False,
 'use_normalization_before_and_after': False,
 'use_split_qkv_input': False,
 'window_size': None
 }

7.3 Training Configuration↩︎

Unless noted otherwise, training choices are held at standard defaults to isolate architectural effects; Adam (constant learning rate) is the baseline optimizer, with a mirror set trained using AdamW (decoupled weight decay). Training configuration templates are shown in (Listing [lst:train_cfg_normal], [lst:train_cfg_normal_wd], [lst:train_cfg_gpt2] and  [lst:train_cfg_gpt2_wd]). We define these templates using TransformerLens HookedTransformerTrainConfig; any hyperparameters not listed use the default values. More details can be found in training code provided in supplementary material.

{
  "batch_size": 5,
  "lr": 0.0001,
  "optimizer_name": "Adam",
  "weight_decay": None,
  "num_epochs": 1,
  "save_every": 66285
}

{
  "batch_size": 5,
  "lr": 1e-4,
  "optimizer_name": "AdamW",
  "weight_decay": 0.1,
  "num_epochs": 1,
  "save_every": 66285
}

{
  "batch_size": 5,
  "lr": 5e-05,
  "optimizer_name": "Adam",
  "weight_decay": None,
  "num_epochs": 1,
  "save_every": 66285
}

{
  "batch_size": 5,
  "lr": 1e-4,
  "optimizer_name": "AdamW",
  "weight_decay": 0.1,
  "num_epochs": 1,
  "save_every": 66285
}

7.4 Prompts Sets & Generation↩︎

We used two prompts resources:

• Primary set (100 prompts): Used for all experiments except 4.4.2. This set was synthesized using ChatGPT (GPT-5.2 Thinking) model; provided at [app:fig:prompts951].

• Length-sweep set (6 x 20 prompts): Six subsets derived from the same 20 base prompts by truncation, yielding target length of 5, 10, 20, 30, 40, and 50 tokens per prompt. This set was used for 4.4.2 and is shown in [app:fig:prompts952].

Actual prompt sets are provided within supplementary material.

7.5 Stability of head \(h_i\) w.r.t other individual pair refit↩︎

12 visualizes per-head stability with respect to 50 individual pair refits for an 8-layer, 8-head MLP model. It shows a series of heatmaps, one per layer, for all heads in the anchor refit (seed = 1). The x-axis indexes pair refits, so each cell reports the stability of a single head against a single pair refit. See 12.

7.6 Within-layer uniqueness of attention heads↩︎

Methodology to calculate commonness (non-uniqueness) of an attention head \(h_i\) (See 4.3):

Setup and notation↩︎

• Let \(\mathcal{P}\) be a fixed set of 100 prompts.

• \(h_{i}\) is a head belonging to layer \(l\) in a refit \(m\).

• \(h_j\) can be any head belonging to layer \(l\) in the same refit \(m\).

• All possible \(h_j\) heads are considered as peer heads to be compared with.

7.6.1 Step 1: Prompt-wise head similarity↩︎

For a head \(h_i\) of layer \(l\) in anchor refit \(m\) and one of its peer heads \(h_{j}\), we define the prompt-wise similarity between \(h_i\) and \(h_j\) as the cosine similarity between their vectorized (flattened) attention score matrices for prompt \(p \in \mathcal{P}\), that is (\(\!\big(A^{m}_{h_{i}}(p)\big), \!\big(A^{m}_{h_{j}}(p)\big)\)), from heads \(h_i\) and \(h_j\) respectively: \[C^{(m)}_{(h_{i},h_{j})}(p) = CosineSim\left( \operatorname{vec}\!\big(A^{m}_{h_{i}}(p)\big), \operatorname{vec}\!\big(A^{m}_{h_{j}}(p)\big)\right) \label{eq:unique:eq1}\tag{8}\]

7.6.2 Step 2: Average similarity score across prompts↩︎

We aggregate over prompts to obtain an average similarity scores between the two heads (\(h_i\) and \(h_j\)): \[\bar{C}^{(m)}_{(h_{i},h_{j})} = \frac{1}{|\mathcal{P}|}\, \sum_{p \in \mathcal{P}} C^{(m)}_{(h_{i},h_{j})}(p) \label{eq:unique:eq2}\tag{9}\]

7.6.3 Step 3: Commonness of a head↩︎

By repeating Step 2 (Eq. 9 ), We compare head \(h_i\) against each peer heads \(h_j \in \{1,\dots,H^m\}\). Averaging these head-to-head comparisons yields per-head redundancy score, quantifying how common (or unique) \(h_i\) is within its layer: \[Commonness^{(m)}_{(h_{i})} =\frac{1}{Number\;of\;heads}\, \sum_{h_j \in \{1,\dots,H^m\}} \bar{C}^{(m)}_{(h_{i},h_{j})} \label{eq:unique:eq3}\tag{10}\]

7.6.4 Step 4: Average redundancy of a layer↩︎

The average of calculated commonness from Step 3 (Eq. 10 ) for all the heads \(h_i \in \{1,\dots,H^m\}\) provides a sense of average redundancy within the layer: \[Commonness^{(m)}_{l} =\frac{1}{Number\;of\;heads}\, \sum_{h_j \in \{1,\dots,H^m\}} \bar{Commonness^{(m)}_{(h_{i})}} \label{eq:unique:eq4}\tag{11}\]

We can now plot average redundancy for all layers of the architecture at hand, which shows that middle layer attention heads are more unique, relative to those in early and later layers (See Fig. 4 and 8.3)

7.7 Stability of residual stream↩︎

We followed the CKA implementation provided in [20] to compare layer counterparts of different refits using their set of activations over the prompt set \(\mathcal{P}\). To compute centered Gram matrices we used an RBF kernel with a threshold of 1.0.

7.8 Layer-wise correlation between the stability and “post-ablation change in perplexity”↩︎

To probe implications of head instability, we conduct a head ablation study on the anchor refit. For each layer \(l\), we:

1. For all the heads \(h_i \in l\), we calculate the head’s stability (\(S_{h_i}\)).

2. For all the heads \(h_i \in l\), we measure the post-ablation change in perplexity.

We define the “post-ablation change in perplexity" as change in perplexity after ablating the head \(h_i\), that is, zero’ing the head \(h_i\) output during the forward pass. The perplexity was calculated based on a 100-prompts set 7.1: \[\Delta \mathrm{PPL_{h_i}}= \mathrm{PPL_{\text{ablated}}}(h_i) - \mathrm{PPL_{\text{baseline}}}\]

For each layer, we then compute the correlation between two sets containing the:

1. \(\{S_{h_{i}}\}_{h_i=1}^H\) : Stability for all the heads \(h_i \in\) layer \(l\)

2. \(\{\Delta \mathrm{PPL_{h_{i}}}\}_{h_i=1}^H\) : “post-ablation change in perplexity" for all the heads \(h_i \in\) layer \(l\)

7.9 Geometric observation of the attention head activations using meta-SNE↩︎

Method for calculating meta-SNE embeddings:

Following the approach laid out in [47], for each head \(h\), we represent its behavior by using its output activations over a fixed prompt set \(\mathcal{P}\). We compute a pairwise distance matrix \(D\) whose \(D_{i,j}\) entry reflect the dissimilarity between pairs of points within the representation space of head \(h\). The points are attention head \(h\) output activation for different prompts in \(\mathcal{P}\). We use Euclidean distance to calculate the distance matrix \(D\). Because meta-SNE consumes distances rather than raw coordinates, it is insensitive to many isometric transformations (common in representation spaces), preserving both local and global neighborhood structure. The resulting distance matrix is then embedded with t-SNE to produce a 2-D embeddings.

Related figures are provided in 4.9.

8 Results↩︎

8.1 layer-wise stability↩︎

Additional architecture-wise plots for [results:head95layer95stab]:

8.2 Cross-layer best-match stability↩︎

Representative plots for few of the architectures 4.2:

8.3 Within-layer uniqueness of attention heads↩︎

Representative plots for few of the architectures 4.3:

8.4 EFFECT OF PROMPT LENGTH ON STABILITY↩︎

Additional architecture-wise plots for 4.4.2:

8.5 Correlation between Query-weight norm and layer-wise stability↩︎

8.6 Stability comparison: Adam vs AdamW↩︎

Additional architecture-wise plots showing improvement in stability when AdamW optimizer is used while pre-training (4.5):

8.6.1 Mechanistic check↩︎

8.7 Performance Parity: Adam vs AdamW↩︎

Additional architecture-wise plots showing Performance Parity between Adam and AdamW variants as mentioned in 4.5. The validation perplexity was calculated using dataset "wikitext-2-raw-v1" [49]:

8.8 Stability of residual stream↩︎

Additional architecture-wise plots for 4.7:

8.9 Correlation between the stability and “post-ablation change in perplexity”↩︎

Additional architecture-wise plots for 4.8:

References↩︎