探究后继头
作者：Emmanuel Ameisen、Joshua Batson；编辑：Jack Lindsey
存在一些机制，例如归纳、前一个 token 追踪和复制抑制，这些机制通常集中在许多不同 Transformer 模型中的少量注意力头上。在此，我们复现了 Gould 等人（2023）的发现：语言模型包含少量后继头，这些头在许多有序序列（数字、星期、月份等）中实现后继功能——将 2 映射到 3，将 Wednesday 映射到 Thursday，将 G 映射到 H 等。在一个 18 层模型中，我们采用了四种互补方法来识别和分析这些注意力头，其中包括一种基于 ICA 的新方法。
权重检查：输出-值电路
我们首先采用了一种紧密遵循 Gould 等人的基于权重的方法。我们整理了一个包含数字（阿拉伯数字、罗马数字和英文单词）、字母表字母、星期几和月份（包括缩写和全称）的有序序列数据集。（所有这些都被分词为单个 token。）我们将每个有序序列输入模型，并在第一个 MLP 层之后捕获残差流（遵循 Gould 等人的做法，我们注意到嵌入层与第一个 MLP 输出之间的表征发生了急剧变化）；我们将其用作每个有序 token 的表征。对于每一层的每个注意力头，我们通过将固定的表征向量输入该注意力头的输出-值（OV）电路，对结果进行解嵌入，并将 logits 限制为其他有序 token，从而具体化一个 n_有序token × n_有序token 的矩阵。然后，我们通过统计输入 token 的有序后继作为最高有序输出的 token 占比，来为每个注意力头打分。尽管我们使用的是从早期层提取的 token 表征，但即使对于较后层的注意力头，这种方法也能给出合理的结果。
在得分最高的注意力头中，约 80% 的序数 token 主要映射到其后继 token。这对应了下方序数矩阵子集中一条显著的超对角线分量（红色表示更高数值）。我们还注意到一些块状结构：在得分上仅次于后继 token 的是其他有序列表中的等价 token（"1" 映射到 "2"，但也映射到 "two"、"second"、"February"），以及同一列表中的其他 token（"3"、"4"）。我们还注意到，同一列表中较早的 token 几乎总是被抑制（每个块的较低对角线通常为蓝色）。
第二个得分较高的注意力头中，约 60% 的序数 token 映射到其后继 token。但错误呈现块状结构，对应相似序数类别中的其他 token（数字形式的数字与文字形式的数字相互映射，日/周/月与日/周/月相互映射）。
我们还随机选取了另一个注意力头作为对比展示。
独立成分分析
受块状结构以及序数继承和类别归属同时出现在两个后继注意力头中的现象启发，我们使用独立成分分析在所有注意力头中寻找共同模式。具体来说，我们将上述序数 token 矩阵展平并堆叠，生成一个 n_head × n_ordinal² 的矩阵，并应用 ICA。我们选择 ICA 是因为我们预期成分载荷是非高斯的（它们应更稀疏且尾部更重），并且具有一定独立性（通过建设性干涉实现注意力头的叠加）。字典学习也是一种可行的方法，它倾向于更稀疏的特征载荷并强制成分权重为正，例如会生成归纳和复制抑制的独立因子（而在此处它们会被合并）。每个成分对应一种"元注意力头"，每个注意力头的行为则是各成分的加权和。
我们发现了一个似乎实现序数继承的成分，在相关类别之间存在一些交叉（如数字与英文数字词，或月份与其缩写），例如将 1 映射到 2/two/second。
我们还发现了一个带有某种类别渗透的归纳成分（将 Monday 映射到 Monday/Mon）。
最后，我们重点介绍了一个将信息投射到某个类别上的成分（在多个成分中之一）：该成分将所有英文数字和阿拉伯数字映射到接近均匀混合的 0/1/2/3/4/one/two/three/four/first/second/third/fourth，并且在负方向上，使每个罗马数字抑制所有其他罗马数字。虽然这些描述并非完全清晰，但 ICA 揭示了跨注意力头反复出现的主题包括序列递进、归纳以及某种类别投射。
上述通过 OV 计数法识别出的排名靠前的注意力头，在序列递进成分上也具有最大的系数。
消融研究
某些注意力头可能具有较高的 OV 计数得分，但在实践中并未对序列递进做出贡献；可能的情况是，该注意力头并未关注序数 token，或者其输出不足以对 logit 产生显著影响，又或者其间接效应远大于直接效应，并且执行了序列递进之外的其他功能。另一种可能是，某个不同的注意力头通过间接效应，其贡献可能大于上述通过直接效应识别出的注意力头。为了探究这一点，我们对完整的序数序列（如“1 2 3 4 5 …”等）进行了前向传播，并计算了均值消融每个注意力头的输出对正确后继 logit 的影响；对所有序数序列和 token 取平均后，即得到该注意力头的“消融效应”。按成分得分排名前三的注意力头（与按 OV 投射得分排名前三的相同）在消融效应上分别排名第 1、第 3 和第 5。因此，仅凭权重识别出的这些注意力头在分布内是重要的。
我们注意到，成分得分较低但消融效应较高的两个注意力头（上述排名第 2 和第 4）位于比排名前三的序列递进注意力头（第 10、11、13 层）更低的层（第 3 和第 5 层）。这些注意力头可能通过与更上层的序列递进注意力头进行 Q 或 K 组合，或者通过影响更上层的 MLP，从而对序列递进做出贡献。
归因分析
我们还在相同的序列数据集上计算了注意力头归因（基于归因修补方法中的方法论），以检验这种更快速的方法是否与较慢的消融方法或基于权重的方法结果一致。虽然归因得分最高的注意力头在所有其他方法中也获得了最高分，但我们惊讶地发现，归因得分与其他三种方法之间的吻合度相对较低。
结论
我们复现了 Gould 等人的发现，即在小型 Transformer 模型中存在多个后继头，这些头通过直接效应促进序数 token 的递进。评估这些头在哪些上下文中被使用将会很有意义，例如，它们是否也参与列表递增、年份递增，或自然散文中的表达（"第二天，星期二，……"）。最后，我们注意到，在注意力头的 ICA 分解中发现的两种成分——一种大致实现了序数递进（但会部分遗忘具体序列），另一种实现了复制序列类别（但会遗忘具体序数）——与 Gould 等人发现的序数值与序列类别的组合线性表示是一致的。总体而言，如果这类注意力头分解方法能够用于识别残差流中的组合线性因子，那将是非常有趣的。
在 SAE 训练集中对某一主题进行过采样，会导致与该主题相关的特征更加细化
Trenton Bricken、Jonathan Marcus、Kelley Rivoire、Thomas Henighan；由 Adam Jermyn 编辑
在探索字典学习在安全相关应用的过程中，我们感兴趣的是在 Claude 3 Sonnet 中能找到哪些与生物武器相关的特定特征。我们发现，即使是我们基于 Sonnet 的最大规模 SAE（3400 万特征）也并未包含很多此类特征。为解决这一问题，我们将合成生成的生物武器数据集（以下简称"生物数据"）纳入 SAE 训练混合数据中。我们的假设是，这将促使 SAE 学习到更多与生物武器相关的特征。
最初，在 SAE 训练数据中未加入生物数据时，最相关的生物学特征主要涉及可用于制造危险或非法物品（如炸弹和违禁药物）的资源。虽然这些特征确实相关，但其范围比我们希望找到的更宽泛。
在将生物数据整合到 SAE 训练数据中后，我们观察到显著的转变。如今，用于预测有害生物提示词的最重要特征，聚焦于讨论病原体改造的文本示例，例如“增强病毒逃避免疫的能力”以及“存活并保持传染性”。另一个重要特征则围绕病毒颗粒的传播展开。
这些结果令人鼓舞，因为它们展示了一种解决字典学习中固有“特征覆盖”问题的潜在方案。我们的 SAE 尚未学会给定层中的所有可能特征，而要实现全面覆盖可能需要超出当前预训练能力的计算资源。然而，我们的发现表明，通过在训练集中策略性地对安全相关行为进行过采样，我们可以引导 SAE 学习更多与安全相关的特征。
Investigating successor heads
Emmanuel Ameisen, Joshua Batson; edited by Jack Lindsey
There are a few mechanisms, such as induction, previous token tracking, and copy suppression, which are often localized to a small number of heads in many different transformer models. Here we replicate the finding of Gould et al. (2023) that language models contain a small number of successor heads which implement succession in many ordinal sequences (numbers, days, months, etc.): mapping 2 to 3, Wednesday to Thursday, and G to H, etc. In an 18 layer model, we employed four complementary methods to identify and analyze these heads, including a novel method based on ICA.
Weight inspection: output-value circuit
We began with a weights based approach closely following Gould et al. We curated a dataset of ordinal sequences of numbers (arabic numerals, roman numerals, and english words), letters of the alphabet, and days of the week and months of the year (abbreviated and not). (All of these are tokenized as single tokens.) We pass each ordinal sequence through the model and capture the residual stream after the first MLP layer (following Gould et al., we noticed a sharp change in the representation between the embedding and the first MLP output); which we use as a representation for each ordinal token. For each head in each layer, we materialize the n_ordinal_token x n_ordinal_token matrix given by passing the fixed representation vectors through the Output Value (OV) circuit of the head, unembedding the result, and restricting the logits to other ordinal tokens. We then score each head by tallying the fraction of ordinal tokens for which the top ordinal output is the successor of the input. This gives sensible results even for heads in late layers, in spite of the fact that we are using token representations drawn from an early layer.
In the top scoring head, about 80% of the ordinal tokens are most mapped to their successor. This corresponds to a strong super-diagonal component in the subset of the ordinal matrix shown below (red denotes higher values). We also note some block structure; behind the successor in score is the equivalent token in other ordered lists ("1" maps to "2", but also "two", "second", "February") as well as other tokens in the same list ("3", "4"). We also note that earlier tokens in the same list are almost always suppressed (the lower diagonals of each block are usually blue).
The second scoring head has a ~60% rate of ordinal tokens being mapped to their successor. But the errors are block structured, corresponding to other tokens in a similar ordinal category (numbers in numerals or text with numbers in numerals or text, days/weeks/months with days/weeks/months).
We also show a randomly selected other attention head for comparison.
Independent Components Analysis
Intrigued by the block structure, and the fact that ordinal succession and category membership appeared in both successor heads, we looked for common motifs across all attention heads using independent component analysis. Concretely, we flattened and stacked the ordinal token matrices above, producing an n_head x n_ordinal**2 matrix, and applied ICA. We chose ICA because we expect component loadings to be non-Gaussian (they should be sparser and heavier tailed) and somewhat independent (allowing for attention-head superposition via constructive interference). (Dictionary learning, which would favor sparser feature loadings and enforce positive component weights, would also be a plausible approach, and would produce e.g. separate factors for induction and copy suppression which here would be merged.) Each component corresponds to a kind of "meta-head", and the behavior of each head is a weighted sum of the components.
We found one component which appears to implement succession, with some bleed between related categories (like numerals and english words for numbers, or months and their abbreviations), e.g. mapping 1 to 2/two/second.
We also find an induction component with some category bleed (mapping Monday to Monday/Mon).
Finally we highlight one of a few components which projects onto a category, here taking all english numbers and numerals to a close-to-uniform mix of 0/1/2/3/4/one/two/three/four/first/second/third/fourth and, in the negative direction, causing each roman numeral to suppress all roman numerals. While these descriptions aren't entirely clean, ICA reveals repeated motifs across heads include succession, induction, and some category projection.
The top heads identified by the OV tally method above also had the largest coefficients on the successor component.
Ablation Studies
Heads might have high OV tally scores without contributing to succession in practice; it could be that the head doesn't attend to the ordinal tokens, or that the output isn't large enough to make a significant difference to the logit, or that its indirect effects are much larger than its direct effects and do something other than succession. It's also possible a different head might, via indirect effects, have a greater contribution than the heads identified via their direct effect above. To examine this, we did forward passes on the full ordinal sequences ("1 2 3 4 5 …" etc.) and computed the effect of mean-ablating the output of each head on the correct successor logits; averaged over all ordinal sequences and tokens this is the "ablation effect" of the head. The top three heads by component score (which are the same as the top three by OV projection score) are ranked 1, 3, and 5 respectively by ablation effect. Thus those heads identified by weights alone are important on-distribution.
We note that the two heads with low component scores but high ablation effects (rank 2 and 4 above) are in lower layers (3 and 5) than the top 3 successor heads (layers 10, 11, 13). It is possible those heads contribute to succession by Q- or K- composition with the later successor heads or by influencing later-layer MLPs.
Attribution analysis
We also computed head attributions on the same sequential datasets (based on the methodology in Attribution Patching), to see if this faster approach agreed with the slower method of ablations or the weight-based methods. While the head with the highest attribution score also had the highest score according to all other methods, we were surprised to find relatively low agreement between attribution scores and all three other methods.
Conclusion
We reproduce the finding of Gould et al. that multiple successor heads exist in a small transformer model, which promote ordinal token succession through direct effects. It would be interesting to evaluate which contexts those heads are used in, e.g., are they also involved in incrementing lists, or years, or in natural prose ("The next day, Tuesday,..."). Finally, we note that two components found in the ICA decomposition of heads, which approximately implemented ordinal succession (but somewhat forgetting which sequence) and copy sequence category (but forget which ordinal) are consistent with the compositional linear representation of ordinal value and sequence category identified in Gould et al. It would be interesting if such head decompositions, in general, could be used to identify compositional linear factors of the residual stream.
Oversampling a Topic in the SAE Training Set Results in More Detailed Features Related to that Topic
Trenton Bricken, Jonathan Marcus, Kelley Rivoire, Thomas Henighan; edited by Adam Jermyn
As part of investigating safety-relevant applications of dictionary learning, we were interested in seeing what features we could find in Claude 3 Sonnet that were specific to bioweapons. We found that even our largest Sonnet-based SAE (34M features) did not contain many such features. To address this, we incorporated synthetically generated bioweapons datasets (hereafter referred to as "bio data") into our SAE training mix. Our hypothesis was that this would encourage the SAE to learn more bioweapons-related features.
Initially, without bio data in the SAE training mix, the most relevant biology features pertained to resources useful for creating dangerous or illegal items like bombs and recreational drugs. While relevant, this was broader than we hoped to find.
After integrating bio data into the SAE training mix, we observed a significant shift. The most important feature for predicting harmful bio prompts now focused on text examples discussing pathogen modification, such as "enhancing the virus's ability to evade" and "survive and remain infectious." Another important feature centered on viral particle dispersal.
These results are encouraging as they demonstrate a potential solution to the "feature coverage" problem inherent in dictionary learning. Our SAEs haven't learned all possible features in a given layer, and achieving comprehensive coverage might require computational resources that exceed current pre-training capabilities. However, our findings suggest that we can guide the SAE to learn more safety-relevant features by strategically oversampling safety-relevant behaviors in our training sets.