[논문 리뷰]FA-GAN: Fused Attentive Generative Adversarial Networks for MRI Image Super-Resolution

1. Introduction

---

이번 논문은 FA-GAN(Fused Attentive Generative Adversarial Network)에 대해서 소개할 것이다. FA-GAN은 enhancement의 일종으로써, MRI image가 주어졌을 때 여러 attention 기법들을 활용하여 low resolution인 부분을 high resolution으로 바꾸어주는 모델이다. 이름에서 볼 수 있듯이 GAN(Generative Adversarial Network)를 활용한다. 

 

해당 논문을 이해하기 위해서 알아야 할 필수 용어들을 몇 가지 소개해보고자 한다.

 

1. Super-Resolution(SR): 저화질(Low-Resolution, LR) 이미지에서 고화질(High-Resolution, HR) 이미지를 재구성하는 것을 의미한다. 구체적으로 MRI image에서 super-resolution을 활용하면 더욱 정교한 해부학적 정보를 주어, 임상 진단 및 의사 결정에 도움을 준다.  Super-Resolution을 얻기 위한 대표적인 방법으로는 interpolation-based, construction-based, machine-learning based가 있다. 본 논문에서는 machine-learning based, 구체적으로 deep learning을 활용하여 super-resolution을 얻는다. Machine-learning based technique가 SR에 있어 SoTA 퍼포먼스를 보여주기 때문이다. 

 

2. SR-GAN: Super-Resolution을 얻기 위하여 GAN과 결합한 모델. 본 모델은 뇌 이미지의 SR을 위하여 제작되었다. 하지만 주로 convolutional network들만 활용하여 이미지 내 멀리 떨어져 있는 정보들을 서로 연결시키기 어렵다는 단점이 존재한다. 해당 문제로 인하여 Self-Attention Generative Adversarial Network(SAGAN) 이라는 모델도 등장했다고 하는데, 추후에 따로 알아보면 좋을 것 같다.

 

SR-GAN 모델의 구조가 궁금하다면 해당 링크를 방문하길 바란다.

https://www.edwith.org/medical-20200327/lecture/63187?isDesc=false

[LECTURE] 5.SRGAN : edwith

 

 

2. Methology

---

먼저 FA-GAN framework(Generator부분)의 input으로는 I^LR, 즉 저화질 이미지가 input으로 들어간다. I^LR은 I^HR, 다시 말해 고화질 이미지를 downsampling하여 확보할 수 있다. I^HR로부터 I^LR로 downsampling은 bicubic interpolation이라는 방법을 통하여 얻었다고 한다. (시간관계상 이에 대한 것까지는 찾아보지 못했다) 어쨌든 이를 수식으로는 아래와 같이 나타낼 수 있다. 

 

 

주석: 여기서 f()는 downsampling과정이라 이해하면 된다. 

 

2.1 SR Network with GAN

---

FA-GAN framework의 output I^SR은 I^LR로부터 upsampling을 통하여 얻을 수 있다. Upsampling은 convolutional layer들을 통과하며, 각 합성곱 층을 통과할 때마다 이미지의 크기를 두 배로 확대한다. 하지만 희망하는 scaling 정도에 따라 upsampling과정에 조금의 변형은 줄 수 있다고 한다. 전체적인 framework를 살펴보면 아래와 같다. 

 

 

 

위 이미지의 LFFB(Local Fusion Feature Block) 모듈은 I^LR로부터 중요한 feature들을 추출해낸다. GFFB(Global Fusion Feature Block) 모듈은 convolutional layer 및 upsampling을 한 결과값을 다양한 attention mechanism으로 결합을 한다. LFFB와 GFFB에 대한 설명은 아래에 상술할 것이다. 

 

 

2.2 Local Fusion Feature Block(LFFB)

---

LFFB는 3개의 다른 convolution kernel을 사용하여 얻은 feature들을 서로 공유한다. 그 과정은 아래와 같다. 

 

 

FA-GAN의 LFFB에서는 3*3, 5*5, 7*7 kernel size를 활용하여 얻은 정보들을 서로 공유하였다. 각각의 kernel size은 2번의 s*s 크기의 covolution 및 그 사이의 ReLU를 진행한다. 그리고 나온 결과값을 합친 다음 다시 한 번 1*1 convolution에 통과시킨다. 이는 각각의 kernel size로 얻은 feature들의 융합을 용이하게 하고 계산적 부담을 줄여주는 역할을 한다. 그리고 그 전 feature에서 얻은 값과 3개의 kernel size를 통과하여 얻은 공유값을 서로 더하여 새로운 feature을 생성한다. 이를 수식으로 나타내면 아래와 같다.

 

F[ ]는 합치는 과정(concatenation)이랑 1*1 convolution을 의미한다.

 

 

 

2.3 Global Fusion Feature Block(GFFB)

---

GFFB를 활용하면 MRI image의 중요한 feature들이 enhance 된다. 그렇게 하여 Super Resolution image가 최대한 현실적이고, 기존 이미지와 비슷한 형태를 띄게 한다. GFFB는 크게 channel attention 모듈, self attention 모듈, fusion operation으로 나눌 수 있다. GFFB의 연산과정은 아래와 같다. 

 

 

(1) Channel Attention Module:

본 논문에서는 channel attention module 중 lightweight channel attention mechanism을 활용한다. 본 메커니즘은 global information이 주어졌을 때 필요한 feature들을 강조하고, 불필요한 feature들에게 제약을 건다. 이 개형은 아래와 같다. 

 

 

먼저 Global Average Pooling을 통하여 global information을 추출하고, 비율 r로 차원 축소를 한다음 ReLU를 거치고, 다시 차원을 증가시키고, 시그모이드 함수르 통과시켜 SR Image를 형성한다. 차원을 축소/증가 시키는 부분들은 FC Layer들로 이루어지고, 마지막에는 input feature들을 rescaling하여 input과 똑같은 차원의 output을 반환한다. 

 

 

(2) Self Attention Module

본 논문에서 self attention module은 전통적인 convolutional feature map을 self-attention feature map으로 변환하는 역할을 한다. 그 개형은 아래와 같다. 

 

 

 

앞선 과정에서 추출된 convolutional feature map이 input으로 들어왔을 때, 1*1 크기의 convolution 구조를 띈 3개의 가지 f(x), g(x), h(x)로 input을 각각 통과시킨다. g(x)는 channel의 개수를 바꾼다. h(x)는 channel의 개수를 그대로 고정시킨다. 여기서 언급된 H, W는 feature map의 높이와 너비이며, C는 채널의 개수이다. f(x)는 channel의 개수를 바꾸고 전치를 시킨다. 그 다음 f(x)와 g(x)를 곱하고 softmax로 정규화하여 [H*W, H*W] attention map을 얻는다. 이 attention map과 h(x)를 곱하면 [H*W, C] feature map을 얻게 되고, 이를 다시 1*1 convolutional layer에 통과시켜 [H*W*C]의 output을 반환한다. 

 

아래는 앞서 설명한 부분의 수식이다. 

 

 

softmax로 정규화를 진행하는 부분이다. 여기서β_j,i는 j번째 영역을 synthesize할 때의 i번째 위치이다.  

 

 

본 attention layer에서의 최종 출력을 하는 부분이다. 여기서 W_g, W_f, W_h, W_v는 모두 학습되어야할 hyperparameter이다.

 

 

앞서 본 개형에서 우리는 한 번의 연산을 더 진행한다. 그것은 바로 attention layer의 output에다 scale parameter γ를 곱한 다음 기존 input과 더해서 self attention module의 최종 결과값을 얻는다. 그 수식은 아래와 같다. 

 

 

γ의 초기값은 0으로 설정한다. 인접한 픽셀에서 정보를 얻기를 권장하기 때문이다. 허나 이런 방법이 불가능하다면 γ를 계속 학습시켜 그 값을 증가시킨다. 이 값을 증가하면 인접하지 않은 픽셀들에 대한 feature들을 더욱 잘 포착할 수 있다. 

 

(3) Fusion Operation

(1), (2)에서 얻은 값들을 융합하는 방법은 크게 두 가지로 나뉜다.

 

A. Direct Connection

Direct Connection은 두 attention 결과값들을 바로 더하는 것이다. 그 수식은 아래와 같다. 

 

 

i는 feature의 인덱스를 의미하고, R은 channel attention의 결과, Y는 self-attention의 결과이다. α와 β는 hyperparameter로써 초기값은 0.5이다. 

 

B. Weighted Connection

Weighted Connection은 channel attention의 결과 R과 self-attention의 결과 Y 사이에 경쟁을 부추긴다. 그 경쟁은 softmax의 형태로 나타낸다. α와 β는 hyperparameter로써 초기값은 0.5이다. 그 외 추가적인 parameter은 없다. 수식은 아래와 같다. 

 

 

(4) Loss Function

Loss Function은 기존의 MRI image와 프레임워크를 통하여 생성된 MRI를 비교한다. Loss function 값이 작을 수록 모델이 더욱 강력한 것이다. Loss의 종류로는 perceptual loss와 adversarial loss가 있다. (pixel loss도 사용했다고 하지만, 이에 대한 설명이 없는 걸로 봐서 perceptual loss를 진행할 때 부속으로 사용한 것 같다.)

 

Perceptual loss를 구하기 위해서 기존의 image와 생성된 image를 각각 VGG Network에 통과시킨다. 그래서 얻은 feature map을 가지고 두 이미지를 Euclidean distance를 활용하여 비교한다.(이곳이 pixel loss) 거기에서 발생되는 loss가 바로 perceptual loss이다. 그 식은 아래와 같다. 

 

 

Φ 부분은 VGG network를 뜻하고, W랑 H는 feature layer의 차원을 의미한다. 

 

Adversarial Loss는 generator로 생성된 이미지를 가지고 discriminator이 average discriminator probability를 출력한 것이다. GAN에서 사용하는 loss라고 생각하면 편하다. 그 식은 아래와 같다. 

 

 

그리고 perceptual loss와 adversarial loss는 아래와 같은 방법으로 더해져서 최종 loss를 출력한다.

 

---

최종적으로, 이 모델은 GAN의 변형이라고 생각해도 좋으며 아래 식에서 Discriminator 값을 최대화, Generator 값을 최소화하는 것이 목표이다. 

 

3. Experimental Results

---

3.1 Datasets and metrics

---

• TeslaV100-SXM2 GPU 사용
• 4개의 MRI dataset에서 각각 50개의 sample 추출하여 40개는 train set, 10개는 test set에 사용
• Adam Optimization 사용 -0.9
• LR: 0.0001
• 학습할 때 걸리는 시간: 약 10시간

모델을 평가하기 위해서는 크게 3가지 evaluation criteria를 사용한다. 

 

1. Peak Signal-to-Noise Ratio(PSNR)

PSNR은 주로 화질 손실 정보를 파악할 때 사용하는 지표이다. 최대 화질(전력)에 비해 잡음 화질(전력)의 비율을 나타내는 것이다. PSNR은 아래와 같은 식으로 이루어져 있다. 

 

 

x는 원본 이미지, y는 SR로 인한 생성 이미지며, i, j는 pixel의 좌표를 나타낸다. M과 N은 이미지의 크기를 나타낸다.  Log scale로 지표를 평가하기에 손실이 적을 수록 더 높은 값이 나온다. 그러므로 이 값이 더 높을 수록 더욱 좋은 성능을 내는 것이다. 

 

2. Structural Similarity Index Measure(SSIM)

SSIM은 구조적 유사 지수를 의미하며, 아래와 같은 식으로 이루어져 있다. 

 

 

μ_x와 μ_y는 x와 y 이미지의 픽셀값의 평균이며  σ_x와 σ_y는 분산을 나타내며, σ_xy는 공분산을 나타낸다. C들은 상수값으로써 안전성을 보장하기 위한 값이라고 한다. 이 값 또한 높을수록 좋은 성능을 내는 것이다. 

 

3. Freshet Inception Distance score(FID)

FID는 feature vector들 간의 거리를 나타내는 지표로써, 아래와 같은 식으로 이루어져 있다. 

 

 

Tr은 행렬의 diagonal 들의 합이며, μ는 평균, Σ는 공분산, x_r은 실제 이미지, g는 생성된 이미지를 의미한다. 이 값은 작을수록 좋은 성능을 내는 것이다. 

 

3.2 Experimental Results

---

4배 해상도의 SR을 진행하고 원본이미지와 FA-GAN 이미지를 비교했을 때의 결과는 다음과 같다. 왼쪽이 원본 이미지, 오른쪽이 FA-GAN 이미지이다. 

 

흉부(측면) 이미지

뇌 이미지

무릎이미지

Multi Modality Whole Heart

 

 

실제로 다른 GAN 기반 모델들과 비교했을 때도 FA-GAN이 우수한 성능을 보인다. 

 

 

 

4. Discussion

---

모델의 우수함을 평가하기 위해서 모델의 각 부분을 제거한 뒤 성능을 평가해보았다. 아래 이미지에서 볼 수 있듯이 기존 FA-GAN 모델이 가장 우수한 성능을 보였지만, 4배 해상도를 SR하는 과정에서는 self attention이 제거된 모델이 가장 우수한 성능을 보였다. (그 이유는 설명되어 있지 않다... 운에 의해서인가 아니면 구체적인 이유가 있었던 것일까...)

 

 

또 Fusion Operation시 direct connection과 weighted connection을 비교해보았는데, weighted connection을 사용했을 때 좋은 결과를 얻었다. 

 

 

추가적으로 discriminator network에 Spectral Normalization(SN)이라는 과정을 추가하였다.  추가 결과 Loss가 SN이 없을 때보다 더욱 빠르고 완만하게 감소하는 모습을 확인할 수 있다. 달리 말해 훈련 과정이 안정적으로 변하였다. 

Spectral Normalization에 대한 설명은 아래의 논문을 참고하길 바란다. 

https://arxiv.org/abs/1802.05957

 

5. Personal Thoughts

---

LR 이미지를 얻기 위해서 HR이미지를 downsampling했는데, 낮은 해상도를 출력하는 MRI image로 한번에 HR 이미지를 얻는 방법이 있었다면 좋을 것 같다!

 

출처:

https://arxiv.org/ftp/arxiv/papers/2108/2108.03920.pdf