의료 인공지능에서의 멀티 태스크 러닝의 이해와 활용

Abstract

최근, 의료 분야에서 인공지능은 많은 발전을 통해 다양한 분야로 확장하며 활용되고 있다. 하지만 대부분의 인공지능 기술들은 하나의 모델이 하나의 태스크만을 수행할 수 있도록 개발되고 있으며, 이는 의사들의 복잡한 판독 과정을 인공지능으로 설계하는데 한계로 작용한다. 멀티 태스크 러닝은 이러한 한계를 극복하기 위한 최적의 방안으로 알려져 있다. 다양한 태스크들을 동시에 하나의 모델로 학습함으로써, 효율적이고 일반화에 유리한 모델을 만들 수 있다. 본 종설에서는 멀티 태스크 러닝에 대한 개념과 종류, 유사 개념 등에 대해 알아보고, 연구 사례들을 통해 의료 분야에서의 멀티 태스크 러닝의 활용 현황과 향후 가능성을 살펴보고자 한다.

서론

최근, 딥러닝을 기반으로 하는 인공지능 기술이 빠른 속도로 발전하면서 의료, 무인 자동차, 음성 인식 등 생활 속 다양한 분야에서 큰 성과를 보이며 개발 및 활용되고 있다(1). 인공지능은 더 이상 미래 기술이 아니며, 우리 일상생활에 밀접하게 스며들며 현실로 다가오고 있다. 이와 같은 상황에서 의료 분야는 인공지능이 가장 활발하게 연구와 개발이 진행되는 분야 중 하나이다. 질병 진단부터 외과 수술, 환자 모니터링, 신약 발굴 등 다양한 분야로 확장하고 있으며, 그중에서도 의료 영상 분야는 인공지능이 가장 활발하게 활용되고 있는 분야이다.

의료 영상 분야에서 초기 인공지능은 병변 분류, 병변 검출, 병변 분할 분야에 국한되어 집중적 연구가 이루어졌으나, 최근에는 영상 정합(image registration), 영상 재구성(image reconstruction), 영상 슈퍼레졸루션(image super-resolution), 영상 합성(image synthesis) 등 다양한 문제 해결을 위한 인공지능 기술들이 보고되고 있다(2,3,4,5,6).

최근, 인공지능 기술을 탑재한 다양한 컴퓨터 보조 진단(computer aided diagnosis; 이하 CAD) 시스템들이 상용화되고 있지만, 대부분의 CAD의 작동 과정은 영상의학과 의사들의 근본적 워크플로우와는 큰 차이가 있다(7,8). 의사들의 판독 과정은 매우 복잡하며, 영상으로부터 여러 장기의 다양한 이상 구조들을 동시에 감지한다. 반면, 대부분의 인공지능 기반 CAD 시스템들은 하나의 장기로부터 단일 병변에 대한 이상 구조를 감지한다. 물론, 하나 이상의 이상 구조를 감지할 수 있는 CAD 시스템도 있으나, 그 범위가 매우 제한적이며 단일 장기 내에서의 이상 구조를 대상으로 하는 경우가 많다. 그 원인은 일반적으로 딥러닝 모델은 하나의 태스크를 위한 데이터셋으로 구성되고 학습되기 때문이다. 여러 종류의 태스크를 수행하기 위해서는 각각의 태스크를 위한 모델을 따로 구축해야 한다. 따라서 수행해야 하는 태스크의 개수가 늘어날 때마다 점진적을 보유해야 하는 모델의 개수가 늘어나는 한계가 있다. 이는 딥러닝 모델의 유지보수도 까다롭고, 저장용량 측면에서도 비효율적이다. 하지만 범용적인 AI 시스템을 위해서는 다양한 태스크들이 동시에 수행될 필요가 있고, 이를 위해서는 여러 태스크들이 동시에 처리되어 빠르게 적용되어야 한다. 최근에는 이러한 한계를 극복하기 위해 경량화와 같이 여러 태스크들이 빠른 속도로 처리되기 위한 다양한 시도들이 이루어지고 있으며, 그중 멀티 태스크 러닝 방법(multi-task learning; 이하 MTL)은 이러한 한계를 극복하기 위한 좋은 대응 방안이 될 수 있다. MTL 방법은 관련 있는 태스크들을 동시에 하나의 모델로 학습시키는 방법으로, 각각의 태스크들이 학습 과정에서 만들어지는 유용한 정보가 서로 공유되어 모든 태스크들의 학습에 도움이 되며, 효율적인 연산, 과적합 방지, 데이터 증강 효과 등의 많은 장점이 있다(9,10,11).

특히, MTL이 가장 활발하게 연구되는 분야로는 자율주행 분야가 있다(12). 자율주행 자동차는 카메라를 통해 입력되는 실시간 도로 상황 영상을 기반으로 도로 및 차선 감지, 표지판 탐색, 사물 분류 등 다양한 문제를 실시간으로 해결해야 한다. 만약 각각의 문제들을 해결하기 위해 각각의 신경망을 만들 경우, 매우 많은 시간과 계산이 소모될 수밖에 없다. 따라서 MTL은 자율주행 분야와 같이 동시에 다양한 태스크를 처리해야 하는 경우 좋은 해결 방안이 될 수 있다. 의료 분야도 자율주행 분야와 마찬가지로 MTL의 활용 가치가 높은 분야라 할 수 있다. 앞에서 언급했듯이, 단일 태스크로는 의사들의 복잡한 워크플로우를 처리하는데 한계가 있으며 다양한 태스크들을 동시에 처리할 수 있어야만 의료 AI 분야가 한 단계 더 발전하여 임상 활용의 가능성이 높아질 수 있다. 이에 본 종설에서는 MTL에 대해 전반적으로 알아보고, 의료 분야에서의 MTL이 어떻게 시도되고 있는지 살펴보고자 한다.

멀티 태스크 러닝(Multi-Task Learning)

MTL은 서로 연관이 있는 태스크들을 동시에 학습함으로써 각각의 태스크에 대한 성능을 전반적으로 향상시키는 학습 방법이다(13). MTL은 하나의 태스크만을 학습하는 단일 태스크 러닝과 비교하여 여러 장점이 있다. 다수의 태스크들이 하나의 심층신경망을 공유하기 때문에 반복적인 계산을 줄여 메모리 양을 줄이면서 추론 속도를 향상시킬 수 있어 학습 효율이 증대될 수 있다. 또한 태스크들 간에 상호 보완적인 정보를 공유하는 경우 서로에 대한 정규화 역할을 하여 각 작업의 예측 성능을 향상시킬 수 있으며, 은닉층들이 하나의 태스크에 과적합 되지 않도록 하여 모델의 일반화 성능 향상이 가능하다(14).

일반적으로 MTL은 공유 네트워크(shared network)를 통해 공통의 feature를 추출하고, 그 이후 분기별로 나누어져 각 태스크를 해결할 수 있도록 추가 학습되는 구조를 가진다(15). 이때 공유 네트워크의 목적, 구성 환경, 접근 방식 등에 따라 다양한 방법으로 구분될 수 있다.

Parameter Sharing Method

MTL은 공유 네트워크의 접근 방식에 따라 크게 hard parameter sharing 방법과 soft parameter sharing 방법으로 구분될 수 있다.

Hard parameter sharing 방법은 MTL에 대한 가장 일반적인 접근 방식으로, 모델의 레이어가 shared layer와 task specific layer로 구성된다(Fig. 1A). Shared layer에서는 데이터로부터 공통적인 feature를 추출하고, task specific layer를 통과하면서 각각의 태스크들에 대한 예측값을 추론하게 된다. Shared layer에서 일반화된 feature를 추출하게 되어 각 태스크별 데이터가 부족한 상황에서도 성능을 향상시킬 수 있는 구조를 가지게 된다(11).

Fig. 1. Multi-task learning structure according to the parameter sharing method.

A. Hard parameter sharing method.

B. Soft parameter sharing method.

C. Mixed parameter sharing method.

Soft parameter sharing 방법은 hard parameter sharing 방법과 달리 각각의 태스크가 고유한 신경망 네트워크를 통해 파라미터를 가진다. 이때 각각의 신경망 파라미터 간의 거리를 줄이기 위한 제약식이 추가되는 형태를 가진다. Shared layer는 존재하지 않으며 constrained layer에서 파라미터 간에 feature 공유가 이루어진다. Fig. 1B를 보면 soft parameter sharing은 동일한 신경망에서 시작하는 것이 아닌 각 태스크마다 별도의 신경망 네트워크에서 시작이 되며, constrained layer의 각 layer끼리 연결되어 서로 간의 feature를 공유하면서 학습되는 구조임을 확인할 수 있다. Soft parameter sharing 방법은 태스크별로 각각의 네트워크를 가지기 때문에 많은 메모리가 소요되며, 태스크별로 성능이 달라 과적합이 특정 태스크에만 일어날 수 있는 단점이 있다(16).

앞서 설명한 hard parameter sharing 방법과 soft parameter sharing 방법을 혼합한 형태로도 구성이 가능하다. Fig. 1C를 보면 soft parameter sharing과 동일하게 shared layer로 구성된 하나의 신경망 네트워크로 시작되는 것을 확인할 수 있다. 이후 각 태스크에 따라 별도의 신경망 네트워크 구조를 가지는데, 이때 hard parameter sharing의 constrained layer와 마찬가지로 각 layer 간에 연결이 되어 feature를 공유하는 구조를 가진다. 이처럼 Hard parameter sharing 방법의 모델 구조에서 task specific layer를 constrained layer로 구성할 수 있으며, 데이터 특성에 따라 shared layer의 크기나 constrained layer의 제약 크기 등을 자유롭게 설정할 수 있다.

현재 대부분의 MTL 논문들은 대부분 hard parameter sharing 방법과 soft parameter sharing 방법을 기반으로 한다. 하지만 최근에는 신경망에서 MTL의 더 나은 메커니즘을 개발하고자 하는 시도들이 점차 증가하고 있다. 2017년 Long 등(17)이 보고한 deep relationship network는 fully connected layer에 행렬 우선순위를 배치하여 모델이 베이지안 모델과 유사하게 태스크 간의 관계를 학습할 수 있도록 하였다. 하지만 이 방식은 기존의 잘 정의된 분야에 적합하며, 새로운 태스크들에 대해선 아직 더 많은 실험이 필요하다. 2017년 Lu 등(18)이 보고한 fully adaptive feature sharing 방법은 얕은 네트워크에서 시작하여 학습 중에 동적으로 네트워크를 확장하는 상향식 접근 방법으로, 전역적으로 최적화되었음을 증명하기 어렵고, 태스크들 간의 복잡한 상호작용을 모델이 학습하게 하기 어렵다는 단점이 있다. 2016년 Misra 등(19)이 보고한 cross-stitch network는 서로 다른 태스크를 위한 아키텍처 간에 십자수 단위 (cross-stitch unit)를 적용하여 서로 다른 태스크에 대한 지식을 활용하는 방식을 결정할 수 있도록 한다. 이와 같이 태스크들 간의 지식 공유를 위한 다양한 MTL 아키텍처들이 연구되고 있다. 하지만 아직까지는 대부분의 모델들이 hard parameter sharing이나 soft parameter sharing에 기반하고 있으며, 새로운 MTL 구조는 기초적인 수준으로, 더 많은 연구와 개선, 증명이 필요한 상황이다.

Transfer Learning과의 비교

MTL과 전이 학습(transfer learning)은 파라미터를 공유하여 성능을 높인다는 면에서 서로 유사한 부분이 있다. 전이 학습은 source 태스크를 통해 학습한 후, knowledge transfer를 통해 target 태스크를 학습시키는데 반해, MTL은 source 태스크와 target 태스크가 동일하게 knowledge를 공유하며 성능을 높이다는 것에 차이가 있다. 즉, 전이학습은 이미 잘 작동하는 source 태스크를 이용하여 target 태스크의 성능을 개선하는 것에 목적을 두고 있고, MTL은 source 태스크와 target 태스크를 동시에 학습하여 성능을 개선하는 것에 목적을 둔다(20,21).

전이 학습은 주로 데이터가 부족한 경우 많이 활용되는 방법으로, 다른 데이터셋이나 목적 함수를 사용해 미리 훈련한 후, 이를 바탕으로 본격적인 학습에서 신경망 가중치 파라미터를 더 쉽게 최적화하고 수렴속도와 성능을 향상시킬 수 있다(22). 대표적으로 영상 처리 분야에서 ImageNet 데이터를 기반으로 학습된 신경망 가중치 파라미터값을 사용하여 다른 문제에 전이 학습을 적용하는 접근 방법이 많이 사용되고 있다(23).

전이 학습을 위해서는 사전 학습된(pre-trained) 모델이 필요하며, 사전 학습된 모델을 새로운 태스크에 적용하기 위해 일부 가중치를 조절하는 미세 조정(fine-tuning) 과정을 통해 전이 학습이 이루어진다(24).

미세 조정은 합성곱 신경망을 구성하고 있는 convolutional base와 classifier 각각에 대해 얼마나 많은 부분을 새로 학습시킬지에 따라 Fig. 2와 같이 크게 3가지 종류로 구분할 수 있다. 여기서 convolutional base는 합성곱층과 풀링층들로 구성된 feature를 추출하는 영역을 의미하고, classifier는 fully connected layer로 이루어져 카테고리를 분류하는 영역을 의미한다. 미세 조정의 첫 번째 방법은 Fig. 2B와 같이 전체 모델을 새로 학습하는 방법으로, 사전 학습 모델의 구조만 사용하면서 새로운 데이터셋에 맞게 전부 새로 학습시키는 방법이다. 두 번째 방법은 Fig. 2C와 같이 convolutional base의 일부분은 고정시키고, 나머지 레이어와 classifier를 새로 학습시키는 방법이며, 세 번째 방법은 Fig. 3C와 같이 convolutional base는 고정시킨 상태에서 classifier만 새로 학습시키는 방법이다(24).

Fig. 2. Fine-tuning methods in various situations.

A. Convolutional neural network’s basic structure.

B. Train the entire model.

C. Train some layers of the convolutional base and freeze others.

D. Train the classifier and freeze the convolutional base.

Fig. 3. Example of a dataset in multi-task learning. It consists of different tasks on different data.

학습시키고자 하는 데이터셋이 사전 학습된 모델의 데이터와 유사성도 높고 규모도 큰 경우에는 어떠한 방법을 사용해도 무방하지만 convolutional base의 높은 레벨 계층 일부와 classifier만 새로 학습시키는 것이 유리하다. 데이터셋이 유사하기 때문에 이전에 학습된 지식을 충분히 활용할 수 있는 방법이다. 데이터셋이 사전 학습된 모델의 데이터와 유사성은 높으나 규모가 작은 경우에는 기존의 convolutional base는 고정시키고 classifier만 새로 학습시키는 방법이 유리하다. 데이터셋이 사전 학습된 모델의 데이터와 유사성은 적으나 규모가 큰 경우에는 전체 모델을 새로 학습하는 방법을 사용하는 것이 좋다. 데이터의 유사성이 적기 때문에 사전 학습된 모델의 구조와 파라미터들을 이용하되, 데이터의 규모가 크므로 처음부터 새롭게 학습하는 것이 유리하다. 데이터셋이 사전 학습된 모델의 데이터와 유사성도 낮고 규모도 작은 경우에는 가장 불리한 상황으로, convolutional base의 좀 더 낮은 레벨 계층부터 classifier까지 새로 학습시키는 방법을 시도하는 것이 좋다. 다만, 너무 많은 계층을 새로 학습시키면 적은 데이터셋에 과적합될 우려가 있고, 너무 적은 계층을 학습시키면 제대로 학습되지 않을 수 있기 때문에 적당한 계층을 고정시켜 학습시키는 것이 중요하다. 이와 같이 효과적인 전이 학습을 위해서는 데이터의 규모와 사전 학습된 모델과의 유사성 정도에 따라 적절한 방법을 선택하여 적용해야 한다.

Multi Label Learning과의 비교

멀티 라벨 러닝(multi label learning; 이하 MLL)과 MTL은 학습 데이터셋의 구성에 따라 구분된다. MTL은 각각의 태스크가 서로 다른 학습 데이터셋으로 구성되는 반면, MLL은 각각의 태스크가 서로 동일한 데이터셋으로 구성된다(25).

MTL은 여러 데이터셋으로 구성되어 동시에 학습되며, 데이터셋에 따라 서로 다른 태스크를 수행할 수 있다. 예로 Fig. 3과 같이 chest X-ray, mammography, abdominal CT 영상으로 데이터셋이 구성되어 chest X-ray에서는 폐 결절 유무 분류, mammography에서는 유방 종괴 검출, abdominal CT에서는 간 분할을 목적으로 동시에 학습하는 방법이다. 반면, MLL은 하나의 데이터셋에서 각 데이터 샘플별로 다수의 레이블이 존재하게 되며, 동시에 여러 레이블들이 학습된다. 예로 Fig. 4와 같이 chest X-ray에서 폐 결절, 기흉, 흉수에 대해 각각의 유무에 대한 레이블을 구성하여 동시에 각 질병의 유무를 동시에 예측할 수 있는 모델을 학습할 수 있다.

Fig. 4. Example of a dataset in multi-label learning. There are three types of labels on one chest X-ray: nodule, pneumothorax, and pleural effusion.

멀티 라벨(multi label)은 멀티 클래스(multi class)와 목적이 동일하기 때문에 혼용해서 사용하는 경우가 많으나, 데이터셋 구성에 차이가 있기 때문에 구분하여 사용해야 한다. 멀티 라벨과 멀티 클래스 모두 2개 이상의 카테고리에 속하는 객체를 대상으로 한다는 공통점이 있지만 멀티 라벨은 하나의 데이터에 여러 객체가 포함되는 반면, 멀티 클래스는 하나의 데이터에 하나의 객체만 포함한다(26).

의료 분야에서의 MTL 활용

데이터 수집이 어려운 의료 분야에서 관련 있는 태스크들을 동시에 학습시킬 수 있는 MTL은 좋은 해결 방안이 될 수 있다. 최근에는 의료 분야에서도 MTL을 접목하여 문제를 해결하고자 하는 시도가 점차 증가하고 있다.

2017년 Chaichulee 등(27)은 병원에서 조산아의 지속적인 비접촉 생체 신호 모니터링을 위해 실시간 프레임 영상에서 환자의 존재를 감지하고 피부 영역을 분할하기 위한 MTL 기반의 convolutional neural network (이하 CNN) 아키텍처를 제안하였다. 30명 조산아의 비디오 녹화 영상을 대상으로 학습 및 검증을 수행한 결과, 환자 검출은 98.75%의 정확도를 보였고 피부 영역 분할은 88.57%의 intersection over union를 보였다.

2020년 Li 등(28)은 안저 영상에서 당뇨병성 망막병증(diabetic retinopathy; 이하 DR)과 당뇨병성 황반부종(diabetic macular edema; 이하 DME)의 등급 예측을 위해 질병 간의 내부 관계를 탐색하여 DR과 DME를 공동으로 등급 예측이 가능한 MTL 기반의 cross-disease attention network (CANet) 아키텍처를 제안하였다. Messidor 공개 데이터셋과 IDRiD 데이터셋을 대상으로 학습 및 검증을 수행한 결과, DR은 96.3%의 area under the curve (이하 AUC), DME는 92.4%의 AUC를 보였다.

2019년 He 등(29)은 흉부 CT 영상에서 식도(esophagus), 심장(heart), 기관(trachea), 대동맥(aorta)의 4종류의 흉부 장기를 분류 및 분할하고자 MTL을 활용한 바 있다. 그들은 MTL 기반의 uniform U-like Encoder-Decoder 아키텍처를 구성하여 40명 환자의 3D 흉부 CT로 학습 후 20명 환자의 데이터로 검증하였다. 그 결과, 식도에서 0.86, 심장에서 0.95, 기관에서 0.92, 대동맥에서 0.94의 Dice similarity coefficient (이하 DSC) 성능을 보였다.

2020년 Zhai 등(30)은 흉부 CT에서 결절의 9개 방향의 2D 단면을 이용하여 2D와 3D에 대한 정보를 구성하고, 각각을 이용하여 양성과 악성의 분류가 가능한 multi-task convolutional neural network (이하 MT-CNN) 프레임워크를 제안하였다. 그들은 각 결절의 후보군을 추출한 후, 서로 다른 9개의 2D 단면을 구성하여 각 단면에서 MTL 기반의 2D MT-CNN 모델을 통해 3D 폐 결절의 특성을 학습하였다. 최종적으로는 9개의 2D MT-CNN 모델의 예측 결과를 가중 융합하여 양성과 악성에 대한 분류 결과를 얻었다. 학습된 모델은 LUNA-16과 LIDC-IDRI 데이터 세트에서 각각 97.3%와 95.59%의 AUC의 성능을 보였다.

2020년 Gao 등(31)은 공개 데이터인 full filed digital mammogram (FFDM) 데이터 세트에서 MTL 기반의 FT-MTL-Net 아키텍처를 통해 양성 및 악성 분류와 종괴 검출, 종괴 분할의 3가지 태스크를 동시에 학습 및 검증하였다. 그 결과, 종괴의 양성 및 악성 분류는 0.92 ± 0.01의 AUC를 보였고, 종괴 검출은 0.91 ± 0.05의 true positive rate과 3.67의 false positive rate per image를 보였으며, 종괴 분할은 0.76 ± 0.03의 DSC를 보였다.

2020년 Amyar 등(32)은 흉부 CT 영상에서 COVID-19의 식별과 분할을 위해 MTL 기반의 CNN 아키텍처를 제안하였다. CNN 아키텍처는 hard parameter sharing 방법으로 설계되었다. 식별 태스크는 정상과 COVID-19, 기타 감염병의 3가지 카테고리 분류를 목표로 하며, 분할 태스크는 COVID-19 영역의 분할을 목표로 하였다. 다수의 공개 데이터베이스로부터 수집된 1219개의 데이터로부터 학습하고 150개의 데이터로 검증한 결과, 정상과 COVID-19, 기타 감염병의 식별에 대해 94.67%의 정확도(accuracy)를 보였으며, COVID-19 영역에 대한 분할은 88.0%의 DSC를 보였다.

이처럼 다양한 의료 데이터에서 MTL을 적용한 다양한 연구들이 보고되고 있으며, 대부분의 연구들에서 싱글 태스크 모델들과 비교하여 대등하거나 향상된 성능을 보이고 있다.

결론

본 종설에서는 MTL에 대해 알아보고 의료 분야에서 MTL이 어떻게 활용되는지 다양한 연구 사례들을 살펴보았다. MTL은 여러 태스크들을 효율적으로 학습시킬 수 있는 최적의 방법임에는 틀림이 없다. 다양한 태스크들을 동시에 학습시킴으로써 복잡한 워크플로우에 대한 모델 설계 시 유용하게 활용될 수 있다. 하지만, MTL이 장점만 있는 것은 아니다. 각각의 태스크들이 서로 무관할 경우 특정 태스크의 성능을 오히려 저해할 수도 있으며, 태스크들 간에 학습 난이도의 편차가 클 경우 모델이 수렴하지 않을 수 있다. 따라서, 아직까지는 태스크들 간의 관계, 계층 구조 및 성능 저하의 원인이 되는 많은 제한 조건들에 대한 많은 연구와 개선이 필요하다. 그럼에도 불구하고, MTL은 현재 대부분의 싱글 태스크 기반 모델들의 한계를 극복하기 위한 좋은 대안이며, 의료 분야에서도 인공지능의 임상 활용을 위한 MTL의 역할은 크다고 할 수 있다. 다만, 의료 분야에서 MTL을 활용할 때에는 모델의 목적과 태스크들 간의 관계와 차이를 정확히 분석하여 적절한 MTL 구조를 설계하는 것이 필요하다. 향후 의료 분야에서도 MTL에 대한 더 많은 연구와 개선이 충분히 이루어진다면 의사들의 복잡한 워크플로우를 인공지능이 이해하고 역할을 수행하는데 많은 도움이 될 수 있을 것으로 기대한다.