arXiv:2605.07588cs.LGcs.AI2026-05

用能量最小化视角重新理解Transformer参数设计,发现共享权重等新结构可行。

Revisiting Transformer Layer Parameterization Through Causal Energy Minimization

  • 将注意力与MLP视为能量函数的优化步骤,推导出权重共享的理论依据。
  • 在亿级参数模型上验证,简化参数的结构仍能稳定训练并匹配基线性能。
  • 适合对模型设计原理感兴趣的研究者,尤其关注架构可解释性与轻量化设计。

Transformer模块通常结合多头注意力(MHA)进行令牌交互与门控MLP进行特征变换,但其参数化设计多为经验性。我们提出因果能量最小化(CEM)框架,将Transformer层视为条件能量函数上的优化步骤,并显式建模层参数化。扩展已有注意力的能量解释,CEM表明权重重用的MHA可视为交互能量的梯度更新,而具有共享上下投影的门控MLP可通过元素级能量解释。该视角揭示了包含层内权重共享、对角加低秩交互、轻量预条件器及递归更新的设计空间。我们在中等规模(约亿级参数)语言模型实验中评估了基于CEM的层结构。尽管参数受限,这些层仍能稳定训练并达到对应基线性能。结果表明,CEM为理解Transformer层参数化提供了有效视角,将架构与能量模型联系起来,激励进一步探索能量引导的层设计。

原文摘要 · Abstract (English)

Transformer blocks typically combine multi-head attention (MHA) for token mixing with gated MLPs for token-wise feature transformation, yet many choices in their parameterization remain largely empirical. We introduce Causal Energy Minimization (CEM), a framework that recasts Transformer layers as optimization steps on conditional energy functions while explicitly accounting for layer parameterization. Extending prior energy-based interpretations of attention, CEM shows that weight-tied MHA can be derived as a gradient update on an interaction energy, and that a gated MLP with shared up/down projections can be viewed through an element-wise energy. This perspective identifies a design space for Transformer layers that includes within-layer weight sharing, diagonal-plus-low-rank interactions, lightweight preconditioners, and recursive updates. We evaluate CEM-derived layers in language-modeling experiments at the moderate hundred-million-parameter scale. Despite their constrained parameterizations, these layers train stably and can match corresponding Transformer baselines. Overall, our results suggest that CEM provides a useful lens for understanding Transformer layer parameterization, connecting Transformer architectures to energy-based models and motivating further exploration of energy-guided layer designs.

Transformer能量模型参数共享

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。