[논문 리뷰] Efficient Image Generation with Variadic Attention Heads

Steven Walton, Ali Hassani외 3명|arXiv (Cornell University)|2022. 11. 10.

Generative Adversarial Networks and Image Synthesis인용 수 8

한 줄 요약

이 논문은 지역적 및 전역적 특징을 효율적으로 캡처하기 위해 주변주의(NA)를 통해 다양한 수신장역할을 가진 주의 헤드의 새로운 분할 방식인 Hydra-NA를 사용하는 트랜스포머 기반 이미지 생성 프레임워크인 StyleNAT을 제안한다. 이는 FFHQ-256에서 FID 점수 2.046, FFHQ-1024에서 4.174의 새로운 SOTA 기록을 달성하며, 기존 GAN 및 트랜스포머보다 28% 적은 파라미터와 56% 높은 추론 속도를 기록한다.

ABSTRACT

While the integration of transformers in vision models have yielded significant improvements on vision tasks they still require significant amounts of computation for both training and inference. Restricted attention mechanisms significantly reduce these computational burdens but come at the cost of losing either global or local coherence. We propose a simple, yet powerful method to reduce these trade-offs: allow the attention heads of a single transformer to attend to multiple receptive fields. We demonstrate our method utilizing Neighborhood Attention (NA) and integrate it into a StyleGAN based architecture for image generation. With this work, dubbed StyleNAT, we are able to achieve a FID of 2.05 on FFHQ, a 6% improvement over StyleGAN-XL, while utilizing 28% fewer parameters and with 4$ imes$ the throughput capacity. StyleNAT achieves the Pareto Frontier on FFHQ-256 and demonstrates powerful and efficient image generation on other datasets. Our code and model checkpoints are publicly available at: https://github.com/SHI-Labs/StyleNAT

연구 동기 및 목표

• 고해상도 이미지 생성에서 표준 트랜스포머의 비효율성과 제한된 적응 가능성 문제를 해결하기 위해.
• Swin 및 NA와 같은 局소화된 주의 메커니즘에서 고정된 수신장역할의 한계를 극복하기 위해.
• 확장된 및 슬라이딩 윈도우 기반 NA를 사용해 헤드를 다양한 수신장역할로 분할하여 데이터 적응형으로도 유연한 주의를 가능하게 하기 위해.
• 모델 크기와 추론 시간을 줄이며 FID에서 최고 성능을 달성하기 위해.
• Swin 및 NA 기반 모델의 국소화된 주의 맵을 해석하기 위한 시각화 방법을 개발하기 위해.

제안 방법

• 주변주의(NA)를 사용해 다중 헤드 주의를 다양한 수신장역할을 가진 별도의 헤드로 분할하는 Hydra-NA 방법을 도입한다.
• 슬라이딩 윈도우 및 확장된(디레이티드) NA 변형을 활용해 헤드가 국소적 이웃 또는 확장된 전역적 맥락에 주의를 기울일 수 있도록 한다.
• StyleGAN에 영감을 얻은 스타일 기반 생성기 설계를 통해 계층적 특징 학습을 가능하게 하며, Hydra-NA 블록을 통합한다.
• 무작위 잠재 벡터를 스타일 코드로 매핑하는 네트워크를 사용해 각 레이어의 특징 맵을 조절한다.
• 잠재 공간 내의 공간적 구조를 유지하기 위해 학습 가능한 위치 인코딩을 적용한다.
• 다양한 해상도 수준에서 주의 맵을 시각화할 수 있는 기법을 개발하여, 워터마크와 같은 데이터셋 특수 아티팩트를 식별한다.

실험 결과

연구 질문

• RQ1 다양한 수신장역할을 가진 주의 헤드로 분할함으로써 트랜스포머 기반 이미지 생성기의 표현력과 효율성을 향상시킬 수 있는가?
• RQ2 변동 가능한 커널 크기와 확장 정도를 가진 국소화된 주의가 고해상도 이미지 생성에서 장거리 의존성 모델링에 어떤 영향을 미치는가?
• RQ3 특정 하이퍼파라미터 조정 없이도 하나의 아키텍처가 다양한 데이터셋에 일반화할 수 있는 정도는 어느 정도인가?
• RQ4 주의 맵을 통해 생성된 이미지에서 워터마크나 ID와 같은 데이터셋 고유 아티팩트가 기억된 것으로 확인될 수 있는가?
• RQ5 기존 GAN 및 트랜스포머에 비해 계산 비용을 줄이며 SOTA 성능을 달성하는가?

주요 결과

• StyleNAT는 FFHQ-256에서 FID 점수 2.046을 기록하며, StyleGAN-XL과 같은 컨볼루션 모델 및 트랜스포머 기반 모델을 모두 능가하는 새로운 SOTA 성능을 달성한다.
• FFHQ-1024에서 StyleNAT는 트랜스포머 기반 GAN 중에서 FID 점수 4.174를 기록하며, 전역적 특징 모델링 능력이 뛰어나다는 것을 입증한다.
• 더 나은 성능를 보임에도 불구하고, StyleGAN-XL 대비 파라미터 수를 28% 감소시키고 추론 속도를 56% 향상시켰다.
• 주의 시각화 결과, 저해상도(예: 16x16)에서 특히 워터마크 및 이미지 ID와 같은 데이터셋 고유 아티팩트를 기억하는 것으로 나타났다.
• 주의 맵 내 워터마크 및 ID의 존재는 그 시각적 존재와 강하게 상관관계를 보이며, 훈련 데이터에 대한 과적합의 가능성을 시사한다.
• 고품질의 이미지 샘플을 생성함에도 불구하고, 교회 데이터셋에서는 FID 점수가 여전히 열악한 편이었으며, 이는 분포 특수 아티팩트에 대한 과적합이 일반화 능력 저하로 이어졌음을 시사한다.