본문 바로가기

딥러닝

UNeXt: 다층 퍼셉트론(MLP)을 활용한 초고속 의료 영상 분할 네트워크 - 무거운 트랜스포머를 버리고 현장 진료(Point-of-Care) 환경을 정복한 비결

반응형


일러두기 (Disclaimer)

본 글은 AI의 도움을 받아 작성되었습니다. 최대한 정확한 정보를 전달하기 위해 검수를 거쳤으나, AI 모델의 특성상 내용에 일부 오류나 부정확한 정보가 포함될 수 있습니다. 따라서 실제 논문의 상세한 수식이나 실험 세부 조건 등은 반드시 원문을 직접 확인하시기 바랍니다.

 

 

논문 링크 : https://arxiv.org/abs/2203.04967

 

UNeXt: MLP-based Rapid Medical Image Segmentation Network

UNet and its latest extensions like TransUNet have been the leading medical image segmentation methods in recent years. However, these networks cannot be effectively adopted for rapid image segmentation in point-of-care applications as they are parameter-h

arxiv.org

 

 

초록 (Abstract)

UNet과 TransUNet 같은 최신 확장 모델들은 최근 몇 년간 의료 영상 분할(Medical Image Segmentation) 분야를 주도해 왔습니다. 하지만 이러한 네트워크들은 매개변수(Parameter)가 지나치게 많고 연산이 복잡하며 구동 속도가 느리기 때문에, 환자 침상 곁에서 즉각적인 진단이 필요한 현장 진료(Point-of-care) 애플리케이션에 도입하기에는 실질적인 어려움이 존재합니다.

이러한 문제를 해결하기 위해, 본 논문에서는 다층 퍼셉트론(MLP, Multilayer Perceptron)과 합성곱(Convolution) 연산을 결합한 새로운 영상 분할 네트워크인 UNeXt를 제안합니다. UNeXt는 네트워크의 초기 단계에는 합성곱 계층을 배치하고, 깊은 잠재 공간(Latent Space)에는 MLP 계층을 배치하는 효율적인 구조로 설계되었습니다. 구체적으로 합성곱 특징(Feature)을 효과적으로 토큰화(Tokenize)하고 투영한 뒤, MLP를 사용하여 데이터의 표현을 모델링하는 '토큰화된 MLP 블록(Tokenized MLP block)'을 제안합니다. 여기에 모델의 성능을 한층 더 향상하기 위해, 입력 채널의 축을 이동(Shifting)시켜 MLP가 지역적인 의존성(Local dependencies)을 학습하는 데 집중할 수 있도록 만드는 기법을 도입했습니다.

잠재 공간에서 토큰화된 MLP를 사용하면 분할에 도움이 되는 훌륭한 특징 표현을 얻어냄과 동시에, 모델의 매개변수 개수와 연산 복잡도를 의미 있게 줄일 수 있습니다. 이 네트워크는 기존 UNet과 마찬가지로 인코더와 디코더의 다양한 계층 사이에 스킵 커넥션(Skip connection)을 포함하고 있습니다. 다양한 의료 영상 분할 데이터셋을 통해 UNeXt를 테스트한 결과, 기존의 최고 성능(State-of-the-art) 의료 영상 분할 아키텍처들과 비교하여 매개변수 수를 72배 줄이고, 연산 복잡도를 68배 감소시켰으며, 추론 속도를 10배 향상하면서도 더 뛰어난 분할 성능을 달성했음을 입증했습니다.

 

1. 한줄 요약 & 3줄 핵심 (Executive Summary)

UNeXt는 기존 인공지능 모델의 무거운 연산량을 덜어내기 위해 합성곱 연산과 다층 퍼셉트론(MLP)을 결합하여, 일반 컴퓨터(CPU) 환경에서도 초고속으로 작동하게 만든 현장 진료용 초경량 의료 영상 분할 모델입니다.

  1. 기존의 문제점 (Pain point): 의료 영상 분할에서 높은 성능을 내는 기존의 모델(UNet, TransUNet 등)은 막대한 매개변수와 연산량을 요구합니다. 이로 인해 고성능 그래픽 처리 장치(GPU)가 없는 실제 병원 진료실 기기(휴대용 초음파, 스마트폰 등)에서 실시간으로 사용하기가 불가능에 가까웠습니다.
  2. 이 논문의 해결책 (Solution): 네트워크의 깊은 층(잠재 공간)에 연산량이 높은 구조를 배치하는 대신, 데이터 특징을 압축하는 토큰화(Tokenization) 기법과 데이터를 물리적으로 밀어서 주변의 지역적 특징을 잡게 해주는 축 이동(Shifting) 기법이 적용된 다층 퍼셉트론(MLP) 블록을 도입했습니다.
  3. 달성한 성과 (Key Result): 이 새로운 구조를 통해 기존 최고 성능 모델 대비 매개변수를 72배, 연산 복잡도를 68배 줄이면서도 추론 속도는 10배 높였습니다. 피부 병변과 유방 초음파 데이터셋에서 기존 모델들을 뛰어넘는 높은 분할 정확도를 달성했습니다.

 

2. 들어가며: 왜 이 연구가 필요했나? (The "Why")

질병을 진단하고 치료 계획을 세우는 데 있어 의료 영상은 필수적인 역할을 수행합니다. 그중에서도 영상 내에서 종양이나 장기의 정확한 경계선을 픽셀 단위로 분류해 내는 '의료 영상 분할(Medical Image Segmentation)'은 컴퓨터 보조 진단(CAD, Computer-Aided Diagnosis)과 영상 유도 수술 시스템의 핵심 기술로 자리 잡고 있습니다. 지난 10여 년 동안 이 분야는 합성곱 신경망(CNN, Convolutional Neural Network)의 발전과 함께 눈부신 성장을 이루어 왔습니다.

 

의료 영상 분할 분야의 가장 대표적인 모델은 2015년에 등장한 UNet입니다. 인코더(Encoder)를 통해 이미지를 압축하며 의미론적 특징을 추출하고, 디코더(Decoder)를 통해 다시 해상도를 복원하되, 잃어버린 위치 정보를 스킵 커넥션(Skip connection)으로 연결하여 보완하는 방식은 의료 영상 분할의 표준이 되었습니다. 이후 UNet의 성능을 개선하기 위해 특징 맵을 더 촘촘하게 연결한 UNet++나 3차원 데이터를 다루는 3D UNet, 잔차 연결(Residual connection)을 더한 ResUNet 등 수많은 변형 모델들이 등장했습니다.

 

최근에는 자연어 처리에서 시작된 딥러닝 트렌드가 컴퓨터 비전으로 넘어오면서, 이미지 전체의 넓은 문맥(Global context)을 이해할 수 있는 트랜스포머(Transformer) 구조가 의료 영상에 도입되었습니다. ViT(Vision Transformer)를 의료 영상에 맞게 변형한 TransUNet, Swin-UNet, MedT와 같은 모델들은 기존 합성곱 기반 모델들의 한계였던 좁은 수용 영역(Receptive field) 문제를 해결하며 예측 정확도를 비약적으로 끌어올렸습니다. 하지만 연구진들은 성능 향상이라는 이면의 치명적인 단점에 주목했습니다.

바로 "현실의 진료 환경에서 이 모델들을 구동할 수 있는가?"에 대한 의문입니다.

 

대부분의 최신 딥러닝 모델들은 성능을 높이는 데만 몰두한 나머지, 모델의 연산 복잡도(GFLOPs), 추론에 걸리는 시간(Inference time), 그리고 모델의 크기를 결정하는 매개변수(Parameter) 수는 크게 고려하지 않았습니다. 대학이나 연구소의 이상적인 환경에서는 고성능 GPU 연산 자원을 제약 없이 사용할 수 있으므로 모델이 무거워도 문제가 되지 않습니다.

 

그러나 실제 의료 시스템의 트렌드는 거대한 영상의학 센터에서 환자 침상 옆으로 다가가는 '현장 진료(Point-of-Care, POC)' 방식으로 진화하고 있습니다. 의사가 스마트폰 카메라 앱을 통해 환자의 피부 병변을 찍어 즉석에서 피부암 여부를 판별하거나, Butterfly iQ와 같은 소형 휴대용 초음파(POCUS) 기기를 환자의 몸에 대고 실시간으로 심장 혈류나 방광의 부피를 계산하는 시대가 도림한 것입니다. 현장 진료 기기들은 휴대성을 위해 배터리로 구동되며, 내부에 고성능 GPU가 탑재되어 있지 않은 경우가 대부분입니다.

 

기존의 무거운 TransUNet이나 거대한 CNN 앙상블 모델을 이런 기기에 탑재하면 이미지를 한 장 분석하는 데 긴 시간이 소요되어 실시간 진료가 불가능해집니다. 트랜스포머 모델의 자가 주의 집중(Self-attention) 메커니즘은 입력 데이터 크기의 제곱에 비례하여 연산량이 증가하는 $O(N^2)$의 복잡도를 가지므로, 고해상도 의료 영상에 적용하기에는 연산 비용이 너무 큽니다.

 

결과적으로, 컴퓨팅 파워가 부족한 일반 CPU 환경에서도 가볍고 빠르게 동작하면서 높은 정확도를 유지할 수 있는 초경량, 초고속 네트워크에 대한 임상적 요구가 극에 달해 있었습니다. 연구진은 이러한 의료 현장의 실질적인 요구사항을 충족하기 위해, 무거운 어텐션 연산이나 합성곱 연산을 과감히 덜어내고 혁신적인 속도를 낼 수 있는 UNeXt 구조를 고안하게 되었습니다.

 

3. 이 논문의 뿌리 (Key Reference)

UNeXt는 기존 컴퓨터 비전과 의료 영상 분야의 중요한 철학을 담은 핵심 논문들의 아이디어를 비판적으로 수용하고 독창적으로 조합하여 탄생했습니다. 이 논문이 뿌리를 두고 있는 세 가지 핵심 선행 연구와 그 관계성은 다음과 같습니다.

  • U-Net (2015): 영상 분할의 기본 골격 제공 본 논문은 U-Net이 제시한 인코더-디코더 기반의 U자형 백본(Backbone) 아키텍처를 기본적인 틀로 사용했습니다. 이미지를 점진적으로 압축하며 의미론적 특징을 추출하는 인코더와, 이를 다시 원래 해상도로 복원하는 디코더 구조를 채택했습니다. 또한 압축 과정에서 손실되는 픽셀의 미세한 위치 정보를 디코더로 직접 넘겨주는 스킵 커넥션(Skip connection) 개념을 계승하여 병변의 테두리를 세밀하게 복원할 수 있도록 했습니다. 다만, 기존 U-Net이 모든 계층에서 연산량이 많은 합성곱 필터를 다수 사용한다는 단점을 보완하기 위해, UNeXt는 모델 초기와 후반부에서만 제한된 수의 필터를 가진 합성곱을 사용하고 중심부는 다른 구조로 대체했습니다.
  • TransUNet (2021) 및 Vision Transformer (ViT): 넓은 시야의 필요성과 그 한계 인식 TransUNet은 U-Net 구조의 인코더 부분에 트랜스포머(ViT)를 결합하여, 이미지 전체의 전역적 문맥(Global context)을 학습할 수 있도록 고안된 모델입니다. 트랜스포머는 멀리 떨어진 픽셀 간의 상관관계를 파악하는 데 탁월한 성능을 보입니다. UNeXt 연구진은 영상 전체의 맥락을 파악하는 것이 분할 정확도에 큰 영향을 미친다는 점을 적극 수용했습니다. 하지만 트랜스포머의 어텐션 메커니즘은 현장 진료 기기에 탑재하기에는 지나치게 무겁다고 판단하여, 어텐션 구조를 배제하고도 전역적 특징을 잡을 수 있는 대안을 탐색하게 됩니다.
  • MLP-Mixer (2021): 다층 퍼셉트론(MLP)의 부활과 경량화의 영감 이 논문이 탄생하는 데 가장 결정적인 '아하(Aha) 모멘트'를 제공한 연구입니다. 딥러닝 초창기에 주로 쓰이던 단순한 다층 퍼셉트론(MLP)만으로도 데이터의 채널(Channel)과 토큰(Token)을 교차로 섞어주면(Mixing) 복잡한 트랜스포머와 유사한 이미지 인식 성능을 달성할 수 있음을 증명한 기념비적인 논문입니다. UNeXt 연구진은 주로 이미지 분류(Classification)에 쓰이던 MLP-Mixer의 아이디어를 픽셀 단위의 분할(Segmentation) 작업에 맞게 변형했습니다. 수천만 개의 매개변수가 필요한 트랜스포머 대신, 가벼운 MLP를 네트워크의 가장 깊은 잠재 공간에 이식함으로써 압도적인 경량화를 이루어냈습니다.

 

4. 핵심 아이디어: 무엇이 다른가? (The "Aha!" Moment)

UNeXt의 가장 핵심적인 혁신은 복잡한 수식이나 무거운 어텐션(Attention) 연산을 완전히 배제하고, 데이터를 극도로 압축한 뒤 단순한 물리적 조작과 기본 신경망 연산만으로 높은 성능을 냈다는 점입니다. 이를 가능하게 한 두 가지 핵심 개념은 '토큰화된 MLP (Tokenized MLP)'와 '축 이동(Shifting) 연산'입니다.

 

이전 방식 (A) vs 발상의 전환 (B)

  • 이전 방식 (A): 이미지의 특정 부분을 잘 파악하기 위해 수많은 필터를 겹겹이 쌓아 올리거나(CNN), 모든 픽셀 사이의 연관성을 계산하기 위해 거대한 곱셈 행렬을 생성(Transformer)했습니다.
  • 발상의 전환 (B): 픽셀 데이터를 작은 덩어리(Token)로 요약한 뒤, 데이터의 배열 순서를 상하좌우로 살짝 밀어버림으로써(Shifting) 억지로 복잡한 계산을 하지 않아도 인접한 정보들이 자연스럽게 섞이게 만들었습니다.

이를 이해하기 쉽게 비유를 들어 설명해 보겠습니다.

 

1. 토큰화(Tokenization): 정예 요원만 선발하기

기존 모델이 책의 모든 단어를 처음부터 끝까지 정독하며 문맥을 파악했다면, 토큰화는 각 문장의 핵심 키워드(정예 요원)만 추출하여 요약본을 만드는 과정과 같습니다. 입력 이미지가 모델의 깊은 곳(잠재 공간)으로 들어오면, 특징(Feature) 맵의 차원을 픽셀 단위에서 작고 추상적인 '토큰(Token)'으로 변환합니다. 이렇게 데이터의 덩치를 대폭 줄인 다음 단순한 다층 퍼셉트론(MLP) 계층에 통과시킵니다. 이 방식은 계산해야 할 파라미터의 수를 획기적으로 줄이면서도 필수적인 문맥 정보를 보존할 수 있게 해줍니다.

 

2. 축 이동(Shifting): 고개를 돌려 주변 풍경 담기

단순한 MLP 구조의 가장 큰 약점은 이미지 픽셀 간의 인접한 공간적 특성, 즉 지역적인 의존성(Local dependencies)을 파악하는 데 취약하다는 것입니다. 기존 연구자들은 이를 해결하기 위해 이미지 내 특정 구역에만 집중하는 지역 주의 집중(Window-based attention)을 도입했지만, 이는 연산량을 크게 증가시키는 원인이 됩니다.

 

UNeXt는 연산량을 전혀 늘리지 않으면서 지역적 정보를 파악하기 위해 데이터가 담긴 채널(Channel)의 차원 축을 물리적으로 이동시키는(Shift) 기법을 택했습니다. 비유하자면, 우리가 풍경을 볼 때 망원경을 고정하고 세밀하게 계산하는 것이 아니라, 망원경 렌즈를 가로로 한 번, 세로로 한 번 살짝 옆으로 밀어서 주변 풍경이 자연스럽게 렌즈 안으로 들어오게 만드는 것과 같습니다.

 

구체적으로 데이터를 $h$개의 파티션으로 쪼갠 뒤, 특정 축(가로 또는 세로)을 따라 데이터를 $j$칸만큼 밀어냅니다. 이 밀어내기 연산은 덧셈이나 곱셈 같은 수학적 연산이 전혀 포함되지 않은 단순한 메모리상의 데이터 자리 바꾸기에 불과합니다. 따라서 모델의 연산 복잡도(FLOPs)는 전혀 증가하지 않습니다. 하지만 이렇게 밀려난 데이터가 토큰화되어 MLP를 통과할 때, 네트워크는 강제적으로 인접한 픽셀 정보를 함께 처리하게 되므로 트랜스포머의 복잡한 어텐션 메커니즘과 유사한 지역 정보 학습 효과를 얻게 됩니다.

 

5. 작동 원리: 어떻게 돌아가는가? (The "How")

UNeXt는 이미지가 입력되어 최종 분할 마스크 결과가 나올 때까지 크게 '합성곱 단계(Convolutional Stage)'와 '토큰화된 MLP 단계(Tokenized MLP Stage)'라는 두 가지 주요 흐름을 거칩니다. 디테일한 수식을 배제하고 데이터가 단계별로 어떻게 처리되는지 그 흐름(Flow)을 살펴보겠습니다. 전체 구조는 입력 이미지의 크기를 줄여나가는 5단계의 인코더와, 다시 크기를 키워나가는 5단계의 디코더로 구성됩니다.

 

Step 1: 얕은 합성곱 단계 (인코더 1~3 블록)

가장 먼저 의료 영상 데이터가 네트워크에 입력되면, 일반적인 합성곱(Convolution) 블록 3개를 차례로 통과합니다. 이 단계의 목적은 이미지의 질감, 테두리, 색상 대비 같은 기본적인 시각적 특징을 추출하는 것입니다.

  • 각 합성곱 블록은 커널 크기 3x3, 패딩 1의 합성곱 계층과 배치 정규화(Batch Normalization), 그리고 활성화 함수(ReLU)로 구성됩니다.
  • 블록을 지날 때마다 풀링(Max Pooling) 연산을 통해 이미지의 가로세로 크기는 절반으로 줄어들고, 특징을 담는 채널(Channel) 수는 증가합니다.
  • 특징적인 점은 연산량을 억제하기 위해 채널의 개수(필터 수)를 32, 64, 128로 기존 모델들보다 매우 적게 설정했다는 것입니다.

 

Step 2: 깊은 토큰화 MLP 단계 (인코더 4~5 블록, 병목 구간)

이미지의 특징이 3번 압축되어 네트워크의 가장 깊고 좁은 잠재 공간(Latent Space)에 도달하면, 연산이 무거운 합성곱 대신 '토큰화된 MLP 블록' 2개를 연속으로 통과합니다. 이 구간이 UNeXt의 핵심입니다. 데이터는 블록 내부에서 다음과 같은 순서로 처리됩니다.

  1. 축 이동(Shifting): 특징 데이터를 가로축(Width) 기준으로 정해진 칸 수만큼 밀어냅니다.
  2. 토큰화(Tokenization): 크기 3의 커널을 사용하여 데이터를 임베딩 차원(Embedding dimension) $E$ 크기의 토큰으로 압축 변환합니다.
  3. 첫 번째 MLP 통과: 압축된 토큰을 MLP에 통과시켜 전역적 특징을 섞습니다.
  4. 위치 정보 각인 (DWConv): 3x3 크기의 깊이별 합성곱(Depth-wise Convolution, DWConv)을 통과시킵니다. 트랜스포머 모델은 픽셀의 위치를 기억하기 위해 복잡한 위치 인코딩(Positional Encoding) 값을 억지로 더해주어야 합니다. 하지만 UNeXt는 파라미터 소모가 극히 적은 DWConv를 중간에 끼워 넣는 것만으로도 패딩(Padding)을 통해 발생하는 가장자리 효과 덕분에 위치 정보를 모델에 효과적으로 각인시킬 수 있습니다.
  5. 활성화 함수: ReLU 대신 더 부드러운 곡선을 가지며 최신 트랜스포머 모델에서 자주 쓰이는 GELU(Gaussian Error Linear Unit) 함수를 통과합니다.
  6. 세로축 이동 및 두 번째 MLP: 이번에는 데이터를 세로축(Height) 방향으로 밀어내고 다시 MLP를 통과시킨 후, 안정적인 학습을 위해 층 정규화(Layer Normalization)를 거쳐 블록 연산을 마무리합니다.

 

Step 3: 다시 토큰화 MLP 단계 (디코더 1~2 블록)

인코더를 통해 극도로 압축된 의미 정보를 바탕으로, 이제 원래 이미지 크기만큼 정답 지도를 복원해 나가는 디코더(Decoder) 단계가 시작됩니다. 디코더의 첫 2개 블록 역시 토큰화된 MLP 블록을 사용합니다. 이를 통해 전역적인 문맥 정보를 잃지 않으면서도 가벼운 연산량으로 이미지의 크기를 점차 키워나갑니다. 이때 해상도를 키우기 위해 파라미터가 늘어나는 전치 합성곱(Transpose Convolution) 대신, 이미지를 부드럽게 확대하는 단순한 이중 선형 보간법(Bilinear interpolation)을 사용하여 파라미터를 극단적으로 절약했습니다.

 

Step 4: 얕은 합성곱 단계 (디코더 3~5 블록) 및 스킵 커넥션

디코더의 마지막 3개 블록은 인코더 초기와 대칭되는 구조로, 다시 일반적인 합성곱 블록으로 구성됩니다. 이 단계에서 가장 중요한 메커니즘은 스킵 커넥션(Skip Connection)입니다. 인코더에서 압축되는 과정 중 동일한 해상도를 가진 계층의 특징 맵을 디코더로 직접 복사하여 넘겨줍니다. 깊은 병목 구간을 거치며 손실되었을지도 모르는 장기나 종양의 미세한 경계선 위치 정보를 얕은 층에서 직접 끌어와 화질을 보정하는 역할을 합니다. 모든 과정을 거친 데이터는 최종적으로 입력 이미지와 동일한 가로세로 크기를 가지며 픽셀별로 병변 여부를 나타내는 예측 지도(Prediction Map)로 출력됩니다.

 

6. 작동 원리: 어떻게 학습하는가? (The "How")

모델이 픽셀 단위로 병변을 칠해내는 법을 배우기 위해서는 정답이 달린 방대한 양의 예시 데이터가 필요합니다. 학습을 위해 입력되는 원본 이미지와, 의사나 의료 전문가가 병변의 위치를 픽셀 단위로 정확히 칠해 놓은 정답 마스크 데이터(Ground Truth)가 하나의 쌍(Pair)으로 구성됩니다. 본 논문에서는 현장 진료 애플리케이션의 실효성을 검증하기 위해 두 가지 주요 의료 데이터셋을 사용했습니다.

 

입력 데이터 및 출력 데이터 구성

1. ISIC 2018 (피부 병변 데이터셋) 스마트폰이나 휴대용 카메라를 이용한 현장 진료 환경을 모사하기 위해 국제 피부 영상 협력(ISIC)에서 제공한 데이터셋을 사용했습니다.

  • 입력 데이터 (Input): 일반 카메라나 더모스코피(Dermoscopy, 피부 확대경)로 촬영한 환자의 피부 병변(흑색종, 모반 등) 이미지입니다. 포맷은 2D 형태의 RGB 컬러 이미지(.png 또는 .jpg)로 제공되며, 총 2,594장의 학습용 이미지가 포함되어 있습니다. 이미지의 원본 해상도는 다양하지만, UNeXt 모델 학습을 위해 모든 이미지를 가로세로 512x512 픽셀 해상도로 규격화(Resize)하여 입력합니다.
  • 출력 데이터 (Label/Mask): 피부암 등 병변이 위치한 픽셀을 명시한 정답 데이터입니다. 동일한 512x512 크기의 2D 흑백 이진(Binary) 마스크 이미지(.png) 형태로 제공됩니다. 병변이 있는 픽셀은 흰색(1), 정상 피부 배경은 검은색(0)으로 칠해져 있습니다.

2. BUSI (유방 초음파 이미지 데이터셋) 저전력 휴대용 초음파(POCUS) 기기에서의 성능을 입증하기 위해, 이집트 Baheya 병원에서 수집한 유방 초음파 데이터셋을 활용했습니다.

  • 입력 데이터 (Input): 25세에서 75세 사이의 여성 환자 600명으로부터 수집된 2D 초음파 흑백 이미지(.png)입니다. 초음파 이미지는 특유의 반점 노이즈(Speckle noise)가 심하고 조직 간의 경계가 흐릿하여 인공지능이 학습하기 매우 까다로운 포맷입니다. 학습의 효율을 높이기 위해 정상 케이스를 제외하고, 양성(Benign) 및 악성(Malignant) 종양 케이스만 선별하여 총 647장의 이미지가 사용되었습니다. UNeXt에서는 연산량을 더욱 줄이기 위해 이 이미지들을 256x256 픽셀 해상도로 크기를 조절하여 사용했습니다.
  • 출력 데이터 (Label/Mask): 의사가 양성이나 악성 종양이 위치한 구역을 세밀하게 마스킹한 256x256 해상도의 흑백 종양 마스크 이미지(.png)입니다. 정상 유방 이미지의 경우 병변이 없으므로 완전히 검은색 마스크가 정답으로 제공됩니다.

 

학습 방식 (Loss Function)

입력 이미지가 모델을 통과해 만들어진 예측 결과와 실제 정답 마스크 사이의 오차를 계산하여 모델을 업데이트합니다. UNeXt는 이 오차를 계산하기 위해 이진 교차 엔트로피(BCE) 손실 함수Dice 손실 함수를 결합하여 사용합니다. 수식으로는 $L = 0.5 \times BCE + Dice$ 와 같이 적용됩니다.

  • BCE 손실 함수는 이미지 내의 픽셀 하나하나를 독립적으로 평가하여 "이 픽셀이 종양일 확률을 얼마나 정확히 맞췄는가?"를 채점합니다.
  • Dice 손실 함수는 "모델이 예측한 전체 종양 덩어리의 모양이 실제 정답 덩어리와 얼마나 완벽하게 겹치는가?(교집합 비율)"를 평가합니다. 픽셀 단위의 미시적 정밀도와 덩어리 단위의 거시적 일치도를 동시에 학습시킴으로써, 불균형한 의료 데이터 환경에서도 뚜렷하고 정확한 분할을 유도합니다.

 

7. 결과: 얼마나 좋아졌나? (Results)

현장 진료(Point-of-Care) 환경을 목표로 한 UNeXt의 가장 핵심적인 결과는 극단적인 경량화와 초고속 추론 속도를 달성하면서도 정확도 면에서 타협하지 않았다는 점입니다.

 

핵심 성능 지표 비교

아래 표는 기존의 대표적 합성곱 모델인 UNet과, 어텐션 메커니즘을 도입한 무거운 최고 성능 모델 TransUNet을 본 논문의 UNeXt와 다각도로 비교한 결과입니다. (ISIC 데이터셋 기준)

비교 항목 UNet (2015) TransUNet (2021) UNeXt (제안 모델)
매개변수 수 (Parameters) 31.13 M 105.32 M 1.47 M (약 72배 감소)
연산 복잡도 (GFLOPs) 55.84 38.52 0.57 (약 68배 감소)
추론 속도 (CPU 기준) 223 ms 246 ms 25 ms (약 10배 향상)
분할 정확도 (ISIC F1 Score) 84.03% 88.91% 90.41% (SOTA 달성)

 

  • SOTA(State-of-the-art) 달성: 성능 향상의 척도로 사용되는 F1 Score(정밀도와 재현율의 조화 평균)에서, UNeXt는 피부 병변 데이터셋(ISIC)에서 90.41%를 기록했습니다. 이는 무려 1억 개가 넘는 매개변수를 가진 무거운 TransUNet(88.91%)이나 기존의 UNet(84.03%) 등 당시의 SOTA 아키텍처들을 가볍게 뛰어넘는 수치입니다.
  • 초고속 추론 및 극단적 경량화: 모바일이나 소형 진단 기기에는 고성능 GPU가 없다는 현실적 제약을 반영하여, 연구진은 인텔 Xeon CPU 환경에서 이미지를 한 장 처리하는 데 걸리는 추론 속도(Inference time)를 측정했습니다. 측정 결과, 기존 모델들이 200ms 이상의 지연 시간을 보인 반면, UNeXt는 불과 25ms(0.025초)만에 처리를 완료했습니다. 이는 10배 이상 향상된 속도로, 의사가 현장에서 초음파 프로브를 움직일 때 버벅거림 없이 실시간 영상 분할 결과를 모니터로 확인할 수 있음을 의미합니다. 또한, 연산 복잡도를 나타내는 GFLOPs는 0.57에 불과하여 모바일 프로세서에서도 부담 없이 구동될 수 있는 경량성을 입증했습니다.

 

8. 이 논문의 계보와 발전 (Impact & Follow-ups)

UNeXt가 2022년에 학계에 발표된 이후, 의료 영상 인공지능 연구의 패러다임은 근본적인 변화를 맞이했습니다. 기존 학계의 트렌드가 "어떻게든 트랜스포머 레이어를 더 쌓고 연산량을 늘려서 벤치마크 점수를 1%라도 올리자"는 성능 지상주의였다면, UNeXt를 기점으로 "어떻게 하면 정확도를 조금이라도 더 보존하면서 모바일 기기에 들어갈 만큼 모델을 쥐어짜 낼 수 있는가?"라는 초경량 네트워크(Lightweight Networks) 최적화 시대가 본격적으로 열린 것입니다.

 

UNeXt는 모바일 의료 인공지능을 평가하는 새로운 '비교 표준(Baseline)'으로 자리매김했으며, 이후 UNeXt의 철학을 계승하거나 한계점을 보완하기 위해 수많은 후속 연구들이 파생되었습니다. 대표적인 후속 발전 트렌드는 다음과 같습니다.

 

후속 모델 진화 방향 및 주요 특징 매개변수 (Params)
UNeXt (2022) MLP와 축 이동을 도입한 경량화의 시발점 (기준 모델). 1.47 M
MALUNet (2022) UNeXt의 크기도 크다고 판단하여, 파라미터를 극한으로 줄임. 4개의 다중 어텐션(Multi-attention) 모듈을 결합해 피부 병변 분할 성능을 유지하면서 모델을 0.17M 단위까지 깎아내는 극한의 다이어트 달성. 0.17 M
EGE-UNet (2023) UNeXt가 초음파 이미지의 흐릿한 경계선 분할에 취약하다는 점을 보완. 가장자리 유도(Edge-guided) 메커니즘과 그룹 향상 주의 집중을 결합. 파라미터 크기를 불과 50KB(0.05M 수준)로 줄이면서도 픽셀 분할 정확도를 개선. 약 0.05 M
CMUNeXt (2023) 초음파처럼 노이즈가 심한 데이터에서는 MLP보다 합성곱(CNN)이 가진 고유한 귀납적 편향(Inductive bias)이 더 유리함을 증명. UNeXt의 병목 구간에 MLP 대신 거대한 커널(Large kernel)을 가진 합성곱 블록을 배치하여 국소적 특징을 보강함. 3.14 M
MK-UNet (2025) 경량화 트렌드의 최신 완성형. 단 하나의 커널 크기에 의존하지 않고 다중 커널 뎁스와이즈(Multi-kernel depth-wise) 합성곱을 교차로 사용. UNeXt 대비 약 5분의 1 크기의 파라미터를 유지하면서도 DICE 스코어(정확도)를 6.7%나 끌어올림. 0.316 M

 

이 외에도, 모델의 성능을 생물학적 관점에서 개선하기 위해 세포의 분열과 사멸(Division and Apoptosis) 알고리즘에서 착안한 최적화 기법을 도입하여 유방 초음파 성능을 보강한 SC-UNext 등의 파생 연구도 활발히 진행되고 있습니다. 

 

9. 마무리

지금까지 살펴본 UNeXt는 딥러닝 연구, 특히 의료 인공지능 산업이 나아가야 할 아주 실용적이고 중요한 방향성을 제시하는 논문입니다. 의사들이 실제 진료실 환경에서 환자에게 곧장 적용할 수 있도록 '현장 진료(Point-of-care)'라는 현실적인 가치에 집중했다는 점에서 실무적으로 엄청난 기여를 한 연구라 평가합니다.   

 

데이터의 차원을 좁은 토큰으로 압축해 다층 퍼셉트론(MLP)에 통과시키고, 복잡한 삼각함수나 어텐션 수식 없이 축을 밀어내는(Shifting) 물리적 조작만으로 주변 문맥을 읽어내는 발상은 딥러닝 최적화가 반드시 복잡한 수학으로만 풀리는 것이 아님을 증명한 혁신적인 아이디어입니다.   

 

하지만 이 모델을 실제 의료 임상 환경에 도입하거나 관련 앱을 개발할 때는 개발자와 의료진이 반드시 염두에 두어야 할 주의 사항들이 있습니다.

  1. 입력 해상도와 배치 사이즈의 실무적 제약: 논문에서는 512x512나 256x256과 같이 규격화된 해상도를 사용했습니다. 모바일 기기의 CPU나 NPU 메모리는 제한적이므로, 스마트폰으로 촬영한 4K 이상의 고해상도 피부 사진을 원본 그대로 네트워크에 밀어 넣으면 작동이 불가능하거나 지연이 발생할 수 있습니다. 기기 환경에 맞춘 이미지 리사이징 전처리 최적화가 실무에서는 매우 중요합니다.
  2. 3D 데이터로의 확장성 숙제: 논문에서 입증된 성능은 2D 이미지 기반의 분할에 국한됩니다. 실제 병원에서 자주 쓰이는 CT나 MRI 등은 수백 장의 슬라이스가 겹쳐진 3D 볼륨(Volumetric) 데이터입니다. 다층 퍼셉트론(MLP)의 구조적 특성상 차원이 높아지면 메모리 요구량이 급증하거나 연속된 슬라이스 간의 공간적 위치 정보 손실이 발생할 수 있으므로, 3D 구조로 적용할 때는 추가적인 구조 확장이 필수적인 과제로 남습니다.   

결론적으로, UNeXt는 무겁고 비싼 대형 하드웨어의 전유물로 여겨지던 의료 인공지능을 가볍고 저렴한 모바일 및 휴대용 디바이스의 영역으로 끌어내리는 데 지대한 공헌을 한 선구적인 연구입니다. 앞으로도 이 논문을 기점으로 한 모델 경량화 연구가 지속되어, 값비싼 진단 장비나 GPU 인프라를 갖추지 못한 자원 부족 국가의 병원이나 긴급 재난 현장에서도 스마트폰 하나로 의사 수준의 진단 보조가 이루어지는 진정한 '의료 인공지능의 평등화(Health equity)'가 실현되기를 진심으로 기대해 봅니다.   



 

반응형