Geometry Matters in Packed 3D Attention

Change the memory. Keep the geometry.

A volume has coordinates. Its tokens also have positions in memory. Move the tokens, keeping their geometry intact, and see how two attention models respond.

The repacking demo

Same cell. Same weights. A new memory order.

Under the repack, about 18 of 262,144 predictions change for PACK3D and about 3,215 for packed RoPE.

PACK3D uses grid coordinates

0predictions changed
0.0069% of positions

Packed RoPE uses memory positions

0predictions changed
1.23% of positions

Training layout: raster order

Drag a cell, swipe sideways, or focus the view and use the arrow keys to turn both.

An illustration of measured results. Counts are the mean changed-prediction rates across six seeds, scaled to 262,144 positions. The cell and the locations of the changes are illustrative; this Space does not run a trained model. The memory strips show sixteen groups moved whole.

Enable JavaScript to interact. The measured rates correspond to about 18 changed predictions for PACK3D and 3,215 for packed RoPE per 262,144 positions. The results below remain available without JavaScript.

The measured result

Does the prediction follow the object?

Both models reach similar accuracy on the training layout. After a pure repack, PACK3D stays nearly unchanged; packed RoPE loses Dice in all six seeds.

View the per-seed results ↗

Allen-15 · mean validation macro Dice
Model	Training layout	Pure repack
PACK3D	0.8247	0.8248
Packed RoPE	0.8229	0.7750

Six matched seeds · 120 validation crops · same checkpoints, no retraining. Values from the released camera-ready results.

Try it in code

Run a real equivariance check.

A pure repack carries each token’s coordinate, neighbors, summary memberships, and within-block slot along with it. The CPU example checks aligned outputs of a randomly initialized attention layer. It is separate from the trained-model results above.

Run the CPU test ↗