Geometry Matters in Packed 3D Attention¶

会议: NeurIPS 2026
arXiv: 2610.38694
论文: https://arxiv.org/abs/2610.38694
领域: 3D表示与场景理解
关键词: 3D注意力机制, 序列紧缩, 轴向RoPE, 局部-全局分解, 体积显微图像分割

一句话总结¶

针对三维高Token体积场展平打包后内存存储顺序破坏物理空间几何的问题，Pack3D提出了物理坐标与内存布局解耦的局部-全局注意力算子，通过轴向网格RoPE在块压缩前旋转保留相对空间位移，在Allen WTC-11显微图像分割中超越同等参数量3D U-Net并在100万Token下实现相比密集注意力125倍加速。

研究背景与动机¶

在活细胞三维荧光显微成像、医学CT/MRI体积扫描以及高光谱遥感等高分辨率3D领域中，单幅体积样本经过分块（Patchify）与线性化展平后往往包含数万至数十万个Token。为了系统层面的显存利用与硬件访存效率，这些Token在内存中通常按照光栅扫描（Raster）、Morton Z序（Morton/Z-order）或块连续（Block-contiguous）方式紧密打包排列。然而，这种存储布局是纯粹的系统工程选择，并非物理坐标系。当模型使用传统的绝对可学习位置编码或基于1D一维序列索引的RoPE旋转位置编码时，自注意力机制直接将内存存储位置误当成样本的物理几何结构，导致模型学到的只是特定的显存排列规律而非真实的物理对象几何。

这种位置与几何错配的矛盾在局部-全局（Local-Global）稀疏注意力范式中尤为致命。密集全自注意力（Dense Attention）计算复杂度随Token数量呈平方级爆炸，无法扩展到长序列；而局部窗口注意力（Window Attention）虽然扩展性好，却受制于感受野孤岛，丢失了细胞或器官尺度的全局空间排布先验。局部-全局注意力是理想的折中方案：近处保持精确交互，远处通过块聚合压缩降低计算量。但远距离压缩操作天然依赖底层位置机制——若位置编码基于序列索引，块池化操作将直接抹除相对坐标关系，使压缩块退化为缺乏空间指向性的匿名内容。此时无论图稀疏结构如何精巧，注意力机制所关注的都是错误的几何流形。

解决这一矛盾的核心切入点在于将“内存存储布局”与“注意力几何坐标”显式解耦。存储顺序可以为了硬件访存连续性自由重排，但查询与键之间的几何交互必须始终绑定在体素网格的真实3D物理坐标上。核心 idea：将3D场展开与坐标表示彻底解耦，构建保真局部邻域与块压缩全局上下文并行的局部-全局算子，并在块级内容聚合前应用物理网格坐标驱动的轴向RoPE，使查询到压缩块的注意力交互严格保持为相对空间位移的加权叠加，确保全局压缩表征在任意打包重排下依然具备几何可解释性与打包等变性。

方法详解¶

整体框架¶

Pack3D是一个兼顾几何一致性与硬件执行效率的高Token三维局部-全局注意力算子。模型输入为展平打包后的三维Token序列，每个Token不仅携带特征向量，还绑定其在原始体素晶格中的真值网格三维物理坐标 $c_i = (z_i, y_i, x_i) \in \mathbb{R}^3$。在主干网络的前向计算中，Pack3D将注意力分解为两条并行分支：精确局部邻域分支负责捕获细胞膜、点状斑等高频形态学细节，坐标感知压缩全局分支负责以较低复杂度汇聚大范围空间拓扑；最后通过逐Token的可学习门控混合模块融合两分支输出，返回残差连接。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入3D体积Token<br/>绑定物理网格坐标 c_i=(z_i,y_i,x_i)"] --> B["线性投影与层归一化<br/>生成 Q, K_loc, V_loc, K_cmp, V_cmp"]
    B --> C["1. 轴向网格RoPE旋转<br/>基于物理坐标 R(c_i) 调制Q与K"]
    C --> D["2. 精确局部邻域注意力<br/>Triton流式核函数计算局部窗"]
    C --> E["3. 块级摘要压缩与全局交互<br/>聚合前旋转保持相对空间位移"]
    D --> F["4. 自适应双分支门控混合<br/>Token级动态权重融合局部与全局"]
    E --> F
    F --> G["残差输出与前馈网络"]

关键设计¶

1. 物理坐标与内存布局解耦：轴向网格RoPE与打包等变性

针对传统位置编码随序列打包顺序变化而崩溃的问题，Pack3D完全废弃了基于序列下标的位置索引，直接采用3D Token网格物理坐标 $c_i = (z_i, y_i, x_i)$ 驱动轴向RoPE（Axial Rotary Position Embedding）。对于每个特征头维度 $d_h$，轴向旋转矩阵 $R(c_i) \in \mathbb{R}^{d_h \times d_h}$ 被分解为 $z, y, x$ 三个坐标轴上平面旋转的正交直和： $$R(c) = \bigoplus_{a \in \{z,y,x\}} \bigoplus_{r=0}^{m_a-1} \begin{bmatrix} \cos(\omega_{a,r} c_a) & -\sin(\omega_{a,r} c_a) \\ \sin(\omega_{a,r} c_a) & \cos(\omega_{a,r} c_a) \end{bmatrix} \oplus I_{\mathrm{rem}}$$ 其中对 $d_h=64$ 采用平衡分配 $m_z=11, m_y=11, m_x=10$，基频设为 10,000。Pack3D严格只对查询 $q_i^R = R(c_i) W_Q \text{LN}(h_i)$ 和键 $k_i^{b,R} = R(c_i) W_K^b \text{LN}(h_i)$ 施加旋转，值向量 $v_i^b = W_V^b \text{LN}(h_i)$ 保持不旋转，因为值向量承载语义内容而非空间位置。利用正交变换的内积不变性 $R(c_i)^\top R(c_j) = R(c_j - c_i)$，注意力Logit仅取决于物理相对空间位移 $c_j - c_i$，与内存索引彻底无关。由此证明了打包等变性定理（Proposition 1）：对于任意保持物理坐标绑定的重排置换 $P$，算子严格满足 $\mathcal{A}(PX, Pc) = P \mathcal{A}(X, c)$。

2. 压缩前旋转与相对几何保持：无损相对位移的块级摘要

在非局部上下文中，为将长序列注意力复杂度从 $\mathcal{O}(N^2)$ 降低到 $\mathcal{O}(NM)$，Pack3D将全域划分为 $M$ 个空间连续的3D块 $\mathcal{B} = \{B_1, \dots, B_M\}$（例如每个块体积 $v = |B_j| = 1 \times 4 \times 4 = 16$）。核心难点在于：常规池化操作（如先均值池化再附加位置）会导致块内相对几何信息的丢失。Pack3D提出了“压缩前旋转（Token-before-compression）”机制：每个块的键摘要是在旋转后的Token键上进行归一化加权平均： $$\tilde{k}_j^R = \sum_{u \in B_j} a_{j,u} R(c_u) W_K^{\mathrm{cmp}} \text{LN}(h_u), \quad \tilde{v}_j = \sum_{u \in B_j} a_{j,u} W_V^{\mathrm{cmp}} \text{LN}(h_u)$$ 由于旋转算子在线性加权下具有分配律，当任意查询 $q_i^R$ 与压缩摘要 $\tilde{k}_j^R$ 计算点积时，满足精确等式（Proposition 2 / Theorem 7）： $$(q_i^R)^\top \tilde{k}_j^R = \sum_{u \in B_j} a_{j,u} q_i^\top R(c_u - c_i) k_u$$ 这意味着查询与全局块摘要之间的注意力交互，精确等价于该块内所有单元相对于查询的空间位移 $c_u - c_i$ 的加权叠加。压缩后的摘要不再是匿名的模糊向量，而是携带细粒度相对空间导向的几何表征。理论分析进一步证明（Theorem 8），相较于坐标盲（Coordinate-blind）的传统压缩，这种几何感知压缩在区分对称或等距不同区域时具备严格的表征优势。

3. 硬件对齐的局部-全局算子实现与流式执行

算法上的稀疏性若依赖显式张量重排（Gather）或大尺寸路由矩阵物化，其理论加速比通常会被显存带宽开销完全吞噬。Pack3D设计了完全对齐现代GPU访存体系的流式融合算子。局部邻域分支 $\mathcal{N}(i)$ 基于物理网格坐标定义（如 $[1,1,1]$ 半径体素窗），采用定制的 Triton 算子，在片上 SRAM 中流式加载局部 K/V Tile，避免在全局显存中物化巨大的局部邻域张量。全局压缩分支仅物化尺寸极小的 $M$ 个摘要向量，查询与摘要间的 Softmax 同样采用流式在线归约（Online Softmax）。系统总算力复杂度分解为： $$T_{\mathrm{LG}} = \underbrace{\Theta(H_q N w d_h)}_{\text{精确局部分支}} + \underbrace{\Theta(H_q N M d_h)}_{\text{查询-摘要分支}} + \underbrace{\Theta(H_{kv} M v d_h)}_{\text{摘要构建}}$$ 在 $N=262,144$、$M=16,384$ 的典型高Token场景下，全局比较次数相比密集注意力减少了 $v=16$ 倍，且执行全程无需在主路径上构造密集路由张量。

损失函数 / 训练策略¶

模型端到端训练采用医学图像分割标准的软 Dice 损失与二值交叉熵（BCE）加权和： $$\mathcal{L}_{\mathrm{seg}}(\hat{y}, y) = \lambda_{\mathrm{dice}} \mathcal{L}_{\mathrm{dice}}(\hat{y}, y) + \lambda_{\mathrm{bce}} \mathcal{L}_{\mathrm{bce}}(\hat{y}, y)$$ 实验中两项损失权重均设为 1.0。在 Allen 显微图像数据集上，输入通道包含荧光目标通道、明场通道（Brightfield）和细胞核 DNA 通道，输入体积尺寸为 $32 \times 512 \times 256$ 体素，经 $1 \times 4 \times 4$ 分块后形成 $32 \times 128 \times 64 = 262,144$ 个 Token，监督区域为中心 $16 \times 256 \times 256$ 体素（对应 65,536 个监督 Token）。模型采用 AdamW 优化器训练 10k 步，学习率及衰减策略在所有对比基线中严格对齐。

实验关键数据¶

主实验¶

论文在 Allen WTC-11 hiPSC 显微镜数据集的15结构联合预测任务（包含 TOMM20、LMNB1、ST6GAL1、GJA1、LAMP1、ATP2A2 等）以及不同位置机制控制实验上进行了全面评估。

表1：Allen-15 联合任务对比（3个随机种子均值±标准差）

模型架构	参数量	Best Macro Dice	Final Macro Dice	架构特性
Pack3D (Ours, Depth-8)	30.1M	0.8228 ± 0.0008	0.8228 ± 0.0007	物理网格坐标轴向RoPE + 局部全局稀疏
TokenGrid 3D U-Net (base128)	31.9M	0.8134 ± 0.0034	0.8093 ± 0.0039	参数对齐的标准三维卷积网络基准

在严格匹配参数量（~30M）的前提下，Pack3D 在整体 15 结构平均指标上超越了卷积架构基准 3D U-Net（Best Dice 提升 +0.0094，Final Dice 提升 +0.0135），验证了坐标感知注意力在处理复杂微观三维胞内空间结构时的有效性。

表2：位置机制对齐对比实验（262k 与 65k Token 尺度，短周期门控测试）

Token尺度 / 深度	位置编码机制	参数量	Best Dice	Final Dice	属性归属
262k / d1	全网格可学习位置编码 (Learned full-grid PE)	137.3M	0.7086	0.6525	显式绝对网格参数表
262k / d1	可分离可学习 3D PE (Separable 3D PE)	3.18M	0.7151	0.7151	紧凑型可学习位置编码
262k / d1	低秩可学习 3D PE (Low-rank 3D PE)	3.07M	0.6891	0.6852	紧凑型低秩分解
262k / d1	坐标傅里叶 MLP (Coordinate Fourier MLP)	3.13M	0.6942	0.6942	坐标驱动的连续隐式编码
262k / d1	打包序列 RoPE (Packed-sequence RoPE)	3.06M	0.7142	0.6351	依赖内存一维排布的RoPE
262k / d1	轴向网格 RoPE (Axial grid RoPE, Ours)	3.06M	0.7375	0.7041	真值物理网格坐标驱动
65k / d1	3D 相对位置偏置 (3D RPB)	3.06M	0.7055	0.6420	加性偏置（耗时高达588s）
65k / d1	3D ALiBI 风格偏置	3.06M	0.7098	0.6691	预设斜率的加性几何偏置
65k / d1	轴向网格 RoPE (Axial grid RoPE, Ours)	3.06M	0.7346	0.7154	乘性无缝正交旋转（耗时仅79.4s）

消融实验¶

表3：合成可控重排实验与推理期压缩分支扰动消融

为证明序列位置编码学到的只是内存布局，论文在合成 3D 几何标注任务上测试不同重排下的泛化性（左），并在训练收敛后的 Pack3D 模型上对压缩分支施加推理期扰动（右）：

测试时打包方式（合成任务Dice）	无位置编码	可学习PE	序列RoPE	轴向网格RoPE
光栅顺序 Raster (训练集)	0.554	1.000	1.000	1.000
Morton 空间填充曲线序	0.554	0.449	0.449	1.000
块连续重排 Block-contiguous	0.554	0.418	0.418	1.000
完全随机打乱 Random	0.554	0.378	0.378	1.000

推理期压缩分支扰动条件（262k / d8 模型）	Macro Dice	相比完整模型 $\Delta$	注意力Logit $L_2$ 偏差
完整模型 Full Pack3D	0.8531	—	0.0000
单独训练的纯局部模型 (Local-only)	0.8387	-0.0144	1.3101
压缩分支门控强制置零 ($g_{\mathrm{cmp}} \to 0$)	0.7648	-0.0883	1.5702
压缩摘要全局随机打乱 (Summary shuffle)	0.6858	-0.1674	1.4834
压缩分支输出强制置零 ($o_{\mathrm{cmp}} \to 0$)	0.6775	-0.1756	1.5444
局部有效门控强制置零 ($g_{\mathrm{loc}} \to 0$)	0.5764	-0.2767	2.9167

表4：单层上下文长度扩展性能对比（Pool-64，NVIDIA GPU测量）

Token 数量	密集注意力 SDPA (ms)	Pack3D (ms)	纯窗口注意力 Window (ms)	Dense / Pack3D 加速比	Pack3D 显存 (GiB)	Window 显存 (GiB)
65k (65,536)	104.55	16.11	13.86	6.49×	1.51	2.04
131k (131,072)	405.62	28.36	26.62	14.30×	2.95	4.02
262k (262,144)	1620.07	54.64	52.49	29.65×	5.82	7.98
524k (524,288)	6641.53	107.49	104.19	61.79×	11.58	15.89
1M (1,048,576)	27034.70	215.51	207.76	125.45×	23.09	31.72

关键发现¶

序列编码学到的是内存索引而非几何：在合成重排测试中，可学习 PE 和序列索引 RoPE 在训练集（Raster）上均能达到 1.000 准确率，但一旦测试时将内存布局改为 Morton 或随机重排，Dice 骤降至 0.378~0.449，甚至显著劣于完全不加位置编码的基线（0.554），说明序列位置信号在重排下变成了严重误导性的假信号；而轴向网格 RoPE 始终稳定在 1.000。
全局压缩分支具有实质性承载作用：在训练期，引入压缩全局分支相比单独训练的纯局部模型使 Best Dice 提升 +0.0144（0.8387 $\to$ 0.8531）；在推理期，直接置零压缩分支输出或打乱压缩摘要会导致 Dice 发生 0.16~0.17 的严重塌陷，证实压缩上下文绝非装饰性结构，而是网络推理的支柱。
近乎纯窗口的硬件执行开销：在系统扩展性测试中，Pack3D 的实际延迟仅比极简的纯局部窗口注意力高出约 3%~16%（如 1M Token 下 215.51ms vs 207.76ms），但远快于密集注意力的 27 秒，成功以极其微弱的系统代价补齐了长程全局上下文。

亮点与洞察¶

“压缩前旋转”的代数自洽性：论文通过正交变换直和的代数性质，巧妙证明了在聚合前对 Token 施加坐标旋转后，查询与摘要的点积天然退化为该摘要块内所有单元物理相对位移的加权和。这一设计以极小的计算开销彻底解决了非局部粗粒度上下文的空间指向模糊问题。
系统工程与几何建模的明晰解耦：论文确立了一个关键范式——Token 打包与重排是属于加速器访存连续性的“系统工程决策”，而注意力距离度量是属于“物理几何建模”的建模决策。二者绝不能混淆。
极其实事求是的论证态度：作者并未盲目吹嘘“SOTA 分割”或“全面超越 U-Net”，而是明确指出在单结构最终收敛 Checkpoint 上，3D U-Net 在 5/6 的结构上依然具备微弱优势；深度为 8 时对可学习 PE 的巨大提升（+0.0978）很大程度上来源于基线参数量过大导致的欠拟合，真正的坐标增益严格由 3.06M 的等参数控制实验给出。

局限与展望¶

作者承认的局限：当前的实验任务集中于 Allen WTC-11 荧光显微结构分割，输入包含目标荧光通道，任务并非纯粹的无标记（Label-free）图像预测，也不代表通用手工金标准分割的最佳水平；跨模态（如 MSD 医疗 CT/MRI）的迁移实验呈现混合结果，泛化性仍有待更深入检验。
几何各向异性的调优空间：显微镜在 Z 轴与 XY 平面往往具有极度各向异性的体素分辨率（如 $1\times 4\times 4$），目前采用统一网格坐标与平衡频率分配虽然简洁稳健，但未充分发掘各向异性专用频域分配的潜力。
未来方向：探索将该几何坐标注意力算子推广至NeRF/3D高斯辐射场等神经体积表征，以及空间转录组学等具有不规则空间物理坐标的高维密集序列建模中。

评分¶

新颖性: ⭐⭐⭐⭐☆ [深刻揭示了序列展平打包与物理几何解耦的关键问题，提出了代数严格自洽的压缩前旋转机制]
实验充分度: ⭐⭐⭐⭐⭐ [控制实验异常严密，合成重排、参数量严格对齐控制、推理期扰动消融与硬件延迟测试层层闭环]
写作质量: ⭐⭐⭐⭐⭐ [论证严谨，边界划定极其诚实清醒，理论证明与系统实现相互呼应]
价值: ⭐⭐⭐⭐☆ [为 3D 显微视觉、高维体数据及高Token几何序列注意力建模提供了极具指导意义的通用算子范式]