[딥러닝]Stable Diffusion에서 너무 어두운 이미지를 학습하지 못하는 이유와 해결 방법(noise offset에 대하여)

동기

Stable Diffusion, Latent Diffusion Model 관련 논문들과 정보글을 찾아다니다가 어두운 이미지를 출력하지 못하고, 학습하지도 못한다는 흥미로운 글을 발견해서 읽어보고 그 글에 나온 원인과 해결 방법까지 알아보게 되었다.

 

 

문제

Stable Diffusion을 사용하여 이미지를 출력하다 보면, midnight, dark를 써줘도 내가 생각한것 만큼 어두운 이미지가 나오지 않는다. 학습 이미지가 부족해서 그런가라는 생각을 가지고 매우 어두운 그림의 이미지를 모델에 학습시켜도 더 밝아져서 나오게 된다.

실험결과,심지어 완전히 검은색인 이미지를 학습시키고 출력해도 아래 그림처럼 검은색 이미지가 나오지 않는다고 하며, 출력되는 이미지 값의 평균이 0.5라고 한다.

A solid black image를 출력한 결과

문제 원인

이 문제는 기존 Stable Diffusion 의 노이즈 프로세스에 있었다. 이전 글에서 설명했듯이, Latent Diffusion Model은 VAE에 의하여 Latent로 변환된 이미지에 여러차례 가우시안 노이즈*를 주고 원본 이미지를 복원하는 것을 반복하여 학습한다.

(* 가우시안 노이즈 : IID 노이즈라고 불리는 노이즈로 정규분포의 형태를 띄는 노이즈이다, independently and identically distributed noise)

 

Diffusion Model에서 원본 이미지에 노이즈를 추가하는 과정을 forward process라고 하고 이미지에서 노이즈를 제거하여 복원하는 과정을 reverse process 라고 한다  그런데 이 forward process는 계속 노이즈를 조금씩 추가하는 과정이기 때문에 원본 이미지의 잔재를 완벽히 없앨 수 없다는 문제가 있다. 또한, Latent vector로 추출된 특징중 long-wavelength 특징은 short-wavelength 특징에 비해 노이즈 프로세스에서 가장 오래 남아있는 특성을 가진다.

 

여기서 long wavelength, short wavelength 특징이란, CNN 구조의 VAE로 3*512*512(3은 RGB 채널)원본 이미지 특징을 추출하면 4*64*64의 latent vector로 생성된다. 이 때 CNN이 여러겹 쌓은 구조의 VAE의 각 레이어에서는 이미지의 특징이 추출된다. long-wavelength 특징은 보통 깊은 레이어에서 추출되며, 이미지의 전체적인 구조나 물체의 일부분 같은 큰 특징을 포함한다. 반대로 short-wavelength 특징은 질감이나 테두리, 모서리 같은 작은 특징을 포함하며 보통 얕은 레이어에서 추출된다. 

 

이미지의 전체적인 구조는 작은 특징보다 노이즈로 가려지기 어려운것은 어떻게 생각해보면 당연하다.

 

실제로도 원본 이미지를 latent로 변환하고 노이즈를 추가하면서 비교하는 아래 영상을 보면 쉽게 알 수 있다. 가장 왼쪽의 long-wavelength 특징이 오른쪽의 short-wavelength 특징에 비하여 거의 변하지 않는 것을 볼수 있다.

결과적으로 우리는 기존 이미지 정보가 조금 남아있는 상태(특히, long-wavelength 특징)에서만 복원하는 것만 학습하고 완전한 노이즈에서 복원하는 것을 학습하지 못하게 된다. 이런 문제로 인해서 같은 노이즈 시드에서 다른 프롬프트를 입력했을때 전체적으로 비슷한 구조를 가지는 이미지가 생성되게 하며, 이미지를 생성 할 때의 reverse process에서는 학습할때와 달리 완전한 noise에서 이미지를 생성하기 때문에 우리가 학습한 것을 완벽히 복원할 수 없게 된다.

 

이 문제는 원본 정보의 dimension이 커질수록 더욱 심화된다. 가우시안 노이즈는 정규분포의 형태이기 때문에 표준편차 값에 따라서 노이즈의 세기가 정해진다.

 

참고: 표준 편차는 $\sigma= \sqrt{ \frac{{(x_1 -m)}^{2} + {(x_2 -m)}^{2} + \cdot \cdot \cdot + {(x_N -m)}^{2} }{N}}$이다.

$\sigma$가 클 수록 양극단의 값이 더 많아진다. 따라서 노이즈의 세기도 강해진다.

longest-wavelength 특징은 이미지 전체의 평균값이라고 할 수 있는데 이는 noise가 표준 정규분포의 형태(평균이 0, 표준편차는 1)이므로 이는 영원히 보존된다. 물론 feature가 전체 이미지의 평균값으로 구해지는 않을 것이므로 극단적인 예시이다.

 

반대로 shortest-wavelength 특징에 noise의 평균은 정규분포의 표본분포의 평균은 모집단의 표본분포와 같으므로 기존과 같다. 분산의 경우에는 $\frac{\sigma ^{2}}{N}$(여기서 $N$은 표본의 개수)이다. 결론적으로 표본이 작을수록 표준편차가 커지므로 노이즈의 영향=표준편차가 이론상 최대 $\sqrt{N}$배 커진다. 

 

Stable diffusion 모델의 경우에는 3*512*512의 이미지를 4*64*64의 latent vector로 변환하여 처리하는데, 이 때의 dimension은  4*64*64=16384이다. 따라서 노이즈의 영향(노이즈의 표준편차)는 최대 128배 (4채널 계산에 미포함시 64배) 차이가 발생하게 된다.

해결 방법

그래서 long-wavelength가 계속 남아있는 문제를 해결하기 위하여 간단한 방법을 제안했다.

바로 새로운 특이한 노이즈인 offset noise를 추가하는 것이다.

# 기존 noise
noise = torch.randn_like(latents)

# 개선된 noise (noise offset 추가)
noise = torch.randn_like(latents) + 0.1 * torch.randn(latents.shape[0], latents.shape[1], 1, 1)

일단 함수와 변수의 의미를 설명해보겠다. torch.randn_like(latents)는 기존 latent와 같은 크기의 가우시안 정규분포 노이즈를 생성하는 함수이다. torch.randn() 함수도 비슷하지만 생성하는 노이즈의 크기를 수동으로 정할 수 있다. latents.shape[0]는 한번에 연산하는 latent의 개수인 배치사이즈를 의미하며 latents.shape[0]는 latent vector의 채널 수인 4이다.

여기까지는 기존 노이즈와 동일하다. 하지만 마지막 부분이 달라진다. 원래 torch.rand_like(latent) 함수는 torch.randn(latents.shape[0], latents.shape[1], 64, 64) 와 같지만 개선된 노이즈에서는 뒤에 64,64가 아니라 1,1 의 랜덤 노이즈를 생성하고 0.1을 곱하여 더함을 알 수 있다. 이럴 경우 이미지 당 노이즈의 dimension이 4 밖에 되지 않아 longest-wavelength와 shortest-wavelength에서 이론상 최대 2배까지 노이즈의 차이(4채널 포함 안하면 차이 없음)가 발생한다.

 

이는 위에서 언급했던 노이즈의 영향이 128배(채널 미포함 64배)까지 차이나는 결과를 줄이기 위한 방법으로, 쉽게 설명하면 각 픽셀별로 계산하는 것이 아니라 채널별로 한번에 같은 노이즈를 추가해버리는 것이다. 하지만 이걸 그대로 적용하면 노이즈가 너무 커져 대부분의 정보가 손실되기 때문에 0.1을 곱하여 노이즈 강도를 낮춘다. 단, 이 값은 정확한 실험을 통해 도출된 값은 아니지만 이 값에서 잘 작동했다고 하니 참고하길 바란다.

 

결과

원 글에 따르면, 이를 실제로 적용한 결과 기존과 다르게 너무 밝거나 어두운 그림도 학습할 수 있으며, 실제로도 출력이 가능했다.

Top right: A dark alleyway in a rainstorm (0.032); Top left: Monochrome line-art logo on a white background (0.974); Bottom left: Snow-covered ski slope on a sunny day (0.858); Bottom right: A town square lit only by torchlight. (0.031)

 

A starry sky before and after offset noise.

Superheroes fighting plant monsters in a dark alley before and after offset noise.

또한, 앞으로의 학습에서 개선된 noise를 사용할 것을 적극 권장했다.

 

개인적인 생각으로는 stable diffusion을 학습할때 문제가 있는 기존 noise를 사용해서 상당히 퀄리티에 나쁜 영향을 미쳤을것으로 생각된다. 앞으로의 모델은 새로운 offset noise를 적용하면 훨씬 개선된 다양한 이미지를 출력할 수 있을것으로 기대된다.

그런데 가우시안 노이즈로 인해 발생하는 근본적인 원인을 해결하는 것은 아니기에 새로운 diffusion process가 나와야 이 문제를 완전히 해결하는 것이 가능할것으로 보인다.

 

참고자료

https://www.crosslabs.org//blog/diffusion-with-offset-noise

https://theailearner.com/2019/05/05/understanding-frequency-in-images/