안녕하세요 콥스랩(COBS LAB)입니다.
오늘 소개해 드릴 논문은 ‘MedSefDiff: Medical Image Segmentation with Diffusion Probabilistic Model’입니다.
해당 내용은 유튜브 ‘딥러닝 논문 읽기 모임'MedSefDiff: Medical Image Segmentation with Diffusion Probabilistic Model’ 영상 스크립트를 편집한 내용으로, 영상으로도 확인하실 수 있습니다. (영상링크:https://youtu.be/gY3iFOIRs48)
본 논문의 구성은 백그라운드에 대해서 이야기를 하고 논문의 순서 그대로 진행하도록 하겠습니다.
처음으로 Diffusion에 대해서 간단하게 이해하고 넘어야 될 부분이 있습니다. 우선 열역학 쪽에서 High Density 상태에서 Low Density로 가고 Equilibrium 상태로 넘어가는 것의 정의를 Diffusion이라고 이야기를 하고 있습니다.
그렇다면 Equilibrium 상태에서 high Density로 모델이 학습할 수 있을까에 대한 아이디어에서 이 모델이 시작되었다고 볼 수 있고 이 부분은 통계학에서도 동일하게 적용이 되어서 complex 한 distribution에서 noise가 굉장히 많은 prior distribution으로 갈 때 모델이 거꾸로 학습할 수 있을까에 대한 아이디어로 시작되게 됩니다.
첫 번째로 Diffusion 모델이 왜 잘되는가에 대해서 좀 궁금증을 가졌습니다. 이 부분의 근거로는 2015년 PMRL에 나왔던 논문을 근거로 준비를 했습니다. 간단하게 내린 결론을 말씀드리면 통계적 근사를 더 정확하게 구현한 것이 아닐까라고 생각을 하게 되었고 크게 두 가지 아이디어를 제시를 하는데 Sequential Monte Carlo 저희가 일반적으로 알고 있는 Markov chain에 대해서 근사화했던 방식과 Non-equilibrium statistical physics에서 수식을 가져와서 모델이 어떤 걸 배우는 게 가장 좋을 것인가에 대한 고민이 있습니다.
첫 번째로 Markov chain을 적용했을 때 Markov stochastic process 같은 경우에는 미래 상태는 과거에 의존하지 않고 현재 상태에 대해서만 영향을 받는다라고 보는 과정입니다. 이 부분에 가장 중요한 부분은 연산을 단순화시킬 수 있다는 장점이 있습니다. 왼쪽 아래 보시면 미래를 표현하기 위해서 과거 데이터를 써야 한다면 딥러닝 모델이라면 당연히 레이어에서 연산을 추가하는 과정이 되고 시간적인 스텝이 굉장히 많아지면 많아질수록 연산이 굉장히 복잡해지는데 그 부분을 오른쪽 아래 그림처럼 단순하게 구현을 해서 연산을 단순화했다고 볼 수 있습니다.
하지만 이상적인 경우가 가장 현실을 잘 표현했다고 볼 수가 있습니다. 근사화로 인해서 정확도가 하락되는 것을 어떻게 커버할 것인가에 대한 궁금증이 있었고 그 두 번째 아이디어로 보는 것은 준정적 과정이라고 보이는 것이 Quasi static process입니다.
매우 천천히 변해서 각 순간마다 평형상태로 간주할 수 있을 만큼 이상적인 과정이라고 볼 수 있는데 그림에서 볼 수 있듯이 noise가 추가되었음에도 변화가 거의 없는 상태를 준정적 과정이라고 볼 수 있습니다.
각 순간에 변화과정에 거의 없는 상태로 구현할 수 있다면 과거에 대한 의존성을 비교적 더 적게 가져갈 수 있기 때문에 Markov stochastic process를 적용했을 때도 근사화로 인한 부정적인 영향을 줄일 수 있을 것으로 보입니다.
마지막으로는 수많은 단계로 이루어진 과정을 수식적으로 표현하는 단계가 필요합니다.
그래서 여기서 적용한 수식 같은 경우에는 Langevin dynamics와 Kolmogorov forward and backward equation입니다. 위에 Langevin dynamics 같은 경우에는 실제 세계에서 분자 시스템이 어떻게 움직이는 가에 대해서 보여주고 있는 식입니다.
식이 조금 복잡하긴 한데 익숙한 식으로 대체를 해놓으면 F=ma 그다음 particle 마다 얼마나 interaction을 할 수 있는가
그다음에 X에 점 하나 있는 부분 같은 경우는 속도에 해당하는 부분이 될 수 있을 것 같습니다. 이 부분에 대해서 준정적 프로세스라고 가정을 하게 된다면 가속도나 particle interaction 하는 것들 그리고 속도에 대한 영향이 굉장히 적어져서 결국에 남는 부분은 Gaussian process만 남게 되고 Gaussian process 같은 경우에 오른쪽 위에 보이는 것처럼 damping function이 들어오게 되면은 그 부분이 점차 Gaussian을 통해서 사라지는 과정을 볼 수 있습니다.
그다음에 결국 각 스텝에 대해서 어떻게 정의할 것이냐에 대한 수식입니다. 이 부분은 Markov process일 때 시간에 따른 Diffusion을 어떻게 표현할 것인가에 대한 수식이고 첫 번째 수식 같은 경우에는 Diffusion에 해당하는 수식입니다.
그리고 아래 수식 같은 경우에는 Diffusion의 역 과정을 표현한 수식인데 각 스텝별로 나온 수식을 보게 되면은 U값과 시그마 값으로 대부분이 표현된 것을 볼 수 있고 이 부분이 두 가지 수식을 통해서 각 스텝별 noise를 Gaussian noise로 적용을 했을 때가 가장 적합하다고 보일 수 있습니다.
다음으로는 complex distribution 데이터의 distribution에서 prior distribution으로 가는 과정을 표현한 방식이고 이 부분을 하나의 트랜스포메이션으로 표현한 수식입니다.
두 번째 라인에 수식을 보면은 Markov chain을 적용했을 때 Stationary Distribution을 표현을 한 수식이고 연속적이기 때문에 적분 값으로 나타나 있지만 세 번째에 수식을 보면은 한 스텝에 distribution을 나타내는 방식은 전 단계가 들어왔을 때 현재 단계를 나타내는 컨디션을 distribution으로 나타낼 수 있습니다.
그런데 이 부분이 현실적으로 들어왔을 때는 굉장히 큰 값으로 스탭이 나눠져 있다면은 이 부분을 거의 연속적이다
라고 근사 시킬 수 있고 네 번째 줄을 보게 되면은 한 단계의 distribution을 굉장히 작은 Gaussian 분포로 나타낼 수 있습니다.
이 과정에서 Gaussian Diffusion이라고 표현을 하고 여기서 Gaussian noise가 계속 추가되는 과정을 Forward Trajectory라고 표현을 하게 됩니다. 이렇게 계속 가우스 noise가 추가되면 마지막에는 Gaussian noise로 해당하는 과정으로 볼 수 있고 하나의 가우스 noise에 해당할 때 베타 값으로 표현을 할 수가 있는데 베타값이 매우 작다면 Diffusion 과정에서 Gaussian noise가 추가되는 분포랑 모델이 noise를 걷어내는 분포가 매우 유사하기 때문에 결국은 모델 입장에서도 noise를 걷어내는 것을 학습할 수 있게 됩니다.
결국에는 정해진 noise를 추가해서 Diffusion을 구현해 놓은 과정이 있고 모델은 그의 역연산으로 모델 파라미터를 통한
학습을 진행한다고 보시면 될 거 같습니다.
본 논문에서는 이 스텝마다 과정을 segmentation에 적용시킨 논문이고 여기서는 dpm을 결국엔 image segmentation, medical image segmentation에 적용하고 그다음에 dynamic conditional encoding과 FF-parser 모듈을 제안해서 성능을 올렸으며 최종적으로는 여러 가지 image modality에 대해서 SOTA 퍼포먼스를 가져온 것을 볼 수 있습니다.
두 가지 메서드에 대해서 설명드리겠습니다.
첫 번째로 dynamic conditional encoding 같은 경우에 적용된 이유 자체가 결국 Medical image segmentation 같은 경우에는 굉장히 애매한 오브젝트를 바탕으로 segmentation을 가져와야 하기 때문에 유니크한 어떤 정보가 오더라도 어느 정도는 noise의 강건한 모델이 필요하다고 보이고 있습니다.
그렇기 때문에 이 부분에서 한 단계 스텝에서 초록색 모듈로 표현된 부분이 dynamic conditional encoding으로 볼 수 있습니다. dynamic conditional encoding을 구현한 arichitecture는 U-Net 구조 어떻게 보면 vanilla DPM에서 가져온 모델이고 이 부분에서 A 부분에 해당하는 부분이 dynamic conditional encoding에서 Attentive- like mechanism을 표현한 수식입니다. 여기서 볼 수 있듯이 layer normalization을 어떻게 보면 전 단계에서 noise를 가지고 있는 segmentation 마스크에서 가져오고 그다음에 input image를 넣어서 두 가지를 활용해서 결과적으로 noise를 걷어내는 과정을 보여 드리고 있습니다.
이 과정을 조금 수식마다 나눠서 표현을 하면 입력 데이터와 이전 noise가 낀 부분을 바탕으로 해서 얼마나 segmentation input 데이터에서 중요한 부분을 고려해야 되나 이 부분을 표현을 하고 있습니다.
결과적으로 중간에 embedding vector로 들어가는 부분은 입력 데이터 피처에서 얼마나 중요한 부분을 남겨서 noise를 걷어 낼까에 대한 과정임을 알 수 있습니다.
두 번째로 FF-Parser 모듈입니다.
FF-Parser 같은 경우에는 저렇게 attentive like mechanism을 사용했을 때라도 high frequency가 negative effects를 가져오는 단점이 있어서 이 부분을 해결하기 위해서 Fourier Transform을 가져왔고 그 부분에서 어떤 high frequency가 중요한지를 걸러낼 수 있도록 구성하였습니다.
feature map이 들어오게 되면은 단순히 Fourier Transform을 거쳐서 거기에 Learnable 한 parameterized 된 attentive map을 그냥 곱하는 형식으로 되어 있습니다. Inverse FFT로 복원시키는 꼴입니다.
이 논문에서 사용하는 데이터셋 같은 경우에는 어떻게 보면 high frequency를 걸러내고 segmentation mask가 low frequency 굉장히 중요한 정보를 많이 담고 있을 확률이 높기 때문에 high frequency를 걷어내는 쪽으로 parameterized map을 그냥 넣어서 학습시켰더라도 high frequency를 걸러내는 효과를 가져오게 된 되는 걸 볼 수 있습니다.
데이터셋 같은 경우에는 여러 가지 modality를 해결했다는 것을 보여 주는 것처럼 광학으로 인한 segmentation 데이터와 뇌 악성종양을 mri로 표현한 데이터셋과 갑상선 초음파 image 이 3가지 modality에 대해서 모두 SOTA 모델보다 더 좋은 퍼포먼스를 가져오게 됐다고 이야기를 하고 있습니다.
위 과정을 그림으로 봤을 때 어떻게 보면은 조금 다른 모델에 비해서 애매하다고 표현될 수 있는 부분을 더 작게
prediction 한 것을 볼 수 있습니다.
결과적으로 dynamic conditional encoding과 FF parser를 넣었을 때 3가지 데이터셋에 대해서도 둘 다 더 넣었을 때
적용된 것이 가장 높은 것을 확인할 수 있습니다.
결과로는 결국에는 DPM 모델을 바탕으로 해서 general image segmentation 태스크에 적용을 했고 그다음에 dynamic conditional encoding과 FF parser 모듈을 적용하였으며 이것을 통해서 이전에 SOTA 모델을 모두 넘었다는 것을 강조하고 있습니다.
그리고 세 가지 Medical image segmentation에서 각각 다른 modality에서 동일하게 퍼포먼스를 보여 준 것을 결론으로 하고 있습니다.
댓글