부분적으로 가역적인 U‑Net을 활용한 메모리 효율적인 3D 영상 분할

본 논문은 3차원(3D) 의료 영상 분할에 널리 사용되는 U‑Net 기반 모델이 GPU 메모리 한계 때문에 깊이와 입력 크기를 제한받는 문제를 해결하고자, RevNet에서 영감을 얻은 ‘부분 가역’ 아키텍처를 제안한다. 기존 3D 분할 네트워크는 전방 전파 시 각 레이어의 활성값을 저장해야 역전파 시 gradient를 계산할 수 있기 때문에, 레이어가 깊어질수록 메모리 사용량이 O(n)으로 급증한다. 특히 전체 필드‑오브‑뷰(FOV)를 한 번에 처리하려면 수백 메가바이트 수준의 메모리가 필요해 일반적인 12 GB GPU에서는 불가능하다. 저자들은 이러한 메모리 병목을 해소하기 위해 ‘가역 블록(Reversible Block)’을 도입한다. 입력 텐서를 두 그룹(x₁, x₂)으로 나누고, 각각에 임의의 연산 F와 G(예: Conv‑GN‑LeakyReLU)를 적용해 y₁ = x₁ + F(x₂), y₂ = x₂ + G(y₁) 로 전파한다. 역전파 시에는 y₁, y₂만 알면 x₁, x₂를 역연산식 x₂ = y₂ – G(y₁), x₁ = y₁ – F(x₂) 로 정확히 복원할 수 있다. 따라서 중간 활성값을 메모리에 보관할 필요가 없으며, 메모리 비용은 블록 내부 파라미터와 가역 시퀀스 끝에서만 저장되는 활성값에 국한된다. 하지만 U‑Net은 다운샘플링·업샘플링과 스킵 연결이라는 복합적인 구조를 가지고 있어, 모든 레이어를 가역화하면 채널 수가 변하는 문제와 트리 구조 자체가 가역성에 부합하지 않는 문제가 발생한다. 이를 해결하기 위해 저자들은 ‘부분 가역’ 전략을 채택한다. 각 해상도 단계(Encoder와 Decoder 각각)에서 전통적인 비가역 연산(다운샘플링, 업샘플링, 스킵 연결)을 그대로 유지하고, 그 내부에 다수의 가역 블록을 연속적으로 쌓은 ‘가역 시퀀스(RSeq)’만을 가역화한다. 이렇게 하면 전체 네트워크는 여전히 U‑Net의 다중 스케일 특성을 보존하면서, 가역 시퀀스 내부에서는 활성값 저장이 필요 없으므로 메모리 사용량이 크게 감소한다. 수식(4)에 기반한 메모리 분석에서는 비가역 네트워크가 각 레이어마다 활성값(M_Al)과 파라미터 상태(M_Pl)를 저장해야 하는 반면, 부분 가역 네트워크는 비가역 레이어의 활성값(M_Nl)과 가역 시퀀스 끝의 활성값(M_Si)만 저장한다. 역전파 시 추가 메모리(M_Bl)는 한계값 하나만 고려하면 되므로 전체 메모리 증가폭이 미미하다. 실제 구현에서는 PyTorch autograd와 호환되는 RevTorch 라이브러리를 개발해 가역 블록과 가역 시퀀스를 손쉽게 정의하였다. 실험은 BraTS 2018 데이터셋(뇌종양 3D MRI)에서 수행되었다. 입력은 160×192×160 크기의 전체 볼륨이며, 전처리로는 각 모달리티별 z‑축을 제외한 표준화와 뇌 외부를 0으로 설정하는 작업이 포함된다. 데이터 증강은 회전·스케일·탄성 변형을 z‑축에 수직인 평면에만 적용해 이미지 품질 저하를 방지하였다. 학습은 배치 크기 1, Adam 옵티마이저(초기 학습률 1e‑4), Dice 손실(세 영역 WT, TC, ET)의 가중합을 사용했으며, 30‑epoch 이동 평균 Dice가 60 epoch 동안 개선되지 않으면 조기 종료하였다. 표 1의 결과를 보면, 비가역 베이스라인(패치 기반 학습, 메모리 128 MB) 대비 가역 모델(패치 기반, 메모리 94 MB)에서 메모리 사용량이 약 30 % 감소하면서 Dice 점수는 거의 동일했다. 전체 FOV를 처리한 가역 모델은 메모리 94 MB(패치와 동일)에서 시작해 가역 블록 수를 1→4개로 늘릴 때 메모리 증가가 94 → 97 MB 수준에 그쳤다. 반면 비가역 모델에서 동일 깊이를 구현하려면 20 GB 이상이 필요했을 것이다. 성능 측면에서는 가역 블록 3개를 사용한 모델이 ET Dice 80.33 %, WT 91.01 %, TC 86.21 %를 기록했으며, Hausdorff 거리(HD95)도 2.58 mm까지 감소해 전반적인 경계 정확도가 향상되었다. 학습 시간은 가역 모델이 비가역 대비 약 50 % 더 오래 걸렸지만, 메모리 절감과 깊이 확장 가능성으로 인해 연구 실무에서는 충분히 감수 가능한 수준으로 평가된다. 또한, 제안된 부분 가역 U‑Net은 별도의 앙상블, 테스트‑타임‑증강, 후처리 없이도 Isensee et al.의 최종 제출 모델(다중 앙상블·테스트‑타임‑증강·후처리 포함)과 비슷한 Dice 점수를 단일 모델로 달성했다. 결론적으로, 본 연구는 가역 연산을 3D 의료 영상 분할에 성공적으로 적용함으로써 메모리 효율성을 크게 향상시키고, 동일 메모리 한도 내에서 네트워크 깊이를 자유롭게 늘릴 수 있는 새로운 설계 패러다임을 제시한다. 향후 연구에서는 가역 블록의 구조 최적화, 혼합 정밀도 연산, 다른 의료 영상 분야로의 확장, 그리고 가역 연산과 기존 메모리 절감 기법(예: 체크포인팅) 간의 시너지 효과 등을 탐색할 여지가 있다.