패킹된 3D 어텐션은 메모리 순서가 아닌 실제 좌표를 따라야 한다

무엇인가

이 논문이 푸는 문제는 3D 볼륨을 패치화해 토큰 시퀀스로 만들었을 때, 메모리 레이아웃이 물리적 기하와 분리된다는 점이다. 볼륨은 시스템 효율을 위해 raster, Morton, block-contiguous 같은 순서로 패킹될 수 있지만, 그 순서는 좌표계가 아니다. 시퀀스 위치를 읽는 위치 인코딩이나 packed-sequence RoPE는 객체의 3D 구조 대신 메모리 배치를 학습할 수 있다. 이는 dense attention의 비용 문제나 window attention의 국소성 문제와 별개로, 압축된 원거리 문맥을 쓰는 local-global attention에서 특히 위험하다. 현미경의 live-cell crop은 수만에서 수십만 토큰을 가지며, Allen WTC-11 hiPSC 데이터셋은 25개 태그 구조에 걸친 20만 개 이상의 3D 세포를 포함한다. 과학적 대상 자체가 공간적인 domain에서는 저장 순서와 위치를 혼동하는 모델이 잘못된 질문에 답하게 된다.

어떻게 동작하나

PACK3D는 이 문제를 메모리 레이아웃과 어텐션 기하를 분리하는 local-global attention 연산자로 푼다. 로컬 분기는 좌표 격자에서 정의한 이웃 N(i)에 대해 정확한 token-to-token 어텐션을 수행한다. 원거리 문맥만 Ω를 연속적인 3D 블록 B_j로 나누고, 블록마다 하나의 요약을 만들어 압축한다. 각 토큰은 모든 M개 블록 요약에 어텐션한다. 위치 지정은 실제 토큰 격자 좌표 c_i에 대한 axial RoPE 하나뿐이며, 패킹된 시퀀스 인덱스는 쓰지 않는다. 쿼리와 키는 R(c_i)로 회전하고, 값은 내용을 담으므로 회전하지 않는다. 이때 ⟨R(ci)q, R(cj)k⟩ = qᵀR(cj−ci)k가 성립해 로짓은 패킹 위치 i, j가 아니라 상대 3D 변위 c_j−c_i에만 의존한다.

무엇과 다른가

압축 분기에서도 기하가 살아남는 이유는 회전을 블록 요약 전에 적용하기 때문이다. 블록 B_j의 키 요약은 각 토큰 키를 R(c_u)로 회전한 뒤 정규화 가중치 a로 풀링한 k̃_j = Σ a R(c_u) W_K LN(h_u)가 된다. 그러면 쿼리-요약 로짓은 Σ a q_iᵀR(c_u−c_i)k_u가 되어, 익명의 풀링된 콘텐츠와의 유사도가 아니라 쿼리에 대한 상대 3D 변위로 가중된 내용 상호작용이 된다. 논문은 요약 중심점만 회전하는 더 싼 근사 변형도 단기 ablation에서 seed noise 안(최고 Dice 0.7378 대 0.7384)이라고 보고하지만, 정확한 항등식을 유지하는 token-before-compression을 canonical form으로 택한다. 로컬/압축 분기는 학습된 tokenwise softmax gate로 섞인다. 구현은 Triton 기반 fused forward/backward, 연속 블록 요약, dense route tensor와 local-neighbor tensor materialization 회피를 강조한다. 비용은 로컬 O(HqNw dh) + query-summary O(HqNM dh) + summary build O(HkvMv dh)이며 M=N/v다. 65k 토큰에서 windowed-attention latency와 맞먹으면서 windowing이 못 하는 장거리 문맥을 제공한다고 보고한다.

어떻게 쓰나

가장 직접적인 증거는 통제된 합성 3D 라벨링 과제다. raster 패킹으로 학습한 뒤 물리 좌표와 라벨 규칙은 고정하고 메모리 순서만 Morton, block-contiguous, random으로 바꾼다. 학습된 positional embedding과 packed-sequence RoPE는 학습 레이아웃에서 Dice 1.000을 달성하지만, 재패킹 후 각각 0.449, 0.418, 0.378로 붕괴한다. 반면 axial grid-coordinate RoPE는 모든 재패킹에서 1.000을 유지한다. 위치 없는 baseline은 모든 패킹에서 0.554로 평평하지만 과제를 풀지 못한다. 시퀀스 인덱스 기반 방식은 재패킹 후 위치 무지 baseline보다도 낮아지는데, 이는 이들이 단순히 이점을 잃은 것이 아니라 재패킹으로 오해를 유발하는 신호를 학습했음을 보여준다. 논문은 이 값을 0.449=22/49처럼 정확한 유리수로 제시하며, Proposition 1에서 좌표 보존 재패킹은 출력을 순열화할 뿐이고 sequence-index RoPE는 일반적으로 packing equivariant하지 않다고 정리한다.

전제와 한계

Allen WTC-11 hiPSC 현미경의 15-structure joint task에서 PACK3D는 파라미터 매칭된 TokenGrid 3D U-Net보다 best macro Dice +0.0094, final macro Dice +0.0135를 기록했다. 262k/depth8의 훨씬 큰 learned-position local-global baseline과 비교해서는 best +0.0978, final +0.1248을 보였지만, 그 baseline이 158.5M 파라미터인 반면 PACK3D는 24.3M이어서 저자들은 이 큰 격차를 좌표계 효과의 깨끗한 측정으로 읽지 않는다. 대신 capacity-controlled ablation으로 이득이 인코더 용량이 아니라 좌표 시스템에서 온다고 주장한다. 구조별로는 learned-PE baseline에 대해 axial RoPE가 여섯 구조 모두에서 best/final Dice를 개선하지만, U-Net과의 비교에서는 final Dice 기준 U-Net이 우세한 구조가 있어 저자들은 경쟁력 있는 결과로 한정하고 구조별 압도를 주장하지 않는다.

압축 분기가 실제로 쓰인다는 증거도 있다. 학습된 axial RoPE W+G는 0.8531로 local-only 0.8387보다 높고, 추론 시 압축 gate를 0으로 고정하면 0.7648, 요약을 섞으면 0.6858, 출력을 0으로 만들면 0.6775로 떨어진다. collar 통제에서는 실제 collar 0.8531, 0으로 치환 0.8447, 섞은 collar 0.8503으로 외부 collar 내용이 이득의 깨끗한 원인은 아니라고 본다. 다만 저자들은 이 ablation이 압축 경로가 쓰인다는 것은 보여도 그 경로가 좌표 인지적이라는 것까지 완전히 분리하지는 못한다고 명시한다. 동일 조건의 좌표 무시 압축 분기와 셔플 강건성 비교가 필요한 통제로 남는다. 외부 전이(MSD Heart, Spleen)는 단일 seed gate로, axial이 best Dice에서 두 데이터 모두 앞서고 Heart의 final Dice에서는 최고지만 Spleen 전체에서는 U-Net이 가장 강해 claim-boundary evidence로만 제시된다. 또한 additive 3D relative bias도 정당한 대안이지만, 262k 토큰에서 naive bias는 head group당 약 4.29B, 두 KV group이면 8.5B를 넘는 query-by-summary 변위 텐서를 materialize해야 해서, axial RoPE가 같은 상대 기하를 메모리 버스 밖에서 주입한다는 시스템적 차이를 강조한다.

개발자에게 이 논문의 실무적 의미는 3D 필드를 토큰으로 패킹할 때 패킹은 시스템 결정, 좌표는 모델링 결정으로 분리해야 한다는 것이다. 의료 볼륨 분할, 현미경, density field/NeRF류 표현, spatial transcriptomics처럼 토큰이 물리 좌표를 가지는 작업에서는 시퀀스 위치 기반 위치 인코딩이 Morton/block-contiguous 배치를 학습하지 않는지 재패킹 스트레스 테스트로 확인해야 한다. 원거리 문맥을 압축할 때는 회전을 요약 전에 적용하거나 좌표 인지 요약을 써서 query-relative geometry를 보존하고, 로컬 고주파 구조는 정확한 이웃 어텐션으로 남겨야 한다. 커널 수준에서는 dense route tensor나 gathered local-neighbor tensor를 materialize하지 않는지, 파라미터 매칭 baseline과 구조별 결과를 함께 보는지 확인할 필요가 있다.

저자들이 밝힌 한계와 전제도 분명하다. Allen 과제는 target fluorescence, brightfield, DNA 채널로 Allen-generated mask를 예측하는 것이며 label-free prediction도 수동 라벨 분할도 아니다. 논문은 segmentation state of the art, U-Net에 대한 보편적 우월성, sparse attention의 일반적 속도 법칙을 주장하지 않는다. 고토큰 H200의 4M·8M 토큰 행은 실행 도달 범위를 보여줄 뿐 실용적 throughput 주장이 아니고, sweep에서는 dense SDPA가 더 낮은 메모리를 보이기도 한다. 요약 중심점 변형의 근사, 주파수 할당, 외부 전이는 추가 multi-seed 검증이 필요하며, label-free prediction, U-Net 대체, 이미징 모달리티 간 전이는 범위 밖으로 남겨져 있다.