跳转至

Parabolic Position Encoding: Vision-Centric, Principled, Extrapolatable, General

会议: NeurIPS 2026
arXiv: 2602.01418
OpenReview: https://openreview.net/forum?id=aHGX6FOsYg
代码: https://github.com/DTU-PAS/parabolic-position-encoding
领域: 其他
关键词: 视觉位置编码, 抛物线编码, 分辨率外推, 视觉Transformer, FlashAttention兼容

一句话总结

针对视觉模态对几何与多维空间的独有要求,本文提出基于抛物线结构的视觉位置编码(PaPE),通过将注意力偏置分解为负二次项距离衰减、一次项方向引导与上下文自适应调制,在原生支持 FlashAttention 算子的同时,实现了在 \(1024^2\) 超高分辨率下领先次优方法 10.5% 的优异外推性能与跨 4 种模态的泛化表现。

研究背景与动机

Vision Transformer 在图像、视频、事件相机流以及 3D 点云等多样化视觉任务中得到了广泛应用。然而,尽管视觉注意力机制在各模态百花齐放,模型所依赖的位置编码体系却大多是从自然语言处理领域直接移植而来的,例如经典的绝对正弦位置编码(Sinusoidal)、旋转位置编码(RoPE)以及线性注意力偏置(ALiBi)。这些方案最初针对一维离散符号序列设计,当研究者尝试通过轴向分解(Axial RoPE)、多维欧氏距离扩展(2D-ALiBi)或独立学习频率(RoPE-Mixed)将其推广到二维或多维视觉空间时,往往只考虑了部分几何特性,难以全面契合视觉模态的本质规律。

视觉空间与一维文本存在本质差异:其一,视觉目标具备平移不变性,局部相对结构比绝对画布坐标更具判别力;其二,视觉感知呈现固有的局部聚集性(距离衰减),但在特定场景下又必须依赖长程全局交互;其三,二维平面与三维空间具备极强的几何方向性(上下左右、对角线与运动矢量),不同朝向的交互往往具有截然不同的语义含义;其四,部分 3D 几何与多模态任务还要求严格的旋转不变性(如空间物体的方向不应改变语义分类)。现有编码方案难以同时兼顾这些原则:Axial RoPE 独立处理各坐标轴而割裂了对角线方向建模;LookHere 通过给各注意力头分配固定视场角引入方向性,却丢失了旋转不变性并破坏了与通用高效算子的兼容;ALiBi 类的密集注意力偏置矩阵则因必须显式实例化 \(N \times N\) 偏置图,无法直接调用 FlashAttention 等低显存优化算子。

为了填补这一底层设计鸿沟,本文系统归纳了平移不变性、旋转不变性、距离衰减、方向敏感性与上下文自适应五大视觉位置编码核心准则。作者突破传统复数旋转与标量距离偏置的思路,将注意力打分函数重构为以相对空间坐标为自变量的抛物线函数簇求和,并推导出了无需显式实例化注意力矩阵的独立 Query-Key 映射形式。核心 idea:将两两 Token 的空间相对位移作为一组可学习凹抛物线的自变量,通过二次项实现距离惩罚、一次项捕获几何方向、并由 Token 内容自适应生成抛物线系数,以可与 FlashAttention 算子精确等价的 Query-Key 投影形式统一了多维视觉位置编码。

方法详解

整体框架

PaPE 的核心目标是在注意力计算中注入具备几何完备性与内容自适应能力的相对位置信息。对于输入 Token 集合 \(X = \{x_1, \dots, x_n\}\) 及其关联的物理空间坐标 \(r_i \in \mathbb{R}^p\)(图像 \(p=2\),点云与时空序列 \(p=3\)),标准注意力打分矩阵由 Query 与 Key 的内积 \(S_{ij} = \langle q_i, k_j \rangle\) 构成。PaPE 将这一打分改写为以相对坐标向量 \(\Delta r_{ij} = W_p (r_j - r_i)\) 为自变量的 \(m\) 个一维通用抛物线 \(y = a x^2 + b x + c\) 的求和,从而将空间相对几何自然分解为二次项、一次项与常数语义项。

为了在训练与推理中摆脱对全量 \(N \times N\) 注意力偏置矩阵的显式存储依赖,PaPE 进一步推导了一套解析变换,将包含二次项与一次项的相对偏置严格重写为 Query 和 Key 的独立绝对位置变换内积 \(\langle f_q(q_i, r_i), f_k(k_j, r_j) \rangle\),使整个计算流程完全兼容 FlashAttention-2 等高效算子。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入视觉 Token 特征 x_i 与空间坐标 r_i"] --> D1["相对坐标投影与自适应系数生成<br/>投影矩阵 W_p 生成 Delta r_ij,Token 动态预测 a_i 与 b_i"]
    D1 --> D2["抛物线分解:距离衰减与方向敏感调控<br/>负 Softplus 约束二次项主导衰减,一次项匹配空间几何朝向"]
    D1 -.->|参数置零与正交对称约束| D4["旋转不变特化 PaPE-RI<br/>消除方向项 b_i 并固定标量权重,满足 SO(p) 对称"]
    D2 --> D3["Query-Key 显式解耦映射<br/>独立升维变换绝对位置,无缝兼容 FlashAttention 算子"]
    D4 -.-> D3
    D3 --> Out["高效注意力计算与多模态表征输出<br/>视频/事件/图像/点云的高精度表征"]

关键设计

1. 抛物线注意力分解与自适应系数生成:将相对几何解耦为二次衰减与线性方向

针对以往位置编码要么仅考虑径向标量距离、要么只能进行一维正弦编码的缺陷,PaPE 从一维抛物线方程 \(y = a x^2 + b x + c\) 出发,将注意力打分矩阵展开为多项式形式。对于任意两个 Token \(i\) 和 \(j\),首先通过可学习矩阵 \(W_p \in \mathbb{R}^{m \times p}\) 对物理相对坐标进行空间线性投影: $\(\Delta r_{ij} = W_p (r_j - r_i)\)$ 其中 \(m\) 为抛物线基底的数量。打分函数展开为: $\(S_{ij} = \sum_{\ell=1}^m \left( a_{i\ell} \Delta r_{ij\ell}^2 + b_{i\ell} \Delta r_{ij\ell} + \frac{\langle q_i, k_j \rangle}{m} \right) = \langle a_i, \Delta r_{ij}^{\odot 2} \rangle + \langle b_i, \Delta r_{ij} \rangle + \langle q_i, k_j \rangle\)$ 式中 \(\odot 2\) 表示逐元素哈达玛平方。抛物线系数 \(a_i, b_i \in \mathbb{R}^m\) 并非全局固定的死参数,而是直接由当前 Token 的语义特征 \(x_i\) 经由可学习线性层 \(W_a, W_b \in \mathbb{R}^{m \times d}\) 动态生成: $\(a_i = -\text{softplus}(W_a x_i), \quad b_i = W_b x_i\)$ 将 \(a_i\) 强制约束在负实数域 \(\mathbb{R}_{<0}^m\),保证了每一个单项抛物线开口始终严格向下(凹函数)。这一设计使得模型天然具备内容自适应性(Context Awareness):当 Token 需要关注长程全局依赖时,网络可以通过调整 \(W_a x_i\) 使得 \(a_i \to 0\) 且 \(b_i \to 0\),抛物线偏置自动退化为标准语义注意力 \(\langle q_i, k_j \rangle\);而在需要细粒度局部感知时,则可通过大负数的 \(a_i\) 施加急剧的距离惩罚。

2. 距离衰减与方向性调控机制:兼具局部偏置与几何指向性

视觉表征既要求距离衰减(近邻 Token 强相关),又要求方向选择性(如检测目标的水平延展或视频中动作的时空轨迹)。PaPE 的二次项与一次项分别精准负责这两项几何属性: - 距离衰减:由于 \(a_i\) 元素皆为负值,当距离 \(\|\Delta r_{ij}\|\) 增大时,二次惩罚项 \(\langle a_i, \Delta r_{ij}^{\odot 2} \rangle\) 会单调降低注意力打分,赋予模型强烈的局部交互归纳偏置; - 方向敏感性:向量内积 \(\langle b_i, \Delta r_{ij} \rangle = \|b_i\| \|\Delta r_{ij}\| \cos(\theta)\),在相对位移方向与系数向量 \(b_i\) 完全平行同向(\(\theta = 0\))时取得极大值。因此,\(b_i\) 实际在特征空间中定义了一个最优偏好方向,使得注意力打分能够沿着特定几何轴向延伸,完美捕捉物体的长宽比特征与连续运动趋势。

3. FlashAttention 兼容的 Query-Key 显式解耦映射:消除 \(N \times N\) 显存瓶颈

以往基于相对位置的 Attention Bias 方法(如 ALiBi、T5 Bias)必须在注意力打分前显式构造 \(N \times N\) 的稠密偏置矩阵,无法使用 FlashAttention 这类基于 IO 感知分块的融合算子。PaPE 通过代数展开,将复杂的两两相对位置偏置完全重构为 Query 与 Key 的独立升维投影。定义对称矩阵 \(M_i = W_p^T \text{diag}(a_i) W_p \in \mathbb{R}^{p \times p}\),对于 Query 变换 \(f_q(q_i, r_i)\) 与 Key 变换 \(f_k(k_j, r_j)\),构造如下高维向量拼接: $\(f_q(q_i, r_i) = q_i \oplus r_i^T M_i r_i \oplus \text{flatten}(M_i) \oplus -2 M_i r_i \oplus W_p^T b_i \oplus -\langle W_p^T b_i, r_i \rangle\)$ $\(f_k(k_j, r_j) = k_j \oplus 1 \oplus \text{flatten}(r_j \otimes r_j) \oplus r_j \oplus r_j \oplus 1\)$ 数学证明表明,两者的直接内积展开后正好精确还原式中的抛物线偏置: $\(\langle f_q(q_i, r_i), f_k(k_j, r_j) \rangle = \langle q_i, k_j \rangle + \langle a_i, \Delta r_{ij}^{\odot 2} \rangle + \langle b_i, \Delta r_{ij} \rangle\)$ 在此变换下,Query 和 Key 的维度仅增加了 \(p^2 + 2p + 2\)(二维图像 \(p=2\) 仅增加 10 维,三维点云/视频 \(p=3\) 仅增加 17 维)。模型可以直接调用未修改的标准 FlashAttention-2 算子进行高效计算,完全规避了中间矩阵的显存开销。

4. 旋转不变特化变体 PaPE-RI:针对 3D 几何与多模态的对称性约束

在部分特定任务中(例如无固定朝向的 3D 点云分类或多传感器融合),模型预测不应随物体的刚体旋转而改变。为了使位置编码严格满足旋转不变性(Definition 3.2: \(f(R r_i, R r_j) = f(r_i, r_j), \forall R \in \text{SO}(p)\)),PaPE 引入了约束变体 PaPE-RI: - 将一次项系数置零:\(b_i = 0\)(牺牲方向性); - 约束二次项系数为各向同性标量:\(a_{ik} = \alpha_i \in \mathbb{R}_{<0}\); - 约束投影矩阵为标量对角阵:\(W_p = w_p I_p\)(\(w_p \in \mathbb{R}\))。

在上述约束下,位置项退化为 \(\alpha_i w_p^2 \|r_j - r_i\|^2\)。由于两点间的欧氏距离在正交旋转变换下严格保持不变,PaPE-RI 在理论上具备严格的 \(\text{SO}(p)\) 旋转不变性,为特定几何场景提供了无偏归纳先验。

损失函数 / 训练策略

PaPE 属于通用底层注意力结构改进,无需定制的辅助损失函数。在所有模态的实验中,均沿用各基准任务原本的端到端损失函数(例如 ImageNet 图像分类的交叉熵损失、YOLOv10 目标检测的多任务损失、点云分割的交叉熵等)。超参数设置上,抛物线数量在 CIFAR-10 验证集上被统一定为 \(m=50\),在后续所有 8 个数据集上均直接沿用而未做任何调优。网络参数均采用 uniform Kaiming 初始化,优化器采用 AdamW,并采用 16-bit 混合精度加速训练。

实验关键数据

主实验

主实验从跨分辨率外推能力与跨多模态通用性两个维度全面评估 PaPE 的性能。

1. ImageNet-1K 分辨率外推能力评估
所有模型统一在 \(224^2\) 分辨率下训练(ViT 架构,不使用额外外推微调),直接在 \(320^2 \sim 1024^2\) 测试分辨率下进行零样本尺寸外推:

位置编码方案 \(224^2\) (训练) \(384^2\) \(512^2\) \(640^2\) \(768^2\) \(896^2\) \(1024^2\) \(1024^2\) 相对次优增益
nD-sincos 79.9 72.6 57.5 41.3 27.2 16.8 9.3 -
RoPE (Axial) 80.3 80.6 78.2 74.0 67.0 56.2 42.3 -
RoPE-Mixed 80.0 79.2 73.6 61.5 43.4 27.1 17.1 -
nD-ALiBi 80.0 77.8 75.3 72.2 68.3 64.1 59.1 -
LookHere 80.5 81.1 79.7 77.5 74.2 69.6 64.4 (次优基线)
PaPE (本文) 80.6 81.7 81.3 80.5 79.2 77.3 74.9 +10.5% (绝对)
PaPE-RI (本文) 80.0 79.2 76.9 73.6 69.7 64.6 58.6 -

在最佳配置(允许基线采用位置插值 Position Interpolation)下,即使 RoPE 和 nD-sincos 借助插值在 \(1024^2\) 下回升至 59.8% 和 56.0%,未经任何插值的原生 PaPE 依然以 74.9% 的高精度超出所有竞争对手 10.5% 以上。

2. 跨 4 种视觉模态的 8 大数据集通用性评测

模态类别 数据集 任务类型 评估指标 nD-sincos RoPE RoPE-Mixed nD-ALiBi LookHere PaPE PaPE-RI
时空视频 UCF101 动作识别 Top-1 准确率 38.7 43.9 41.2 41.6 - 49.5 42.2
事件流 DvsGesture 手势分类 Top-1 准确率 93.4 93.1 83.3 89.2 - 93.4 90.6
事件流 GEN1 自动驾驶目标检测 mAP@[0.5:0.95] 27.3 33.6 31.7 28.8 - 34.8 28.5
图像 ImageNet-1K 图像分类 Top-1 准确率 79.9 80.3 80.0 80.0 80.5 80.6 80.0
图像 COCO 目标检测 mAP@[0.5:0.95] 34.7 38.8 38.0 33.9 37.3 38.9 34.3
点云 ScanNet 点云语义分割 mIoU 72.6 71.7 71.1 70.6 - 71.0 69.3
点云 ModelNet40 3D 物体分类 Top-1 准确率 92.2 93.0 93.0 93.2 - 93.3 92.7
多模态 nuScenes 自动驾驶 3D 检测 mAP 67.4 68.3 66.3 68.5 - 68.7 68.9
平均 8 任务总览 - 平均分 63.3 65.3 63.1 63.2 - 66.3 63.3

消融实验

在 ImageNet-1K 上以 ViT-S 架构训练 150 个 epoch,系统性验证各项核心设计对最终精度的独立贡献:

消融配置 准确率 (%) 相对完整模型性能变化 机制解析说明
Full PaPE 完整模型 79.2 基准 完整包含自适应抛物线衰减、方向项与空间投影
移除二次项 \(\langle a_i, \Delta r_{ij}^{\odot 2} \rangle\) 77.2 -2.0 剥离距离惩罚机制,注意力失去局部先验,精度显著下滑
移除一次项 \(\langle b_i, \Delta r_{ij} \rangle\) 77.0 -2.2 失去方向指引敏感性,各向同性导致几何结构判别受阻
移除上下文调制 (\(W_a, W_b\)) 77.2 -2.0 退化为与内容无关的静态参数,模型无法动态切换局部/长程交互
移除 \(a_i\) 负激活 (Softplus) 77.5 -1.7 开口方向未受限导致距离增加时注意力可能错误放大
移除线性投影 \(W_p\) (\(m=p\)) 78.9 -0.3 降维导致基底数量受限,拟合复杂曲面能力轻微减弱

关键发现

  • 分辨率外推优势极大:PaPE 在训练分辨率 \(224^2\) 放大到 \(512^2\) 时准确率甚至反超训练精度(提升至 81.3%),直到 \(640^2\) 以上才轻微下降;在 \(1024^2\) 极端尺度下依然保持 74.9%,比 LookHere(64.4%)和 RoPE(42.3%)展现出压倒性的鲁棒性,主要得益于负二次抛物线提供了平滑连续的距离衰减边界。
  • 时空动态感知尤为突出:在快速运动密集的 UCF101 视频数据集上,PaPE 达到了 49.5% 的准确率,相较 RoPE(43.9%)取得多达 5.6% 的绝对提升,三折交叉验证统计显著性达到 \(p=0.011\);证明几何方向向量与自适应系数对快速帧间位移具有极其敏锐的捕捉能力。
  • 绝对位置在部分场景下不可替代:在 ScanNet 室内点云语义分割任务中,纯相对位置的 PaPE(71.0)略逊于经典绝对正弦编码 nD-sincos(72.6)。论文指出这是因为室内家具具备绝对高度物理先验(例如桌面通常高出地面 1 米),纯相对编码在局部窗口未完全跨越全局坐标系时会丢失此类绝对空间锚点。

亮点与洞察

  • 抛物线数学构造与高效内核的精妙结合:通常引入相对距离二次方与方向偏置的注意力改进都需要显式计算大尺寸矩阵,而作者通过巧妙利用张量外积与展平向量的线性代数恒等式,将抛物线形式等价分解为独立的 Query-Key 映射,仅付出极其轻微的维度扩展(\(p^2+2p+2\))就直接复用了 FlashAttention 算子。
  • 上下文感知解开“局部偏置与长程关联”的博弈矛盾:传统的 ALiBi 固化了注意力随距离线性衰减的斜率,强行压制了长程依赖;PaPE 将抛物线曲率与顶点方向交由输入 Token 的特征实时动态生成,允许模型在需要长程语义检索时自主通过 \(a_i \to 0\) 消除距离惩罚,赋予网络巨大的表示弹性。

局限与展望

  • 可学习参数量微增:相比于完全零参数的 RoPE 或 sincos,PaPE 引入了 \(W_p, W_a, W_b\) 线性映射矩阵,当基底维度 \(m=50\) 时带来了一定程度的模型参数量增长(例如 ViT-B 在 UCF101 上参数从 86.3M 增至 97.4M,显存开销增加约 7%~20%)。
  • 超参数 \(m\) 的轻量化探索:未来的改进方向包括探索更高效的低秩投影方案,或者直接约束 \(m=p\) 以彻底消除投影层参数,同时探究在底层 CUDA 级别将抛物线映射直接融合编译进 FlashAttention 内核的更深层优化。

相关工作与启发

  • vs RoPE / Axial RoPE: RoPE 通过复数旋转在一维序列中表现优异,但多维推广时往往割裂坐标轴;PaPE 通过多维联合投影 \(W_p\) 与显式二次距离衰减,原生自然地刻画了多维欧氏几何结构与对角线交互。
  • vs 2D-ALiBi & LookHere: 2D-ALiBi 具备旋转不变性却缺乏几何方向性感应;LookHere 引入了方向性却打破了旋转对称且仅能应用于 2D 图像,且无法使用 FlashAttention。PaPE 兼顾了方向性、上下文感知与内核高效性,并通过 PaPE-RI 提供了即插即用的旋转对称变体。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次将多维抛物线注意力偏置理论体系化引入视觉模态,并给出了与高效内核完全相容的代数解耦证明。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 4 种模态、8 个跨领域主流数据集、分辨率达 \(1024^2\) 的极端外推实验以及充分的参数匹配消融对比。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨自洽,原理图示与设计动机阐述条理清晰。
  • 价值: ⭐⭐⭐⭐⭐ 解决了 Vision Transformer 长期套用语言编码的根本矛盾,为高分辨率推理与跨模态通用视觉架构奠定了坚实的基础组件。