논문 링크:
Deep Joint Source-Channel Coding for Wireless Image Transmission with Adaptive Rate Control
We present a novel adaptive deep joint source-channel coding (JSCC) scheme for wireless image transmission. The proposed scheme supports multiple rates using a single deep neural network (DNN) model and learns to dynamically control the rate based on the c
arxiv.org
출처: M. Yang and H.-S. Kim, “Deep joint source - channel coding for wireless image transmission with adaptive rate control,” arXiv:2110.04456, Oct. 2021.
요약
JSCC에 관한 최근 연구들은 대부분 모델을 고정된 rate로 학습시킨다. 즉 하나의 모델에 하나의 rate만 대응되기에 여러 rate가 필요한 경우 여러 모델을 훈련시켜야 하는 비효율이 발생한다. 이러한 문제를 해결하고자 하나의 모델에 여러 rate를 훈련시키는 방법도 있지만 이 방법 또한 하나의 SNR에서만 훈련되게 채널 환경에 잘 적응하지 못하는 한계를 갖는다. 이에 따라 이 논문에선 채널 SNR과 이미지 내용의 조건에 따라 적응하는 JSCC를 제안한다. 이렇게 adaptive하게 rate를 선택하기 위한 policy network도 설계된다. 이 논문에서 제안하는 JSCC 구조도는 아래 그림과 같다.
먼저 소스 인코더는 이미지에서 특징들을 추출해 이를 \( G_s + G_n \)으로 분리한다. \( G_s \)는 policy network에 따라 활성화여부가 결정되며 \( G_n \)은 항상 활성화된다. 이렇게 feature 선택이 이루어진 후에 전력 정규화가 이루어지고 신호는 채널을 통해 전송된다. 또 수신단에선 이 신호를 받아 채널 디코더와 소스 디코더를 순차적으로 적용한다. SNR값은 채널 인코더와 디코더 그리고 policy network에 전달되어 모델이 채널 환경에 적응할 수 있도록 돕는다. 그리고 이 논문에선 전송률을 channel usage per pixel(CPP)로 나타낸다.
채널 인코더와 디코더의 구조는 아래와 같다. 가장 아래 그림은 SNR-adaptive 모듈을 나타내며 여기서 입력 요소들은 각 채널을 통해 average pool 되며 SNR과 결합된 후 두 MLP층을 통과하게 된다. 이에 따라 이 모듈에선 채널 값을 반영한 스케일링 요소들을 얻을 수 있다. 채널 인코더 \( E_c \)에선 입력 요소들이 ResNet과 방금 정리한 SNR-adaptive 모듈들을 통과한다. 그 이후 2D 컨볼루션을 통해 목표하는 크기로 바뀐다. 채널 디코더 \( D_c \)에선 앞서 다룬 활성화된 특징들만 수신받기에 비활성화 된 요소들의 공간을 zero-pad 시켜 입력 크기를 동일하게 유지한다.
다음으론 policy network의 구조에 대해서 다루는데 이는 이미지 요소를 먼저 입력으로 받은 후에 pool, concat, 2 MLP 계층을 거친다. 그리고 Gumbel Softmax를 거쳐 확률에 관한 값을 변환된 후에 thermometer-coded 벡터 W로 최종 변환된다. 이 W는 연속된 feature 그룹들이 활성화되도록 보장하기에 활성 여부를 알려주는 추가적인 제어 메시지가 필요하지 않다.
이산적인 값을 갖는 sampling 과정을 거치기에 policy network는 미분이 불가능하기에 back-propagation도 할 수 없다. 그리하여 이에 대한 해결책으로 보상 함수를 사용할 수 있으나 이는 수렴하는데 오래 걸리고 편차가 높다. 그렇기에 대안으로 Gumbel-Softmax 함수를 사용하여 미분 불가능한 문제를 해결한다.
\[
\hat{P} = \arg\max_k \big( \log p_k + g_k \big), \quad k \in \{0,1,\ldots,G_s\}
\]
\[
\tilde{P}_k = \frac{\exp\left((\log p_k + g_k)/\tau\right)}{\sum_{j=0}^{G_s} \exp\left((\log p_j + g_j)/\tau\right)}, \quad k = 0,1,\ldots,G_s
\]
\(p_k\)는 카테고리 k가 선택될 확률, \(g_k\)는 Gumble 분포를 통한 난수, \(\tau\)는 분포의 이산성을 결정해주는 parameter다. \(\tau\)가 무한대로 향할 때 분포는 uniform 분포에 수렴하고, \(\tau\)가 0에 가까울 땐 이산적인 one-hot 벡터와 근접한다. 위 두 식들은 각각 순서대로 forward, backward propagation에 사용된다. 첫번째 식은 policy를 위한 이산적인 분포로 사용되고 one hot vector를 만들며 이는 다시 thermometer encoding으로 바뀌어 마스크 W를 만든다. 두번째 식은 gradient 추정을 위해 사용되어 네트워크를 학습할 수 있게 해준다.
\[
\mathbb{E}_{X \sim D_{\text{train}}} \left[ \| X - Y \|_2^2 \;+\; \alpha \sum_{i=1}^{G_s} W_i \right]
\]
또한 이 모델의 loss function은 위 식과 같은데 이는 이미지 reconstruction을 반영하는 항과 활성화된 feature group수에 비례하는 채널 사용량을 반영하는 두 항으로 이루어져 있다. 즉 이 loss function은 이미지를 잘 복원하면서 채널 사용량도 줄이는 방향으로 학습을 유도한다.
CIFAR-10 데이터셋에 모델을 훈련시킨 후 고정 rate 모델을 baseline으로 활용하고, 제안된 방법의 α를 바꿔가며 위 그림과 같이 비교를 진행했다. SNR이 낮을 때 설계된 모델은 더 높은 rate를 갖고 SNR이 증가함에 따라 더 적은 rate로 수렴한다. 이는 채널 환경이 나쁠 땐 채널 자원, CPP를 더 많이 사용하고 채널 환경이 좋을 땐 채널 자원을 더 적게 써서 더 적은 CPP를 갖는다는 의미와 같다. α값을 증가시킬수록 loss function은 이미지 복원보다 rate에 더 집중하기 때문에 SNR 증가에 따른 rate 감소율이 크다.
다음으론 제안된 모델과 최신 기술인 (1)Shannon capacity기반 BPG와 (2)고정 rate 모델들과의 비교를 진행했다. 비교 결과 더 적은 α에 해당하는 (a)에서는 (b)와 비교했을 때 상대적으로 더 높은 CPP를 선택한다. 또한 (a)에서는 낮은 SNR에서 BPG의 성능을 뛰어넘는다. (b)의 경우 CPP가 사억으로 적으며 모든 SNR에서 BPG보다 더 높은 성능을 달성한다. 고정 rate로 학습된 모델들과 비교했을 땐 거의 유사한 성능을 보이나, 제안한 방법은 오직 하나의 모델만 학습시키에 더 강점을 갖는다.
다음으론 SNR과 α를 고정시키고 이미지 종류를 변경시키면서 이미지 종류에 따른 rate와 PSNR을 비교했다. 그래프를 보면 이미지 종류마다 rate가 다른 것을 확인할 수 있는데, 더 복잡한 특징을 갖고 있는 이미지 종류일수록 더 많은 CPP를 사용하기 때문이다. 또한, 제안한 모델은 재복원 품질,PSNR의 편차가 기존의 방법보다 더 적다. 즉 이는 평균 CPP가 같을 이미지 복원을 더 잘한다는 뜻이다.
결론
이 논문에선 새로운 deep JSCC 모델을 제안했다. 이 모델은 policy network를 통해 rate를 SNR과 이미지 내용에 따라 동적으로 달리한다. 또한 이진 마스크를 사용해 이미지 feature들을 선택적으로 반영한다. 이 policy network를 미분가능하게 하기 위하여 Gumbel-softmax 기법이 사용되었다. 이 기법을 통해 SNR이 높을 때, 즉 채널 환경이 좋을 땐 주파수 대를 적게 사용할 수 있게 되었으며 여러 rate에 학습시키는 기존의 번거로운 방법과 비슷한 성능을 얻을 수 있었다.