TransGAN: Two Transformers Can Make One Strong GAN
1. Abstract
•
현재 Computer Vision Task에서 transformer가 좋은 성능으로 많은 관심을 받고 있으나 classification이나 detection에만 한정된 경우가 많음
•
Vision에서 어렵기로 유명한 GAN에서도 이것을 적용해보고자 하는 의지를 보임
•
기존 ConvNet 기반 GAN에 견줄수 있는 높은 성능을 보여 transformer의 가능성을 보임
2. Our Contributions
•
오직 transformer만을 사용한 최초의 completely free of convolution 모델
◦
기존의 인코더 블록에만 적용한 방식과는 완전히 다르다 주장
•
기존 transformer 기반 CV모델들은 classification 혹은 detection에만 주로 사용되었기에
구조나 색상, 질감 등의 공간적인 일관성에도 효과가 좋은지 분명하지 않음
•
이미지를 출력하는 몇 안되는 transformer 기반 모델들이 존재하나
하나같이 convolution 기반의 encoder를 사용함
◦
visual transformer 학습은 무겁고 매우 어렵기로 알려져 있음
◦
GAN 자체가 학습이 불안정하고 mode collapse가 쉽게 일어남
◦
이 두가지를 동시에 해결해야 가능하기에 어렵다고 함
•
이것을 해결하기 위한 방법을 제시
◦
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 x W x C 크기를 가진 vector로 변경 후 learnable 한 positinal encoding 을 추가
◦
upsampling을 위해 매 스테이지 뒤에 pixelshuffle 모듈을 사용
▪
우선 1D로 된 sequence를 2D로 변경 후 pixelshuffle를 적용해
2배로 upsample을 하면서 dimension 을 1/4로 줄인 후
다시 1D sequence로 변화시킴 (sequence가 변하지 않음)
▪
최종적으로 목표 해상도에 도달했을때 channel을 3으로 조정
▪
이방식으로 memory and computation explosion을 완화함
•
Tokenized-Input For Discriminator
◦
픽셀을 정확히 예측 해야 하는 Generator 와는 다르게 discriminator는 오직 real / fake 의 분류만 하면 된다
▪
이것은 입력 이미지를 patch 레벨로 tokenizing 하는것을 허용하게 해준다
◦
이미지를 8 x 8패치로 나누고 각각 flatten을 거쳐 마치 word 인것마냥 [cls] 토큰을 추가
◦
이후 leanable token embedding을 거치고 [cls] 토큰의 출력만으로 real / fake 분류를 한다
•
Evaluation of Transformer-based GAN
◦
transformer 기반 모델의 성능을 파악하기위해 기존 SOTA 모델인 AutoGAN과 비교
▪
4가지 설정으로 비교
•
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을 위해 더 파고듬1
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 데이터 전체상태에서의 성능을 비교
◦
Conv에선 전체적으로 별 효과가 없으나 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 과제를 co-training 해봄
◦
downsampling 한 이미지를 정답으로 쓰면 되므로 데이터 생성이 쉬움
◦
2번째 스테이지에 input으로 넣어 최종 스테이지에서 mse loss를 사용함
•
이 방식을 사용하여 약간의 성능향상을 이룸
3.4 Locality-Aware Initialization for Self-Attention
•
Transformer 기반의 문제
◦
CNN 기반 모델은 이미지를 natural image smoothness 하는 기능이 있으나
Transformer 에서는 이부분이 매우 빈약했음
◦
하지만 기존 연구에서 transformer는 convolution의 구조를 학습하는 경향을 발견함
◦
그래서 이 convolution의 inductive bais와 transformer의 flexibility 를
같이 encoding 할수 있는지에 대해 의문을 가짐
•
Locality-Aware Initialization
◦
기존연구에선 convolution을 살짝 섞어 해결했지만
이것을 warm-starting만으로 비슷한 효과를 얻어냄
◦
마스크되지 않은 픽셀과만 상호작용이 가능하도록 설정
◦
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
•
Conclusion
◦
pure transformer한 모델로 CNN 기반 모델에 근접하는 모델을 만들었다
◦
transformer 모델의 가능성 제시
◦
특화가 아닌 범용성 있는 구조로 여러 task 에 사용된 기법들을 사용할 수 있음
•
Discussion
◦
generator와 discriminator 에 대한 정교한 토큰화 (e.g. semantic grouping )
◦
pretext tasks 를 사용하여 pretrained 된 모델을 사용하면 MT-CT보다 좋아질수 있음
◦
more efficient self-attention forms
▪
모델 성능만이 아닌 메모리 효율까지 이득을 볼수있음
◦
Conditinal image generation