研究3D注意力打包中的几何信息保持问题
原文：Geometry Matters in Packed 3D Attention
将3D体数据打包为token序列不仅仅是长序列建模。一旦体数据被打包，序列顺序就不再编码样本的几何结构——因此，以序列位置为索引的位置机制可能学到的是内存布局而非对象本身。这是对打包几何数据做注意力时的一种普遍失效模式。我们提出 PACK3D，一种局部-全局注意力算子，将内存布局与注意力几何分离。PACK3D 保持局部邻域注意力精确，仅将远距离上下文压缩为连续的3D块摘要，并用基于真实token网格坐标（而非打包位置）的轴向 RoPE 为每个查询-键交互建立索引。由于旋转在块被摘要之前应用，查询到摘要的 logit 仍然是查询相对3D位移的加权和；因此压缩后的全局上下文在打包后仍保持空间可解释性，而不会坍缩为匿名的池化内容。我们直接测量了该失效模式。在一个受控合成任务上——在一个打包上训练，仅改变token的内存顺序，保持物理标注规则固定——学习的位置嵌入和打包序列 RoPE 能完美解决训练布局，但在 Morton、块连续和随机重打包下均崩溃，而轴向网格坐标 RoPE 始终保持精确。同一坐标系统可迁移到真实数据：在 Allen WTC-11 hiPSC 显微图像上，PACK3D 在15结构联合任务上优于参数匹配的3D U-Net，并优于匹配的学习位置局部-全局基线，对照实验将增益来源隔离为坐标系统而非编码器容量。融合 Triton 实现在65k token时达到窗口注意力的延迟，同时提供窗口化无法给出的长程上下文——因此该算子不仅正确，而且可大规模训练。作者：Amil Khan、Anantajit Subrahmanya、B. S. Manjunath。类别：cs.CV、cs.AI、cs.LG