Transformer 的每一层被读成一次 EM 迭代:注意力做置信传播,层间跳变充当参数更新
导语
Transformer 的注意力被解释为在逐样本变化的马尔可夫随机场上做最大积置信传播,层与层之间的参数跳变则充当 EM 的参数更新步,这一视角把自注意力、专家混合、门控单元和选择性状态空间模型归入同一类由输入生成参数的元模型。
正文
Transformer 的内部计算被重新描述为一种概率推断过程,而不是一组彼此独立的架构部件。 此前的解释各自只覆盖一部分行为,例如把上下文学习对应到类梯度下降或贝叶斯更新,或把注意力对应到图上的消息传递,但没有把自注意力、深度、隐变量细化、图结构选择和逐层参数变化放进同一个概率图景。 这一对应关系建立在推导之上:把编码器 Transformer 底层的贝叶斯网络经道德化和变量消元化为全连接马尔可夫随机场,再假设高阶因子可由成对因子近似,从而得到注意力与最大积置信传播之间的映射。
自注意力被写成在动态结构马尔可夫随机场上的最大积置信传播的一步近似。 置信传播原本用于结构和参数已知的部分观测模型,EM 用于结构已知而参数未知的模型,而真实数据往往结构未知、参数也未知,单独使用其中任一算法都难以奏效。 把对数消息在各自极大值处做二阶泰勒展开后,信念极大值可写成各消息极大值的精度加权平均;在负 Hessian 各向同性的假设下,该式与自注意力的查询、键、值计算形式一致,其中注意力权重对应可学习的结构亲和度,值向量对应消息极大值。
Transformer 的深度被解释为一次 EM 迭代:前向传播更新隐变量,层间参数变化充当参数更新。 此前把 Transformer 深度解释为迭代概率拟合的工作,或绑定在特定建模选择上,或是针对特定模型类的表达能力结论,而这里把算法框定为概率隐变量推断。 在监督学习设定下,第 l 层前向传播计算隐变量,对应 E 步,但只跑一次近似最大积置信传播得到估计的极大值;下一层的参数取为使目标最大的参数,对应 M 步,并由 Gibbs 不等式给出目标改善与对数似然改善之间的关系。
由 EM 视角引出元模型这一设计范式:输入先生成该次前向传播要用的参数。 超网络、专家混合、门控卷积与循环模型、选择性状态空间模型通常被当作彼此分离的设计家族来研究,而这里把它们视为同一原则的实例或近亲。 自注意力和专家混合都可写成先由输入经元前向传播确定参数、再执行前向传播的两段式;在由高斯混合模型生成的二维合成数据上,元线性层在多个子分布的数据集上持续优于普通线性层,其中 SwiGLU 变体在最高复杂度下取得 0.6842 的评估准确率,而普通线性层退化到 0.5171。
接下来
这一解释面向编码器 Transformer,其底层模型被设定为部分观测的马尔可夫随机场,并依赖两条假设:元数大于二的因子可由有限个成对因子之积近似,以及编码器 Transformer 隐式地在逐样本变化的成对因子马尔可夫随机场上做推断。实验部分在多个电影评论和来自 X 的真实文本上,用 Huggingface 上的多个 Transformer 取最后一层各 token 的输出特征,去掉标点、换行等非语义或特殊 token 后计算余弦相似度矩阵,再跑单链接凝聚层次聚类并画成树状图。文中提出的后续检验方向包括:在合成图模型任务上把注意力图与已知依赖图比较,追踪逐层表示与参数是否改善代理目标,以及设计阶数与稀疏度可控的参数生成架构。
文中把该解释表述为一种视角,并称经验结果支持这一视角;两条关于成对因子近似和逐样本变化马尔可夫随机场的假设是推导的前提。元线性层的参数量是线性层的若干倍,广义元线性层的参数量随阶数指数增长。实验中的元线性层在最高复杂度下评估准确率为 0.6842,普通线性层为 0.5171,而把专家数设为真实子模型数的 Ground-truth MoE 在该设定下取得 0.9400、0.9022 和 0.8079,被文中作为这些实验里分类准确率的上界。
