Stable Diffusion, Latent Diffusion Model은 논문에서 컨셉은 잘 설명하고 있으나, 명확한 구조를 알려주진 않는다.
그 구조 역시 인터넷에 좋은 자료가 없는 것 같아 코드를 이해하면서 구조를 만들어보았다. 전체적인 구조를 다 그릴 수 없어서 UNet에 집중해서 작성하였다. 여기서 예시는 클래스 조건부 모델이다.
모델의 시작은 Latent 공간에서 시작하며, Down/Mid/Up layer가 각각 2개라고 가정한다.
전체적인 구조는 다음과 같다.
Down Layer와 Up Layer는 각각 ResNet 블럭과 Self Attention 블럭으로 구성된다. 반면 Mid Layer의 경우 ResNet 블럭이 등장 한 후, Self Attention과 ResNet 블럭의 쌍으로 2개 layer로 구성된다. 공간 활용을 위해 이렇게 그렸지만, 엄밀히 따지면 Mid Layer의 경우
ResNet 블럭 + Mid Layer1 (Self Attention + ResNet) + Mid LAyer2 (Self Attention + ResNet)
로 이해할 수 있다.
또한 Down Layer의 ResNet Block에서 Up Layer의 ResNet Block의로 Skip Connection 동작을 하는데(UNet의 기본 구조), 이 때 ResNet Block 시작 지점의 텐서들이 이동하게 된다. 즉 예를 들면 Down Layer2의 ResNet Block 시작점의 텐서가 Up Layer1의 ResNet Block 시작점으로 Skip connection이 발생한다. 다르게 표현하면 Down Layer1의 Self Attention의 output이 Up Layer1의 ResNet Block에 더해진다고도 볼 수 있다.
하나의 ResNet Block 구조를 보면 다음과 같다.
우선 가장 처음에 있는 텐서는 마지막 Convolution을 통과 한 후 skip connection으로 더해지게 된다. 이 때, 경우에 따라서는 첫 텐서의 차원이 마지막 convolution을 통과한 후의 차원과 다른 경우가 발생할 수 있다. 이런 경우를 대비하기 위해서 채널 값을 맞추기 위한 1x1 convolution이 이어진 형태로 skip connection 동작을 한다.
또한 조건부 모델이기 때문에 모든 조건이 ResNet Block에 더해지게 되는데, 이 때 첫 번째 Counvolution 통과 후 조건이 주입된다. 조건의 경우 MNIST 데이터 같은 클래스 조건부 모델이라고 가정하면, 시간 정보(diffusion model의 timestep에 대한 정보)와 클래스 정보가 함께 임베딩된 벡터를 주입해준다. 앞선 그림과 같이 이 조건 정보는 모두 같은 값들이 각각의 layer에 주입된다.
마지막으로 전체 구조를 보면 다음과 같다.
여기서는 채널이 첫 레이어에서 128 → 256으로 변경된 후, 변화 없이 마지막까지 진행된 후 마지막 레이어에서 다시 원복된다. 따라서 변경되는 부분이 Down Layer1의 첫 번째 ResNet Block과 Up Layer2의 마지막 ResNet Block이 된다. 이 부분에 대해 Skip connection 수행 시 채널의 차원이 잘 맞는지 확인할 필요가 있다.
예시는 down layer 에서 채널을 처음에만 확장하고 계속 유지하다가, 마지막 up layer에서만 복원하는 형태를 가지고 있어서 위와 같은 그림으로 표현하였다. 하지만 지속적으로 channel을 확장하였다가(down layer에서), 이를 동일하게 복원하는 형태로도 만들 수 있다.
퍼가시는건 좋으나 꼭 출처를 남겨주세요.
댓글