论文

通过因果能量最小化重新审视 Transformer 层参数化

Revisiting Transformer Layer Parameterization Through Causal Energy Minimization

模型架构Transformer

摘要

Transformer 块通常将用于词元混合的多头注意力 (MHA) 与用于词元特征转换的门控 MLP 结合起来,但其参数化中的许多选择在很大程度上仍然是经验性的。我们引入了因果能量最小化(CEM),这是一个框架,它将 Transformer 层重新定义为条件能量函数的优化步骤,同时明确考虑层参数化。 CEM 扩展了先前基于能量的注意力解释,表明权重关联的 MHA 可以作为交互能量的梯度更新导出,并且可以通过元素能量查看具有共享上/下投影的门控 MLP。该视角确定了 Transformer 层的设计空间,其中包括层内权重共享、对角线加低秩交互、轻量级预处理器和递归更新。我们在中等亿参数规模的语言建模实验中评估 CEM 衍生层。尽管参数化受到限制,但这些层可以稳定地训练并且可以匹配相应的 Transformer 基线。总体而言,我们的结果表明,CEM 为理解 Transformer 层参数化、将 Transformer 架构连接到基于能量的模型以及激发对能量引导层设计的进一步探索提供了有用的视角。