论文

Transformer 的黎曼注意力机制:理论框架和架构设计

Riemannian Attention Mechanisms for Transformers: A Theoretical Framework and Architecture Design

模型架构Transformer

摘要

所有基于 Transformer 的 大语言模型 通过欧几里德内积计算注意力,这是 Dong 等人的架构选择。 (2021)证明,在纯自注意力堆栈中,表征排名会随着深度呈双指数衰减。我们开发了一个理论框架,通过用学习的每个词元黎曼度量代替平面欧几里得度量来针对数学层面上的这种结构限制。我们的贡献是三重的。 (1) 我们证明具有异构每个标记度量的黎曼注意力分数是非 Gram 的——它们不能被分解为分解维度为 O(d) 的 QK^T。我们明确表示这是一个结构性观察,而不是排名保持的证明。 (2) 我们建立了低秩度量因子使所有几何运算变得易于处理:每个标记的测地距离为 O(d*r),通过 Woodbury 恒等式为 O(d*r^2) 的度量反演——两者都远低于一般矩阵的 O(d^3) 成本——使得黎曼注意力在十亿参数规模上可行,且开销可以忽略不计。 (3) 我们提出了 Fiber Bundle Transformer,这是一个完整的架构规范,其中每个词元位置都带有自己的黎曼度量,注意力是测地距离计算,前馈更新使用度量预处理步骤,并且连接带有显式曲率和扭转代理。我们得出了关于正确实现的几何架构的正式预测,并确定了中心开放问题:证明或反驳异构黎曼度量可以防止行随机注意力矩阵导致的排名崩溃。本文提出了理论分析和架构设计;实证验证是未来工作的主题。