DEVOCEAN

안녕하세요! 저희는 'AI 모델을 이용한 X-ray 영상 자동 진단'이라는 주제를 연구하는 6조 beyond label팀입니다.

이번 포스트에서는 저희가 지난 3개월동안 어떤 연구를 해왔는지 공유하려고 합니다.

우선 저희는 이번 중간발표를 기점으로 phase1과 phase2를 나눠서 프로젝트를 진행하였는데요.

Phase1에서는 어떤 연구를 해왔으면, phase2에서는 어떤 계획을 가지고 있는지에 대해 공유드리겠습니다.

주제 소개

아래의 그래프는 논문 출판 년도와 그 논문에서 사용한 학습 데이터 양을 보여줍니다.

최근의 논문일수록 많은 수의 학습 데이터를 사용하는 것을 볼 수 있는데요.

이는 supervised 방법으로는 data labeling에 상당한 부담이 된다는 것을 의미합니다.

실제로 아래의 자료들을 보시면 데이터 라벨링에 투여하는 비용이 5년전에 비해 2배 이상이 증가했음을 확인할 수 있습니다.

또한 인공지능의 연구에서 데이터 라벨링이 차지하는 시간은 무려 1/4이라고 합니다.

특히 의료 분야는 다른 분야보다 라벨링을 진행하는 데에 많은 리소스가 투입됩니다. 숙련된 전문의가 필요하고, 고비용과 엄청난 시간을 필요로 합니다.

이런 문제를 해결하고자, labeling 데이터를 따로 필요로 하지 않는 self supervised learning 방법론이 최근 각광을 받고 있습니다.

또한 self-supervised로 학습하는 모델들이 supervised 방식으로 학습된 모델보다 데이터의 representation을 더 잘 추출한다고 알려져 있습니다.

저희는 이러한 self-sueprvised 방법론을 의료분야에 적용하여 고비용의 labeling cost문제를 해결하고, 최고의 성능을 내는 모델을 개발하고자 합니다.

Phase1 연구

ICCV 2023 CVAMD Challenge

위에서 언급했다시피, 저희는 전체 프로젝트를 이번 중간발표를 기점으로 phase1과 phase2로 나누어 진행하고 있습니다.

Phase1에서는 ICCV 2023 CVAMD(Computer Vision for Automated Medical Diagnois) Challenge에 참여하였습니다.

ICCV 2023 CVAMD challenge는 x-ray 이미지에서 26개의 질병에 대한 multi-label classification task를 수행하는 것이었습니다.

아래의 그림은 저희가 챌린지에서 사용한 데이터셋이 한 이미지당 가지고 있는 class 개수를 histogram으로 표현한 것입니다.

한 이미지는 최대 11개의 class를 가지고 있음을 확인할 수 있습니다.

또한 long tailed problem을 해결하는 것도 본 챌린지에서 매우 중요한 요소 중 하나였습니다.

Long-tailed data란 적은 수의 head class에 대다수의 데이터가 속해있고, 많은 수의 tail class에 적은 수의 데이터가 속해있는 것입니다.

이렇게 데이터의 class distribution이 일정하지 않으면, 모델은 head class에 bias를 가지고 학습이 되기 때문에 좋은 성능을 낼 수 없습니다.

아래의 그림은 저희가 챌린지에서 사용한 데이터셋의 class별 데이터 개수를 histogram으로 표현한 것입니다.

Most class 데이터와 least class ㄷ이터의 개수가 최대 180대 1까지 벌어지는걸 확인할 수 있습니다.

챌린지 1차 연구 방향

[Self-supervised Multi-label Classification을 위한 최적의 Model 탐색]

저희 팀의 메인 연구주제는 Self Supervised Learning이었기 때문에, 가장 먼저 Self-supervised Learning을 이용하여 multi label classification 문제를 해결하는 방법을 모색하였습니다.

다양한 Self Supervised Learning framework에 대한 조사를 진행해본 결과, chest x-ray 480만장에 대해 MOCO-V2 방식으로 학습한 CheSS라는모델이 가장 높은 성능을 보였습니다.

[Long-tailed Problem 완화를 위한 최적의 Loss 탐색]

또한 대회의 핵심인 long tail problem을 해결하기 위해 loss 단에 변화를 주는 방향으로 접근하였습니다.

Object detection에서 흔히 쓰이는 focal loss와, focal loss를 multi label 상황에 특화시킨 asymmetric loss에 대한 실험을 진행했습니다.

실험 결과

[Self-supervised Multi-label Classification을 위한 최적의 Model 탐색]

가장 먼저 ImageNet에 대해 supervised로 pretrained ResNet50과 chest x-ray에 대해 self-supervised learning으로 pretrained 된 ResNet50인 CheSS를 비교해봤습니다.

그 결과 CHESS를 사용한 모델이 더 좋은 성능을 내는 것을 확인할 수 있었습니다.

저희는 그 원인을 pretrained data와 target data의 domain이 일치했기 때문과 비지도학습을 통한 좋은 representation learning의 효과로 생각하였습니다.

[Long-tailed Problem 완화를 위한 최적의 Loss 탐색]

다음으로 전체 데이터셋을 downsampling하여 focal loss와 Asymmetric loss (ASL)을 비교하는 실험을 진행했습니다.

표에 보이시는 것처럼 focal loss를 사용했을 때보다 Asymmetric Loss를 사용하였을 때 약 1.5%의 의미 있는 성능 향상이 있음을 확인할 수 있었습니다.

1차 연구 성과

이 과정에서 저희가 찾은 best model은 ASL로 CheSS 모델을 학습시킨 것이었습니다.

하지만 이는 전체 15등 정도밖에 못 미치는 수치로 다소 실망스러웠기에, 저희는 진행방향을 틀어 supervised learning 으로 2차 연구를 진행하기로 결정하였습니다.

챌린지 2차 연구 방향

[Supervised Learning (SKT 자체 모델)을 이용한 Multi-label Classification]

2차 연구 방향에서는 지도학습을 사용하는 SKT 자체 모델을 이용하여 multi-label classification 을 진행했습니다.

이는 CNN multi-label에서 SOTA를 달성하고 있는 TResNet 을 base 모델로 사용하여 Feature pyramid network 구조를 결합한 것으로, 다양한 크기의

질환을 포착할 수 있도록 개발된 모델이었습니다.

[Long-tailed Problem 완화를 위한 데이터 증강]

또한 Long-tailed 문제를 해결하기 위해 데이터를 증강하였습니다. 이는 tail class에 해당하는 데이터들을 증강하여 head class와의 ratio를 맞춰주는 것으로 진행됐습니다.

이를 위해 저희는 tail class에 해당하는 질병 patch를 정상 이미지에 붙여 인공 데이터를 제작했습니다.

이때 원본 데이터에 존재하는 질병간 상관관계를 유지하기 위해 질병의 correlation 또한 유지되도록 data를 증강하였습니다.

아래의 그림에서 왼쪽(MIMIC correlation)이 원본 데이터의 질병간 상관관계를 나타내고, 오른쪽(Augmented MIMIC Correltation)은 증강 데이터의 질병간 상관관계를 나타냅니다.

두 colormap이 상당히 유사한걸 확인할 수 있습니다.

2차 연구 결과

연구 결과, TResNet에 Feature Pyramid network를 결합한 모델은 기존 baseline model인 TResNet보다 최대 1%p 성능 향상을 보였습니다.

이는 다양한 feature map을 사용함으로써 TResNet보다 최대 1%p 성능 향상을 보였습니다.

이는 다양한 feature map을 사용함으로써 더 많은 representation 정보를 사용했기 때문이라고 생각합니다.

또한 데이터 증강을 통해 class imbalance 문제를 완화하였는데요. 아래의 그림에서 original의 class distribution보다 augmented의 class distribution이 더 고른 것을 확인할 수 있습니다.

위와 같이 증강한 데이터를 모델 학습에 사용함으로써 증강 데이터를 사용하기 전보다 약 0.3%p의 성능 향상을 확인하였습니다.

ICCV 2023 Challenge 성과

지난 3개월동안 Phase 1에 몰두하며 저희는 많은걸 배우고 성장할 수 있었습니다.

챌린지에 참여하는 것부터 논문 작성까지 일련의 과정들을 경험해볼 수 있어 좋았습니다.

또한 챌린지 전체에서 top 5 등수를 기록하고, ICCV 2023 workshop에 2저자로 accept 되는 큰 성과를 기록할 수 있었습니다.

10월에 열리는 ICCV workshop에 저희 멘토님들께서 SKT 이름으로 해당 연구에 대한 포스터 발표를 앞두고 있어 상당히 기대가 되기도 합니다.

마지막으로, 이번에 진행했던 Data augmentation 방식을 잘 정리하여 특허 출원을 준비중에 있습니다.

이번 phase 1의 성과가 SKT의 기술력을 해외 학회에 선보이는 등 회사에도 기여할 수 있어

상당히 의미있는 시간이었다고 생각합니다.

Phase2 연구 계획

phase1에서 진행했던 Self-supervised learning 실험을 바탕으로 phase2에서 저희만의 novel modeling을 진행하려고 합니다.

목표는 다소 추상적이지만 SSL 방식을 통해 좋은 representation을 추출하고, 이를 기반으로 다양한 task에서 좋은 성능을 내는 것입니다.

아래 그림은 저희 연구에 참고할 모델의 아키텍처입니다(Hyun Seok Seong, CVPR 2023).

좋은 representation을 추출하는 feature extractor, 추가적인 정보를 추가할 수 있는 module, 그리고 task-specific head로 구성되어 있습니다.

먼저 여러 task에 강인한 좋은 feature를 추출할 수 있도록 SSL 방식의 backbone을 탐색하고, pretrained weight을 의료 데이터에 맞게 domain shift하는 방법에 대해 연구해보고자 합니다.

다음으로 Feature extractor에서 추출한 representation에 pixel-level부터 image-level까지의 정보를 제공하는 모듈을 추가하여 self-supervised learning 학습을 진행할 것입니다.

마지막으로 간단한 head를 부착하여 task 별 성능을 검증하고 phase 1에서 검증한 ASL 등 loss function을 활용하여 의료 데이터의 long-tail 문제를 완화하는 실험을 할 것입니다.

추가 연구 방향

앞서 말씀드린 모델이 완성된다면 추가적인 연구도 진행할 것입니다.

일부 x-ray 데이터셋은 이미지뿐만 아니라 clinical report를 포함하고 있습니다.

Report가 포함된 dataset의 경우 이미지와 report 정보를 함께 학습하면 질환 classification 및 detection의 성능이 향상된다는 연구들이 있습니다(Yuhao Zhang, MILHC 2022).

x-ray 이미지와 해당 report의 유사도를 최대화하는 방향으로 학습을 진행하여 풍부한 visual representation을 추출하는 연구도 진행할 것입니다.

Phase2 기대효과

저희는 phase 2의 연구 성과도 phase 1과 마찬가지로 저희의 개인적인 성장과 회사의 기술 성장을 모두 잡기를 목표하고 있습니다.

Phase2의 연구가 phase 1에서처럼 논문 혹은 특허로 이어지기를 기대하고 있습니다.

또한 SKT의 supervised 기반 수의 진단 AI인 X-caliber가 저희의 연구를 발판삼아 비지도학습으로 나아가기를 기대하고 있습니다.

이를 통해 의료 AI 시장에서 비지도학습 도입을 열고, 사용자의 labeling 리소스를 낮추기를 기대합니다.

연구 Timeline

저희가 지난 3개월동안 진행했던 phase 1과 앞으로 진행할 phase2의 연구 timeline입니다. 참고해주시면 좋을 것 같습니다.

맺으며..

짧은 기간동안 ICCV 2023 challenge를 준비하고, top-5 등수를 달성하고, 논문을 작성하며 저희 팀원들 모두 한층 더 성장할 수 있었습니다.

매주 온라인/오프라인 회의를 하며 방향성을 잘 잡아주시고, 연구적인 조언을 아끼지 않아주신 멘토님들께 다시한번 감사하다는 말씀을 드리고 싶습니다!

남은 두 달, phase2를 힘차게 이끌어 가보겠습니다. 이상 labeling 없는 의료 AI의 시장화를 꿈꾸며 원 팀으로 연구하는 beyond label이었습니다. 감사합니다