동적 패치 스케줄링으로 효율적인 디퓨전 트랜스포머

초록

DDiT는 이미지·비디오 생성에 사용되는 디퓨전 트랜스포머(DiT)의 고정된 패치 토크나이징을 동적으로 조절한다. 시점별·내용별 복잡도에 따라 패치 크기를 크게 혹은 작게 바꾸어 초기 단계에서는 거친 패치로 전역 구조를, 후반 단계에서는 세밀한 패치로 디테일을 복원한다. 이를 위해 기존 패치 임베딩에 LoRA 브랜치를 추가하고, 잠재 공간의 변화율(Δz) 기반의 무학습 스케줄러를 도입한다. 실험 결과 FLUX‑1.Dev와 Wan‑2.1에서 각각 최대 3.52×·3.2× 속도 향상을 달성하면서 이미지 품질과 프롬프트 일치도는 유지한다.

상세 분석

DDiT는 디퓨전 트랜스포머(DiT)의 핵심 병목인 어텐션 연산의 토큰 수를 동적으로 제어함으로써 연산 효율을 크게 개선한다. 기존 DiT는 VAE 잠재를 고정된 크기(p×p) 패치로 나누어 토큰화하고, 전체 denoising 단계에서 동일한 토큰 수를 유지한다. 그러나 디퓨전 과정은 초기 timesteps에서 전역적인 구조를 형성하고, 후반 timesteps에서 세밀한 디테일을 추가한다는 사실을 이용한다. DDiT는 이러한 단계별 특성을 반영해 “큰 패치(코스 그레인)”와 “작은 패치(파인 그레인)”를 교대로 적용한다.

핵심 기술은 두 가지이다. 첫째, 기존 패치 임베딩 레이어를 다중 패치 크기를 지원하도록 확장하고, 각 새로운 패치 크기(p_new)마다 별도의 임베딩 가중치와 bias를 추가한다. 여기에는 LoRA(Low‑Rank Adaptation) 브랜치를 각 트랜스포머 블록에 삽입해, 새로운 패치 크기에 대한 미세 조정을 최소한의 파라미터만으로 가능하게 한다. 또한, 기존 위치 임베딩을 bilinear interpolation으로 스케일링하고, 패치 크기 식별자를 위한 학습 가능한 토큰을 도입해 모델이 현재 사용 중인 패치 크기를 인식하도록 한다.

둘째, 학습 없이 동작하는 “동적 패치 스케줄러”를 설계했다. 스케줄러는 현재 timestep t에서의 잠재 변화율 Δz_t = z_t – z_{t‑1}를 계산하고, 이를 패치별 표준편차 σ_{p_i,t}와 비교한다. 변화율이 낮으면(잠재가 천천히 변함) 전역 구조만 필요하다고 판단해 큰 패치(p·2, p·4 등)로 전환하고, 변화율이 높으면 세밀한 디테일이 생성되고 있다고 보고 작은 패치로 전환한다. 이 판단은 간단한 임계값 기반 혹은 비율 기반 정책으로 구현되며, 전체 denoising 루프 동안 매 timestep마다 재평가된다.

실험에서는 FLUX‑1.Dev(텍스트‑투‑이미지)와 Wan‑2.1(텍스트‑투‑비디오) 모델에 DDiT를 적용하였다. 이미지 품질 평가지표인 ImageReward와 CLIPScore, 비디오 품질을 측정하는 VBench에서 원본 모델과 거의 동일한 점수를 유지하면서, 1.88×, 2.2×, 3.5× 등 다양한 가속 비율을 달성했다. 특히 3.52× 가속 시에도 ImageReward가 1.0182, CLIPScore가 0.3117로 미세한 감소에 그쳤다. 비디오에서는 3.2× 가속 시 VBench 점수가 80.53으로 원본 81.24와 큰 차이가 없었다.

DDiT의 장점은 (1) 사전 학습된 DiT를 그대로 활용하면서 최소한의 파라미터(LoRA)만 추가해 적용이 용이하고, (2) 테스트 시점에 연산량을 직접 제어할 수 있어 다양한 하드웨어·예산 상황에 맞춤형 배포가 가능하다는 점이다. 반면, 현재 스케줄러는 변화율 기반 임계값 설정에 민감하며, 복잡한 프롬프트(예: 다중 객체, 복잡한 텍스처)에서는 최적 패치 전환 시점을 놓칠 위험이 있다. 또한, 큰 패치로 전환할 때 발생할 수 있는 경계 효과를 완화하기 위한 오버랩 패치나 멀티스케일 피드백 메커니즘이 추가적으로 필요할 수 있다.

전반적으로 DDiT는 디퓨전 트랜스포머의 연산 효율성을 크게 향상시키는 실용적인 접근법이며, 특히 실시간 생성이나 저전력 디바이스에서의 적용 가능성을 열어준다.