📊 statistics

Revisiting Transformer Layer Parameterization Through Causal Energy Minimization

本文介绍了因果能量最小化（CEM），这是一个将 Transformer 层重新诠释为条件能量函数上的优化步骤的框架，旨在推导出在语言建模任务中与标准基线相匹配的受约束且参数高效的架构。

原作者： Jin Xu, Camille Couturier, Victor Rühle, Saravan Rajmohan, James Hensman

发布于 2026-05-11

📖 1 分钟阅读☕ 轻松阅读

原作者： Jin Xu, Camille Couturier, Victor Rühle, Saravan Rajmohan, James Hensman

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一个现代大型语言模型（就像那些写故事或回答问题的模型）是一座巨大的多层工厂。这座工厂的每一层都是一个"Transformer 层”。在每一层上，都有两台主要机器协同工作：

注意力机器（Attention Machine）： 它审视句子中的所有单词，找出它们之间的关联（例如将“狗”与“吠叫”联系起来）。
特征机器（MLP）： 它接收这些关联，并将它们转化为新的、更复杂的概念。

多年来，工程师们通过试错法来构建这些机器——不断调整旋钮和刻度盘，直到工厂产出良好的结果。他们并不总是知道为什么某个特定设置有效，他们只知道它确实有效。

这篇论文介绍了一种看待这些机器的新方法，称为因果能量最小化（Causal Energy Minimization, CEM）。以下是简单的分解：

核心理念：“山丘与山谷”类比

作者建议我们停止将这些机器仅仅视为“数学运算”，而是开始将它们视为试图寻找山谷最低点的徒步者。

能量景观： 想象一个凹凸不平的地形，高耸的山丘代表“糟糕”或“混乱”的状态，而深邃的山谷代表“良好”或“清晰”的状态。
目标： 机器的任务是将一块数据（一个单词或句子）滚下山坡，直到它 settle 在最深、最稳定的山谷中。
步骤： 在旧的方法中，机器一步跨到底部。而 CEM 框架则说：“让我们观察坡度，小心翼翼地迈出一小步。”

他们如何重造这些机器

作者意识到，标准的机器（注意力和特征 Transformer）实际上只是沿着这个能量山丘迈出了一步。他们问道：如果我们专门设计这些机器，使其更擅长这种“滚下山坡”的过程，会怎样？

他们发现了改进工厂楼层的两种主要方法：

1. 共享工具（权重绑定，Weight Tying）

在标准工厂中，注意力机器使用一套工具来寻找关联，并使用另一套不同的工具来输出结果。

CEM 的洞察： 数学表明，如果你使用相同的工具来寻找关联并输出结果，这就好比机器沿着能量山丘上完美、自然的轨迹下行。
结果： 他们构建了一个共享这些工具的机器版本。它使用了更少的部件（更少的参数），但效果与标准、更庞大的机器一样好。这就像使用瑞士军刀，而不是携带整个工具箱。

2. 多走几步（递归，Recursion）

标准机器沿着山丘迈出一大步，然后说：“好的，我们完成了。”

CEM 的洞察： 一步往往不足以到达山谷的最底部。有时你需要再迈出一小步，甚至是第三步才能到达那里。
结果： 他们在机器内部添加了一个“循环”。数据不再在迈出一步后就离开，而是留在这一层，沿着同一个能量山丘再迈出第二步。
益处： 这使得机器能够找到“更好”（能量更低）的状态，而无需添加更多工具或扩大机器规模。它只是在同一层上多思考了一会儿。

他们测试了什么

团队构建了这些新的“能量最小化”机器，并在从小型到中型（约 1 亿到 1.6 亿参数）的语言模型中进行了测试。

发现：
- 效率： 共享工具（权重绑定）的新机器使用了显著更少的部件（注意力机器的部件减少了约一半），但表现与标准模型一样好。
- 改进： 当他们添加“多步”（递归）时，模型的表现实际上比标准模型更好，尽管它们的规模更小。
- 稳定性： 新设计训练稳定，没有崩溃或出现异常行为。

他们没有声称什么

重要的是要坚守论文实际所说的内容：

他们没有声称这使模型在编码或数学等特定任务上变得更聪明（尽管可能是这样，但论文并未测试这一点）。
他们没有声称这解决了 AI 幻觉的问题。
他们没有声称这已准备好立即用于大规模、万亿参数模型（他们测试了约 1.6 亿参数，并指出更大规模需要更多研究）。

核心启示

将这篇论文想象为一位工程师意识到工厂楼层的构建存在不必要的冗余。通过将过程视为“让球滚下山坡”，他们找到了如何：

移除重复的工具（节省空间和成本）。
让球多滚动几次（提高质量），而无需扩大工厂规模。

他们尚未建造一座新工厂，但他们提供了一张新蓝图，展示了如何更高效、更有效地建造同样的工厂。

技术摘要：通过因果能量最小化重新审视 Transformer 层参数化

问题陈述

Transformer 架构，特别是多头注意力（MHA）与门控多层感知机（MLP）的组合，构成了现代大语言模型（LLM）的基础构建模块。然而，关于层参数化的具体选择（如权重共享、投影结构和更新规则）主要受实证研究驱动，而非明确的计算原理。尽管基于能量的模型（EBM）此前提供了一种将计算视为搜索低能量状态（例如在 Hopfield 网络中）的视角，但现有工作尚未将这些基于能量的解释与标准 Transformer 层的具体参数化完全联系起来，也未系统探索这一视角如何指导设计更高效或更有效的层变体。

方法论：因果能量最小化（CEM）

作者引入了因果能量最小化（CEM），这是一个将 Transformer 层重新框架化为条件能量函数上优化步骤的框架。其核心前提是：一个层将输入序列 $h_{1:J}$ 映射到输出 $h'_{1:J}$ ，方法是把每个位置 $i$ 视为优化变量 $x_i$ （初始化为 $h_i$ ），并通过过程 $\mathcal{A}$ 更新该变量，以最小化依赖于因果历史的条件能量 $\epsilon(x_i | h_{1:i})$ 。

本文从该框架推导出了标准 Transformer 组件，并提出了扩展：

权重共享注意力的推导：
通过定义基于对数 - 求和 - 指数形式（类似于现代 Hopfield 网络）的交互能量，作者表明单步梯度下降可恢复标准 MHA 更新。关键在于，该推导自然地需要特定的权重共享约束：
- 键投影（ $W^K$ ）与值投影（ $W^V$ ）共享权重。
- 查询投影（ $W^Q$ ）与输出投影（ $W^O$ ）共享权重。
- 温度参数 $\tau$ 与头维度相关。
权重共享门控 MLP 的推导：
通过定义逐元素能量，独立地为每个 token 的特征向量分配能量，梯度更新恢复了门控 MLP 的结构。这一解释意味着 MLP 中的“上”和“下”投影矩阵应当共享权重（具体而言，下投影的转置等于上投影）。
提出的扩展：
基于能量优化视角，作者提出了三种具体的架构增强：
- 对角加低秩（D+LR）参数化： 针对交互矩阵 $A_k$ ，作者不再使用纯低秩投影，而是提出 $A_k = \text{diag}(d_k) + W^{Q\top}_k W^K_k$ 。这捕捉了低秩矩阵所遗漏的交互，同时保持计算高效。他们还提出在所有头之间共享对角项以减少参数量。
- 可学习的轻量级预条件器： 为了在单步梯度之外改善收敛，他们引入了具有 D+LR 结构的预条件矩阵 $P$ （ $P = \text{diag}(\text{softplus}(d)) + UV^\top + VU^\top$ ）。这些矩阵作为曲率信息的可学习代理，在不产生完整二阶方法成本的情况下对梯度进行重缩放。
- 层内递归更新： 框架允许在单层内进行 $T$ 次递归更新（ $x^{(t+1)}_i = x^{(t)}_i - \eta \nabla \epsilon$ ），而非单次更新。这使得层能够更深入地最小化其能量函数，而无需添加新参数，这与重用整个块的“循环”Transformer 不同。

主要贡献

理论统一： 本文确立了具有特定权重共享方案（注意力中的 Q/O 和 K/V 共享；MLP 中的上/下共享）的标准 Transformer 层，在数学上等价于特定交互能量和逐元素能量函数上的单步梯度更新。
设计空间扩展： CEM 视角为 Transformer 层识别了一个新的设计空间，包括 D+LR 交互、轻量级预条件器和递归更新，这些均由优化理论而非启发式方法驱动。
实证验证： 作者在中等规模（数亿参数）的语言建模任务上评估了 CEM 导出的层，证明这些受约束的参数化方案可以达到或接近标准 Transformer 基线的性能。

实验结果

实验在 SlimPajama 数据集上进行，模型规模从约 8600 万到约 1.62 亿参数不等，采用计算最优的 token 预算（Chinchilla 缩放）进行训练。

参数效率： 使用权重共享的 CEM 模型（使用约一半 MHA 参数和三分之二 MLP 参数，相对于 Llama 基线）实现了具有竞争力的测试困惑度。对于 CEM 注意力，尽管参数大幅减少，性能下降极小。
扩展的影响：
- 递归： 将递归步数从 $T=1$ 增加到 $T=2$ ，一致地提升了 CEM 注意力和 MLP 的性能。值得注意的是，递归 CEM 注意力变体在参数更少的情况下优于 Llama 基线。
- 预条件器： 与递归相比，可学习的预条件器仅提供了边际增益，表明递归结构是性能提升的主要驱动力。
- 对角项： 在交互矩阵中包含对角项（具体为跨头共享的对角项）对于良好性能至关重要，其表现优于纯低秩或无对角项的变体。
稳定性： 完全由 CEM 导出层（ $T=2$ 且带有 D+LR 预条件器）组成的 Transformer 能够稳定地进行端到端训练，以较少的参数实现了与 Llama 基线相当的困惑度。
合成数据： 在合成高斯过程数据上的实验证实，更深的递归（ $T=4, T=8$ ）能进一步降低均方根误差（RMSE），尽管在标准训练设置下的完整 Transformer 环境中，增益仅限于 $T=2$ 。

意义与主张

本文主张 CEM 为理解 Transformer 层参数化提供了一个有用的理论视角，有效地将 Transformer 架构与基于能量的模型联系起来。其意义在于：

原则性设计： 它将层设计从经验性的试错转变为基于能量最小化和优化动力学的框架。
效率： 它证明了由此框架导出的权重共享和递归更新，能够构建在保持竞争力的同时参数更高效的模型。
未来方向： 作者指出，这一视角激发了对以下领域的进一步探索：
- 自定义内核： 针对结构化权重共享和递归更新优化运行时（例如融合操作）。
- 推理时扩展： 利用层内递归作为推理时扩展计算量的额外维度。
- 硬件协同设计： 针对新型硬件架构专门设计层参数化方案。

作者保持谦逊，指出其评估仅限于数亿参数规模，需要更大规模的研究和下游评估才能充分评估这些优势向最先进模型迁移的可行性。他们还承认，当前的实现尚未针对运行时性能进行优化。