Abstract

• 현재 보편적인 CV 태스크에서 transformer는 매우 각광받고있다
• 하지만 더나아가 GAN 같은 어려운 태스크에서도 잘될것인가
• 컨볼루션에서 완전히 벗어난 GAN을 생성함
• TransGAN이라 부르며 memory-frendly transformer 기반 generator 를 사용하여 progressively 증가하는동안 임베딩 차원은 줄어든다
• 그리고 패치단위 discriminator 또한 gransformer 기반이다
• 아방식은 데이터 augmentation, generator의 multy task co-training, 부분적인 이미지의 평활함에 강점이 있다
• 효과적으로 큰모델과 함께 고해상도 이미지 데이터에 적용할수 있다
• 현재 ConvNet 기반 SOTA 모델들과도 큰 차기아 나지않는다
• 이방식의 한계점과 가능성에 대해 탐구한다

1.1 Our Contributions

• 오직 transformer 만을 사용한 최초의 completely free of convolution 모델
  • 기존의 인코더 블록에만 적용한 방식과는 완전히 다름
• 기존 transformer 기반 CV모델들은 classification 혹은 detection에만 주로 사용되었기에 다른 문제가 발생
  • 구조나 색상, 질감 등의 공간적인 일관성에도 효과가 좋은지 분명하지 않음
• 이미지를 출력하는 몇안되는 transformer 기반 모델은 하나같이 convolution 기반의 encoder를 사용함
  • visual transformer 학습은 무겁고 매우 어렵기로 알려져 있다
  • GAN은 일반적으로 불안정하고 mode collapse가 쉽게 일어남
  • 두가지 문제가 더해져 매우 어려운 도전임

• 이것을 해결하기 위한 노력과 혁신적인 방법들을 공개함

• Contribution
  • Model Architecture
    • Purely transformers and no convolution
    • Memory friendly generator and a patch-level discriminator
    • Effiectively scaled up to larger models
  • Training Technique
    • Data Augmentation
    • Multi-task co-training for generator with self-supervised auxiliary loss
    • localized initailization

3. Technical Approach

• A Journey Towards GAN with Pure Transformers
  • GAN은 Generator 와 Discriminator 로 구성되고 이 두가지를 transformer로 변경하는 것으로 시작
  • 둘다 메모리 효율적인 구조로 대체하여 vanilla TransGAN을 구성
  • 이후 약점을 보완할 테크닉들을 하나씩 도입하여 deeper / wider model을 만들고 높은 퀄리티의 이미지를 생성함

    3.1 Vanilla Architecture Design for TransGAN

• Transformer Encoder as Basic Block
  • Image Transformer (Vaswani et al., 2017) 구조를 채용
    • 인코더는 multy-head self-attention 과 feed-forward with GELU 로 구성됨
  • layer normalization을 두가지 파트양쪽 앞에 적용하고 residual connection 을 사용
• Memory Frendly Generator
  • 이미지는 작은 이미지도 픽셀단위 시퀸스로 변화하면 매우 커질수 있어 self-attention에서 매우 큰 비용이 발생한다
  • PGGAN에서 영감읋 받아 점진적으로 증가하며 해상도가 높아질수록 Embedding 사이즈를 줄이는 방식을 사용
    • multiple stages로 구성되며 매 stage마다 목표 해상도가 될 때까지 두배씩 커짐
    • 매 stage에는 encoder block 이 수회 반복됨
  • noise input을 바로 사용하지않고 MLP를 통과 후 $H \times W \times C$ 크기를 가진 vector로 변경 후 learnable 한 positinal encoding 을 추가
  • upsampling을 위해 매 스테이지 뒤에 pixelshuffle 모듈을 사용
    • 우선 1D로 된 sequence를 2D로 변경 후 pixelshuffle를 적용해 2배 upsample을 하고 dimension 을 1/4로 줄인 후 1D sequence로 변화시킴 (sequence가 변하지 않음)
    • 최종적으로 목표 해상도에 도달했을때 channel을 3으로 조정
    • 이방식으로 memory and computation exposion을 완화함
• Tokenized-Input For Discriminator
  • 픽셀을 정확히 예측해야하는 Generator 와는 다르게 discriminator는 오직 real / fake 의 분류만 하면 된다
    • 이것은 입력 이미지를 patch 레벨로 tokenizing 하는것을 허용하게 해준다
  • 이미지를 $8 \times 8$패치로 나누고 각각 flatten을 거쳐 마치 word 인것마냥 [cls] 토큰을 추가
  • 이후 leanable token embedding을 거치고 [cls] 토큰의 출력만으로 real / fake 분류를 한다
• Evaluation of Transformer-based GAN
  • transformer 기반 모델의 성능을 파악하기위해 기존 SOTA 모델인 AutoGAN과 비교
    • AutoGAN G, D
    • Transformer G AutoGAN D
    • AutoGAN G, Transformer D
    • Transformer G, D
    • Generator는 스테이지에 5,2,2 블록을 사용하였고 Discriminator에는 7개의 block을 가진 1개의 stage만 사용
  • Cifar-10 에서 IS 및 FID를 비교
    • 기존모델에 공개된 최고의 성능이 나오도록 hyperparameter를 튜닝
  • result
    • Transformer G 가 매우 강력한 성능을 가짐 (AutoGAN 이상의 성능)
    • Transformer D는 매우 성능이 좋지않고 Transformer와 함께했을때 약간의 성능향상을 보임
    • 하지만 결과적으로 CNN기반 discriminator에는 미치치 못하는 성능을 가짐
  • discriminator는 training 이후엔 사용되지 않으므로 이미 transformer 기반 모델에선 SOTA 성능을 가진 모델이나 completely free of convolution을 위해 더 파고듬

3.2 Data Augmentation is Crucial for TransGAN

• 앞에 결과에서 discriminator가 잘 훈련되지 않았다는 벽에 대해 고민하게함
  • transformer 기반의 분류기는 human designed bias를 제거하여 매우 datahungry 함을 깨달음
    • 데이터가 매우 많아지기 전까진 convolution 기반에 밀렸었음
  • 이것을 해소하기 위해 data augmentation을 연구
    • GAN에서는 기본적으로 augmentation이 잘 언급되지 않았지만 최근 few-shot 모델들이 각광 받고 있음
    • DiffAug 방식을 사용하여 CNN기반 모델들과 CIFAR 10 데이터 전체상태에서의 성능을 비교
    • 전체적으로 별 효과가 없으나 StyleGAN-v2만 눈에 보이는 성능향상을 보임
    • 반면 TransGAN은 압도적인 성능 향상을 보이는것으로 datahungry 한 모델임을 확인하여 큰 도움이됨

      3.3 Co-Training with Self-Supervised Auxiliary Task

• transformer는 NLP에서 multiple pre-training task 가 좋은 결과를 보임
• GAN의 학습에서도 회전을 예측하는 auxiliary co-training 이 좋은 결과를 보인 연구가 있음
• 그래서 super resolution 과제를 적용함
  • downsampling 한 이미지를 정답으로 쓰면 되므로 데이터 생성이 쉬움
  • 2번째 스테이지에 input으로 넣어 최종 스테이지에서 mse loss를 사용함 ($\lambda$는 50을 사용)
  • 이 방식을 사용하여 약간의 성능향살을 이룸

    3.4 Locality-Aware Initialization for Self-Attention

• CNN 기반 모델은 이미지를 natural image smoothness 하는 기능이 있다
  • Transformer 에서는 이부분이 매우 빈약했음
  • 하지만 기존 연구에서 transformer는 convolution의 구조를 학습하는 경향을 발견함
  • 그래서 이 convolution의 inductive bais와 transformer의 flexibility 를 같이 encoding 할수 있는지 궁금함
  • 기존연구에선 convolution을 살짝 섞어 해결했지만 이것을 warm-starting만으로 비슷한 효과를 얻어냄
• Locality-Aware Initialization
  • 마스크되지 않은 픽셀과만 상호작용이 가능하도록 설정
  • stage 마다 mask를 확대하여 최종 스테이지에선 receptive field가 global 해지게 만듬
    • 초기에 큰 도움이 될 수 있으나 이후 훈련에선 문제가 발생할 수 있음
  • 초기 훈련에만 사용되고 점차사라지는 정규화로 간주하여 적용
  • 결과가 조금더 좋아짐

    3.5 Scaling up to large Models

• 지금까지의 학습기법들은 전부 TransGAN을 안정화 시키는 역할에 사용되었으나 이제 모델을 확장에 대해 연구함
• 다른 튜닝 없이 Embedding dimension을 늘리는것 만으로 성능이 향상되는 것을 보임
• Depth와 Dimension을 늘렸을때도 효과적으로 증가함 (왜 Depth만 늘린 실험은 없는걸까?)

4. Comparison with SOTA GANs

• CIFAR-10
  • Style-GAN v2 (augmentation) 버전과 견줄정도의 성능이 나옴
  • 눈으로봐도 만족스로운 품질과 다양성이 나왓다고 함
• STL-10
  • 데이터가 CIFAR 10보다 두배 커서그런지 SOTA 모델보다 잘나옴
  • 근데 StyleGAN은 이데이터에 대해 점수가 없어서 모르겟
• Higher Resolution
  • ConvNet 기반 모델에 비해 성능이 좋진 않지만 나쁘지않은 점수가 나왔다함
  • 근데 점수는 잘나왔어도 분명 체리픽킹 했을텐데 그런거치곤 사진이 좀..
  • 아무튼 계속 이방향으로 연구하겟다함

5. Conclusion, Limitations, and Discussions

• pure transformer한 모델로 CNN 기반 모델에 근접하는 모델을 만들었다
  • transformer 모델의 가능성 제시
  • 특화가 아닌 범용성 있는 구조로 여러 task 에 사용된 기법들을 사용할 수 있음
• 더 발전할 수 있는 내용
  • generator와 discriminator 에 대한 정교한 토큰화 (e.g. semantic grouping )
  • pretext tasks 를 사용하여 pretrained 된 모델을 사용하면 MT-CT보다 좋아질수 있음
  • more efficient self-attention forms
    • 모델 성능만이 아닌 메모리 효율까지 이득을 볼수있음
  • Conditinal image generation