StyleGAN-XL: Scaling StyleGAN to Large Diverse Datasets 논문 리뷰

2월 공개된 StyleGAN-XL는 현재 이미지 생성 부분의 ImageNet의 모든 size를 포함해 많은 데이터셋에서 SOTA를 차지하고 있다.

StyleGAN-XL은 Nvidia에서 공개한 정식 StyleGAN 시리즈는 아니고, Projected GAN 연구에 참여했던 Axel Sauer가 다른 사람들과 함께 Projected GAN의 아이디어를 StyleGAN3에 특화시킨 모델이다.

Project Page

stylegan-xl

Abstract

StyleGAN은 제어 가능성을 위해 설계되었기 때문에 ImageNet과 같은 대규모 비정형 데이터 세트에서 성능이 심각하게 저하된다. 대조적으로, 연구진은 최근에 도입된 Projected GAN 패러다임에 따라 기존의 강력한 신경망과 점진적인 성장 전략을 활용하여 ImageNet에서 최신의 StyleGAN3을 성공적으로 훈련시킨다. 또한 최종 모델인 StyleGAN-XL는 대규모 이미지 합성에서 SOTA를 기록하며 이러한 데이터 세트 규모에서 1024²의 해상도로 이미지를 생성하는 최초의 모델이다. 실험에서 모델이 초상화 또는 특정 객체 클래스의 좁은 영역을 넘어 이미지를 반전하고 편집할 수 있음을 보여준다.

Introduction

ImageNet과 같은 대규모 비정형 데이터셋에서, StyleGAN은 아직 만족스러운 결과를 얻지 못했다. 그리고 또 다른 문제는 모델을 높은 해상도로 확장할수록 계산 비용이 기하급수적으로 늘어난다는 것이다. ImageNet 이미지 합성에서 StyleGAN보다 더 좋은 성능의 BigGAN은 성능이 훈련 실행마다 다르고 이미지 편집에 필수적인 중간 잠재 공간을 사용하지 않기 때문에 StyleGAN2와 비슷한 위치에 도달하지 못했다.

2021년, 샘플을 사전 훈련된 feature 공간에 투영함으로써 훈련 안정성, 훈련 시간 및 데이터 효율성을 크게 향상한 Projected GAN이 도입되었다. 그러나 Projected GAN은 단일 데이터 세트에서의 부분적인 확장에 불가했기 때문에, 이번 연구에서는 이를 해결하기 위한 아키텍처 변경을 제안한다. 결과를 더욱 개선하기 위해, Projected GAN의 사전 훈련 네트워크를 분석하고 CNN와 ViT가 공동으로 사용될 때 성능을 크게 향상시킨다는 것을 발견했다. 마지막으로, 추가 클래스 정보를 주입하기 위해 확산 모델에서 도입된 기술인 분류기 지침을 활용한다.

본 논문의 기여는 대규모 이미지 합성에서 기존 GAN 및 확산 모델을 능가하여 SOTA를 달성했다. 또한 ImageNet 클래스에 대한 반전 및 편집을 소개하고 새로운 반전 패러다임인 PTI가 모델과 잘 결합되며 학습된 잠재 공간에 도메인 밖의 이미지까지 부드럽게 임베딩한다는 것을 발견했다.

본 논문의 효율적인 훈련 전략을 통해 표준 StyleGAN3의 피라미터를 3배 늘리면서도 이전 SOTA인 확산 모델의 성능을 그들의 훈련 시간 내에서 달성했다. 또한 1024²의 해상도로 ImageNet 규모에서 이미지 합성을 최초로 시연한다.

Background

StyleGAN3 : 개념편, 아키텍처편

Alias-Free GAN (StyleGAN3) 리뷰 - 개념

Projected GAN

Projected GANs Converge Faster 논문 리뷰

위의 글을 다 읽어야 이해가 편함. 특히 StyleGAN3은 무조건 알아야 함.

Scaling StyleGAN to ImageNet

본 논문의 목표는 ImageNet에서 StyleGAN3을 성공적으로 훈련하는 것이다.
다음 표과 같이 vanilla StyleGAN3을 기준으로 개선사항을 점진적으로 도입한다.

StyleGAN-XL는 StyleGAN3보다 3배 많은 피라미터를 가지면서도 성능과 훈련 시간, 계산 비용의 측면에서 뛰어나다.

 

Adapting Regularization and Architectures

회전 등변량 개선 모델인 StyleGAN3-EQ-R은 지나치게 대칭적인 만화경 같은 이미지를 생성하기 때문에 StyleGAN-EQ-T를 기본으로 한다.

Regularization

• 정규화는 FFHQ와 같은 단일 모드 데이터셋에서 결과를 개선할 수 있지만 다중 모드 데이터셋에서는 위험할 수 있기 때문에 가능한 정규화를 피하는 것을 목표로 한다.
• Path length regularization, style mixing은 원래 StyleGAN3에서 사용하지 않는다(StyleGAN3 아키텍처 편 D참고 = StyleGAN3-D로 표기). 하지만 path length 정규화는 고품질 반전을 가능하게 하기 때문에 200k 이미지 훈련 이후에 적용한다.
• 판별기의 경우 gradient penalty 대신에 스펙트럼 정규화 사용.
• 훈련 시 처음 200k 이미지에 대해 가우시안 필터를 사용하여 모든 이미지를 흐리게 한다. (StyleGAN3-R)

Low-Dimensional Latent Space
Projected GAN은 StyleGAN 보다 FastGAN에서 더 잘 작동한다. 주요한 이유는 StyleGAN의 잠재 공간이 고차원이기 때문인데, 자연 이미지의 고유 차원은 상대적으로 낮으며 ImageNet의 차원 추정치는 약 40 정도이다. 따라서 StyleGAN 잠재 코드 z의 크기를 64로 줄이고 Projected GAN과 함께 훈련한다. (config-B)

매핑 네트워크 G_m의 용량을 제한하지 않기 위해 스타일 코드 w의 차원은 512로 유지.

Pretrained Class Embeddings
연구진은 config-B가 클래스마다 유사한 샘플을 생성하여 높은 IS(Inception Score)를 초래하는 경향이 있음을 관찰했으며, Recall 메트릭을 활용하여 config-B가 0.004의 낮은 recall을 달성함을 확인했다.

그들은 Projected GAN이 클래스 임베딩을 붕괴시킨다고 가정했다(Projected GAN에서 개별 판별기 D_i가 보는 분포는 실제 분포가 아니기 때문에). 따라서 사전 훈련을 통해 임베딩 최적화를 용이하게 하는 것을 목표로 한다.

 

ImageNet의 각 클래스에서 Efficientnet-lite0(가벼운 분류 모델)의 가장 낮은 해상도 feature를 추출하고 풀링해 클래스당 평균을 계산한다.(네트워크는 config-B 따라 임베딩 차원을 작게 유지하기 위해 채널 수가 적다.) 그러고 나서 z의 크기와 일치하도록 선형 투영을 통과한다. 매핑 네트워크 G_m과 개별 판별기 D_i는 이와 같은 임베딩을 조건으로 받는다. GAN훈련 동안 임베딩 및 선형 투영은 전문화를 허용하도록 최적화된다.

이 구성을 사용하여 모델이 클래스당 다양한 샘플을 생성하고 재현율이 0.15로 증가한다. (config-C)
Config-C에 대한 훈련 시간은 15 V-100 days로 제한한다. 따라서 절대 recall은 완전히 훈련된 모델에 비해 현저히 낮다.
(사전 훈련 조금만 하고 실제 훈련에서 최적화된다는 말인가?)

Reintroducing Progressive Growing

PGGAN은 낮은 해상도의 레이어부터 순차적으로 훈련시키는 방법을 사용하며, StyleGAN에서 사용했었다. 하지만 texture sticking 아티팩트를 유발한다는 이유로 StyleGAN2에서는 빠졌었는데, StyleGAN3에서 texture sticking의 유발 요인인 aliasing을 거의 다 제거했기 때문에 다시 사용한다.

StyleGAN3에서는 해상도에 따라 고정된 레이어를 사용하지 않는다. 출력 해상도에 상관없이 무조건 14개의 레이어를 사용하며, 정교하게 설계된 flexible layer 체계를 통해 출력 해상도에 따라 피라미터가 자동 조정된다. (StyleGAN3-T)

무조건 마지막 두 계층은 critical 샘플링(StyleGAN3-G)을 위한 임계 샘플링 계층인데, 이 계층은 앨리어싱을 도입하기 때문에 무조건 마지막에 있어야 한다. 그래서 점진적으로 다음 단계의 레이어를 추가할 때는 임계 샘플링 계층을 제거한 후에 추가하며, flexible layer 체계를 준수하도록 피라미터를 조정한다.

PGGAN에서는 G와 D를 동시에 성장시키지만, Projected GAN에서는 일반적인 D를 사용하지 않기 때문에 D는 성장하지 않으며, Projected GAN의 사전 훈련 네트워크 F를 위해 작은 이미지로 훈련할 때는 G의 출력 이미지를 투영 전에 224² 픽셀로 업샘플링한다.

 

• StyleGAN-XL에서는 16²의 해상도와 11개의 레이어로 시작한다. 
• 해상도가 높아질 때마다 임계 샘플링 계층을 잘라내고 7겹의 레이어를 추가한다. 1024²의 마지막 단계에서는 5개의 레이어만 추가한다. 총 레이어는 39개에 달한다.
• 고정된 성장 일정 대신 FID가 떨어지지 않을 때까지 하나의 단계를 훈련시킨다.

 

16², 32²의 해상도에서는 경험적으로 2048의 큰 배치를 사용한다. 새로운 레이어가 추가되면 이전의 레이어는 고정된다.

 

이번 구성(config-D)에서는 FID 이득은 크지 않지만, 주요 장점은 고해상도에서 훈련 시간의 엄청난 단축이다. 512² 해상도에서 2 V-100 days 만에 SOTA(FID=3.85)에 도달한다.

EQ-T(Translation Equivariance)는 config-C에서 55, config-D에 48이다. EQ-T가 조금만 떨어진 것은 아직 앨리어싱이 남아있다는 것을 의미한다.

Exploiting Multiple Feature Networks

Projected GAN의 ablation study에서, 사전 훈련 네트워크 F는 네트워크 아키텍처에 관계없이 유사하게 수행된다는 것을 발견했다. 하지만 여러 F를 사용하는 것이 유리한지는 검증하지 않았기 때문에, 다른 네트워크 F2를 추가하여 실험해본다.

CNN과 ViT를 결합하면 성능이 크게 향상된다. 두 아키텍처는 서로 다른 표현을 학습하기 때문에, Projected GAN에 대한 보완 효과가 있는 것으로 보인다.

네트워크를 더 추가했을 때는 의미 있는 개선이 없었다. 그래서 CNN + ViT 조합을 채택했다. (config-E)

 

Classifier Guidance for GANs

어떤 논문에서 확산 모델에 클래스 정보를 주입하기 위한 분류기 지침을 도입했다.
확산 모델이란, 어떤 이미지에서 점점 노이즈를 섞고, 그 반대로 노이즈에서 이미지를 복원할 수는 없을까? 하는 모델이다.

분류기 지침은 사전 훈련된 분류기를 통해 각 확산 단계에서 분류기의 gradient를 추가하여 예측을 수정한다. 여기서 우리는 모델이 이미 임베딩을 통해 클래스 정보를 수신했더라도 분류기 지침으로부터 지속적인 이익을 얻을 수 있음을 알 수 있다.

생성된 이미지 x를 사전 훈련된 분류기 CLF를 통해 클래스를 예측한 다음 생성기 손실에 추가한다. 분류기로는 DeiT-small을 사용한다. 이는 훈련에 많은 오버헤드를 추가하지 않으면서 강력한 분류 성능을 보여준다.

이 변경(config-F)으로 IS가 상당히 개선되고 생성 품질이 좋아진다.

분류기는 32² 보다 큰 해상도에서만 작동하며 그렇지 않으면 모드 붕괴로 이어질 수 있다.

Limitations and Future Work

현재 이미지 생성 부분의 많은 벤치마크에서 SOTA를 차지하고 있을 정도로 성능은 보장되어 있고, 다중 모드 데이터셋에 약했던 기존의 GAN의 한계를 극복했다.

한계점은 피라미터가 너무 많고 StyleGAN3의 등분산 보존 설계 때문에 의미론적 제어 가능성을 희생했다는 것이다. (근데 후자는 연구자들의 충분한 관심이 있다면 쉽게 극복 가능해 보인다.)

Results

Image Synthesis

분류기 네트워크도 활용하는 FID와 IS에 대해 모델이 부주의하게 최적화되지 않도록 하기 위해, random FID(rFID)를 제안한다. rFID는 무작위로 초기화된 inception 네트워크의 pool_3 계층에서 Frechet 거리를 계산한다. 생성 모델 평가를 위한 무작위 feature의 효과는 이 논문에서 입증되었다. 또한 공간 구조를 평가하기 위해 sFID를 보고한다. 마지막으로 샘플 충실도와 다양성은 precision과 recall을 통해 평가된다.
(-U : 업샘플링, -G : 분류기 지침)

같은 잠재 코드에서 해상도별 생성된 이미지

 

Interpolation

전혀 다른 도메인의 이미지끼리도 잘 보간되는 것을 볼 수 있다.

Image Manipulation

위에서부터 "smile", "no stripes", "big eyes".

Inversion

StyleGAN-XL의 flexible layers