한 번에 끝내는 뇌종양 분할 다중작업 네트워크와 교차작업 주의 메커니즘
본 논문은 기존 모델 캐스케이드 방식이 안고 있던 복잡성 및 작업 간 상관관계 무시 문제를 해결하기 위해, 하나의 네트워크 안에서 세 가지 세분화 작업을 동시에 수행하는 One‑pass Multi‑task Network(OM‑Net)를 제안한다. 공유 파라미터와 작업별 전용 파라미터를 결합하고, 온라인 데이터 전이와 커리큘럼 학습 전략을 도입해 학습 효율을 높였다. 또한 이전 작업의 예측 결과를 활용해 채널‑별 주의를 조정하는 Cross‑tas…
저자: Chenhong Zhou, Changxing Ding, Xinchao Wang
뇌종양, 특히 교모세포종과 같은 악성 종양은 조기 진단과 정확한 부피 측정이 치료 계획에 핵심적인 역할을 한다. 기존 연구에서는 클래스 불균형 문제를 완화하기 위해 모델 캐스케이드(MC) 방식을 널리 사용해 왔다. MC는 coarse‑to‑fine 전략으로 여러 개별 네트워크를 순차적으로 실행해 종양 영역을 점진적으로 정밀화한다. 그러나 이 접근법은 (1) 여러 모델을 별도로 학습·저장해야 하는 시스템 복잡성, (2) 각 모델이 독립적으로 학습돼 작업 간 상관관계를 활용하지 못함, (3) 추론 시 GPU‑CPU 간 번갈아 가며 실행돼 연산 효율이 떨어진다는 단점을 가진다.
본 논문은 이러한 한계를 극복하고자 “One‑pass Multi‑task Network”(OM‑Net)라는 경량화된 단일 네트워크를 설계하였다. OM‑Net은 전체 종양, 종양 핵심, 증강 종양이라는 세 가지 서브태스크를 하나의 3D FusionNet 기반 아키텍처에 통합한다. 네트워크는 초기 인코더에서 모든 작업이 공유하는 공통 특징을 추출하고, 이후 디코더 단계에서 작업별 브랜치를 통해 특화된 특징을 학습한다. 이렇게 하면 파라미터 수가 기존 3‑stage 캐스케이드의 약 1/3에 불과하면서도, 각 작업이 서로에게 제공하는 정보(예: ROI)를 실시간으로 활용할 수 있다.
학습 효율을 높이기 위해 두 가지 전략을 도입하였다. 첫째, “온라인 트레이닝 데이터 전이”는 앞선 작업에서 얻은 ROI를 기반으로 뒤쪽 작업의 학습 패치를 확대한다. 예를 들어, 전체 종양 마스크가 넓게 추출되면 이를 이용해 종양 핵심과 증강 종양 작업의 샘플링 영역을 넓혀, 희소 클래스에 대한 데이터 부족을 보완한다. 둘째, 커리큘럼 학습 기반의 단계적 작업 추가 방식이다. 초기에는 전체 종양 분할만 수행하고, 학습이 안정화되면 종양 핵심, 마지막으로 증강 종양 작업을 순차적으로 네트워크에 삽입한다. 이 과정은 손실 함수의 급격한 변동을 방지하고, 모델이 점진적으로 더 어려운 서브태스크에 적응하도록 돕는다.
또한, 기존 채널‑주의 메커니즘인 SE 블록은 전체 채널에 대해 하나의 스칼라 가중치만을 부여해, 다중 라벨이 혼재하는 의료 영상에 적합하지 않다. 이를 개선하기 위해 “Cross‑task Guided Attention”(CGA) 모듈을 설계하였다. CGA는 이전 작업의 예측 결과를 이용해 각 라벨별 채널 통계(예: 클래스별 평균 활성화)를 계산하고, 이를 채널‑주의 네트워크에 입력한다. 결과적으로 채널 가중치가 라벨별로 다르게 조정되어, 특히 증강 종양과 같이 매우 작은 비중을 차지하는 클래스에 대한 감도가 크게 향상된다.
추론 단계에서는 전체 뇌 MRI를 32×32×16×4 크기의 3D 패치로 슬라이딩하면서 겹치는 영역을 평균화하는 오버랩‑타일 전략을 사용한다. 마지막으로, 간단한 후처리 단계에서 작은 고립 영역을 제거하고, 연결된 후보 영역을 정제한다. 이 후처리는 복잡한 CRF 기반 방법보다 연산량이 적으며, 전체 파이프라인의 속도를 유지한다.
실험은 BraTS 2015, 2017, 2018 데이터셋을 대상으로 수행되었다. OM‑Net은 동일한 하드웨어 환경에서 기존 3‑stage 캐스케이드 대비 약 3배 적은 파라미터와 추론 시간을 보이며, Dice 점수에서는 전체 종양, 종양 핵심, 증강 종양 모두에서 기존 최고 기록을 넘어섰다. 특히 증강 종양에 대한 Dice 점수는 0.02~0.04 정도 상승했으며, 이는 CGA 모듈의 기여가 크게 작용한 결과로 분석된다. 또한, 온라인 데이터 전이와 커리큘럼 학습을 각각 제거한 실험에서도 성능 저하가 확인되어, 제안된 학습 전략이 모델 수렴과 일반화에 중요한 역할을 함을 입증하였다.
코드는 GitHub(https://github.com/chenhong-zhou/OM‑Net)에서 공개되어 재현 가능성을 높였으며, 제안된 OM‑Net과 CGA는 다른 장기·단기 의료 영상 분할 문제에도 적용 가능할 것으로 기대된다.
원본 논문
고화질 논문을 불러오는 중입니다...
댓글 및 학술 토론
Loading comments...
의견 남기기