[Page]
[arXiv](2024/03/05 version v1)
텍스트 생성 능력이 엄청나게 향상되었다. SD3
주요 변경 사항: Transformer backbone, Rectified flow
+ 자잘한 내용들 많음
Simulation-Free Training of Flows
Flow matching의 개념에 대해 설명하는 챕터. 확률분포가 아니라 벡터 필드를 회귀하는구나~ 정도만 알아도 된다.
Flow matching에 대한 ChatGPT의 직관적인 설명:
z의 probability path:
u는 역 ψ매핑에 대한 벡터 필드이다.
일반적인 flow matching 목표는 확률적으로 다루기 어렵다.
ϵ에 대한 조건부 벡터 필드를 이용한 flow matching은 결정적이므로 다루기 쉽다.
ψ, λ를 t에 대해 미분하여 대입하면 u를 다음과 같이 나타낼 수 있다.
이를 CFM loss에 대입하고 ϵ으로 reparameterization 하면
고전적인 diffusion simple loss의 형태로 작성할 수 있다.
Flow Trajectories
Flow matching의 다양한 변형들.
논문에서는 여러 가지 변형을 소개하고 있지만 가장 성능이 좋은 Rectified Flow만 소개함.
Rectified Flow
RF는 다음과 같이 forward process를 매우 단순한 직선 경로로 정의하며
두 분포 간의 직선 방향을 네트워크 vΘ를 통해 직접 예측한다.
Tailored SNR Samplers for RF models
t = 0에서 최적 예측은 π1의 평균으로의 방향, t = 1에서는 π0의 평균으로의 방향을 예측하면 되므로 쉽다. (RF는 역방향 샘플링도 가능하다.)
어려운 것은 중간의 t에 대한 예측이다. 따라서 훈련 시 중간의 t를 더 많이 샘플링하도록 해야 하며 이는 손실의 가중치를 변경하는 것으로 달성될 수 있다.
이 부분 또한 논문에서는 π(t)의 여러 가지 변형을 소개하고 있지만 하나만 소개함.
Logit-Normal Sampling
실제로 m, s를 통해 정규 분포의 평균과 분산을 조절하여 timestep에 대한 가중치를 조절한다.
Text-to-Image Architecture
Diffusion transformer인 DiT를 기반으로 하며 전체 아키텍처는 아래 그림과 같다.
특이한 점은 텍스트와 이미지에 대해 각각 별도의 parameter를 통과하며 attention 시에만 합쳐진다는 것이다.
Experiments에도 내용 완전 많으니까 여기까지만 보고 나가지 마세요.
Experiments
Improving Rectified Flows
Rectified Flow / lognorm(0.00, 1.00)의 조합이 가장 성능이 좋았다.
Improving Modality Specific Representations
Improved Autoencoders
VAE의 잠재 채널 수를 늘리면 성능이 크게 향상된다.
Improved Captions
최첨단 MLLM인 CogVLM을 통해 생성된 캡션을 사용하며 해당 모델에 없는 개념을 잊지 않기 위해 원본 캡션을 1:1 비율로 사용.
Improved Text-to-Image Backbones
본문에서 제안한 MMDiT (2sets)와 다른 아키텍처 간 비교.
3 sets MMDiT는 추가되는 계산량에 비해 성능 이득이 크지 않으므로 채택하지 않았다.
Training at Scale
Data Preprocessing
불건전한 컨텐츠, 시각 및 의미적으로 중복되는 컨텐츠 필터링.
전체 image, text dataset에 대한 임베딩을 미리 계산.
장점:
- 총 훈련 시간이 줄어듦
- VAE 인코더를 사용하지 않으므로 메모리 요구량이 낮아짐
단점:
- 각 데이터에 대한 무작위 증강이 불가능
- 추가 저장 비용이 발생하고 로딩 시간이 길어짐
Finetuning on High Resolutions
QK-Normalization
모든 모델을 256×256 해상도에서 사전 훈련한 뒤 다양한 종횡비의 더 높은 해상도로 fine-tuning 한다.
이 때 혼합 정밀도 훈련이 불안정해진다. 완전 정밀도를 사용하면 성능이 거의 2배 가까이 떨어졌다고 한다. (완전 정밀도가 더 성능 좋은 거 아니었어?)
한 논문에서 이러한 문제의 원인이 attention entropy의 폭증임을 관찰했고 attention 직전 Q, K의 정규화를 제안했다. 이를 따라 스케일링 가능한 RMSNorm을 통해 Q, K를 정규화한다.
Positional Encodings for Varying Aspect Ratios
낮은 해상도에서 사전 훈련 후 다양한 종횡비의 높은 해상도에 적응하기 위해 position grid의 확장과 보간의 적절한 조합으로 주파수를 임베딩한다. (왜 RoPE 안썼지? 요즘 RoPE로 해상도 확장한 논문이 많이 나오는데 RoPE가 더 좋을 듯 ㅇㅇ)
Resolution-dependent shifting of timestep schedules
고해상도의 경우 저해상도와 동일한 수준으로 이미지를 파괴하려면 더 많은 노이즈를 추가해야 하므로 timestep schedule을 변경.
언급한 기술들과 함께 1024×1024 해상도까지 fine-tuning 후 LoRA를 추가하여 Diffusion-DPO를 수행한다.
Results
당연하지만 큰 모델일수록 성능이 좋다.
또한 큰 모델의 rectified flow가 더 직선적이다.
Flexible Text Encoders
3개의 text encoder는 각각의 dropout 확률이 46%이다.
가장 많은 메모리를 차지하는 T5를 제거해도 미적 품질에는 거의 영향이 없지만 정확한 텍스트 생성 능력은 떨어짐.
