Latent Diffusion Model과 Cross-Attention을 구현 관점에서 읽기

Stable Diffusion의 구현 내용을 볼 때는, '이미지를 노이즈에서 생성한다'는 설명만으로는 모듈 간 연결 구조를 파악하기 어렵습니다. Latent Diffusion Model (LDM: 잠재 공간에서 확산 과정을 다루는 생성 모델)은 처리를 다음 3 영역으로 나눕니다.

pixel image
-> VAE Encoder
-> latent
...

본 글에서는 High-Resolution Image Synthesis with Latent Diffusion Models와 Stable Diffusion v1의 공식 configuration을 기반으로, 텐서 shape, noise prediction objective, Cross-Attention, scheduler의 경계를 구현 관점에서 정리합니다.

Diffusion의 forward/reverse process나 Steps・Guidance Scale의 의미부터 순차적으로 확인하고 싶다면, Stable Diffusion의 기본: Latent Diffusion이란?에 다이어그램을 포함하여 요약했습니다.

LDM은 denoiser의 계산 위치를 바꾼다

Pixel-space Diffusion과 Latent Diffusion의 핵심적인 차이점은 noise를 예측하는 U-Net이 어떤 표현(representation)을 처리하느냐입니다.

관점	Pixel-space Diffusion	Latent Diffusion
Denoising 대상	RGB 등의 pixel tensor	Autoencoder가 만든 latent tensor
...
512×512 pixel gridと64×64 latent gridで反復denoising範囲を比較

LDM은 denoising step 수 자체를 0으로 만드는 방식이 아닙니다. 반복 횟수가 많은 U-Net 처리를 pixel보다 작은 latent grid로 옮긴 것입니다. VAE (Variational Autoencoder: 이미지와 잠재 표현을 상호 변환하는 모델)의 encode/decode는 남아있기 때문에, '입력 요소 수가 줄어든 비율'과 'end-to-end 속도 향상률'은 같지 않습니다.

512×512 이미지와 4×64×64 latent를 분리하다

Stable Diffusion v1에서는 512×512의 RGB 이미지를 예로 들면, VAE 전후의 shape는 다음과 같이 표현됩니다.

각 공간 축은 8분의 1, 공간 위치 수는 64분의 1입니다. 단순한 스칼라 수(scalar number)는 다음 비율을 가집니다.

48은 입력 텐서의 스칼라 수 비율입니다. U-Net 내부의 channel 폭, down/up block, Attention, VAE, sampling step을 포함하는 FLOPs 비율이 아닙니다. 구현 프로파일링(profile)에서는 적어도 다음을 분리하여 측정합니다.

• Text Encoder의 실행 시간
• 1 step당 U-Net 실행 시간
• Scheduler update
• VAE decode
• Peak VRAM과 activation memory

또한, 4 channels,
downsampling factor 8,
latent scaling factor 0.18215는 Stable Diffusion v1의 configuration에 연결된 값입니다. LDM 전반의 상수가 아닙니다.

Random timestep의 training sample을 만들기

Denoising Diffusion Probabilistic Models에서 사용되는 forward process에서는 clean sample

LDM에서는 같은 작업을 VAE가 만든 latent

import logging
import torch
from torch import Tensor
...

Training sample마다

LDM objective를 읽다

Text 조건

각 기호를 모듈에 대응시킵니다.

| 기호 | 구현상의 대상 |
|---|---|---|
| 학습 이미지 | |
| ... |

식만 보면 VAE도 U-Net과 동시에 업데이트하는 것처럼 보이지만, LDM은 pre-trained autoencoder의 latent space 위에서 Diffusion Model을 학습하는 구성입니다. 어떤 파라미터를 freeze 할지는 training recipe를 통해 확인합니다.

구현 시에는 objective의 표기에도 주의가 필요합니다. Checkpoint나 라이브러리에 따라 U-Net의 target이 noise prediction_type이 학습(training) 시와 스케줄러 측에서 일치하지 않으면, 텐서 모양(tensor shape)이 올바르더라도 업데이트식의 의미가 달라집니다.

Cross-Attention의 Q는 이미지, K/V는 텍스트

LDM 논문은 컨디셔닝(conditioning)을 U-Net에 넣기 때문에 Cross-Attention (서로 다른 시퀀스 간 참조 관계를 계산하는 Attention)을 사용합니다. Text Encoder의 출력을

context_dim은 Text Encoder 출력과 Cross-Attention projection의 인터페이스입니다. Stable Diffusion v1의 공식 config에서는 context_dim: 768이 지정되어 있습니다. 다른 Text Encoder로 교체할 경우, 단순히 임베딩 폭(embedding width)을 맞추는 것뿐만 아니라, 토크나이저(tokenizer), 시퀀스 길이(sequence length), 학습된 projection의 대응까지 확인해야 합니다.

Training과 Sampling에서 U-Net 호출 방식이 달라짐

학습(Training)에서는 깨끗한 latent부터 임의의 [빈칸]을

Sampling update를 추상화하면 다음 형태가 됩니다.

Step은 DDPM, DDIM, Euler 계열 등의 스케줄러(scheduler)에 따라 다릅니다. UI에 표시되는 sampling steps는 training timesteps와 같지 않습니다. 예를 들어 1000단계의 noise schedule로 학습한 모델을, inference scheduler가 선택한 수십 개의 timestep으로 샘플링하는 구성이 있습니다.

Steps를 늘린 효과는 스케줄러, 모델, 해상도(resolution), 가이던스(guidance)와의 조합으로 측정합니다.

• Seed: 초기 latent 노이즈를 결정합니다.
• Steps: U-Net과 scheduler 업데이트의 반복 횟수를 결정합니다.
• Guidance scale: Conditional 예측과 unconditional 예측의 차이를 얼마나 강하게 만들지 결정합니다.

같은 seed라도, scheduler, precision, library version, hardware별로 비결정적인 연산이 달라지면 완전히 일치하지 않을 수 있습니다.

체크포인트를 컴포넌트 계약(component contract)으로 검증하기

Stable Diffusion 파이프라인은 여러 개의 체크포인트와 구성(configuration)으로 이루어져 있습니다. 파일이 로드될 수 있다는 것과, 학습 시의 조합을 재현할 수 있는 것은 별개의 문제입니다.

input image
-> VAE latent shape / dtype / mean / std
-> noisy latent shape
...

게다가, 모델 카드(model card)나 리포지토리의 라이선스를 컴포넌트별로 확인합니다. CompVis의 Stable Diffusion v1 리포지토리는 코드, 구성, 가중치(weight)의 이용 조건을 각각 다르게 안내하고 있습니다. 리포지토리 코드의 라이선스만 보고, 모델 가중치나 학습 데이터의 조건까지 같다고 판단하지 않도록 합니다.

요약

Latent Diffusion Model은 Diffusion의 반복 계산을 픽셀 공간에서 autoencoder의 latent 공간으로 옮깁니다. Stable Diffusion을 구현 관점에서 읽으려면 다음 경계가 중요합니다.

• VAE가 픽셀과 latent를 변환하고, U-Net은 latent 위에서 노이즈를 예측한다.
• Stable Diffusion v1의 4×64×64나 0.18215를 LDM 일반의 상수로 취급하지 않는다 - 학습 시에는 무작위한 1 timestep을 직접 만들고, 샘플링에서는 U-Net을 반복한다.
• Cross-Attention에서는 이미지 feature가 Q(Query), text embeddings가 K/V(Key/Value)가 된다.
• Scheduler, prediction type, latent scaling을 체크포인트 호환성의 일부로 검사한다.

'Stable Diffusion model'이라는 하나의 파일로 보는 것이 아니라, VAE・U-Net・Text Encoder・Scheduler 간의 계약으로 추적하면, shape은 맞는데 이미지가 무너지는 문제를 구분하기 쉬워집니다.

참고 자료

Discussion