data에서 noise로 가는 forward path를 정의해서 좋은 학습 결과들을 보여왔지만, 여전히 어떤 경로를 선택해야할지는 확실하지 않았다.
예를 들어, 완전히 noise가 되지 않고 원래 data의 정보가 남아있는 상태로 학습하게 되면 test 환경에서의 분포에 차이가 생겨 아티팩트가 생길 수 있다.
여기서는 data와 noise를 직선으로 연결하는 Rectified Flow를 선택했다.
이전에는 Rectified Flow를 작은 사이즈가 class-condition에서만 실험했었는데 SD3에서는 large size + text condition으로 진행한다.
Cross attention 대신 서로 정보가 흐를 수 있는 아키텍처도 선택한다.
validation loss가 scale에 따라 확실히 변하며, 낮은 valid loss = evaluation metric도 좋더라
Simulation-free Training of Flows
기존의 neuralODE와는 달리, 시간 t에서의 vector field를 직접 regress하기 때문에 학습 시에 ODE solve가 필요하지 않다.
과 사이의 probability path p_t는 아래처럼 정의된다.
DDPM에서 노이즈가 추가되는 방식은 variance preserving
3.1 Tailored SNR Samplers for RF models
RF loss 학습과정에서 timestep은 [0, 1]에서 uniform하게 샘플링된다.
하지만 직관적으로, t=0 근처나 t=1 근처에서는 예측이 [0, 1] 중간보다는 쉽다.
그래서 중간의 타임스텝이 더 자주 샘플링되도록 함.
→ 어떻게?
Logit-Normal Sampling
실제로는 정규 분포에서 u를 샘플링한 다음 그걸 standard logistic function에 통과시켰다 → [0 ~ 1]
cosine schedule을 쓸 수도 있다.
t가 0과 1부근에서는 학습이 크게 어렵지 않은 이유는, 노이즈가 너무 심하면 어짜피 뭘 하든 크게 의미가 없고,
노이즈가 거의 없어도 그 다음 예측이 쉽기 때문이다.
어려운건 노이즈와 시그널(중요한 정보)가 비슷하게 섞여 있을 때(logSNR이 0부근)이다. 따라서 그 때의 timestep이 더 많이 샘플링 되도록 한다.
전부 실험해봤고, 결론적으로는 log norm (m=0, s=1)이 제일 좋았다고 함.
아키텍처
DiT 활용 - timestep과 class condition을 둘다 adaLN modulation 방식으로 넣는다.
SD3에서는 텍스트를 사용하고, 텍스트에 대한 전체 벡터 뿐만 아니라 더 디테일한 정보를 담고있는 전체 sequence embedding(token수, dim)도 사용한다. (전자는 CLIP, 후자는 T5-XXL)
이미지에는 2*2 patch 적용하고 flatten
복잡해 보이지만.. 우선 CLIP의 두 임베딩은 timestep embedding과 합쳐지고, adaLN module에 사용된다.
각 레이어마다 Norm → scale&shift → Layer → scale(gate)
text와 image가 같은 block에 들어간다. adaLN parameter와 FeedForward는 각자 계산되기 때문에 신경쓰지 않아도 된다.
block마다 residual connection과 QK norm(by RMS Norm)이 있다.
Q, K, V를 각각 계산한다음 concat한 뒤 self attention 연산을 하고, 다시 분리한다.(= no cross attention 근데 파라미터는 더 많을듯)
scaling은 depth와 head 수를 조절해서 실험한다. 각 헤드당 dimension은 64로 고정하기 때문에 사실상 model_dim과 같다.
FID는 CLIP feature로 계산
- timestep sampling 방식을 lognorm(0, 1)로 변경한게 항상 좋은 성능
- 기존의 LDM의 eps/Linear 보다 좋았던 성능은 rectified flow with modified timestep sampling밖에 없다.
- 타임스텝이 적을 때 성능이 대체적으로 rf가 좋았다.
5.2.1 Improved AudoEncoders
LDM은 pixel space가 아니라 latent space에서 처리해서 성능 + 효율을 보였다. H*W*3 → h*w*d
이 autoencoder의 reconstruction quality가 생성 모델의 upper bound를 결정한다.
여기서 latent의 dimension d를 높일수록 recon 퀄리티가 증가한다(당연한게 아닌가..?)
대신 d가 클수록 거기 충분히 정보를 담는 task는 어렵기 때문에, 모델 자체도 더 커져야 한다(AE).
SD3에서는 latent dim이 16인걸 만들어서 사용했다.
5.2.2 Improved Captions
DALL-E3에서는 synthetic caption을 쓰는게 모델을 향상시킨다고 했다. human caption은 정확하긴 하지만, 어쩔 수 없는 bias와 경향성을 가진다 - 덜 디테일하거나 중요한 object에 집중하는 등
그래서 CogVLM쓰고, original-synthetic 반씩 사용했다.
original만 쓰는 것과 실험 비교해봤을 때 확실히 좋았다.
5.2.3. IMPROVED TEXT-TO-IMAGE BACKBONES
DiT, Cross attention with DiT, MM-DiT, UViT를 비교해봤다.
MMDiT가 제일 좋았고, 그다음 Cross attention → UViT → DiT 순이었다.
실험은 적당히 수렴성을 보일 때 끊는 것 같다. 50만 step이나 하기는 하네
이렇게 다 하고나서 8B로 스케일을 올려서 학습함.
학습할 때 Image latent와 text embedding이 필요한데 둘다 frozen network를 쓴다 → 학습전에 전부 뽑아뒀다.
장점
- encoder를 GPU에 올려둘 필요가 없어 메모리가 감소하고 → 학습에 더 쓸 수 있다.
- 1 epoch 이상인 경우에는 시간이 linear하게 감소
단점
- dataloader에서 random augmentation이 불가능하다.
- 미리 embedding을 뽑아둔뒤 저장하면 용량이 크다 → half precision으로 저장함
5.3.2 High resolution finetuning
보통 256으로 먼저 학습 → mixed aspect ratio로 학습.
이 때 mixed precision을 쓰면 학습이 불안정하고 loss가 발산하는 경우가 있었다 → full precision을 쓰면 낫다. 근데 2배 더 비효율
ViT에서, 이런 경우 attention entropy가 커져서 = softmax 값이 다 비슷해져서 = 일 수 있다. attention값이 만약 너무 작다면, 2배가 차이나더라도 softmax 값은 비슷하다. (e0.0001≈1.0001, e0.0003≈1.0003)
그래서 Q, K를 연산 전에 RMSNorm with learnable scale parameter로 정규화한다.
이건 pretrained model이 QK Norm을 사용하지 않고 학습했더라도 적용가능하다.
일반적인 방법인지는 아직 모른다. 잘 사용해야함
Positional Encoding for Varying Aspect Ratios
학습한 비율이 아닌 비율로 inference하고 싶을 때 interpolation하면 비율이 깨진다. 대신 큰 사이즈의 PE Grid를 만들고 필요할 때 잘라서 쓰는 느낌으로 활용
Results
CLIP말고 T5를 써야 확실히 글씨 이미지가 잘 나옴. 보통 글씨는 많이 못하는 경우가 커서 가장 차이가 크다.
Text encoder 3개나 썼는데 46.8%로 drop하면서 학습.
제일 큰건 depth를 38까지..
Scaling result
항상 이렇게 논문을 보면 내가 믿고 사용해도 될만큼 좋은건지가 헷갈린다. 이 논문에서는 특히 MMDiT 구조가 그렇다.
Share article
Subscribe to our newsletter
