论文

元注意力:用于高效 Transformer 推理的贝叶斯每词元路由

Meta-Attention: Bayesian Per-Token Routing for Efficient Transformer Inference

模型架构Transformer

摘要

标准 Transformer 架构在所有标记和序列位置上统一应用单一注意力机制,而不考虑本地上下文或计算预算。我们提出了元注意力(Meta-Attention),这是一个通过贝叶斯元控制器动态地将每个标记路由到最合适的注意力策略的框架——完整的 softmax 注意力、线性(内核)注意力或滑动窗口局部注意力。与使用确定性或无先验学习路由的先前路由方法不同,元控制器将每个词元机制选择视为计算感知狄利克雷先验下的后验推理:路由权重是使用证据下界(ELBO)目标训练的摊销变分后验 q(alpha | x_t; phi) 的输出,该目标联合编码任务性能和注意机制成本。这种设计产生有原则的路由不确定性估计,控制软到硬路由转换,减轻路由崩溃而不会造成临时负载平衡损失,并以可忽略的开销产生比确定性或无先验学习路由更好的计算性能权衡。 Tiny LM 基准的第一阶段实证结果证实了核心预测:贝叶斯控制器的学习路由分布意味着硬路由下的预计归一化 FLOP 成本为 25.1%,而无先验基线 (-34.2 pp) 的预计标准化 FLOP 成本为 59.3%,并将路由熵从 55.8% 降低到 43.3% (-12.5 pp),这表明狄利克雷先验可以防止路由崩溃,而非贝叶斯模型默认为完全关注。我们提出了贝叶斯架构、ELBO 训练目标以及验证前向传递正确性、后验多样性以及针对先验自由基线的受控消融的第一阶段 PyTorch 原型。代码位于:https://github.com/KFEAL/meta-attention