本文是一份初步更新，描述了我们仍在开发中的一种方法。因此，我们在此呈现的许多结果尚不完整。我们将其公开，是希望它能激发外部研究社区的后续工作和改进。
在发布此更新后，我们了解到同期有一项关于高度相关架构（“混合解码器”，Oldfield 等人）的研究。我们已在相关工作中更新了文章，描述了我们方法与该方法之间的异同。我们建议任何对此工作感兴趣的人也阅读他们的论文。
在最近的工作中，我们训练了转码器——稀疏、超宽的多层感知机——作为模型原始 MLP 层的更具可解释性的替代方案。我们利用转码器神经元（“特征”）作为理解模型计算的基础。我们使用“归因图”来描述计算，这些图展示了产生模型输出的特征之间的因果交互。
尽管这种方法已被证明非常有用，但我们认为转码器特征存在关键局限性。转码器将模型计算“打碎”成许多极其细粒度的片段。这可能导致它们以与底层模型不同且效率更低的方式表示计算，并可能引入特征分裂和吸收等病态问题。
举一个具体例子，我们发现转码器使用“查找表”特征来分解加法电路，这些特征代表了个别的、高度特定的计算，例如“6 + 9 = 5（模 10）”。其他研究表明，Transformer 架构将其数字表示嵌入到几何结构中，从而可以利用简单的变换来计算算术运算。虽然查找表对加法的描述在某种意义上很准确，但它遗漏了模型中的关键结构，并且效率极低。这种低效在实践中是一个严重问题——为了训练足够大的转码器以捕获模型中的所有计算，我们可能需要天文数字般的特征数量，使用一个参数比底层模型多得多的转码器。
在这份初步更新中，我们描述了一种我们一直在研究的方法，旨在解决其中一些担忧。我们的方法是用稀疏混合线性变换（MOLT）来替代 MLP。与转码器不同，MOLT 不会学习嵌入在模型激活空间中某个向量方向上的稀疏激活特征。相反，它学习的是稀疏激活的变换，这些变换对残差流施加一个线性变换，从而对 MLP 输出做出贡献。与转码器特征（同时充当计算单元和表征单元）不同，MOLT 变换是纯粹的计算对象，在各层表征之间起到“桥梁”作用。因此，它们不旨在被单独研究，而是与 SAE 等其他表征分解算法结合使用。
到目前为止，我们发现 MOLT 在替代 MLP 方面比转码器更具计算效率，且在机制上更忠实。我们发现 MOLT 变换被激活的条件与转码器特征具有相似的可解释性。我们的初步实验表明，这些变换可用于理解某一层的特征如何被转换为后续层的特征。MOLT 变换可以通过用“承载”该边的变换集合来注释图边，从而被纳入归因图中。我们还推测（但尚未证明）MOLT 变换的组合可用于理解模型表征中的组合结构。
实现
MOLT 的参数化如下：
f(x) = \sum_t [\phi(\mathbf{e_t} \cdot \mathbf{x} - b_t) * (U_t V_t \mathbf{x})]
其中
- \mathbf{x} 是输入的残差流激活值
- t 是变换的索引
- \mathbf{e_t} 和 b_t 是变换 t 的编码器向量和偏置，它们决定该变换是否被激活以及激活程度。\phi 是一个非线性函数，例如 ReLU 或 JumpReLU
- \(U_t\) 和 \(V_t\) 分别是维度为 \(d_{model} \times k_t\) 和 \(k_t \times d_{model}\) 的矩阵。因此，\(U_t V_t\) 是一个秩为 \(k_t\) 的矩阵。不同的变换可以有不同的秩（我们将秩的分布视为待优化的超参数）。
MOLT 的训练方式与转码器相同，旨在模仿模型 MLP 层的输出。我们对变换的激活值施加稀疏惩罚（例如 L1 或 tanh），并按变换矩阵的 Frobenius 范数进行缩放（即惩罚 \(\|U_t V_t\|_F \cdot \phi(\mathbf{e_t} \cdot \mathbf{x} - b_t)\)）。这与对转码器特征的激活值施加稀疏惩罚的方式类似，后者按解码器向量的范数进行缩放。
相关工作
混合解码器（MxD）架构（Oldfield 等人）与我们的架构非常相似——它用稀疏混合的线性变换替代了 MLP 层。MxD 与 MOLT 之间存在一个关键实现差异。使用许多独立学习的全秩变换的稀疏混合在计算上是不可行的，因为这会需要过多的参数。MOLT 通过使线性变换为低秩（具有不同秩的分布；见下文）来规避这个问题。MxD 允许变换为全秩，但通过一种在不同变换之间共享参数的方式对其进行参数化。
具体来说（将 MxD 论文的符号转换为我们所用），每个 \(U_tV_t\) 变换都是全秩的，但第 \(t\) 个变换的 \(U_t\) 矩阵的第 \(h\) 列被参数化为向量 \(\mathbf{d}_h\) 与 \(\mathbf{c}_t\) 的逐元素乘积，其中 \(\mathbf{d}\) 向量在不同变换之间共享，而 \(\mathbf{c}\) 向量在不同列索引之间共享。我们目前尚不确定 MxD 与 MOLT 策略各自的优缺点。我们推测，我们的低秩变换参数化方式会促使变换执行更具可解释性、范围更紧凑的计算角色，但需要更多工作来比较这两种方法。
我们还注意到，我们为 MOLT 设想的用例（将在后文描述）与 MxD 论文的侧重点有所不同——具体而言，我们尤其感兴趣的是将 MOLT 解释为在残差流特征（例如来自 SAE 的特征）之间实现（可能具有组合性的）变换，并将其整合到归因图中，将 MOLT 变换视为一种类似于注意力头的 MLP 对应物。
MOLT 还与跳跃式转码器（skip transcoders）有几分相似，后者通过线性变换增强了转码器；然而，两者在以下几个方面存在差异：
- MOLT 变换仅在条件满足时激活，并且会被惩罚以使其不频繁激活。
- MOLT 变换不是满秩的。
- MOLT 拥有多个线性变换，而非仅有一个。
- MOLT 没有任何“常规”的转码器特征。
直觉理解
为什么稀疏激活的变换可能是表示计算的一种好方法？其直觉理解是什么？我们提供几个视角。
- 首先，我们注意到 MOLT 与 MLP 进行计算的方式大致对应。考虑一个使用 ReLU 神经元的 MLP 层。如果你以激活的神经元集合为条件，MLP 就只是对其输入应用一个线性变换。因此，激活神经元的模式对应于线性变换，其秩等于该模式中的神经元数量。这些可以粗略地类比为 MOLT 变换。
- 稀疏激活的线性变换可以捕捉模型利用残差流几何结构执行计算的某些方式。如上所述，转码器需要“查找表特征”来选择特定的输入对以执行加法运算。相比之下，MOLT 可以更高效地实现加法。假设数字 0–9 被模型表示在一个圆上。一个秩为 2 的“加 3”变换可以通过绕圆旋转 8/10 圈来表示这一操作，并可以将一个“1”特征变换为下一层中的“说 4”特征，将“2”特征变换为“说 5”特征，依此类推。“加 3”变换的编码器会检查上下文中的文本线索，例如“+3”或“加三”。
- 另一种视角是，MOLT 是一种以组合方式参数化特征的方法。如果我们在第 L 层的残差流输入中发现了 100 万个特征，并在第 L 层的 MLP 中发现了 100 万个变换，那么我们就有效地描述了第 L+1 层中的 1 万亿个“伪特征”，这些伪特征对应于特征与变换的配对（当一个伪特征对应的特征和变换都处于激活状态时，该伪特征即为“活跃”，其“方向”是将变换矩阵应用于特征解码器的结果）。此外，如果我们考虑通过多个变换链传播的特征，那么我们所捕获的“伪特征”数量将随链长呈指数级增长。
实验结果
秩的最优分配
训练 MOLT 时的一个重要超参数是变换的秩分配。这是一个高维超参数空间，我们尚未对其进行充分探索。然而，在 Claude 3.5 Haiku 上训练 MOLT 时，我们通过使用从 32 到 512 不等的秩分布，获得了最佳性能（以 MSE/L0 帕累托前沿衡量）。具体来说，我们使用了一组 N 个秩为 512 的变换、2N 个秩为 256 的变换、4N 个秩为 128 的变换、8N 个秩为 64 的变换以及 16N 个秩为 32 的变换。为了扩大运行规模，我们改变 N 的值，但保持比例不变。我们发现，在控制总参数量的情况下，使用可变秩的变换优于使用全部相同秩的变换。
机器学习性能
我们在 Claude 3.5 Haiku 的中间层上训练了 MOLT（使用上述秩分配）和转码器，并改变了运行中使用的计算量。我们将训练步数与特征数量按比例缩放，并使转码器和 MOLT 运行的参数量相匹配。因此，每次 FLOPs 增加 4 倍，都对应参数量和训练步数各增加 2 倍。最大规模的（“1024 倍 FLOPs”）转码器运行包含约 1000 万个特征。
我们发现，在给定的 L0 下，控制参数量时，MOLT 的重建误差（MSE）显著低于 transcoder。这里最小的 MOLT 运行在帕累托效率上优于使用 1024 倍 FLOPs 的 transcoder 运行。此外，transcoder 的性能在更高的计算规模下似乎趋于饱和（尽管这种平坦化可能仅仅是由于机器学习调参不佳），而我们在批量运行中并未观察到 MOLT 有这种饱和现象。
我们还评估了 MOLT 与 transcoder 相比的机制忠实度——即 MOLT（或 transcoder）对输入扰动的响应方式与底层 MLP 层的匹配程度。在扰动大小趋近于无穷小的极限情况下，忠实度可以通过比较替换层与底层在给定数据点上的雅可比矩阵，并在多个数据点上取平均来计算。我们发现，在相同的 L0 下，MOLT 的雅可比相关性（展平后的雅可比矩阵的余弦相似度）远高于 transcoder；此外，transcoder 的忠实度似乎随规模增大而下降，而 MOLT 的忠实度则更为稳定。这种更高的忠实度是合理的，因为 transcoder 的雅可比矩阵被限制为低秩（秩的上限为 L0），而 MOLT 的雅可比矩阵的秩可以远高于其 L0。
请注意，以下结果来自与上述不同的模型（我们电路追踪论文中使用的 18 层模型，而非 Claude 3.5 Haiku）；我们尚未对 Haiku 进行这种雅可比分析。具有给定数量“特征等价物”的 MOLT 运行，其参数量与具有相同数量特征的 transcoder 运行相同。
Transformer 可解释性
Transformer 由两个特性来刻画：
- 它们在什么条件下被激活（以及激活程度如何）
- 它们在激活时执行什么功能
为了理解第一部分，我们可以使用与 SAE 和转码器特征相同的可视化策略——高亮显示激活该变换的数据集示例。当我们这样做时，我们发现变换在性质上与特征相似——我们看到在较早层中选择 token 级别信息的变换，以及在中间层和较后层中选择更抽象上下文信息的变换。
我们还观察到，更高秩的变换倾向于更高密度（即更频繁地被激活）。然而，这些更高密度、更高秩的变换条件在可解释性上似乎仍与较低密度的变换（以及转码器特征）相当。例如，我们观察到一个高秩变换会在句号 token 上激活，另一个则会在西班牙语文本上激活。
解释变换的功能则更为困难。我们最初的策略是在 MOLT 层之前和之后的残差流上训练 SAE，并识别通过给定变换相互作用最强的特征-特征对（\mathbf{e}_t UV \mathbf{d}_s，其中 \mathbf{e}_t 是目标特征编码器，\mathbf{d}_s 是源特征编码器）。然而，我们发现这些信息难以解释，推测是由于与干扰权重相同的问题，使得转码器特征之间的原始交互强度也难以解释。
我们在归因图的上下文中解释变换方面取得了更多成功，如下所述。
将 MOLT 集成到归因图中
方法
在我们最近的论文中，我们构建了基于转码器特征的归因图。特征之间的边是通过计算源特征（通过其解码器方向）对目标特征编码器方向（直接通过残差连接，或通过注意力头）施加的影响来确定的。
MOLT 变换没有固定的解码器方向用于写入，因此相同的归因图策略无法应用于它们。然而，我们在另一种归因图策略上取得了初步成功：
- 在模型的每个残差流层上训练 SAE（或采用跨层变体，如弱因果交叉编码器）。我们发现，图谱受益于跨层字典，但为简化说明，我们将以逐层 SAE 的情况为例进行解释。
- 对于每一对源特征与目标特征，将这两个特征之间的归因分解为两类项之和：
- 由 MOLT 变换（通过 MLP 层进行的归因）介导的项。对于解码器为 \mathbf{d}_i、激活值为 a_i 的源特征 i，编码器为 \mathbf{e}_j、激活值为 a_j 的目标特征 j，以及激活值为 a_t 的变换 UV，该项为 \mathbf{e}_j UV \mathbf{d}_i \cdot (a_i a_t)。
- 由注意力头（通过注意力层进行的归因）介导的项。对于解码器为 \mathbf{d}_i、在 token 位置 q 的激活值为 a_i 的源特征 i，编码器为 \mathbf{e}_j、在 token 位置 k 的激活值为 a_j 的目标特征 j，以及一个在 token 位置之间具有注意力模式 a_h(q, k) 且包含 OV 矩阵的注意力头，该项为 \mathbf{e}_j OV \mathbf{d}_i \cdot (a_i a_h(q, k))。
- 在归因图中的每条边上标注出对该边中介作用最强的 MOLT 变换和注意力头的列表。悬停在 MOLT 变换上会显示其“可视化”面板——一个展示其最高激活数据集示例的面板。（点击注意力头则会显示一个界面，展示其“QK 归因”，这将在后续更新中描述，待定）。
- 请注意，MOLT 变换本身也有输入边，其计算方式与上述相同，但使用该变换的编码器向量作为归因目标。
定性发现
利用此类归因图，我们发现了 MOLT 变换执行可解释计算的实例，例如：
- 在提示词“The Spanish word for hot is "calor"”中，我们观察到，“说出一个以‘cal’开头的单词”这一特征，通过一个在西班牙语相关语境中活跃的变换，接收到了来自“hot”特征的输入。
- 有趣的是，我们还看到一些来自“Spanish”特征的输入，这些输入是通过在“hot”相关语境中活跃的变换传递的！
- 在提示词“3 + 5 = 8”中，我们观察到，一个“说 8”特征通过一个在数字“5”近期出现过的上下文中被激活的变换，接收来自“加 3”特征的输入。
然而，我们也看到变换扮演着一些不那么容易解释的角色。例如，许多边似乎是由那些针对“is”或“Assistant”等关键词进行选择的变换所中介的。我们还看到变换扮演着与其所承载特征看似冗余的角色，例如一个“巴黎”特征通过一个“法国”变换接收来自“法国”特征的输入。
结论
我们认为 MOLT 是跨编码器的一种有前景的替代方案，并且可能能够以参数效率更高的方式捕获 MLP 计算，从而更忠实地反映底层模型所执行的计算。我们推测，要在一个合理的计算成本下捕获前沿模型 MLP 层的全部方差，一种类似 MOLT 的解决方案将是必要的。我们看到了 MOLT 变换能够执行可解释计算的迹象，即“变换”输入特征为输出特征。
我们感到兴奋但尚未探索的一个方向是，利用 MOLT 变换来理解那些未被我们的 SAE 捕获的组合表示。具体来说，我们 SAE 在每一层的重构误差可以重写为对应来自上一层的（特征，变换）对的项之和。这些项可能对应于那些因过于罕见而无法被我们有限大小的 SAE 捕获的概念，但它们是由一个相对常见的特征与一个相对常见的变换组合构建而成的。在归因图中，以这种方式将 SAE 误差分解为（特征，变换）对，会表现为某些图边由连续层中的变换链所中介。
要断定 MOLT 严格优于跨编码器，还需要做更多工作。包含变换信息的归因图比基于跨编码器的归因图要更笨重一些，并且并非所有由变换中介的计算都是清晰可解释的。未来在扩展和改进 MOLT 以及相关的归因图逻辑和用户界面方面的工作，可能会解决其中一些问题。
This work is a preliminary update describing a method we are still developing. As such, many of the results we present here are incomplete. We present them in the hopes that it inspires follow-up work and improvements from the external research community.
Upon publishing this update, we became aware of contemporaneous work on a highly-related architecture (“Mixture of Decoders,” Oldfield et al.). We have updated the post to describe the similarities / differences between our methods in the Related Work section. We recommend that anyone interested in this work read theirs as well!
In our recent work, we trained transcoders – sparse, extra-wide MLPs – as more interpretable replacements for a model’s original MLP layers. We used the transcoder neurons (“features”) as a basis for understanding model computation. We described computations using “attribution graphs,” which depict the causal interactions between features that give rise to the model’s outputs.
While this approach has proved very useful, we believe that transcoder features have key limitations. Transcoders “shatter” model computation into many extremely granular pieces. This can lead them to represent computations in a different, less efficient way from the underlying model, and can introduce pathologies like feature splitting and absorption.
To take one specific example, we found that transcoders decomposed addition circuits using “lookup table” features that represented individual, highly specific computations, like “6 + 9 = 5 (mod 10)”. Other work suggests that transformers embed their representations of numbers in a geometric structure such that simple transformations can be used to compute arithmetic operations. While the lookup table description of addition is accurate in some sense, it misses key structure in the model, and is highly inefficient. This inefficiency is a serious issue in practice – in order to train transcoders that are sufficiently large to capture all the computation in a model, we would likely require an astronomical number of features, using a transcoder with many more parameters than in the underlying model.
In this preliminary update, we provide a description of a method we have been working on that is intended to address some of these concerns. Our approach is to replace MLPs with a sparse mixture of linear transforms (MOLT). Unlike a transcoder, a MOLT does not learn sparsely active features that are embedded along a vector direction in the model’s activation space.Instead, it learns sparsely active transforms, which apply a linear transformation to the residual stream to give their contribution to the MLP output. Unlike transcoder features, which double as both computational and representational units, MOLT transforms are purely computational objects that “bridge” representations between layers. As such, they are not intended to be studied on their own, but in combination with another representation decomposition algorithm like SAEs.
So far, we have found that MOLTs are a more compute-efficient and mechanistically faithful way of replacing MLPs than transcoders are. We find that the conditions in which MOLT transforms are active are similarly interpretable to transcoder features. Our preliminary experiments suggest that these transforms can be used to understand how features in one layer are transformed into features in a subsequent layer. MOLT transforms can be incorporated into an attribution graph by annotating graph edges with the set of transforms that “carried” the edge. We also suspect (but have not yet shown) that compositions of MOLT transforms can be used to understand compositional structure in model representations.
Implementation
A MOLT is parameterized as follows:
f(x) = \sum_t [\phi(\mathbf{e_t} \cdot \mathbf{x} - b_t) * (U_t V_t \mathbf{x})]
where
- \mathbf{x} is the input residual stream activations
- t indexes the transforms
- \mathbf{e_t} and b_t are the encoder vector and bias for transform t, which determine whether the transform is active, and how active it is. \phi is a nonlinearity like ReLU or JumpReLU
- U_t and V_t are matrices of dimensions d_{model} \times k_t and k_t \times d_{model}, respectively. Thus, U_t V_t is a rank-k_t matrix. Different transforms may have different ranks (we treat the distribution of ranks as a hyperparameter to be optimized).
A MOLT is trained to mimic the output of a model MLP layer, just like a transcoder is. We apply a sparsity penalty (e.g. L1 or tanh) to the activations of the transforms, scaled by the Frobenius norm of the transform matrix (i.e. we penalize \|U_t V_t\|_F \cdot \phi(\mathbf{e_t} \cdot \mathbf{x} - b_t)). This is similar to how one applies a sparsity penalty to the activations of transcoder features, scaled by the norm of their decoder vector.
Related Work
The mixture of decoders (MxD) architecture (Oldfield et al.) is quite similar to ours – it replaces MLP layers with a sparse mixture of linear transforms. There is one key implementation difference between MxDs and MOLTs. Using a sparse mixture of many independently learned, full-rank transforms is intractable because it would require too many parameters. MOLTs get around this issue by making the linear transforms low-rank (with a distribution of different ranks; see below). MxDs allow the transforms to be full-rank, but parameterize them in a way that shares parameters across different transforms.
Specifically, (translating the MxD paper notation into ours), each U_tV_t transform has full rank, but the hth column of the U_t matrix for the tth transform is parameterized as an elementwise product of a vector \mathbf{d}_h \ast \mathbf{c}_t, where the \mathbf{d} vectors are shared across transforms and the \mathbf{c} vectors are shared across column indices. We are not yet confident in the pros and cons of the MxD vs. MOLT strategies. We suspect that our low-rank transform parameterization incentivizes transforms that perform more interpretable, tightly-scoped computational roles, but more work is needed to compare the two approaches.
We also note that our intended use-cases for MOLTs (described later in the post) are somewhat different than the focus of the MxD paper – in particular, we are especially interested in interpreting MOLTs as implementing (potentially compositional) transformations between residual stream features (e.g. from SAEs), and integrating them into attribution graphs, treating MOLT transforms as a kind of MLP analog to attention heads.
MOLTs also bear some resemblance to skip transcoders, which augment transcoders with a linear transform; however, they differ in several ways:
- MOLT transforms are only conditionally active, and penalized to be active infrequently
- MOLT transforms are not full-rank
- MOLTs have multiple linear transforms, rather than just one
- MOLTs do not have any “regular” transcoder features
Intuition
What’s the intuition for why sparsely active transforms might be a good way to represent computation? We offer a few perspectives.
- First, we note that MOLTs have a rough correspondence to how MLPs do computation. Consider an MLP layer that uses ReLU neurons. If you condition on the set of active neurons, the MLP is just applying a linear transformation to its input. Patterns of active neurons thus correspond to linear transformations, with rank equal to the number of neurons in the pattern. These can be thought of as loosely analogous to a MOLT transform.
- Sparsely active linear transforms can capture some ways in which the model leverages geometry of the residual stream to perform computations. As noted above, transcoders require “lookup table features” that select for specific pairs of inputs in order to perform addition. By contrast, a MOLT can implement addition more efficiently. Suppose the digits 0–9 are represented by the model on a circle. A rank-2 “plus 3” transform, could represent this by performing a rotation 8/10ths of the way around the circle, and could transform a “1” feature into a “say 4” feature in the next layer, a “2” feature into a “say 5” feature, and so on. The “plus 3” transform’s encoder would check for textual cues in the context like “+3” or “add three.”
- Another perspective is that MOLTs are a way to compositionally parameterize features. If we have uncovered 1 million features at the residual stream input to layer L, and uncovered 1 million transforms in layer L MLP, we have effectively described 1 trillion “pseudofeatures” in layer L+1, corresponding to pairs of features and transforms (a pseudofeature is “active” when its corresponding feature and transform are active, and its “direction” is the result of applying the transform matrix to the feature decoder). Moreover, if we consider features propagated through chains of multiple transforms, the number of “pseudofeatures” we capture grows exponentially in the chain length.
Results
Optimal allocation of ranks
An important hyperparameter in training a MOLT is the allocation of ranks to transforms. This is a high-dimensional hyperparameter space, and we have not explored it fully. However, when training MOLTs on Claude 3.5 Haiku, we have obtained our best performance (in terms of the MSE/L0 pareto frontier) using a distribution of ranks, varying from 32 to 512. Concretely, we use a collection of N transforms of rank 512, 2N of rank 256, 4N of rank 128, 8N of rank 64, and 16N of rank 32. To increase the scale of runs we vary N, but keep the proportions the same. We have found that using transforms of variable ranks outperforms using transforms of all the same rank, controlling for the total number of parameters.
ML performance
We trained MOLTs (using the rank allocation given above) and transcoders on the middle layer of Claude 3.5 Haiku, varying the amount of compute used in the run. We scaled the number of training steps proportionally to the number of features, and matched the number of parameters between transcoder and MOLT runs. Thus each 4× increase in FLOPs reflects a 2× increase in both number of parameters and training steps. The largest (“1024x FLOPs”) transcoder runs contain approximately 10 million features.
We find that at a given L0, the reconstruction error (MSE) is significantly lower for MOLTs than transcoders, controlling for the number of parameters. The smallest MOLT runs here Pareto-dominate transcoder runs that use 1024× as many FLOPs. Moreover, transcoder performance appears to be saturating at the higher compute scales (though it is possible this flattening is simply due to poor ML tuning), while we observe no such saturation for bulk runs.
We also evaluated the mechanistic faithfulness of MOLTs compared to transcoders – that is, the degree to which the MOLT (or transcoder) responds to input perturbations in the same way as the underlying MLP layer. In the limit of infinitesimal perturbation sizes, faithfulness can be computed by comparing the Jacobians of the replacement layer to the underlying layer on a given datapoint, and averaging over datapoints. We find that MOLTs have a substantially higher Jacobian correlation (cosine similarity of the flattened Jacobian matrices) than transcoders do, at the same L0; moreover, the faithfulness of transcoders appears to deteriorate with scale, whereas that of MOLTs is more stable. The greater faithfulness makes sense, given that the Jacobians of transcoders are constrained to be low-rank (rank upper-bounded by the L0), whereas MOLT Jacobians can have rank much higher than their L0.
Note that the results below are from a different model than those above (the 18-layer model used in our circuit-tracing paper, rather than Claude 3.5 Haiku); we have not yet performed this Jacobian analysis on Haiku. MOLT runs with a given number of “feature-equivalents” have the same number of parameters as a transcoder run with that many features.
Transform interpretability
Transforms are characterized by two properties:
- Under what conditions are they active (and how active)
- What function do they perform when active
To understand the first part, we can use the same visualization strategy we use for SAE and transcoder features – highlighting dataset examples that activate the transform. When we do so, we find that transforms appear qualitatively similar to features – we see transforms that select for token-level information in earlier layers, and transforms that select for more abstract contextual information in middle and later layers.
We also observe that higher-rank transforms skew higher-density (i.e. are active more often). However, the higher-density, higher-rank transform conditions still appear comparably interpretable to lower-density transforms (and to transcoder features). For instance, we observed a high-rank transform that activates on period tokens, and another that activates on text written in Spanish.
Interpreting the function of transforms is more difficult. Our initial strategy was to train SAEs on the residual stream prior to and immediately after the MOLT layer, and identify feature-feature pairs that most strongly interact via a given transform (\mathbf{e}_t UV \mathbf{d}_s, where \mathbf{e}_t is the target feature encoder and \mathbf{d}_s is the source feature encoder). However, we found this information difficult to interpret, presumably due to the same problem of interference weights that makes raw interaction strengths between transcoder features difficult to interpret.
We have had more success interpreting transforms in the context of attribution graphs, as described below.
Integrating MOLTs into attribution graphs
Method
In our recent paper, we constructed attribution graphs built on top of transcoder features. Edges between features were computed by determining the influence that a source feature (via its decoder direction) exerted on a target feature’s encoder direction (either directly, via residual connections, or via attention heads).
MOLT transforms have no fixed decoder direction that they write out to, so the same attribution graph strategy cannot be applied to them. However, we have had preliminary success with another attribution graph strategy:
- Train SAEs on each residual stream layer of the model (or a cross-layer variant, like weakly causal crosscoders. We find graphs benefit from cross-layer dictionaries, but we will explain the per-layer SAE case for simplicity).
- For each pair of source and target features, decompose the attribution between the features into a sum of terms of two kinds:
- Terms mediated by MOLT transforms (attributions via MLP layers). For source feature i with decoder \mathbf{d}_i and activation a_i, target feature j with encoder \mathbf{e}_j and activation a_j, and transform UV with activation a_t, the term will be \mathbf{e}_j UV \mathbf{d}_i \cdot (a_i a_t).
- Terms mediated by attention heads (attributions via attention layers). For source feature i with decoder \mathbf{d}_i and activation a_i at token position q, target feature j with encoder \mathbf{e}_j and activation a_j at token position k, and attention head with an OV matrix with attention pattern a_h(q, k) between the token positions, the term will be \mathbf{e}_j OV \mathbf{d}_i \cdot (a_i a_h(q, k)).
- Annotate each edge in the attribution graph with the list of MOLT transforms and attention heads that most strongly mediated that edge. Hovering over MOLT transforms surfaces their “viz” – a panel displaying their top-activating dataset examples. (Clicking on attention heads reveals an interface that displays their “QK attributions,” to be described in another update, TK).
- Note that MOLT transforms themselves have input edges, computed the same way as above, but using the transform’s encoder vector as the attribution target
Qualitative findings
Using such attribution graphs, we have uncovered instances of MOLT transforms performing interpretable computations, such as:
- In the prompt The Spanish word for hot is "calor”, we observed that “say a word beginning with ‘cal’” features received input from “hot” features via a transform that is active in Spanish-language-related contexts.
- Interestingly, we also see some inputs from “Spanish” features via transforms that are active in “hot”-related contexts!
- In the prompt 3 + 5 = 8, we observed a “say 8” feature receiving input from a “plus 3” feature via a transform active in contexts in which the number “5” recently appeared
However, we also see transforms playing roles that are less clearly interpretable. For instance, many edges appear to be mediated by transforms that select for key words like “is” or “Assistant.” We also see transforms playing roles that appear redundant with the features they carry, such as a “Paris” feature receiving input from a “France” feature via a “France” transform.
Conclusion
We believe MOLTs are a promising alternative to transcoders and may be able to capture MLP computation in a more parameter-efficient way that more faithfully reflects the computations performed by the underlying model. We suspect that a MOLT-like solution will be necessary to capture all the variance of frontier model MLP layers at a reasonable computational cost. We see signs of life that MOLT transforms can perform interpretable computations, “transforming” input features into output features.
One direction we are excited about, but have not explored yet, is using MOLT transforms to understand compositional representations that are not captured by our SAEs. In particular, the reconstruction error of our SAEs at each layer can be rewritten as a sum of terms corresponding to (feature, transform) pairs from the previous layer. These might correspond to concepts that are too rare to be captured by our finite-size SAEs, but are built out of composing a relatively common feature with a relatively common transform. In an attribution graph, decomposing SAE errors into (feature, transform) pairs in this fashion would manifest as having some graph edges that are mediated by chains of transforms in consecutive layers.
More work needs to be done to conclude that MOLTs are strictly preferable to transcoders. Attribution graphs that include transform information are somewhat more unwieldy than transcoder-based attribution graphs, and not all transform-mediated computations are clearly interpretable. Future work on scaling and improving MOLTs, and the associated attribution graph logic and UI, may address some of these issues.