BigDatasetGAN: Synthesizing ImageNet with Pixel-wise Annotations

안녕하세요 콥스랩(COBS LAB)입니다.

오늘 소개해 드릴 논문은 ‘BigDatasetGAN: Synthesizing ImageNet with Pixel-wise Annotations’입니다.

해당 내용은 유튜브 ‘딥러닝 논문 읽기 모임' 중 ‘BigDatasetGAN: Synthesizing ImageNet with Pixel-wise Annotations’ 영상 스크립트를 편집한 내용으로, 영상으로도 확인하실 수 있습니다. (영상링크:https://youtu.be/VgYlBHXKAs0)

일단 Synthetic 데이터가 무엇인지 간단하게 보겠습니다. 위 사진은 Gartner라는 곳에서 가져온 사진입니다. 올해 2022년 주목할 기술로 Generative AI가 언급되었습니다. Generative AI를 간단하게 설명하고 있는데 오리지널과 비슷하지만 똑같이 따라 하지는 않는 데이터를 생성하는 것을 말합니다.

 

그래서 좀 떠오르고 있는 분야입니다.

 

Synthetic 데이터로는 왼쪽과 같이 어떤 물리엔진을 사용해서 렌더링 한 데이터가 될 수도 있고 오른쪽 같이 어떤 GAN 네트워크를 사용해서 이미지를 생성하는 데이터가 될 수도 있습니다.

 

BigDatasetGAN이 무엇인지 살펴보도록 하겠습니다. 이 그림은 BigDatasetGAN이 아니라 DatasetGAN이라는 논문에서 가져온 그림인데 두 논문 다 말하고자 하는 것이 똑같고 살짝 scale up만 했기 때문에 이 그림을 봐도 충분히 이해가 잘 될 것 같습니다.

 

1번 2번이 나눠져 있긴 한데 사실 크게 세 단계로 나눌 수 있습니다. 일단 DatasetGAN는 학습하는 부분과 그리고 DatasetGAN으로 Synthetic Dataset을 생성하는 과정과 그리고 이렇게 생성한 Dataset을 가지고 원하는 task의 모델을 학습하는 과정이 있습니다.

 

 

DatasetGAN을 학습하는 과정을 조금 더 자세히 보겠습니다.

이미지 대상이기 때문에 이미지를 생성하는 GAN 모델 아무거나 상관없고 GAN의 Generator를 생각을 해보면 입력으로는 Random Latent Vector z가 들어가고 출력으로는 이 Generator를 학습할 때 봤던 이미지들과 비슷한 어떤 이미지가 나오게 됩니다.

 

그래서 Generator가 생성한 이미지를 사람이 annotation을 합니다. 일반적으로는 실제 이미지를 annotation 하는데 GAN이 생성한 이미지를 annotation 한다는 것이 조금 특이합니다.

 

 

그래서 방금 본 Generator가 이미지를 생성하는 부분과 Generator 뒷부분에 label Branch를 하나 만들어서 방금 사람이 annotation 한 이 mask를 가지고 이 label branch를 학습을 합니다. 그래서 label branch의 입력은 GAN Generator의 중간 피처들을 입력으로 받습니다. 이 label branch의 target을 마스크로 해서 이 branch만을 학습을 합니다. 이때 이 GAN Generator 부분은 업데이트하지 않고 그냥 그대로 사용을 합니다.

위 사진은 전체적인 과정입니다. 이렇게 학습된 Dataset GAN을 가지고 무한한 이미지와 label pair를 만들 수 있다고 말하고 있습니다.

이 Dataset GAN에서 이미지를 생성하는 어떤 GAN이 필요하다고 했는데 간단하게 설명을 드리겠습니다.

 

첫 번째로 Dataset GAN에서는 StyleGAN을 사용을 했는데 StyleGAN는 기존의 입력이 하나의 Gaussian distribution에서 샘플링됐기 때문에 어떤 feature들이 좀 얽혀 있다고 합니다. 그래서 안경을 쓴 사람을 생성을 했을 때 안경만을 제어할 수 없고 안경을 바꾸게 되면 아예 성별이 바뀌기도 하고 그렇게 돼 있는 거를 Entangled 되어 있다고 표현을 합니다. 그래서 StyleGAN은 그거를 조금 해결하고자 Random Latent z를 그대로 입력으로 받는 게 아니라 어떤 non linear mapping을 통해서 원하는 Style 개수만큼 Random vector를 생성을 해서 이 사이사이에 넣어줍니다.

 

high level feature에 들어가는 Style은 이미지에 굵직굵직한 특징들을 제어를 할 수 있게 되고 출력과 가까워지는 low level 쪽으로 갈수록 머리카락에 굵기나 색깔등이 제어가 되게 됩니다.

 

그래서 Generator의 중간 피처들이 어떤 semantic 한 정보를 잘 담고 있다고 볼 수 있습니다.

BigDatasetGAN에서는 BigGAN을 사용했습니다.

 

이거는 기존 GAN이 Scalability와 Stability가 조금 떨어진다는 점을 개선을 했습니다. GAN 하면은 떠오르는 게 오른쪽 아래처럼 아예 이상한 이미지들이 생성되는 것을 보신 적이 있을 겁니다. 그래서 이거를 좀 보완하기 위해서 Truncation trick이라는 trick을 사용을 합니다. 이 trick은 이 Gaussian distribution에서 학습을 할 때 Gaussian distribution 샘플링을 그대로 사용을 하는데 test시에는 자주 sampling 되지 않는 이 끝값들을 사용하지 않고 자주 sampling 되는 가운데 값들만 사용하겠다는 것입니다.

 

그렇게 해서 ability와 그 버라이어티를 trade off 해서 사용자가 그 조절할 수 있게 됩니다. 그래서 Truncation trick을 좀 약하게 주면 다양한 이미지들이 생성되지만 이런 artifact들이 좀 생기게 되고 Truncation trick을 조금 강하게 줄수록 다양성은 떨어지지만 퀄리티가 높은 이미지들이 나오게 됩니다.

BigDataset GAN에서 VQGAN도 사용을 했습니다.

 

VQGAN는 워낙 성능이 좋다 보니까 트랜스포머를 GAN에서도 사용을 하고 싶은데 computational cost가 너무 커가지고 사용하지 못하는 문제가 있었습니다.

 

그리고 VQGAN 이전 논문인 VQVA라는 논문에서 언급된 것은 이미지도 사실 연속적이긴 하지만 discrete 하게 표현이 가능하고 지금까지는 continuous 한 Latent Vector 들을 사용했는데 그것이 오히려 auto encoder 구조에 모델에서 어떤 overfitting을 유발할 수 있다고 주장합니다.

 

그래서 VQGAN에서는 Convolution과 트랜스포머를 결합해서 사용을 했고 encoder-decoder 부분에서는 CNN을 사용을 하고 가운데 Latent Vector를 조작하는 부분에서는 임의로 이런 discrete 한 codebook을 만들어서 Vector quantization을 수행을 하고 각 Latent 간의 어떤 글로벌한 interaction 같은 거는 트랜스포머를 학습을 합니다.

 

그래서 high resolution 이미지를 생성할 수 있는 모델입니다.

이렇게 GAN들을 먼저 살펴본 이유는 이 논문에서 직관적으로 설명한 부분이 있기 때문입니다. 그 부분을 살펴보면 최신 GAN들의 Generator가 좋은 품질의 데이터를 만들 수 있는 이유는 이 Generator의 중간 feature들이 풍부한 semantic knowledge를 가지고 있기 때문이라고 말하고 있습니다.

 

그래서 그 잘 학습된 GAN의 Generator는 이미 좋은 피처가 학습되어 있기 때문에 아까 말씀드린 label branch를 학습할 때는 많은 데이터가 필요하지 않다고 설명을 하고 있습니다.

 

결론은 이미 잘 학습된 GAN과 아주 약간의 annotation만 있다면 이 무한한 이미지 label쌍을 만들 수 있는 Dataset GAN을 학습시킬 수 있다고 합니다.

조금 더 자세히 보겠습니다. 이 Dataset GAN을 학습시키기 위해서 labelIng을 하는 과정입니다. 실제 데이터인 ImageNet과 이 BigGAN이 생성한 한 이미지에서 카테고리별로 10 장씩 랜덤 하게 선택을 합니다. 그래서 한 명의 annotator가 리얼 ImageNet과 Synthetic 이미지를 Segmentation 레벨로 annotation을 합니다.

 

그래서 10장을 다 사용한 건 아니고 좋은 거를 SELECT 해서 real annotated는 8k 그리고 Synthetic annotated는 5k 해가지고 8000장 5000장을 만들었고, 이 데이터로 DatasetGAN을 학습을 합니다. 그렇게 해서 밑에 BigGAN-sim이랑 VQGAN-sim은 이 학습된 Dataset GAN으로 생성한 이미지 annotation pair 개수입니다.

 

그래서 이 정도 개수를 생성을 해놨고 뒷부분은 이미지 퀄리티, label 퀄리티를 real과 Synthetic과 비교를 한 것입니다. 당연히 Synthetic 이 조금 더 떨어집니다.

GAN을 사용하다 보니까 앞서 말씀드린 것처럼 어떤 trick을 쓰지 않으면 이런 이상한 이미지들이 나올 수가 있습니다. 그래서 이런 걸 방지하기 위해서 여러 가지 장치를 마련을 해 두었습니다. 그래서 BigGAN을 쓸 때는 앞서 말씀드린 Truncation trick과 그리고 이미 학습된 Image classifier를 사용을 해서 confidence score를 측정을 합니다. 그래서 confidence score가 만족할 만큼 나오지 않으면 이미지를 버리는 식으로 합니다. 

 

앞에서 말씀드린 것처럼 매우 적은 이미지를 가지고도 label branch를 학습시킬 수 있다고 주장을 하고 실제로도 labeling을 조금만 했기 때문에 label branch를 하나만 학습을 하면 이게 overfitting이 될 수 있다고 합니다. 그래서 16개 head를 만들어서 Ensemble을 했다고 얘기합니다. 그리고 head 간의 Uncertainty를 구해서 filtering을 했다고 합니다.

 

그리고 VQGAN을 사용할 할 때는 Nucleus sampling을 사용을 해서 이상한 이미지들을 걸러내는 작업을 했다고 설명합니다.

앞서 말씀드린 label branch입니다. 이 부분은 논문에서 Feature interpreter라고 표현을 하고 있습니다. 이 왼쪽 그림은 BigGAN의 Generator 구조입니다. 아까 말씀드린 것처럼 Feature Interpreter의 입력으로는 이미지가 들어오는 것이 아니라 Generator의 중간 피처들이 입력으로 들어갑니다. 그리고 아웃풋은 segmentation 마스크가 됩니다.

DatasetGAN의 목적 자체가 어떤 태스크를 하는 모델을 학습하기 위한 Dataset을 만드는 것이기 때문에 온라인 샘플링, 오프라인 샘플링이 있습니다.

 

오프라인 샘플링은 DatasetGAN을 사용해서 데이터를 왕창 만들어 놓고 그 데이터를 가지고 모델을 계속 학습하는 것입니다.

 

온라인 샘플링은 트레이닝하는 중에 계속 생성을 해서 모델이 똑같은 데이터를 두 번 보지 않게, 그러니까 훨씬 더 많은 데이터를 계속 보게 해서 학습을 한 것입니다.

 

예상하다시피 온라인이 당연히 성능이 더 잘 나오고 그 대신에 온라인은 학습 시간이 좀 더 오래 걸린다는 단점이 있습니다.

실험 부분입니다.

 

사실 리서치적이나 그런 게 없고 Ablation study도 이 사람들이 그냥 해봤으니까 되었다는 식이어서 딱히 크게 중요한 부분은 없지만 간략하게 짚고 넘어가겠습니다.

 

Training Dataset으로는 Synthetic annotated Dataset 5k를 사용을 했고, 평가용 Dataset으로는 real annotated Dataset을 사용했습니다.

 

그리고 Segmentation method로는 DeepLabv3 with Resnet 사용을 했습니다. 여기서 Segmentation은 DatasetGAN과는 상관이 없이 DatasetGAN으로 생성된 Dataset을 학습한 method을 의미합니다.

 

Metric은 mIoU를 사용해서 비교하고 있습니다.

BigDatasetGAN에서 DatasetGAN으로 생성한 데이터를 Pretrain으로 사용하고 downstream task에 전이 학습을 하는 식으로 해서 성능 표를 뽑았습니다.

 

BigGAN-sim과 VQGAN-sim은 10만 장 labeled image입니다. 그래서 supervised pretrain으로는 랜덤 하게 하면 각 테스크들을 설명드리면 좋은데 안 나왔습니다.

 

그러니까 dog랑 bird는 Segmentation 했고 FG/BG, 뒤에도 다 비슷합니다. 그래서 여러 가지 task로 실험을 해봤는데 random initialize 했을 때는 56.4인데 supervised로 pre train 했을 때 82%가 나왔고 아까 말씀드린 오프라인 샘플링이랑 온라인 샘플링해서 성능이 이 정도로 많이 올랐다고 합니다.

 

그리고 Supervised Contrastive learning이나 self도 다 비교를 했는데 Synthetic을 추가로 사용했을 때 더 성능이 잘 나왔다고 합니다. 그리고 self supervised method에도 Synthetic 이미지를 사용을 한 게 더 도움이 되었다고 합니다.

사람이 labeling 한 2000장 가지고 supervised 한 거랑 BigGAN이 생성한 거 가지고 학습을 한 거랑 비교를 했습니다.

 

사람이 labeling 한 걸로 했을 때 빨간 선 정도 성능이 나왔는데 BigGAN으로 생성한 거는 똑같이 2000장 정도 될 때는 성능이 좀 낮지만 한 10배가 넘어가는 시점부터는 실제 이미지 2000장의 성능을 넘었다고 하고 있습니다. 그런데 무한한 데이터를 만들 수 있기 때문에 계속 늘려가면서 실험을 해봤는데 200만 장까지 늘어났을 때 성능도 보이고 있습니다.

 

그리고 backbone 사이즈에서도 Supervised로 했을 때는 이 빨간색이 성능입니다. 모델 파라미터가 늘어남에 따라서 증가하기는 하지만 여기서는 좀 많이 증가했다가 살짝 줄어드는 느낌이 있는데 BigGAN을 사용해서 데이터를 훨씬 더 많이 넣었을 때는 모델 파라미터가 늘어남에 따라 좀 더 많이 증가하는 추세를 보인다고 말하고 있습니다.

 

그리고 DenseCL과 BigGAN을 데이터를 같이 써서 pre train 한 것을 비교했는데 그 수렴이 훨씬 더 빨랐다고 합니다. 최종 성능은 크게 다르지 않은 것 같은데 BigGAN까지 사용했을 때 5000 iteration에서 55% 나왔는데 BigGAN을 사용을 안 했을 때는 12000 iteration까지 가야 한 55%가 나옵니다. 그래서 수렴 속도에서도 차이가 좀 많이 나왔다고 합니다.