SANA: EFFICIENT HIGH-RESOLUTION IMAGE SYNTHESIS WITH LINEAR DIFFUSION TRANSFORMERS 논문 리뷰
SANA by NVIDIA paper review 논문 리뷰
Oct 24, 2024
요약
Nvidia에서 낸 Text-to-image 모델, 4096의 resolution의 이미지 생성.
결과물부터
성과
- 매우 빠른 속도
4096 기준 100배, 1024는 25배 빠른 속도
- 매우 작은 모델
위에서 비교한 Flux-dev는 12B, SANA는 1.6B랑 0.6B
- 속도, 크기 대비 좋은 성능
일반적으로 이미지 생성에서 보는 5가지 metric
FID 좋고, CLIP은 4,5위
GenEval - Object에 집중, DPG - 긴 프롬프트
제가 생각하는 contribution
32x compression Auto Encoder
기존의 압축에 영향을 주는 값들이라고하면
- Down-sample factor
- Channel
- DiT의 Patch size
결론적으로는 SDXL, Flux 등에서 보통 사용하는 F8C4P2보다 2배 더 압축하는 AE 구현 : F32C32P1
DiT에 Patchify를 시키지 않고, 채널은 32로 하는게 제일 좋았다.
Decoder-only LLM Text Encoder
기존에 사용하는 Text encoder : T5, CLIP
SANA에서 사용 : Gemma-2
- image-text alignment와 reasoning 능력이 더 좋고, 인코딩 속도도 동일 성능을 보이는 T5-XXL대비 6배 더 빠르다.
- prompt를 LLM에 바로 태워서 마지막 layer의 feature를 썼더니 불안정했다(loss NaN)
그래서 봤더니 T5의 분산이 훨씬(10, 100배?) 작았다 → feature를 분산 1로 normalize하고, 학습가능한 scaling factor를 0.01로 초기화해서 사용
- prompt를 바로 LLM에 태우는게 아니라, instruction을 설계한다.
from Exploring the role of large language models in prompt encoding for diffusion models
이 프롬프트에 맞는 이미지를 자세히(위치, 컬러, 구도 등) 묘사해봐 → 출력값을 사용
Linear Attention : →
2020년 논문(Transformers are rnns: Fast autoregressive transformers with linear attention)에서 제시한 개념
아주 간단하게 이해한 바로는..
- 기존에는 attention 연산을 위해 Q와 K를 matmul하고 softmax를 적용해서 weight score를 구한다음 V에 곱한다.
여기서 Q, K 연산이 time complexity N제곱
- 근데 꼭 softmax를 적용해야하나? 가중합인게 중요한거지
→ Q와 K에 각각 비선형 함수를 적용하고(ReLU) 둘을 곱하는 형태로 분해
- 그럼 ReLU(Q) ReLU(K) V가 되는데, 이럼 matrix 결합법칙 때문에 뒤의 K, V 계산부터 해도 된다. 그리고 이건 dot product, O(N)
그거랑 앞의 Q와의 연산도 O(N)
최종적으로 O(N)이 된다. 2N인건 안중요함?
약 1.7배 속도가 빨라졌다.
궁금한 것 : 이게 그렇게 좋다면 왜 진작에 다른 연구들에서 안썼을까? 분명히 이유는 있을 것.
여기서 말하는건 결국 softmax가 없어짐으로 인해 sub-optimal 성능을 보이고 최종적으로 성능이 비슷하더라도 수렴 속도가 오래 걸린다고 함.
→ 대체하더라도 토큰간의 상관관계를 파악하는 능력이 softmax보다는 떨어진다.
Mix-FFN
→ 그렇다면 FeedForward 연산에서, linear layer말고 토큰끼리 서로의 정보를 참고할 수 있게 설계해서 그 단점을 보완하자 : Conv layer 사용
위의 문제를 해결하기위해 Mix-FFN 사용. linear layer대신에 3*3 convolution을 Gated로 적용한다. local 정보를 파악하는 능력이 올라가니까.
→ 근데 이렇게 Conv를 적용해서 토큰끼리 상관관계도 학습할 수 있게하니, position 정보도 학습해서 positional encoding을 쓰지 않아도 성능하락이 없었다.
원래도 여기에 관한 연구가 있음(How much position information do convolutional neural networks encode?)
예상해볼 수 있는 단점이라고하면 길이가 계속해서 달라질 때 성능이 유지되지 못할 것 같다는 점인데, 최근의 LLM쪽 연구(The impact of positional encoding on length generalization in transformers)에 따르면 generalization ability도 좋다고 합니다.
Flow Matching objective, Flow-DPM-Solver sampling
Flow-based training, sampling을 사용해서 속도향상(SD3에서 차용)
v-objective, EDM도 테스트 해봤다고 하는듯
짧게 이해한 바로는
두개만 비교 : DDPM의 noise prediction과 Rectified flow - flow matching
- DDPM
를 사용한 공식에 따라서 timestep t에서의 noise-perturbed x_t, t, x를 사용해서 이 때 더해진 noise를 예측
즉, 모델의 출력값(예측하고자 하는 것)은 timestep t에서 더해진 noise
- Rectified flow
이 원본, 이 noise (t는 0~1)
noise 가 이 되기위해서 어느방향으로 얼마나 이동해야하는지를 나타내는 벡터 그 자체 = flow?
그리고 이 flow를 직선이라고 단순하게 모델링한다. 즉 target은 이고, t와 무관하게 구해진다.
는, 당연하게도 x0과 1을 사용하여 순수한 interpolation 계산으로 얻는 값.
= t + (1-t)
모델은 (, t, condition)을 입력으로 받아서 위 target을 예측하도록 학습된다.
학습은 위의 objective로 하고,
샘플링은 Flow-DPM-Solver를 사용
Flow-DPM-Solver를 사용하니 sampling step이 14~20이어도 성능이 좋았다. 기존의 Flow-Euler는 28~50. 약 2배 빨라졌다고 볼 수 있을듯
추가
- 데이터 처리 : 4개의 VLM모델로 이미지를 captioning하고, 그중에 CLIP score가 높은 텍스트가 학습 과정에서 더 많이 선택되도록 했다.
- DiT를 썼을 때 Block layer 수가 20 ~ 30 사이인게 제일 좋다. 모델을 키울 때도 depth를 늘려서 키우지 않음.
Share article
Subscribe to our newsletter
