Paper Review 4 - DiT: Scalable Diffusion

DiT(Diffusion Transformer)는 2022년 Meta에서 발표한 논문(https://arxiv.org/abs/2212.09748)으로, 최근 OpenAI의 비디오 생성 모델인 Sora에 큰 영향을 미친 것으로 알려져 다시금 주목받았다. OpenAI의 전략은 대형언어모델부터 음성인식까지 다양한 AI 분야에 걸쳐 일관되게 나타나는데, 그것은 바로 "귀납적 편향이 적은 모델(즉, 트랜스포머)에 방대한 데이터와 연산 자원을 투입하면 획기적인 성능 향상을 이룰 수 있다"는 것이다. 이러한 맥락에서 기존 디퓨전 모델의 컨볼루션 기반 U-Net을 트랜스포머로 대체하려고 한 DiT 논문이 Sora의 주요 참고 자료가 된 것은 그다지 놀랍지 않다.

DiT 아키텍처는 그 이름에서 알 수 있듯이 디퓨전 모델 중 가장 널리 사용되는 DDPM 구조에 ViT(Vision Transformer)와 유사한 방식을 적용한 것이다. 2020년에 발표된 ViT 논문은 2차원인 이미지를 분류하는 모델에 1차원 시퀀스를 처리하는 트랜스포머를 접목하기 위해 각 이미지를 조각조각 나눠 입력으로 사용했다. 이와 같은 맥락에서 DiT는 이미지(정확히는 이미지의 잠재적 표현)를 패치로 분할하고 선형으로 정렬하여 토큰 시퀀스로 변환한다. 이렇게 생성된 토큰 시퀀스에 위치 임베딩을 적용하여 최종적인 모델 입력을 만든다. 이 외에도 이미지 생성을 위한 조건인 레이블과 타임스탬프 정보를 추가적인 입력 값으로 사용할 것이다.

DiT 블록은 U-Net을 대체하는 역할을 하며, 트랜스포머의 표준적인 구성을 차용했다. 다만, 레이블과 타임스탬프 정보를 적절히 반영하기 위해 저자들은 4가지 대안을 고려하고 실험을 진행했다.

첫번째 안은 레이블과 타임스탬프의 벡터 임베딩을 이미지 입력 시퀀스의 끝에 단순히 추가하는 형태로, ViT의 CLS 토큰과 유사하다. 마지막 블록을 통과한 후에는 해당 조건부 토큰을 시퀀스에서 제거한다.

두번째 안은 레이블 및 타임스탬프 임베딩 벡터를 병합한 다음, 멀티 헤드 크로스 어텐션 레이어를 통해 이미지 정보와 융합하는 것이다. 이 접근 방식은 LDM, 즉 스테이블 디퓨전의 접근 방식과 유사하며, 연산량 측면에서 가장 큰 오버헤드를 가져온다.

세번째 안(adaLN)은 레이어 정규화 시 차원 별 스케일링 및 시프트 파라미터를 학습하는 대신, 레이블과 타임스탬프 임베딩 합계의 회귀 값을 적용하여 해당 정보를 반영하게 한다. 이 방식은 연산량 측면에서 오버헤드가 가장 적다.

네번째 안(adaLN-Zero)은 adaLN에 더하여 DiT 블록 내 잔차 연결 직전에 적용되는 차원 별 스케일링 파라미터에 해당 회귀 값을 사용한다.

저자들은 실험을 통해 네번째 안인 adaLN-Zero 전략이 가장 우수한 성능을 보인다는 사실을 밝혀냈다.

마지막으로 최종 DiT 블록의 출력 토큰 시퀀스를 레이어 정규화(또는 adaLN)한 다음, 노이즈 및 공분산 예측(대각 행렬)을 통해 디코딩한다. 이는 선형의 토큰 시퀀스를 원래 입력과 동일한 공간적 레이아웃으로 재정렬하는 것을 의미한다.

DiT 논문은 이와 같이 이미지 생성을 위한 새로운 아키텍처를 제안했다. 이 아키텍처는 연산량이 증가할수록, 즉 트랜스포머 모델의 크기가 커지고 이미지 패치를 더 세밀하게 분할할수록 성능이 지속적으로 향상되는 특징을 보인다. 또, 대형 모델일수록 계산 효율성 측면에서도 더 우수한 결과를 나타낸다. 이러한 특성들은 DiT가 현재의 디퓨전 모델 구조를 뛰어넘는 차세대 이미지 및 비디오 생성 기술로 발전할 가능성을 암시한다.

#Diffusion #Transformer