规模化可解释性的工程挑战¶

2024 年 6 月 13 日 · 原文: https://www.anthropic.com/research/engineering-challenges-interpretability

在这篇文章以及上面的圆桌视频中，我们的研究人员反思了科学进步与工程进步之间的紧密联系，并讨论了他们在将可解释性研究扩展到规模大得多的 AI 模型时遇到的技术挑战。

去年 10 月，Anthropic 可解释性团队发表了 Towards Monosemanticity，这篇论文将字典学习技术应用到了一个小型 transformer 模型上。今年 5 月，我们又发表了 Scaling Monosemanticity，把同样的技术应用到了大几个数量级的模型上。我们在 Claude 3 Sonnet 中发现了数千万个"特征"——与语义概念相关的神经元组合——这代表着我们在理解 AI 模型内部运作方面迈出了重要的一步。

要继续取得这样的进展，我们需要更多的工程师。

如果你只读过我们早期的论文（例如 Frameworks 和 Toy Models of Superposition），这个说法可能会让你感到意外，因为那些工作对工程的要求相对较少。但读一读更新的研究会让你清楚地看到我们所面临的工程挑战的规模。

下面，我们分享两个我们最新研究中涉及的技术工程问题的例子。它们体现了我们的工程师眼下正在解决的那类问题，也有助于解释为什么我们认为工程将是 AI 可解释性研究——乃至最终的 AI 安全研究——取得进展的主要瓶颈之一。

如果你是一名工程师，这篇文章就是写给你的。如果你受到下面讨论的工程问题示例的启发，我们强烈建议你申请我们的研究工程师职位。

工程问题一：分布式混洗（shuffle）¶

我们的稀疏自编码器（sparse autoencoder, SAE）——用来研究"特征"的工具——在 transformer 的激活值上训练，而这些激活值需要被打乱，以防止模型学到虚假的、依赖顺序的模式。我们刚开始训练稀疏自编码器时，训练数据可以放进单个 GPU 的内存，打乱起来轻而易举。但后来，我们想要扩展到超出内存容量的规模（想象一下，从洗一副牌这样的简单任务，扩展到洗堆满整个仓库的牌——这是一个困难得多的问题）。

这时，我们本可以实现一个能够扩展到 PB 级的分布式混洗方案。但我们最终选择了一种可以快速实现、却不太能扩展的方案。我们把混洗拆分成 K 个任务，每个任务负责 1/K 的混洗后输出数据。我们生成一个置换（permutation），让每个任务对流式读取全部训练数据，然后写出自己那一份输出。这让我们得以进一步扩展，但缺点也很明显：每个任务都必须读取全部训练数据。起初这要花几个小时，后来要花几天。等我们着手做 Towards Monosemanticity 时，训练数据已经达到 100TB（1000 亿个数据点，每个 1KB），混洗成了一个让人头疼的大问题。

实现一个可扩展的分布式混洗并不是什么新奇或前沿的问题。但它只是我们为了取得科学进展而必须快速解决的众多工程问题之一。

在这个问题上，我们找到了一篇有用的博客文章，并把它的方法扩展到了多趟执行。一趟中我们有 N 个任务。每个任务读取 1/N 的数据集，将其打乱，然后写成 K 个文件，每个文件包含 1/NK 的数据。每个任务写出的第一个文件的内容，对应最终混洗数据的前 1/K，但它仍需打乱。第二个文件同理，以此类推。一趟下来，我们把对整个数据的一次混洗，变成了 N 次各自小 K 倍的混洗。现在，如果这些混洗能放进单台机器的内存，我们就可以直接打乱收工。如果放不下，就再跑一趟。

假设每个任务能在内存中保留 100GB 的数据，我们写出 100 个 1GB 的文件。每一趟会把所需混洗的规模缩小 100 倍。一趟可以混洗 100GB 的数据，两趟可以混洗 10TB，三趟 1PB，四趟 100PB，以此类推。

自从实现这个方法以来，我们就不再为混洗操心了。现在它执行得又快又稳。当然，有比我们的方案更好的方法和更快的实现。但这个方案解决了我们的瓶颈，让我们腾出手来解决下一个问题。

工程问题二：特征可视化流水线¶

另一个工程挑战是为我们的特征可视化生成底层数据。特征可视化让用户可以看到单个特征中激活最强的特定词元（token），以及其他信息（见《Towards Monosemanticity》论文的特征浏览器（Feature Browser），见此链接）。

对于每个特征，我们都想找到一系列以不同程度激活它的数据集样本，从而探索它的完整分布。对数百万个特征高效地完成这件事，是一个有趣的分布式系统问题。最初，这一切都运行在单个任务里——但我们很快就超出了它的规模。下面是我们当前方案的草图。

我们的可视化数据集有 1 亿词元，需要处理数百万个特征。首先我们按数据集和特征进行"分片"（shard），把它们拆成许多部分。每个任务遍历自己的数据集切片，并针对自己的特征切片，记录每个特征激活最强的 K 个词元，以及 10*K 个激活该特征的随机词元（我们已把 transformer 激活值缓存到 S3 中，因此无需重新计算）。

接下来，我们按特征分片，汇总上一趟的结果。这样就得到了每个特征在整个数据集上激活最强的词元，以及一组激活该特征的随机词元。这些就是我们要在特征可视化中展示的样本。

对于这些样本中的每一个，我们都需要计算特征在其周围词元上的激活情况。我们的第一个方案是按特征分片。每个任务加载它所负责的特征样本对应的 transformer 激活值。问题在于这些样本随机分布在数据集中：没有简单的方法只读取任务需要的那部分数据。

为了改进这一点，我们增加了一趟按数据集分片的处理。在这个设置中，每个任务处理一份 transformer 激活值切片，并把每簇特征所需的激活值保存到单独的文件中。然后，我们可以跑一趟按特征分片的处理，轻松访问我们恰好需要的数据。我们计算出特征在周围词元上的激活程度，然后以我们的前端网站能够读取和展示的格式写出所有相关数据。

我们在寻找什么样的人¶

当我们刚开始研究稀疏自编码器时，并不知道这个方法是否行得通。随着实验越做越多，我们对这项研究的信心也越来越足。这促使我们加大基础设施投入，以便运行更大规模的实验。这个过程一直持续到我们最新的论文 Scaling Monosemanticity。

这个过程也指明了我们在可解释性团队所做的那类工程工作——以及我们把研究与工程视为不可分割的事实。我们的团队成员常常在研究与工程之间来回切换：先从当前系统中挤出更多规模以启动一个新实验，然后再回到研究上来。由于许多研究想法最终行不通，我们在看到成功之前不会对它们的基础设施投入过多。

研究是一项团队工作，它既关乎反复琢磨想法，也关乎把想法付诸实现。我们不只提出假说；我们还会测试、构建、迭代和扩展。

正因为如此，我们特别想招聘能够在不同领域之间灵活工作的通才型工程师——无论是搭建流水线、运行机器学习实验，还是优化 GPU 使用。

如果你是符合条件、又对 AI 安全充满热情的工程师，我们非常期待看到你的申请。请看看我们研究工程师职位的职位描述，也可以在我们的招聘页面上查看可解释性团队的其他几个空缺职位。

常见问题（FAQ）¶

你们计划招多少人？ 可解释性团队目前有 18 名成员，并且正在快速壮大。目前我们打算在多个地点至少招聘五位高级工程师和两位团队经理。职位列表请见我们的招聘页面。
可解释性团队与其他团队合作吗？ 合作。我们与其他研究团队（尤其是对齐团队）合作非常紧密。Anthropic 是一个"团队科学"型组织，团队之间的边界以最好的方式保持着开放——这让我们能够快速完成很多事情。
你们如何选择要做的项目？ 我们根据 Anthropic 的负责任扩展政策来思考研究路线图，该政策承诺公司在开发或部署超过相应能力水平的模型之前，先达成各项安全里程碑。在思考研究方向时，我们会考虑可解释性的研究版图是什么样的、我们处于什么位置能够解决哪些问题，以及可解释性团队的工作如何影响这些安全里程碑。在日常层面，我们的研究往往非常具有探索性，但这种探索是由上述考量引导的。
可解释性团队的成员都有什么样的背景？ 团队成员来自各种各样的职业背景，包括神经科学、数学、生物学、物理学、数据可视化和软件工程。
你们接受湾区以外的候选人吗？ 团队目前有成员分布在旧金山、波士顿、纽约、西雅图和伦敦。旧金山的人数最为集中，这些成员每周有几天会来办公室。我们也接受远程工作，但要求每年约有 25% 的时间到 Anthropic 办公室。更多关于所有空缺职位的信息，请参见我们的招聘页面。