RefineFormer3D 경량 3D 의료 영상 분할을 위한 적응형 멀티스케일 트랜스포머

초록

RefineFormer3D는 GhostConv3D 기반 패치 임베딩, 저랭크 MixFFN3D, 그리고 크로스 어텐션 융합 디코더를 결합한 경량 3D 트랜스포머 모델이다. 전체 파라미터는 2.94 M에 불과하면서도 ACDC와 BraTS 데이터셋에서 각각 93.44 %와 85.9 %의 평균 Dice 점수를 달성한다. 추론 시간은 8.35 ms(볼륨당)로, 메모리와 연산 효율성 면에서 임상 현장 적용이 가능하도록 설계되었다.

상세 분석

RefineFormer3D는 3차원 의료 영상 분할에서 전통적인 CNN‑U‑Net 구조가 갖는 국소성 한계를 극복하고, 기존 트랜스포머 기반 모델이 안고 있는 파라미터·메모리 과다 문제를 해결하기 위해 세 가지 핵심 모듈을 설계하였다. 첫 번째인 GhostConv3D 기반 패치 임베딩은 일반 Conv3D와 달리 기본 채널을 적은 수의 “프라임” 필터로 추출한 뒤, 깊이별 경량 DWConv3D를 통해 “고스트” 특성을 생성한다. 이 과정은 채널당 파라미터를 r = 2 정도로 절반 수준으로 감소시키면서도 공간 연속성을 유지한다. 두 번째인 MixFFN3D는 기존 Feed‑Forward Network의 완전 연결 레이어를 저랭크 선형 변환과 3D depthwise convolution으로 대체함으로써, 전역적인 표현력은 유지하되 연산량과 파라미터를 크게 줄인다. 특히 SILU 활성화와 LayerNorm을 결합해 학습 안정성을 확보한다. 세 번째인 크로스 어텐션 융합 디코더는 인코더의 다중 스케일 스킵 연결을 단순 연결이 아닌 쿼리‑키‑밸류 구조의 크로스 어텐션으로 재구성한다. 디코더 특징이 현재 재구성 단계에서 필요한 정보를 동적으로 질의함으로써, 서로 다른 해상도에서 추출된 특징을 의미론적으로 선택적으로 결합한다. 추가로 SE‑Attention을 적용해 채널 차원의 중요도를 재조정한다. 이러한 설계는 전체 파라미터를 2.94 M로 압축하면서도, ACDC(93.44 %)와 BraTS(85.9 %)에서 기존 nnFormer, SegFormer3D, UNETR 등을 능가하거나 동등한 성능을 보인다. 실험 결과는 파라미터 수 대비 Dice 점수와 추론 속도에서 뛰어난 효율‑정확도 트레이드오프를 확인한다. Ablation study에서는 GhostConv3D 없이 일반 Conv3D를 사용했을 때 파라미터는 1.8배 증가하고 Dice가 1.2 % 감소했으며, MixFFN3D를 표준 MLP로 교체했을 때 메모리 사용량이 2.3배 늘고 연산 시간이 1.6배 증가함을 보여준다. 전체적으로 RefineFormer3D는 경량화와 고성능을 동시에 달성한 3D 의료 영상 분할 모델로, 제한된 하드웨어 환경에서도 실시간 임상 적용이 가능하도록 설계된 점이 가장 큰 강점이다.