Geometry Matters in Packed 3D Attention

Amil Khan; Anantajit Subrahmanya; B. S. Manjunath

cs.CVOct 28, 2026

Geometry Matters in Packed 3D Attention

Authors: Amil Khan, Anantajit Subrahmanya, B. S. Manjunath

Organizations: Department of Electrical & Computer Engineering University of California, Santa Barbara

Abstract

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.

Figures & tables

Model	Best Dice	Final Dice	Params
Pack3D (ours), d8	0.8228 $\pm$ 0.0008	0.8228 $\pm$ 0.0007	30.1M
TokenGrid 3D U-Net, base128	0.8134 $\pm$ 0.0034	0.8093 $\pm$ 0.0039	31.9M

Table 1: Allen-15 joint task: Pack3D is competitive with a parameter-matched 3D U-Net, using a coordinate-indexed attention operator in place of convolution. Best and final macro Dice, mean

\pm

std over three seeds. The

+0.0094

+0.0135

gap is consistent in sign but modest against seed variation; see text for the significance annotation.

Setting	Position mechanism	Best	Final	Params	Role
262k/d1	Learned full-grid PE	0.7086	0.6525	137.3M	full-grid learned
262k/d1	Learned separable 3D PE	0.7151	0.7151	3.18M	compact learned
262k/d1	Learned low-rank 3D PE	0.6891	0.6852	3.07M	compact learned
262k/d1	Coordinate Fourier MLP	0.6942	0.6942	3.13M	coordinate MLP
262k/d1	Packed-sequence RoPE	0.7142	0.6351	3.06M	memory-order RoPE
262k/d1	Axial grid RoPE	0.7375	0.7041	3.06M	physical-coord. (matched)

Table 2: At matched parameter count, physical-coordinate RoPE beats sequence-index RoPE and compact learned alternatives. Short Allen gates; best/final macro Dice. The matched-capacity comparisons are the 3.06M rows — in particular packed-sequence vs. axial RoPE at identical 3.06M — which isolate the coordinate system. The depth-8 pair differs in capacity (158.5M vs. 24.3M) and is shown only to confirm the effect persists at depth.

Condition	Macro Dice	$\Delta$ vs. full
Full Pack3D	0.8531	—
Trained local-only	0.8387	$-0.0144$
Compressed gate $\to 0$	0.7648	$-0.0883$
Compressed summary shuffle	0.6858	$-0.1674$
Compressed output $\to 0$	0.6775	$-0.1756$
Local gate $\to 0$	0.5764	$-0.2767$

Table 3: The trained compressed pathway is load-bearing: corrupting it at inference causes large Dice drops. Inference-time perturbations of trained 262k/depth8 Pack3D checkpoints, mean over three seeds. These show reliance on the compressed branch, not that the branch is coordinate-aware (see text).

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Item	Value
Dataset	Allen WTC–11 hiPSC microscopy, multistructure preprocessed FOV crops.
Structures	TOMM20, LMNB1, ST6GAL1, GJA1, LAMP1, ATP2A2.
Task	Target fluorescence + brightfield + DNA to Allen-generated structure segmentation.
Not claimed	Label-free prediction; manual labels; segmentation state of the art.
Split	FOV/crop/target-path disjoint train/validation; plate/well overlap possible.
Samples	192 train, 48 validation for paired 262k context-collar rows.

Appendix

Table 4: Reproducibility details for the main Allen 262k rows.

Method	Macro Dice	FG recall	Pred. FG frac.
Per-structure tuned threshold	0.5582	0.7089	0.0892
Otsu per crop	0.5021	0.8108	0.1239
Otsu + CC cleanup	0.5004	0.7882	0.1197
Gaussian + Otsu	0.3565	0.8763	0.2079
Adaptive local	0.2171	0.8967	0.3802

Appendix

Table 5: Classical target-fluorescence controls on the six-structure Allen validation set. These rows show that the audited task is not reducible to thresholding.

Variant	Best Dice	Final Dice	Local gate	Cmp. gate
Token-before-compression	0.7378	0.7048	0.5054	0.4946
Summary-center RoPE	0.7384	0.6947	0.4942	0.5058
Compressed branch unrotated	0.7333	0.6830	0.5134	0.4866

Appendix

Table 6: RoPE placement in the compressed branch (262k/depth1). Token-before-compression is the exact formulation of Theorem 7 ; summary-center RoPE is a strictly cheaper coordinate-aware approximation, competitive in this short gate and to be read as a promising simplification, not a beaten baseline.

Positional method	Best Dice	Final Dice	Time (s)	Peak GiB
Learned full-grid PE	0.6939	0.6939	71.8	2.60
Packed-sequence RoPE	0.7099	0.6670	68.3	2.22
Axial grid-coordinate RoPE	0.7346	0.7154	79.4	2.27
3D relative position bias	0.7055	0.6420	588.0	3.35
3D ALiBI-style bias	0.7098	0.6691	74.3	2.97

Appendix

Table 7: Feasible-scale coordinate-aware positional alternatives (65k tokens). Axial grid-coordinate RoPE leads 3D RPB and 3D ALiBI on best and final Dice; note also the order-of-magnitude time cost of the materialized-bias baseline.

Test	Verified property
Packed-order invariance	Repacking the same 3D coordinates leaves attention logits unchanged up to permutation.
Sequence-index failure diagnostic	Packed sequence position changes under repacking, confirming it is not a valid 3D coordinate.
Relative-displacement identity	$\langle R(c_{i})q,R(c_{j})k\rangle=q^{\top}R(c_{j}-c_{i})k$ holds numerically.
Values unrotated	RoPE rotates Q/K only; values pass through unrotated.
Branch-key rotation	Local, compressed, and selected branch keys are RoPE-rotated when axial RoPE is enabled.

Appendix

Table 8: Geometry-level implementation tests for axial RoPE. Each row is a numerical check that an identity from § 3 holds in the executed operator.

	Dice under test-time packing
Positional mechanism	Raster (train)	Morton	Block-contig.	Random
No positional encoding	0.554	0.554	0.554	0.554
Learned PE	1.000	0.449	0.418	0.378
Packed-sequence RoPE	1.000	0.449	0.418	0.378
Axial grid-coordinate RoPE	1.000	1.000	1.000	1.000

Appendix

Table 9: Task-level synthetic repacking benchmark. Each model is trained on raster-packed tokens and evaluated on the same physical 3D labeling rule after changing only the token memory order. Sequence-indexed mechanisms solve the training layout and collapse under repacking — below the position-blind baseline — while axial grid-coordinate RoPE is invariant.

Encoding	Packing change	Behavior
Axial grid-coordinate RoPE	Raster/block/Morton permutation	Invariant up to permutation.
Packed-sequence RoPE	Same coordinates, changed memory order	Logits change with the sequence index.
Learned absolute table	Changed coordinate-to-index assignment	Not invariant unless repacked consistently.
3D RPB / ALiBI	Same coordinate deltas under repacking	Coordinate-aware in principle; high-token comparison needs streaming bias-aware kernels.

Appendix

Table 10: Synthetic packing-invariance diagnostics. Axial grid-coordinate RoPE is invariant up to permutation under repacking; packed-sequence RoPE is not.

Structure	Axial W+G	Learned W+G	Axial local-only	U-Net
TOMM20	0.892/0.876	0.783/0.729	0.857/0.860	0.892/ 0.893
LMNB1	0.921/0.893	0.796/0.769	0.909/0.906	0.903/0.907
ST6GAL1	0.855/0.852	0.807/0.794	0.834/0.839	0.861/ 0.864
GJA1	0.742/0.740	0.683/0.646	0.759/0.762	0.761/ 0.765
LAMP1	0.889/0.881	0.802/0.765	0.865/0.866	0.883/ 0.886
ATP2A2	0.819/ 0.815	0.661/0.605	0.808/0.794	0.801/0.778

Appendix

Table 11: Per-structure best/final Dice on the six Allen structures (entries are best/final). Axial RoPE improves over learned PE on all six; against U-Net the comparison is even, with U-Net winning or tying on final Dice for several structures. Bold marks the higher final Dice between Axial W+G and U-Net.

Condition	Macro Dice	$\Delta$ vs. full	Logit L2
Full W+G	0.8531	—	0.0000
Trained local-only	0.8387	$-0.0144$	1.3101
Central summaries only	0.8350	$-0.0181$	0.6454
Collar summaries only	0.8458	$-0.0074$	0.3585
Compressed gate $\to 0$	0.7648	$-0.0883$	1.5702
Compressed summary shuffle	0.6858	$-0.1674$	1.4834

Appendix

Table 12: Inference-time compressed-branch ablations on trained 262k/depth8 axial RoPE W+G checkpoints. Large Dice drops with large logit shifts establish that the compressed pathway is load-bearing; they do not isolate coordinate-awareness (see text). Mean over three seeds.

Group	Best Dice	Final Dice	Best $-$ final gap
Axial RoPE W+G, real collar	0.8531	0.8429	0.0103
Axial RoPE W+G, zero collar	0.8447	0.8422	0.0025
Axial RoPE W+G, shuffled collar	0.8503	0.8417	0.0086
Axial RoPE local-only, real collar	0.8387	0.8379	0.0008
Learned PE W+G, real collar	0.7553	0.7181	0.0372

Appendix

Table 13: Collar and stability controls. Zeroing or shuffling the external collar leaves Dice close to the real-collar row, so collar content does not cleanly explain the gain.

Task	Method	2k Dice	10k Best	10k Final	Peak GiB
Heart	TokenGrid U-Net	0.8067	0.8307	0.8157	0.204
Heart	Window+decoder	0.8231	0.8231	0.7767	3.357
Heart	LG+decoder, learned PE	0.8155	0.8155	0.8026	2.717
Heart	LG+decoder, packed RoPE	0.8231	0.8624	0.8538	2.342
Heart	LG+decoder, axial 3D RoPE	0.8072	0.8713	0.8713	2.388
Spleen	TokenGrid U-Net	0.9156	0.9360	0.9292	0.204

Appendix

Table 14: Decoder-controlled external 3D transfer on MSD Heart and Spleen. Single-seed gates; mixed outcomes, treated as claim-boundary evidence. Bold marks the best value in each column-block.

Tokens	Route granularity	Fwd+Bwd ms	Peak GiB	Tokens/sec
4,194,304	query-block shared	45,113	51.50	92,973
4,194,304	token-exact	113,675	51.50	36,897
8,388,608	query-block shared	176,190	102.04	47,611
8,388,608	token-exact	452,953	102.04	18,520

Appendix

Table 15: Representative high-token reach rows. These rows show that the implementation can execute at multi-million-token scale, but they are not practical-throughput claims.

Tokens	Dense ms	Pack3D ms	Window ms	Dense/ Pack3D	Dense GiB	Pack3D GiB	Window GiB
65k	104.55	16.11	13.86	6.49	1.19	1.51	2.04
131k	405.62	28.36	26.62	14.30	2.32	2.95	4.02
262k	1620.07	54.64	52.49	29.65	4.57	5.82	7.98
524k	6641.53	107.49	104.19	61.79	9.08	11.58	15.89
1M	27034.70	215.51	207.76	125.45	18.10	23.09	31.72

Appendix

Table 16: Full-layer context scaling in the pool64 setting. Pack3D stays close to Window latency and far below dense SDPA as context grows.

Spatial Attention

See all 14 papers

RelFlexformer: Efficient Attention 3D-Transformers for Integrable Relative Positional Encodings
May 11, 2026Byeongchan Kim, Arijit Sehanobish, Avinava Dubey +23D ViTs Self-Attention
1Seoul National University · 2Independent · 3Google Research +2
ZipTok3D: High-Fidelity 3D Tokenization with Compact Token Prefixes
Sep 1, 2026Mingda Lin, Weijie Wang, Zeyu Zhang +73D Reconstruction 3D Representation Learning
1Zhejiang University · 2Monash University · University of Adelaide
Attention Itself Could Retrieve.RetrieveVGGT: Training-Free Long Context Streaming 3D Reconstruction via Query-Key Similarity Retrieval
May 10, 2026Zichen Zou, Xiaosong Jia, Zuxuan Wu +1Memory-Efficient Attention 3D ViTs
Institute of Trustworthy Embodied AI (TEAI), Fudan University · 2Shanghai Key Laboratory of Multimodal Embodied AI

May 11, 2026J/K move · Enter open · S save

RelFlexformer: Efficient Attention 3D-Transformers for Integrable Relative Positional Encodings

Byeongchan Kim, Arijit Sehanobish, Avinava Dubey, Min-hwan Oh+1

1Seoul National University · 2Independent · 3Google Research+2

3D ViTs Self-Attention Relative Positional Encoding Efficient ViTs Positional Encoding

Open paper AskPDF arXiv

We present a new class of efficient attention mechanisms applying universal 3D Relative Positional Encoding (RPE) methods given by arbitrary integrable modulation functions

f

. They lead to the new class of 3D-Transformer models, called RelFlexformers, flexibly integrating those RPEs, and characterized by the

O(L \log L)

time complexity of the attention computation for the

L

-length input sequences. RelFlexformers builds on the theory of the Non-Uniform Fourier Transform (NU-FFT), naturally generalizing several existing efficient RPE-attention methods from structured settings with tokens homogeneously embedded in unweighted grids into general non-structured heterogeneous scenarios, where tokens' positions are arbitrarily distributed in the corresponding 3D spaces. As such, RelFlexformers can be applied in particular to model point clouds. Our extensive empirical evaluation on a large portfolio of 3D datasets confirms quality improvements provided by the NU-FFT-driven attention modulation techniques in the RelFlexformers.

Abstract

Spatial Attention

Explore similar work

RelFlexformer: Efficient Attention 3D-Transformers for Integrable Relative Positional Encodings

ZipTok3D: High-Fidelity 3D Tokenization with Compact Token Prefixes

Attention Itself Could Retrieve.RetrieveVGGT: Training-Free Long Context Streaming 3D Reconstruction via Query-Key Similarity Retrieval