Geometry Matters in Packed 3D Attention

1 Introduction

• Dense attention preserves global access but scales quadratically, whereas window attention scales but sees only locally; reducing the graph and fixing its geometry are independent requirements.Source: §1 · pp. 1–2
• PACK3D keeps local neighborhood attention exact, compresses distant context into contiguous 3D block summaries, and indexes interactions by true token-grid coordinates rather than packed positions.Source: §1 · p. 2
• PACK3D’s fused forward/backward implementation uses Triton local attention and contiguous summaries, aligning its sparse operator with memory movement and training execution.Source: §1 · p. 2
• Applying RoPE before block compression preserves query-relative 3D displacement in query-to-summary logits, keeping compressed summaries spatially interpretable.Source: §1 · p. 2
• On Allen WTC–11, PACK3D improves over a parameter-matched TokenGrid 3D U-Net and a matched learned-position local–global baseline, with controls attributing the gain to coordinates rather than position-encoding parameters.The comparison concerns the 15-structure joint task.Source: §1 · pp. 2–3
• Learned positional embeddings and packed-sequence RoPE fit the training layout but collapse under Morton, block-contiguous, and random repacking, while axial grid-coordinate RoPE stays exact.Source: §1 · p. 2
• Label-free prediction, U-Net replacement, and transfer across imaging modalities are outside the studied scope.Source: §1 · p. 3

Attention for 3D vision

• Prior vision and biomedical models use attention or hybrid architectures to couple local evidence with long-range context, but PACK3D does not claim to replace U-Nets.Source: §2 · pp. 3–4

Efficient and sparse attention

• Efficient attention methods reduce cost through IO-aware implementations, approximations, or sparse connection patterns; PACK3D adds the constraint that positions must represent physical geometry under packing.Source: §2 · p. 3

Position encodings and geometry under packing

• Learned, sinusoidal, relative, and rotary encodings provide positional information, while PACK3D distinguishes physical 3D coordinates from useful raster, Morton, or block-contiguous memory layouts.Source: §2 · p. 3

Microscopy, Allen cellular organization, and task scope

• Allen WTC–11 hiPSC microscopy tests packed attention using target fluorescence, brightfield, and DNA to predict Allen-generated masks.Source: §2 · p. 4

Memory layout and systems realization

• PACK3D uses packing and block-contiguous summaries for memory movement while indexing attention logits by physical coordinates.Source: §2 · p. 4

3.1 Packed token fields and coordinate geometry

• PACK3D indexes attention by true 3D coordinates, so coordinate-preserving repacking permutes the computation without changing its geometry.Sequence-index RoPE changes its rotary coordinate with memory order and is not packing equivariant in general.Source: §3.1 · pp. 4–5

3.2 Geometry-aware local–global factorization

• Exact coordinate-defined local attention preserves nearby detail, while contiguous 3D block summaries compress only nonlocal context.RoPE is applied before pooling, so each query-to-summary logit retains a weighted sum of query-relative displacements rather than comparing against anonymous pooled content.Source: §3.2 · pp. 5–6

3.3 Branch mixing

• A learned tokenwise mixer combines exact local morphology with lower-resolution global organization, and the segmentation objective and main comparisons hold other training components fixed.The shared objective combines Dice and binary cross-entropy losses.Source: §3.3 · p. 6

3.4 Systems realization of the local–global graph

• The systems design co-develops coordinate-defined neighborhoods, contiguous summaries, and a backward pass that follows the forward pass’s streamed access pattern.The main Allen rows use raster_block layout; native_block_morton is also supported for layout and systems diagnostics.Source: §3.4 · p. 6

3.5 Complexity and systems contract

• PACK3D replaces N remote tokens per query with M = N/v summaries while retaining exact nearby token interactions; it is not linear attention or approximate local attention.Streamed online-softmax reductions avoid gathered local-neighbor tensors and dense token-to-summary score matrices.Source: §3.5 · pp. 6–7

4.1 Allen contract

• The Allen contract evaluates six-structure runs and an Allen-15 joint comparison; at 262k tokens it reports best and final macro Dice on the fixed central target region over three seeds.The strongest tuned threshold baseline reaches 0.5582 macro Dice, far below every neural row.Source: §4.1 · p. 7

4.2 Main Allen results

• At 262k/depth8, PACK3D reaches 0.8531 ± 0.0005 best and 0.8429 ± 0.0136 final Dice; its compressed branch adds +0.0144 best and +0.0050 final over coordinate-aware local-only attention.The +0.0978 best gap to the learned-position model at 0.7553 ± 0.0042 best combines the compressed-branch contribution with a weak baseline, so it is not a clean coordinate-geometry measurement.Source: §4.2 · p. 7, §4 · p. 7
• 0.8228 ± 0.0008 best and 0.8228 ± 0.0007 final macro Dice for PACK3D, versus 0.8134 ± 0.0034 best and 0.8093 ± 0.0039 final for the parameter-matched TokenGrid 3D U-Net.The gains are +0.0094 best and +0.0135 final, consistent in sign across three seeds but modest against seed variation.Source: §4.2 · p. 7, Table 1 · p. 8

4.6 Claim boundaries

• PACK3D wins the Allen-15 joint average and best-checkpoint comparisons, but the U-Net wins final Dice on five of six structures; the claim is competitiveness, not per-structure dominance.The Allen task predicts Allen-generated masks from target fluorescence, brightfield, and DNA, not label-free predictions or manual-label segmentation; the paper also disclaims segmentation state of the art and universal U-Net superiority.Source: §4.2 · p. 8, §4.6 · pp. 9–10

4.3 Positional controls

• At matched 3.06M parameters, axial RoPE reaches 0.7375 best Dice versus 0.7142 for packed-sequence RoPE at 262k, and leads 0.7346 to 0.7099 at 65k.Compact learned 3D positional encoding remains competitive on final Dice, bounding the conclusion to physical-coordinate geometry rather than one uniquely superior rotary scheme.Source: §4.3 · p. 8, Table 2 · p. 8

4.4 Compressed-summary ablations

• Adding the compressed branch raises best Dice from 0.8387 to 0.8531 and final Dice from 0.8379 to 0.8429; inference-time gate clamping, summary shuffling, and output zeroing drop Dice to 0.7648, 0.6858, and 0.6775.These perturbations show reliance on a functional compressed pathway, but do not by themselves establish that its contents are coordinate-aware; the matched coordinate-blind control remains unresolved.Source: §4.4 · pp. 8–9, Table 3 · p. 9

4.5 Systems scaling

• At 65k tokens, PACK3D is 6.49× faster than dense SDPA but slower than Window attention, which cannot provide the compressed global context associated with the Allen Dice gains.Dense attention completes through 1M tokens and remains lower-memory in this sweep, so the systems result is not universal memory superiority.Source: §4.5 · p. 9

• Memory layout is a systems decision, not an attention coordinate system; packed geometric attention should index compatibility by physical coordinates.Source: §5 · p. 10
• PACK3D combines exact local attention with compressed global summaries and coordinate-aware query-relative geometry.Source: §5 · p. 10
• Under synthetic repacking, learned or sequence-index positions fit one layout and fail under another, while coordinate-indexed attention remains invariant.Source: §5 · p. 10

• The claim is limited to coordinate-aware local–global attention for packed 3D fields, not segmentation state of the art, label-free prediction, broad transfer, or universal speed and memory gains.Source: A Experimental Contract and Claim… · p. 11
• +0.0978 best and +0.1248 final macro Dice over matched learned-position local–global attention; +0.0094 best and +0.0135 final over a parameter-matched TokenGrid 3D U-Net.These are the main 262k-token depth-8 comparison and the updated Allen-15 joint comparison, respectively.Source: A Experimental Contract and Claim… · p. 11
• Synthetic repacking preserves axial-RoPE logits up to permutation, whereas sequence-index RoPE fails when layouts change.Source: A Experimental Contract and Claim… · p. 11

A.1 Reproducibility contract

• The main Allen 262k reproducibility contract uses a split disjoint by FOV, crop, and target path, while allowing plate or well overlap.It is not claimed as a strict plate-disjoint benchmark.Source: §A.1 · p. 11, Table 4 · p. 11

A.2 Classical threshold baselines

• Classical threshold controls remain below the neural rows, and higher-recall thresholds over-predict foreground relative to the Allen-generated masks.Source: §A.2 · p. 11, Table 5 · p. 12

B Mathematical Details and Theory for Geometry-Aware Packed 3D Attention

• The theory distinguishes exact repacking invariance from approximate coordinate awareness and develops identities, symmetry, equivariance, separation, and bounds.Source: B Mathematical Details and Theory for… · p. 12

B.1 Setup and notation

• A packing maps physical lattice sites to memory positions; coordinate-defined neighborhoods and blocks remain tied to the lattice rather than packed-index intervals.Source: §B.1 · p. 12

B.2 Axial RoPE as an orthogonal representation of displacements

• Axial RoPE’s orthogonal displacement identity makes logits depend on the difference between physical coordinates, so a common translation leaves them unchanged.The global coordinate origin is therefore unobservable; the results also hold for fixed anisotropic or normalized coordinate maps.Source: §B.2 · p. 13, §B.9 · p. 17

B.3 Packing equivariance of the full operator

• The full operator is packing-equivariant when neighborhoods and blocks are coordinate-defined, summary weights are order-independent and symmetric, and pointwise maps are shared.Source: §B.3 · p. 14

B.4 Why sequence position is not a geometric coordinate

• Sequence-index RoPE is not packing-equivariant: the same physical token pair can receive different logits under different packings.Source: §B.4 · pp. 14–15

B.5 Relative geometry survives block compression

• Rotating token keys before block compression makes each query-to-summary logit a weighted sum of query-relative token displacements, rather than similarity to a positionless pooled key.Source: §B.5 · p. 15

B.6 Coordinate-awareness is strictly more expressive

• In a two-block construction with identical pooled keys and opposite values, coordinate-blind pooling incurs squared error 1, while coordinate-aware logits can approximate the target arbitrarily closely without a fixed norm bound.Source: §B.6 · pp. 15–16

B.7 RoPE placement in the compressed branch: exact vs. approximate

• Token-before-compression preserves the exact identity, while summary-center RoPE is a cheaper approximation that is competitive in the short gate; the evidence does not show it is empirically beaten.Source: §B.7 · p. 16, Table 6 · p. 16

B.8 How far the summary-center approximation can drift

• The summary-center approximation’s error is controlled by block diameter, within-block key variation, and rotary smoothness; larger blocks or higher frequencies can worsen the bound.Source: §B.8 · pp. 16–17

B.10 Attention-graph complexity; B.11 Comparison with additive 3D relative biases

• With block volume v, PACK3D uses Nw exact local interactions and N²/v query-to-summary interactions; at 65k tokens, axial RoPE leads tested 3D RPB and 3D ALiBI on best and final Dice.Axial RoPE injects relative geometry through rotation without materializing the query-by-summary displacement tensor required by the described additive-bias approach.Source: §B.10 · p. 17, Table 7 · p. 18, §B.11 · p. 18

B.12 Summary of the theory

• Repacking a physical 3D field only permutes the output of coordinate-indexed local–global attention.The operator is well-defined independently of the packing layout.Source: §B.12 · p. 18
• Sequence-indexed positional attention is not geometric on packed 3D fields because changing memory layout changes its positional coordinate.The empirical lesson is to tie positional geometry to physical coordinates rather than memory order.Source: §B.12 · p. 18
• Coordinate-aware compressed summaries preserve query-relative geometry: query-to-summary logits combine content interactions modulated by relative 3D displacement.This makes the summary more than a content-similarity score to a pooled token.Source: §B.12 · p. 18
• Coordinate-aware compression can distinguish identical-content remote blocks at different relative positions, unlike content-only compression.The theory’s claims are independent of Allen microscopy and segmentation; packing is for memory movement, coordinates for attention geometry.Source: §B.12 · p. 18

C Diagnostics, Supporting Results, and Claim Boundaries

• The following evidence chain checks operator identities, reports Allen results with their unevenness against U-Net, and marks the claim boundaries.The diagnostics test code against the math; the Allen tables test real-data performance; boundary controls identify where the claim stops.Source: C Diagnostics, Supporting Results, and… · p. 19

C.1 Geometry diagnostics: does the implementation match the theory

• Table 8 reports numerical checks that the executed axial-RoPE operator satisfies identities from §3.The diagnostics are intended to confirm that the implementation realizes the theoretical operator before task-level results are interpreted.Source: Table 8 · p. 19, §C.1 · p. 19

C.2 Task-level synthetic repacking benchmark

• Learned embeddings and packed-sequence RoPE share the same repacking failures because both key on sequence position, which becomes misleading when memory order changes.Their repacked Dice values fall below the position-blind baseline, while the task itself requires positional information.Source: §C.2 · p. 19
• 1.000 Dice on raster falls to 0.449 (Morton), 0.418 (block-contiguous), and 0.378 (random) for learned embeddings and packed-sequence RoPE; axial grid-coordinate RoPE stays at 1.000.The no-positional baseline remains at 0.554 across packings but does not solve the task.Source: §C.2 · p. 19, Table 9 · p. 20

C.3 Synthetic packing-invariance diagnostics

• Axial grid-coordinate RoPE logits remain invariant up to permutation under coordinate-preserving repacking, whereas packed-sequence RoPE logits change with memory order.These are operator-level diagnostics, distinct from the task-level Dice measurements.Source: Table 10 · p. 20, §C.3 · p. 20

C.4 Per-structure Allen results

• Axial RoPE improves over learned positional encoding on all six Allen structures, while U-Net wins or ties on final Dice for several structures.The table describes the U-Net comparison as even, rather than showing per-structure dominance by axial RoPE.Source: §C.4 · p. 20, Table 11 · p. 21
• Compressed-branch ablations show that the pathway is load-bearing, but do not isolate coordinate-awareness as the cause of the effect.The ablations use trained 262k/depth8 axial RoPE W+G checkpoints and average results over three seeds.Source: Table 12 · p. 21

C.5 Compressed-branch inference ablations

• At inference, clamping the compressed gate to zero drops Dice to 0.7648, shuffling summaries to 0.6858, and zeroing the output to 0.6775; the logit changes show the pathway is load-bearing.The ablations establish pathway use, not that it is coordinate-aware.Source: §C.5 · p. 21

C.6 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.The reported best Dice values are 0.8531 for real collar, 0.8447 for zeroed collar, and 0.8503 for shuffled collar.Source: §C.6 · p. 21, Table 13 · p. 22

C.7 Decoder-controlled external transfer

• External transfer is mixed: axial grid-coordinate RoPE is the best attention row by best Dice on Heart and Spleen, but U-Net remains strongest overall on Spleen.These are single-seed gates, and the results are treated as claim-boundary evidence rather than proof of generalization.Source: §C.7 · p. 22

D Optional Selected Refinement and Systems Extensions

• Selected refinement, query-tiled execution, multi-atom summaries, and high-token reach are implementation extensions, not main quality claims.The selected sparse-refinement branch is retained for systems reach and appendix analysis, not the main Allen quality path.Source: D Optional Selected Refinement and… · p. 23, §D.1 · p. 23

D.1 Selected sparse refinement

• Selected sparse refinement pools queries, scores compressed block summaries, then lets each query attend to raw K/V tokens from selected blocks.For GQA/MQA, selected block indices are shared within each KV group to keep selected K/V access blockwise and hardware-aligned.Source: §D.1 · p. 23

D.2 High-token systems reach

• Query-tiled paths run to 4M and 8M tokens, but are slow and skip some diagnostics above the auxiliary-materialization threshold, so these are execution-reach results, not practical-throughput claims.The representative high-token rows show multi-million-token execution reach, not practical throughput.Source: §D.2 · p. 23, Table 15 · p. 24
• On full layers, PACK3D stays close to Window latency and well below dense SDPA as context grows, while Dense SDPA remains lower-memory in the sweep.The evidence supports measured latency scaling, not universal memory superiority.Source: §D.2 · p. 23, Table 16 · p. 24

NeurIPS Paper Checklist

• The checklist says data are public, code will be released with final publication, and full experimental details are in the appendices.It also reports H200 hardware and error bars and accuracy information where applicable and practical.Source: NeurIPS Paper Checklist · pp. 30–31