PACK3D 提出针对打包 3D 数据的局部-全局注意力算子
Geometry Matters in Packed 3D Attention
做 3D 视觉或体数据建模的可以看看：这篇指出了 packed attention 学到内存布局而非几何的坑，PACK3D 用真实坐标的轴向 RoPE 解决，还有 Triton 实现，65k token 下速度和窗口注意力一样快。
arXiv 论文指出，3D 体数据打包成 token 序列后，序列位置不再对应真实几何结构，基于序列位置的编码会让注意力学到内存布局而非物体本身。作者提出 PACK3D：局部邻域注意力保持精确计算，远距离上下文压缩为 3D 块摘要，并用基于真实 token 网格坐标的轴向 RoPE 索引所有 query-key 交互。在合成实验中，学习式位置编码和打包序列 RoPE 在更换 Morton、块连续、随机等重打包方式时全部崩溃，轴向网格坐标 RoPE 保持精确。在 Allen WTC-11 hiPSC 显微成像的 15 结构联合任务上，PACK3D 超过参数量匹配的 3D U-Net 和学习式位置的局部-全局基线；其 Triton 融合实现在 65k token 时延迟与窗口注意力持平。
Geometry Matters in Packed 3D Attention
Packing a 3D volume into a token sequence is not merely long-sequence modeling. Once the volume is packed, sequence order no longer encodes the geometry of the sample---so a positional mechanism indexed by sequence position can learn the memory layout rather than the object. This is a general failure mode for attention over packed geometric data. We introduce PACK3D, a local-global attention operator that separates memory layout from attention geometry. PACK3D keeps local neighborhood attention exact, compresses only distant context into contiguous 3D block summaries, and indexes every query-key interaction with axial RoPE on true token-grid coordinates rather than packed positions. Because the rotation is applied before a block is summarized, the query-to-summary logit remains a weighted sum of query-relative 3D displacements; compressed global context therefore stays spatially interpretable after packing, instead of collapsing to anonymous pooled content. We measure the failure mode directly. On a controlled synthetic task---train on one packing, change only the memory order of the tokens, hold the physical labeling rule fixed---learned positional embeddings and packed-sequence RoPE solve the training layout perfectly yet collapse under Morton, block-contiguous, and random repacking, while axial grid-coordinate RoPE stays exact throughout. The same coordinate system carries to real data: on Allen WTC-11 hiPSC microscopy, PACK3D improves over a parameter-matched 3D U-Net on the 15-structure joint task and over a matched learned-position local-global baseline, with controls isolating the coordinate system, not encoder capacity, as the source of the gain. A fused Triton implementation matches windowed-attention latency at 65k tokens while supplying the long-range context windowing cannot---so the operator is not only correct but trainable at scale.