Stable Diffusion 3 논문 읽어보려고 한다.
CVPR 2024 이다. 이걸 지금 읽다니. 나왔을 때 읽어볼 걸!
https://arxiv.org/pdf/2403.03206
이번에도 취사선택해서 읽어봤다.
0. Abstract
주요 성과
1. Rectified Flow 개선
기존 diffusion 대신 Rectified Flow 기반 생성 방식을 사용한다.
학습 과정에서 시각적으로 관련이 있는 것들에 집중하는 noise sampling 기법 제안했기 때문에 고해상도 text-to-image 생성에서 기존 diffusion 방식보다 우수한 성능을 보인다.
2. MMDiT 아키텍처 제안
text, image 각각이 다른 weight를 사용하면서도 transformer 안에서는 양방향으로 정보를 주고받을 수 있는 구조이다.
텍스트 이해력, 글자 생성 품질, 사람 선호도가 높아졌다.
그리고 이렇게 하니 val loss가 낮아질 수록 성능이 높아지는 고런 관련성이 높게 나왔다! 기존에는 그러지 못했음
제일 큰 모델 성능이 SOTA보다 좋다.
3. Flow Trajectories
기존 생성모델들이 사용하는 다양한 flow trajectory들을 같은 수학 틀 안에서 비교해봤다.
Recified Flow
Rectified Flow(RFs)는 원본 이미지랑 노이즈를 선형보간하는 flow 방식이다. 수식으로 표현하면 아래와 같다.
$x_0$: 원본 이미지, $\epsilon$: 랜덤 노이즈
그리고 *CFM loss를 기반으로 한다. 가중치는 $w_t^{RF} = \frac{t}{1-t}$이고 모델이 velocity $v_{\theta}$를 직접 예측한다.
* CFM loss: conditional flow matching. 현재 점(z_t)에서 얼마나 이동해야 하는지(v)를 학습하는 loss
EDM
Euclidated Diffusion Models은 원본 이미지에 노이즈를 점점 많이 더하는 diffusion 방식이다. 수식으로 표현하면 아래와 같다.
(뭐 수식 어쩌구저쩌구 많이 나오긴 하는데 복잡해서 생략하겠다. 나중에 돌아와서 다시 봐야겠다.)
아무튼 노이즈를 직접 샘플링하지 않고 먼저 정규분포에서 $\lambda_t$를 뽑고 $b_t = e^{\lambda_t}$로 변환해서 샘플링 함!
그리고 가중치 $w_t$가 들어간 MSE loss를 사용한다.
Cosine
Cosine 방식은 아래와 같은 수식이다.
t=0일 때는 원본이미지 그대로이고 t=1일 때는 완전 노이즈가 된다.
직선이 아니라 곡선이라는 느낌만 가져가면 될 듯
$\epsilon$-prediction을 사용하면 $w_t = sech(\frac{\lambda_t}{2})$가 되고
$v$-prediction을 사용하면 $w_t = e^{\frac{-\lambda_t}{2}}$가 된다. (이것도 왜 이렇게 되는지는 나중에 다시보겠다)
(LDM-)Linear
LDM은 DDPM을 변형해서 사용한다. 둘 다 $a_t^2$ + $b_t^2$ = 1이다.
DDPM은 timestep마다 $\beta_t$를 선형 증가시키는데 LDM은 $\sqrt{\beta_t}$를 선형 증가시키고 제곱을 한다. (얘도 나중에)
그래서 여기서 하고 싶은 말이 무엇이냐!
-> RF가 trajectory 때문에 좋은 게 아니라 timestep 샘플링 방법이나 loss weight 때문에 좋은 것일수도 있기에
다양한 방법을 같은 수학 틀(z_t = a_t x_0 + b_t \epsilon)에서 비교해는 것임
4. Text-to-Image Architecture
Text를 조건으로 image를 샘플링하기 위해서는 text랑 image 정보를 모두 사용해야 한다.
이미지는 VAE를 사용해서 latent representation으로 변환하고, 전체 구조는 LDM(Latent diffusion model)을 따른다.
텍스트는 pretrained, frozen된 text encoder를 사용해서 text embedding으로 변환한다.
이것들을 diffusion backbone에 입력해서 latent 공간에서 학습시키는 것이다.
Multimodal Diffusion Backbone
DiT 아키텍처를 기반으로 한다. DiT는 조건부 이미지 생성만 하기 때문에 diffusion timestep하고 클래스 라벨만 사용한다.
이거랑 비슷하게 우리도 timestep하고 cvec(한줄요약본)을 사용한다. 근데 cvec만 사용하면 너무 coarse하기 때문에 cctxt(모든 토큰)도 같이 사용한다.
e.g. A red sports car parked in front of a blue house에서의 cvec과 cctxt는 아래와 같다.
cvec: car and a house
cctxt: A, red, sports, car, parked, in, front, of, a, blue, house
그리고 text랑 image의 임베딩으로 구성된 하나의 시퀀스를 만든다.
Latent image에 positional encoding을 추가한 뒤 2x2로 나눠서 flatten시킨 다음 cctxt랑 연결하는 것이다.
이 다음에는 DiT처럼 attention하고 MLP를 반복한다.
Text랑 image 임베딩이 본질적으로 꽤 다르기 때문에 각각 서로 다른 가중치를 사용한다.
이거는 text랑 image에 각각 다른 transformer을 두는 것과 비슷한 상황이다.
근데 attention 연산에서는 text랑 image를 같이 사용한다.
따라서 text랑 image가 자기 특징을 유지하면서도 서로를 참고할 수 있게 된다.
🌟 여기서 짚고 넘어가보는 cross attention하고 MMDiT의 차이!!
Cross attention은 text -> image 단방향임. image가 text를 보면서 생성되지만 text는 image를 볼 수 없음
MMDiT는 text <-> image 양방향임. 서로서로 참고할 수 있음
그래서 cross attention이 MMDiT에 비해서 계산량이 적고 구현이 쉽지만 MMDiT가 cross attention에 비해서 성능이 좋음
구조는 아래와 같다.
구조 하나하나 설명해줬더라면 좋았을텐데 아쉽게도 없어서 지금은 한번에 이해하기 어려울 것 같고 계속해서 보면서 이해해야 할 것 같다.
하나하나 보자면 먼저 input으로 caption, timestep, noised latent가 들어간다.
Timestep
지금 이미지가 얼마나 노이즈 상태인지 알려주는 숫자이다.
이게 sinusodial encoding 돼서 MLP layer까지 통과한다. ··········· ①
Caption
CLIP-L/14, CLIP-G/14는 각각 이미지-텍스트 모델이고 G가 L보다 조금 더 큰 버전이다. T5 XXL은 더 디테일한 텍스트를 잡는다.
Pooled 사각형은 은 문장 전체가 요약된 벡터들이고
저 CLIP-G/14 토큰 77개, CLIP-L/14 토큰 77개, T5 XXL끼리 합쳐진 사각형은 세부 텍스트 정보이다.
Pooled 사각형은 ①이랑 합쳐져서 벡터 y가 되고 ··········· ②
CLIP-G/14, CLIP-L/14, T5 XXL끼리 합쳐진 사각형은 linear layer로 크기 통일돼서 합쳐져서 행렬 c가 된다. ··········· ③
Nosie latent
VAE Encoder로 압축된 latent를 transformer에 들어갈 수 있도록 patch로 나눠준다.
이거 linear layer로 차원 맞춰준 다음 positional embedding 해줘서 행렬 x로 만든다. ··········· ④
이 ②, ③, ④가 MMDiT 안에서 쿵짝쿵짝 어쩌구저쩌구 되어서 무언가가 나온다.
얘가 modulation 거치면서 y(현재 timestep + 현재 프롬프트)에 의해서 한번 더 보정된다.
그러고 linear layer로 차원 조정되고 unpatching 되어서 output으로 나온다.
Output은 이미지가 아니다. Latent를 어떻게 수정해야하는지에 대한 속도값 v_{\theta}이다!
이 속도값을 가지고 latent를 수정한 다음 VAE Decoder로 디코딩해서 이미지를 만들어내는 거임!!
6. Conclusion
본 연구는Text-to-image rectified flow 모델의 스케일링 분석을 했고
새로운 timestep sampling 방법과 MMDiT 구조를 제안했다.
그리고 모델 규모를 확장했다. 8B이다! (불과 2년 전까지만해도 8B가 큰 모델이었구나)
아직 더 연구할 게 많다는 말과 함께 끝~
---
아 모델 구조 세세하게 모르겠다. SD1 보면서 뭐가 어떻게 달라졌는지 확인해봐야겠다.
이미지 생성 도메인 생각보다 재밌다. 연구할 task가 진짜 많아서 고민할 수 있는 거리가 많아보인다.
잘 해봐야겠다 화이팅