Augmented Equivariant Mesh Networks for Anatomical Mesh Segmentation

Motivation

• Known: Triangle mesh는 수술 계획, 치과, 혈관 재건 파이프라인의 native 출력으로 풍부한 surface topology를 담고 있으나, 기존 point cloud나 volumetric grid 기반 방법들은 이 구조를 직접 활용하지 못한다. 기존 task-specific mesh segmentation 방법들(Fast-TGCN, MeshGraphCNN 등)은 equivariant하지 않아 test-time에서의 자세 변화나 mesh resolution 변화에 취약하다.
• Gap: Robust mesh segmentation을 위해서는 임의의 patient pose와 다양한 mesh resolution을 모두 존중하는 모델이 필요하지만, 현재의 task-specific 방법들은 이러한 equivariance/robustness 속성을 갖추고 있지 않으며, equivariant mesh network를 여러 임상적으로 상이한 task에 걸쳐 체계적으로 평가한 연구도 부재하다.
• Why: 해부학적 mesh는 의료 ML에서 실질적으로 중요한 structured-data regime이며, task-specific architecture 재설계 없이도 compact equivariant model이 강건성을 확보할 수 있음을 보이는 것은 다양한 임상 mesh 응용(수술 계획, 치과, 혈관 재건)에 걸쳐 일반화 가능한 프레임워크 설계에 중요한 함의를 갖는다.
• Approach: EMNN(Equivariant Mesh Neural Network)을 mesh-native equivariant backbone으로 채택하고, 여기에 HKS·dihedral angle·local area 등 intrinsic feature와 PCA 기반 anatomical frame을 결합했으며, 수용 영역을 넓히기 위해 soft regional aggregator(SRA)와 virtual node(VN)라는 두 가지 equivariance를 깨지 않는 long-range context module을 추가했다.

Achievement

세 가지 임상적으로 상이한 task(vertex-level IntrA 동맥류, face-level 치아 segmentation, edge-level 간 표면 segmentation)에서 EAMS 변형들이 DGCNN, PTv3, Fast-TGCN 등 강력한 task-specific baseline과 canonical 조건에서 경쟁력 있는 성능(예: IntrA에서 SRA+EAMS 95.81% AN Dice)을 보였고, 특히 rotation 등 test-time geometric perturbation 하에서 baseline 대비 성능 저하가 거의 없거나 전혀 없는(E(3)-invariant) 강건성을 입증했다.

How

• EMNN 기반 equivariant message passing: invariant scalar edge/face message와 equivariant coordinate update(식 1, 2)를 통해 face normal과 edge displacement 정보를 손실 없이 전파
• Anatomy-aware featurization: HKS(8차원), pointwise local area, dihedral angle 등 intrinsic descriptor와 PCA 기반 anatomical frame(치아/간 대상, cylindrical coordinate 표현)을 결합해 SE(3)-invariant에 가까운 feature 구성
• SRA(soft regional aggregator): K개의 region prototype에 대한 soft assignment로 region token을 생성하고 소형 transformer로 mixing 후 vertex에 다시 broadcast, diversity·equipartition 정규화 적용
• VN(virtual node): 그래프 수준의 global memory bank를 유지하며 실제 노드와 메시지 교환, Gaussian kernel 기반 에너지로 collapse 방지 및 표면 근접 유지
• 학습 손실: boundary-aware bDoU prediction loss + contrastive boundary loss(Lcbl)를 공통으로 사용하고, SRA는 diversity/equipartition term, VN은 Gaussian kernel energy, 간 데이터셋은 local continuity loss를 추가

Originality

• Equivariant mesh network(EMNN)를 여러 임상적으로 상이한 anatomical segmentation task(vertex/face/edge-level)에 걸쳐 체계적으로 평가하는 구조화된 데이터 강건성 벤치마크를 최초로 제시
• SE(3) 하에서 mesh와 함께 co-rotate하는 PCA 기반 anatomical frame을 도입해 intrinsic feature와 task-aware global prior를 equivariance를 깨지 않고 결합
• 기존 local one-ring 기반 EMNN의 제한된 수용 영역을 확장하기 위해 SRA(구조화된 regional bottleneck)와 VN(유연한 global memory)이라는 상호 보완적인 두 가지 경량 long-range context 모듈을 equivariance 보존 방식으로 설계

Limitation & Further Study

• EGNN과 EMNN을 직접 비교하지 않아 equivariance 자체의 이점과 explicit mesh structure의 이점을 분리하지 못했음을 저자도 인정하며 향후 연구로 남김
• PCA 기반 anatomical frame의 sign-resolution heuristic이 완전한 SE(3) invariance를 보장하지 않는 유일한 근사 지점으로 남아 있어, 극단적인 mesh 형태나 노이즈 하에서 실패 가능성 존재
• 평가 데이터셋(IntrA 103개, 3D-IOSSeg 180개 등)의 규모가 비교적 작아 대규모 임상 적용 시의 일반화 성능에 대한 추가 검증이 필요
• 다양한 K(region prototype 수), V(virtual node 수) 등 하이퍼파라미터에 대한 민감도 분석이 부족해 실제 배포 시 튜닝 부담이 있을 수 있음

Evaluation

총평: Equivariant mesh network에 anatomy-aware feature와 경량 long-range context module을 결합해 다양한 임상 mesh segmentation task에서 robustness와 성능을 동시에 달성한 실용적이고 잘 설계된 연구로, EGNN 대비 EMNN의 순수 이점 분리 등 추가 ablation이 보완되면 더욱 완성도가 높아질 것이다.