Meta华人研究员提出轻量级结构Canon层，提升模型推理深度和广度

“大模型设计迎来伽利略时刻。正如比萨斜塔实验推动了现代物理学的发展，我们开发的受控合成预训练平台揭示了大模型架构的真正极限。这可能是一个转折点，将大模型的研究划分为‘之前’和‘之后’。” 清华本科校友、美国麻省理工学院博士毕业生、Meta 研究员朱泽园在 X 上写道。

这则推文介绍了一个名为“语言模型物理学”的长期项目。在这则推文发布之前，当地时间 5 月 2 日关于这一长期项目的一篇论文上线 [1]。

在这篇论文中，朱泽园和同事提出一种名为“Canon 层”的轻量级结构。

“Canon”原本是一个音乐术语，它指的是一种复调音乐的作曲技法，其特点是让多个声部以相同的旋律和不同时间进入，从而形成层叠交错的效果。

而本次提出的“Canon 层”和上述音乐技法有着异曲同工之妙，它能促进相邻 tokens 之间的水平信息流动，提升推理深度、推理广度、知识容量和操控能力，并能在跨多种架构的层次结构学习上实现突破。

通过开展这些合成任务，研究团队实现了架构性能的严格分离，让有着 13 亿参数和 1000 亿 tokens 数据量的学术级预训练实验中的噪声与随机干扰因素，得到了有效消除。

研究人员在论文中表示，Canon 层能够显著提升传统弱势架构的性能，例如它能将无位置编码（NoPE，No Positional Encoding）的 Transformer 提升至旋转位置编码（RoPE，Rotary Positional Encoding）水平，使线性注意力（GLA，Generalized Linear Attention）达到 Mamba 级表现，同时大幅降低对旋转位置编码的依赖，从而有效改善长程泛化能力。（注：Mamba，是此前由其他研究团队提出的一种高效序列建模架构，下文的 Mamba2 是 Mamba 的升级版。）

据介绍，Canon 层能够计算附近 tokens 表示的加权组合，并能无缝集成到 Transformer、线性注意力、状态空间架构或任何通用序列模型中。

本次研究还表明，线性模型推理深度不足的原因在于 tokens 的压缩过程和检索过程效率低下。同时，本次研究中的合成任务也验证了受控可预测基准在架构评估中的核心价值。

与此前经典的学术级别实验场景相比，本次研究中的合成实验平台能够更清晰地揭示模型的真实能力和扩展行为。

凭借无限高质量的数据，研究团队希望该框架能够预测架构如何随着训练流程的改进而演变。例如，预测该如何通过更好的数据整理或基于强化学习的后训练，解锁更深层次的推理能力和层次化推理能力。

研究团队在论文中表示，他们更建议将智能分解为原子级组件，比如分为推理深度和推理广度，并通过构建高度可控的合成预训练任务，来独立地分离和评估这些能力。

这样一来，在理想化的实验环境之下，就能清晰地界定不同架构的性能优势与扩展潜力。

据研究人员介绍，本次方法通过实现单一技能的评估，消除了真实世界预训练数据中的混杂因素。例如，本次方法既可以针对架构 A 在推理深度上是否优于架构 B 开展严谨对比，还能确保修改不会降低其他能力。

通过隔离固有的架构偏差，他们在合成预训练任务中揭示了那些在传统设置中往往会被噪声掩盖的特性。

研究中该团队打造了一个合成训练环境，并引入了五项合成预训练任务，以便系统地评估语言模型中的关键认知能力。

通过这一实验平台研究人员揭示了在小规模场景下的架构趋势：其一，线性注意力模型在各类任务中表现持续滞后；其二，Mamba2 这样的递归架构在依赖记忆的任务中优势显著；其三，标准 Transformer 在复杂推理任务中占据主导型的优秀表现。

对于分析不同架构下的固有偏差和可扩展性来说，这些受控基准测试能为其提供一个较为严谨的框架。研究人员指出，Transformer 在层内缺乏水平信息流，即使在关联回忆这样的简单任务上也会出现效率低下的情况。

另据研究人员在论文中表示，Canon 层可以灵活地嵌入网络的不同位置，其中：

• Canon-A：插入注意力机制之前；
• Canon-B：嵌入注意力机制内部；
• Canon-C：置于多层感知机（MLP，Multilayer Perceptron）之前；
• Canon-D：整合至多层感知机内部。

虽然 Canon 层的实现方式有很多种，即便使用简单的随机平均操作也能取得显著效果。但是，本次研究聚焦于采用可训练的一维线性卷积核（核大小为 4），这种方法只需使用轻量计算、而且易于集成，只需少量代码即可适配任意架构。

在研究团队的实验平台中，Canon 层分别能将推理深度提升 200%-400%、将推理广度提升 30%、将知识操作长度提升 30%，并且只需要极小的计算开销能够实现这些提升，而这主要依赖于更优的层次化学习动态。

将 Canon 层集成到模型中，能够提高无位置编码模型的性能，其性能与“旋转位置编码模型+Canon 层”的组合相当，甚至更胜一筹。在性能上，Canon 层优于像线性偏置注意力（ALiBi，Attention with Linear Biases）或混合线性偏置注意力（H - Alibi，Hybrid Attention with Linear Biases）这样的位置修正方法，并且能够降低或消除旋转位置编码的使用，从而显著提升模型的长程泛化能力。

研究人员在论文中表示，Canon 层通过不同子层位置的累积式作用机制来实现性能提升，其效果独立于注意力或多层感知机组件。残差连接提高了训练效率，在不影响稳定性的前提下，几乎不需要进行参数调整。

通过提高训练效率和稳定性，Canon 层还能恢复在门控多层感知机或混合专家架构中损失的部分知识承载能力。

研究中，该团队还将线性注意力与 Canon 加以结合。通过此，Canon 层能够显著提升门控线性注意力的推理能力。其中，推理深度能从单跳（1-hop）被扩展至四跳（4-hop），推理广度与知识操作长度能够实现双倍增长。最终可以使门控线性注意力达到与状态空间模型架构 Mamba2 相当的性能，并能在 Brevo 等任务中实现反超。

据介绍，Mamba2 的性能得益于其内置的一维卷积（conv1d）机制，该机制相当于一个应用于选定坐标的非线性 CanonB 层。移除一维卷积机制之后，其性能下降到与门控线性注意力模型相当的水平。而使用完整的 Canon 层替代它则能进一步提升效果，这凸显了在结构化状态空间模型设计中水平信息流的重要性。

通过消融研究，该团队发现 Canon 层在集成点和残差链接上的选择会影响到 Mamba2 的性能。作为一种使用学习线索来初始化 Transformer 自注意力权重的方法，模仿初始化（Mimetic initialization）虽然能在长度泛化上实现优化，但却会给短上下文任务造成损害，而这恰恰凸显多样化预训练环境的必要性。

研究中，该团队还开展了架构比较。他们分别在旋转位置编码、无位置编码、Mamba2 和门控线性注意力中同时使用完整的 Canon 层。结果表明，在层次推理任务中，完整 Transformer 的表现优于线性模型，其推理深度达到了线性模型的两倍。

研究人员指出，线性模型在推理深度方面存在一定的局限性，这种局限性源于压缩与检索过程中累积的误差，而非记忆容量的不足。而将滑动窗口 Transformer 与线性模型相结合的混合架构，则能为深度推理任务提供可扩展的解决方案。

与此同时，他们还开展了学术级别的真实预训练。具体来说，他们在 1000 亿个 tokens 上训练具有 4096 个上下文长度的 13 亿参数模型，结果显示噪声较大且分辨率有限，这使得大多数架构差异在统计上并不显著。

尽管如此，课题组也收获了一些发现：

首先，在检索任务繁重的任务上，线性架构即使配备了 Canon 层，也始终落后于完整的 Transformer 架构。

其次，Canon 层能够显著提升门控线性注意力和无位置编码，使门控线性注意力能够达到与 Mamba2 相当的水平，以及使无位置编码能够达到与旋转位置编码相当的水平。然而，移除 conv1d 会使 Mamba2 性能下降到与门控线性注意力相当。

再次，即使在只有 100 个 tokens 的短文本中，所有模型在处理两跳推理任务时都面临一定的困难，这凸显了学术级别预训练的局限性。

最后，通过减少或去除旋转位置编码这一做法，尤其是在添加了 Canon 层的情况下使用这一做法，能在不影响整体性能的前提下，提升长上下文的泛化能力。

总之，Canon 层从根本上改善了跨不同架构的水平信息流，从而实现了更深入的推理和高效的可扩展性。未来，该团队希望借助无限高质量的数据，通过本次成果来预测未来架构的演进方向，例如通过更完善的数据策展或基于强化学习的训练后优化，从而释放更深层次的推理与分层推断能力，进而推动大模型训练流程的持续进步。

参考资料：

1.论文地址：https://papers.ssrn.com/sol3/papers.cfm?abstract_id=5240330

https://www.linkedin.com/in/zeyuan-allen-zhu/details/education/

https://x.com/ZeyuanAllenZhu

http://zeyuan.allen-zhu.com/index.php

运营/排版：何晨龙