Alias-Free GAN (StyleGAN3) 리뷰 - Architecture

지난 리뷰 StyleGAN3 - 개념편 에서 이어집니다.

지난 리뷰를 대충이라도 읽고 나서 이번 리뷰를 보는 걸 권장합니다.

 

 

 

translation equivariance EQ-T, rotation equivariance EQ-R은 PSNR와 유사한 metric으로 평가된다. PSNR

I_max는 이미지의 최대범위를 의미한다. 논문에선 -1 ~ +1 이므로 2.

z₀은 G의 시작점인 상수 Z₀을 연속표현으로 변환한 것. 

g는 generator, t_x는 translation 함수, w는 네트워크의 가중치, c는 각 채널, v는 픽셀 공간, p는 각 픽셀을 나타낸다.

분모인 g(t(z))와 t(g(z))의 차이가 적을수록 EQ-T가 높아진다. 

 

 

이제부터 metric을 이용해 각각의 구성요소에 대해 알아볼 것이다.

stylegan3 architecture

 

B. Fourier features

stylegan3에서는 모델의 최초 입력 상수로 4x4x512의 텐서를 사용하지 않는다.

[channels, 2] 크기의 frequencies, [channels] 크기의 phases를 상수로 사용한다. (첫 단의 channels = 512)

매핑 네트워크의 출력을 받아 FC layer(w_dim -> 4)를 통과해 (c, s, x, y)형태로 만든다.

c, s는 rotation에 대한 정보이고 x, y는 translation에 대한 정보를 가지게 한다.

이 정보를 이용해 freqs와 phases에 반영하고 affine_gird변환으로 [B, C, H, W] 형태로 만든 후 훈련 가능한 피라미터를 행렬곱한 값을 다음 레이어에 넘긴다.

 

정리하자면 모델의 최초 입력 상수로 주파수와 위상을 사용하고, 매핑 네트워크에서 받은 rotation, translation에 대한 정보를 반영한 후 익숙한 이미지의 형태로 바꿔서 출력한다.

 

fourier feature를 사용함으로써 자연스럽게 공간적으로 무한하게 정의되는 이점도 있다고함. (파동에는 시작과 끝이 존재하지 않으니까.)

 

 

C. No noise inputs

Texture sticking을 유발하기 때문에 노이즈를 입력받지 않는다. 하지만 모델이 전혀 노이즈를 입력받지 않는 것은 분명히 좋은 방향은 아닌데, equivariance를 해치지 않으면서 노이즈를 입력하는 방법을 찾아야 한다고 함.

 

 

D. Simplified generator

 

1. mapping network 깊이 축소

매핑 네트워크의 깊이를 8에서 2로 줄입니다. ADA의 논문에서 적은 데이터로 학습하기 위해 얕은 매핑 네트워크를 처음 제안했고, stylegan3에서 얕은 매핑 네트워크를 사용해도 딱히 문제가 없었던 듯 합니다.

 

2. mixing regularizaion(style mixing) 제거

Making Anime Faces With StyleGAN에서 스타일 믹싱이 사람 얼굴과 같은 단순한 데이터셋에서는 효과가 있지만 복잡하고 multi-modal인 데이터셋에서는 악영향을 끼친다고 언급.

 

3. path length regularization 제거

잠재공간의 변환에 페널티를 부여하여 texture sticking을 유발함. rotation과 translation에 대해 공간이 유연하게 변해야 하는 equivariance의 개념과 상충된다.

 

4. skip connection 제거 

skip connection 대신 훈련 중 모든 모든 픽셀과 feature map에 대한 EMA를 추적하고 각 레이어의 컨볼루션 연산 전에 반영한다. 이 변경은 generator를 단순화해서 높아진 FID를 원래수준으로 되돌리고 equivariance를 향상시킨다.

 

 

E. Boundaries & upsampling

1. boundaries

sampling_rate는 연속 공간에서의 샘플링 주기, 픽셀 공간에서의 픽셀 수를 나타낸다.

하지만 실제로 각 레이어에서 출력되는 feature map의 크기는 가로세로로 20씩 큰데, stylegan3 개념편에서 설명했던 테두리의 위치정보를 배우는 것을 방지하기 위해 큰 크기의 feature map을 연산 후 crop해서 다음 레이어의 입력으로 사용한다.

 

2. upsampling

기존의 방법이 아니라 Kaiser window 기반의 직접 제작한 필터를 이용해 업샘플링한다.

오른쪽과 같은 형태의 함수를 각 픽셀에 2d convolution 한다고 생각하면 된다.

필터의 기본 kernel 크기는 6이며, 더 크면 성능은 좋아지지만 연산비용이 기하급수적으로 늘어난다.

 

이 config E(boundaries, upsampling)는 equivariance는 향상되지만 FID는 높아짐.

 

 

F. Filtered nonlinearities

앞서 개념편에서 언급했듯이, nonlinearity는 완벽하게 alias-free 하게 만들수가 없어서 upsampling->nonlinearity->downsampling의 순서를 거친다.

이 과정이 엄청나게 많은 계산비용을 필요로 하기 때문에 연구진은 custom CUDA kernel을 제작하여 사용했다. 실제로 pytorch operation보다 10배가 넘는 속도 차이를 보였다.

커널의 자세한 작동 원리가 궁금하다면 stylegan3 논문의 24p 부록 D 참고.

 

 

G. Non-critical sampling

critical sampling은 cutoff가 차단하려는 주파수에 정확히 위치하는 것을 의미한다. (빨간 선)

x축: 주파수, y축: 감쇠 정도(0일 때 완전 차단)

하지만 실제로 cutoff 주위에 신호를 서서히 감쇠시키는 transition band가 있어서 cutoff를 정확히 내가 차단하려는 1.4 지점에 둔다면 높은 주파수의 정보가 남아서 조금의 앨리어싱이 발생하게 된다.

보통의 이미지 처리 응용 프로그램에서는 이 transition band가 적당한 높은 주파수의 정보를 유지하고 적당한 안티-앨리어싱의 균형을 잘 유지하게 해 주기 때문에 이상적이다. 하지만 이 모델의 목적은 alias-free를 이뤄내는 것이기 때문에 stopband의 주파수를 완전히 차단하도록 cutoff를 half-width만큼 낮춘다.

(초록색 선은 cutoff를 위치시킨 뒤 f_h를 증가시켰을 때의 그래프. transition band가 넓어진다.)

 

이걸 최종 임계 샘플링 레이어을 제외한 모든 레이어에 적용한다. 최종 레이어를 제외하는 이유는 주파수를 낮추면 디테일이 떨어지는데, 모델의 최종 출력은 세부사항이 구현된 선명한 이미지여야 하기 때문이다.

그리고 feature map의 수를 이미지의 크기인 sampling rate에 반비례 하는 것이 아닌 cutoff와 반비례 하도록 조정한다. cutoff가 낮으면 이미지의 주파수가 떨어지고, 낮은 주파수로 세부적인 특징을 표현하려면 많은 feature map이 필요하기 때문이다.

 

이 config G에서부터 FID가 기존의 stylegan2보다 낮아진다.

 

 

H. Transformed Fourier features

사실 config B에서 config B, H를 같이 설명했었다.

주파수와 위상을 입력 상수로 사용하는 게 config B

매핑레이어로부터 받은 잠재벡터로 기하학적 변환 정보를 얻어 입력 상수에 반영 하는 게 config H

변환 과정을 자세하게 알고 싶다면 논문 28p 부록 F와 affine_gird, Input Layer참고

 

 

T. Flexible layers

여기까지의 과정으로 equivariance가 많이 개선되었지만 아직도 일부 눈에 보이는 아티팩트가 남아 있었고, 최저해상도에서 여전히 필터의 감쇠가 부족했다고 한다. 낮은 단계의 레이어일수록 높은 주파수 콘텐츠가 필요없기 때문에 non-critical sampling을 했음에도 낮은 레이어에서는 더 강한 주파수 감쇠(더 낮은 cutoff, 더 큰 half-width)를 적용해야 했다.

 

그러면 낮은 레이어 일수록 높은 감쇠를, 최종 레이어에서는 낮은 감쇠를 적용시키고 싶은데 그 계수를 각각의 레이어마다 일일이 지정해놓지 말고 몇개의 인자만 입력하면 자동으로 모든 레이어의 cutoff, stopband가 지정되게 하면 어떨까?

sampling late s(=이미지 크기), 레이어 갯수 N, 임계 샘플링 레이어의 층수, start cutoff, start min_stopband, 임계 샘플링 레이어에서 cutoff와 stopband의 차이를 인자로 넘기면 모든 레이어의 f_c, f_t, f_h가 자동으로 정해진다.

 

위 문장에서 이상한 것을 느꼈는가? 바로 이미지의 크기와 레이어의 갯수 N을 따로 받는다는 것이다. 이미지의 크기에 따라 N을 정했던 기존의 모델과는 다르다. 그리고 더욱 신기한건, N의 고정된 선택이 여러 출력 해상도에서 일관되게 작동하고 학습률과 같은 다른 하이퍼파라미터가 더 예측 가능하게 동작한다는 것이다. 이미지의 최종 출력 크기와는 상관없이 N=14(본 논문에서)를 사용한다는 것. 이건 좀 신기하다.

 

 

R. Rotation Equivariance

EQ-R을 위한 두 가지 변경사항이 있다.

 

1. convolution

모든 레이어에서 3x3 conv를 1x1 conv로 교체하고 feature map을 두배로 늘려 감소된 용량을 보상한다.

모든 conv를 교체했기 때문에 픽셀간의 정보를 확산시키는 커널은 업샘플링과 다운샘플링 뿐이다.

 

2. downsampling

업샘플링에 새로운 필터를 사용했던 것과 마찬가지로 kaiser window 기반 + 방사형 대칭인 커스텀 필터를 이용하여 다운샘플링을 한다. 마지막 임계 샘플링 계층에서는 equivariance보다 최종 출력의 세부묘사를 생성하는 게 더 중요한 부분이기 때문에 적용하지 않는다.

 

이 두 변경을 적용하면 각 계층에 훈련 가능한 피라미터가 56% 더 적음에도 불구하고 FID, EQ-T, EQ-R 모두 향상된다.

이것도 신기하네..

 

또한 초기 훈련에서 가우시안 필터를 사용하여 판별자가 보는 모든 이미지를 흐리게하는 안정화 트릭을 사용한다.

이것은 판별자가 초기에 고주파수에 너무 많이 초점을 맞추는 것을 방지한다. 그리고 이 트릭이 없으면 생성기가 때때로 작은 지연으로 고주파수를 생성하는 방법을 학습하여 판별자의 작업을 trivializing하게 만들기 때문에 config R 모델이 조기 붕괴되는 경향이 있다.

 

 

 

Config B~R 

 

 

 

 

Result

stylegan2는 config B(+fourier features)

FFHQ 데이터셋에 대해 훈련시간 SG2-1106,  SG3T-1576(+42%),  SG3R-2248(+103%)

 

 

 

stylegan2 ablation study, dataset : FFHQ-U

 

 

 

 

 

 

영상들을 포함한 더 많은 결과는 https://nvlabs.github.io/stylegan3/

 

 

 

 

 

Limitations, discussion, and future work

이 논문에서는 generator만 고려했지만 discriminator 또한 equivariance하게 만들면 더 좋을 것 같다고 한다.

예를 들면, FFHQ 모델에서 얼굴이 회전 할 때 치아가 완전히 equivariance하게 움직이지 않는데,  판별자가 특정 위치에서 앞니를 보는 것에 익숙하기 때문으로 의심된다고 한다.

 

stylegan3은 훈련 데이터에 대한 암시적 가정을 포함하고 있다.

예를 들어, 모든 테두리에 검정색 마진이 있는 흑백 만화 데이터셋이 있다고 하자.

등장인물의 얼굴이 회전하거나 이동한다고 해서 테두리의 마진은 변하지 않는다. 하지만 이 모델에서는 지역 고정적인 특성을 전혀 고려하지 않기 때문에, 훈련에 어려움을 겪을 수도 있다.

 

texture sticking을 유발하지 않는 더 나은 방법으로 노이즈와 path length regularization을 추가할 수 있다면 흥미로울 것.

 

신호처리에서 안티-앨리어싱은 원래 톤 매핑 전에 처리해야 하지만 이 모델을 포함한 지금까지의 모델은 톤 매핑 후인, 그러니까 convolution을 적용할 수 있는 실제 이미지와 유사한 형태의 feature map에서 작동했다. 

 

최근에 많이 나오는 어텐션 기반 GAN는 equivariance의 개념과 상충함. 어텐션 자체가 픽셀들간의 위치 연관을 깨트리기 때문에 어텐션을 equivariance 하게 만들 수 있을 지는 미지수.