Stable diffusion 같은 Text-to-Image 모델들은 학습에 어마어마한 시간과 비용이 들어간다. 이를 해결하기 위해 저자들은 transformer 기반의 생성 모델 아키텍처를 제안하는데, 어떤 방법을 통해 어떤 문제를 해결하는지 알아보겠다.
기존 연구의 한계점
서두에서 밝힌 것처럼, stable diffusion, Imagen, Midjourney 같은 기존의 text-to-image 모델을 학습시킬 때는 시간과 비용이 굉장히 많이 발생한다. Stable diffusion 같은 경우, A100 한대로 6,250 시간의 학습시간이 발생하는데, A100 의 경우 클라우드 서버 호스팅을 이용하면 시간당 약 $5 정도이니, 비용은 약 3천5백만원 정도가 발생하게 된다. (저자들은 논문에서 이를 약 $320,000, 원화 3.5억 정도로 산정했는데, 계산이 잘못되지 않았나 싶다.)
저자들은 이러한 한계점이 트레이닝 파이프라인과 데이터에서 발생한다고 보았다.
특히, T2I 태스크는 3가지로 구분될 수 있는데, 1) 원본 이미지에서의 pixel 분포 학습, 2) 텍스트와 이미지 사이의 관계(alignment) 학습, 3) 이미지 품질 제고 이다.
기존의 방법들은 이러한 프로세스가 분리되지 않아 다량의 데이터에서 한꺼번에 학습하는 비효율로 인해 학습 시간이 오래 걸린다고 주장한다.
또한, 데이터셋에서의 이미지 캡션 역시 문제로 지적했는데, 몇몇 이미지와 맞지 않는 캡션이 들어가 있어 이를 LLaVA를 이용해 refine 하였다.
아이디어
따라서, 위 세가지 프로세스를 disentangle해서, 한 번에 학습하는 것이 아니라, 세 가지 프로세스로 학습하는 방법을 제안했는데, 1) 원본 이미지에서 pixel 분포를 학습하고, 2) 텍스트와 이미지 사이의 관계(alignment)를 학습하고, 3) 마지막으로 이미지의 품질을 높인다.
이를 위해 저자들은 Diffusion Transformer (DiT) 아키텍를 사용해 이미지 정보를 학습했고, ImageNet에서 pretrained된 모델을 사용해 학습 속도를 높였고, 텍스트-이미지 관계를 위해 기존 데이터셋을 다듬었다. 마지막으로 high-quality aesthetic 데이터를 이용해 고해상도 이미지를 생성하도록 파인튜닝을 마쳤다.
또한, 효과적인 모델 학습을 위해, DiT 전체 파라미터의 약 27%를 차지하는 AdaLN 모듈을 줄이기 위해 모든 block에서 동일한 값을 사용하는 AdaLN-single 아키텍처를 제안한다.
느낀점
사실 논문만 봤을 때는 DiT를 T2I 태스크에 적용하면서 몇몇 파라미터들을 효율화 한 정도로 보이지만, 보여주는 결과들이 굉장히 디테일이 잘 재현된 것 같다.
생성 모델의 성능은 사실 논문에 기재된 metric보다는 직접 reproduce를 해봐야 알지만, 논문에 기재된 fine-tuning된 결과물 등은 좋은 성능을 보이고 있다.
