[논문 리뷰] HyperStyle
Introduction
이번에 소개드릴 논문은 CVPR 2022에 올라온 따끈따끈한 논문입니다 :)
본론으로 들어가기 전에, Hyperstyle로 어떤 재미있는 실험들을 할 수 있는지 구경해 볼까요?
마블 영화, 롤 게임 캐릭터의 어린 시절 모습, 먼 미래의 모습 또는 웃는 모습이 궁금했던 적 있지 않으신가요?
HyperStyle을 이용하면 실제 사람이 아닌 도메인에서도 사물의 ID를 유지하면서 이미지를 복원할 수 있어, 원하는 방향으로 실시간 수정이 가능합니다!
HyperStyle을 자세히 설명하기 전에 StyleGAN의 latent 공간 내에 존재하는 몇 가지 상충관계(trade-off)들을 짚고 넘어가겠습니다.
1) Time-accuracy Trade-off
현재, 목표 이미지를 생성하기 위해 좋은 latent를 찾는 GAN inversion에 대한 연구가 많이 진행되고 있습니다. 그리고 보다 더 정확한 이미지 재건을 위한 GAN inversion 방법은 크게 두 가지로 분류할 수 있습니다.
- optimization-based :
- 임의의 latent로부터 생성된 이미지가 목표 이미지에 가까워지도록 per-image 거리를 최소화하는 방향으로 latent를 업데이트하는 방법입니다. 정확도는 높지만 매우 느리다는 단점이 있습니다.
2. encoder-based :
- 여러 장의 이미지를 latent 공간으로 매핑하는 encoder를 학습하는 방법입니다. per-image가 아니기 때문에 빠르지만 정확도는 떨어집니다.
최근 발표된 논문 중, 이미지별로 좋은 latent를 찾은 뒤 그 latent를 입력으로 받는 generator를 fine-tuning 하는 논문[1] 이 있습니다. Per-image 학습이 진행되기 때문에 optimization-based로 분류할 수 있지만 latent를 업데이트하는 것이 아닌, latent 이 고정된 상태에서 generator를 튜닝하는 것이 특징입니다.
아래의 그림처럼 생성된 이미지와 주어진 이미지 간의 차이를 줄이기 위해, 다음 loss (eq (1))로 generator weight를 학습합니다.
Hyperstyle 역시 pre-trained StyleGAN generator의 가중치를 업데이트하며, 빠른 inference time을 위해 hypernetwork를 도입하여 encoder-based 방법처럼 빠른 inference time 을 가능하게 합니다.
실제로 기존 파인튜닝 방법들은 한번 업데이트하는데 이미지 별로 10초의 시간이 걸린 반면 hyperstyle은 2초 이하의 시간만 걸려, real-time inference 가 가능한 매우 큰 장점이 있습니다.
2) Distortion-Editability Trade-off
GAN inversion으로 출력된 latent는 W 또는 W+ 공간에 존재한다고 볼 수 있습니다.
W 공간은 StyleGAN에서 mapping network를 통과한 뒤 generator의 입력으로 존재하는 공간이고, W+는 generator의 다른 레이어마다 다른 latent를 할당한 공간이라고 볼 수 있습니다.
W 공간으로 반전된(inverted) 이미지들은 latent 공간에서의 적은 수정으로도 이미지를 부드럽게 변경할 수 있는 editability가 좋은 반면, 다양한 범위의 이미지를 표현할 수 있는 표현력이 부족합니다. 따라서, 많은 연구들은 W+ 공간으로 이미지를 반전시켜 distortion을 줄이는 대신 약간의 editability를 희생하는 방법을 택하고 있습니다.
e4e [2]는 이러한 특성을 이용하여 W 가까이에 있는 W+를 예측하도록 하는 encoder를 설계하였습니다. Hyperstyle 은 pre-trained e4e를 encoder로 사용하여 W공간에 가까운 latent를 사용하였습니다.
Method
- 목표
저자들의 목표는 이미지 x에 대하여 eq(1) 을 만족하는 generator weight를 찾는 것임을 기억합시다.
그러기 위해서는 먼저 W 공간에 가까운, 적당한 초기 latent 코드 w_init 을 찾아야 합니다. 이는 앞서 소개한 pre-trained encoder e4e[2] 를 통해 구합니다. 그리고 새로운 weight θ ̂를 예측하기 위한 hypernetwork H를 도입하게 되는데, 아래의 loss와 같이 여러 장의 이미지 x 와 처음 재건된 이미지 ŷ_init=G(ŵ_init ; θ) 간의 distortion을 줄이는 방향으로 학습이 됩니다.
즉, 예측된 weight 가 θ ̂=H(y ̂_init; x) 일 때, 최종적인 이미지 재건 결과는
y ̂=G(w ̂_init; θ ̂)가 됩니다. 여기서 특징적으로 저자들은 바로 weight를 예측하지 않는 대신, 점차적으로 weight를 업데이트하기 위해 ‘iterative refinement’라는 방법을 사용합니다.
요약하자면, hyperstyle은 5번 정도의 적은 수의 step으로 layer (ℓ=1, …, M) 별 weight offset Δ_ℓ 를 구하여, 점진적으로 generator weight를 튜닝합니다 (eq(3)) . 이러한 방법은 동일한 저자들의 기존 논문 Restyle [3]에서 먼저 소개되었는데, 더욱 정확한 이미지 재건에 효과적이라고 합니다.
Hypernetwork Design
StyleGAN2의 파라미터 수가 30M 개 정도 있는 상황에서 각각의 offset을 모두 구하는 것은 매우 고가의 비용이 듭니다. 따라서 학습 경량화를 위하여 세 가지 방법을 사용합니다.
1) 각 파라미터 별로가 아닌 channel-wise로 offset을 구하며 2) 같은 크기의 합성곱 블록끼리 refinement block의 fc layer를 공유하고 3) 마지막으로 모든 layer가 아닌 특정 layer의 가중치만 업데이트하는 방법을 사용합니다.
1. channel-wise로 offset
StyleGAN의 ℓ번째 합성곱 레이어의 weight (θ_ℓ^ij)를 i 번째 합성곱 필터의 j 번째 channel 단위로 업데이트를 하게 됩니다. (i = 1, …, C_ℓ^out, j = 1, …, C_ℓ^in)
Restyle[3] 에서 소개된 weight update 방법처럼 T (=5) step을 수행한다고 했을 때 각 step t마다 offset Δ_t=H(y ̂_(t-1), x) 을 구하여, output 이미지 y ̂_t=G(w ̂_init; θ ̂_t) 를 업데이트하기 위해 weight를 다음과 같은 방식으로 튜닝합니다.
그림으로 다시 설명하자면 hypernetwork는 그림 2의 회색 박스 부분에 해당하고, ResNet34 backbone을 통과한 16×16×512 feature map 이 다시 refinement block의 입력으로 들어가서 offset을 예측합니다. channel-wise offset의 크기가
1×1×C_ℓ^in×C_ℓ^out 기 때문에 θ_ℓ 와 크기를 맞춰주기 위해서 kernel 사이즈인 k_ℓ×k_ℓ번 복사되어 eq(3)에 사용됩니다.
2. sharing refinement block :
refinement block 은 다음 그림과 같이 여러 개의 downsampling layer들과 fc layer로 이루어져 있는데, 가장 많고 크기가 큰 3×3×512×512 합성곱 블록들끼리 같은 fc layer를 공유하도록 해줍니다.
3. Selective layer refinement :
toRGB layer는 픽셀 별로 그 텍스처와 색을 조절하기 때문에 toRGB layer의 weight는 변경하지 않는 것이 좋습니다. 따라서 non-toRGB layer의 weight만 업데이트하는데, 이때도 course, medium, fine layer들 중 medium 과 fine에 해당하는 layer들만 튜닝합니다. 이에 대한 ablation은 논문에 나와있습니다.
Results
- Dataset : 학습을 위해 FFHQ, 테스트를 위해 CelebA-HQ 데이터 셋을 사용합니다.
- Reconstruction Quality
encoder-based 방법들(pSp, e4e, Restyle) 과 비교했을 때, 훨씬 높은 정확도로 이미지를 재건하며 optimization-based 방법들 (PTI, latent vector opimization)에 비해 훨씬 적은 시간과 계산량으로도 비슷한 수준으로 이미지를 재생성합니다.
2. Editability Quality
- 정확한 이미지 생성만큼이나 높은 editability를 갖게 되는 것은 중요한데, 앞서 다뤘듯이 distortion 과 editabilty 간의 상충관계가 존재하기 때문에 둘 다 좋은 성능을 보이기는 쉽지 않습니다. 여기서 editability quality를 측정하기 위해 이미지의 ID를 유지하면서 변경 가능한 범위를 살펴보았습니다.
- 기존 encoder-based 방법들의 한계를 먼저 살펴보면 pSp 와 Restyle_pSp는 W+ space로 이미지를 반전시키기 때문에 editibility가 좋지 않으며, e4e 와 Restyle_e4e의 경우 editability는 좋지만 얼굴의 ID를 유지하지 못합니다. Hyperstyle은 optimization-based 방법들에 비해 훨씬 빠른 속도로도 비슷한 editabilty 성능을 보입니다.
3. Domain Adaptation
- 다른 유사한 도메인으로 적응(adaptation) 시키기 위해 사전에 학습된 StyleGAN 을 파인튜닝할 수 있습니다. 그러나 수염이나 머리 색과 같은 몇 가지 얼굴 특성들이 사라질 수 있기 때문에 Hyperstyle은 파인튜닝된 target generator에 source generator로 학습된 offset을 적용합니다.
- 아래 결과에서 보시다시피 Hyperstyle만이 ID를 잃지 않고 이미지를 다른 도메인으로 성공적으로 이전할 수 있습니다. 심지어는 out-of-domain으로도 좋은 generalization adaption 성능을 보이는데, 그 실험 결과들은 논문을 참조해 주시기 바랍니다.
Conclusion
지금까지 StyleGAN inversion을 위한 방법으로 성능과 실용성을 모두 갖춘 Hyperstyle에 대해 살펴보았습니다. 실제 산업에서 이미지 재건 및 수정을 하려면 실시간으로 할 수 있어야만 실효성이 있습니다. 특히 저희 Klleon은 이미지뿐만 아니라 video를 많이 다루는데, 한 장 한 장의 이미지 frame마다 시간을 단축할 수 있어야 합니다. 이런 측면에서 기존 상충관계를 해결하여 빠른 시간으로 좋은 이미지 성능을 유지할 수 있는 Hyperstyle은 KLleon에서도 유용하게 사용할 수 있을 것 같습니다. 다음에도 stylegan inversion을 적용하여 다양한 task을 해결한 논문들을 소개하겠습니다.
References
[1] Roich, Daneil et al., ‘Pivotal Tuning for latent-based editing of real images’, arXiv 2021
[2] Tov, Omer et al., ‘Designing an encoder for stylegan image manipulation’, 2021
[3] Alaluf, Yuval., ‘Restyle: A residual-based stylegan encoder via iterative refinement’., ICCV 2021
