모바일용 고성능 디퓨전 트랜스포머 SnapGen++

초록

SnapGen++은 모바일·엣지 디바이스에서 1K 해상도 이미지를 실시간으로 생성할 수 있도록 설계된 경량 디퓨전 트랜스포머이다. 전역‑지역 스파스 어텐션, 엘라스틱 슈퍼네트워크 기반 서브‑DiT 공유 학습, 그리고 few‑step 교사 모델을 활용한 Knowledge‑Guided Distribution Matching Distillation(K‑DMD)이라는 세 가지 핵심 기술을 결합해 0.4 B 파라미터 모델이 1.8 초 안에 1024 × 1024 이미지를 생성한다.

상세 분석

본 논문은 고해상도 이미지 생성에 필수적인 전역 컨텍스트와 지역 디테일을 동시에 보존하면서도 연산량을 크게 줄이는 ‘Adaptive Global‑Local Sparse Attention(AGLSA)’을 제안한다. 구체적으로, 초기 저해상도 단계에서는 KV‑Compression 기반 전역 어텐션을 적용해 토큰 수를 압축하고, 이후 고해상도 단계에서는 블록‑와이즈 Neighborhood Attention을 이용해 인접 토큰 간의 정밀한 관계만을 계산한다. 이때 어텐션 스파싱 비율은 입력 이미지의 복잡도와 하드웨어 메모리 상황에 따라 동적으로 조정된다.

두 번째 기여인 ‘Elastic Training Framework’는 하나의 슈퍼네트워크 안에 다양한 깊이·폭을 가진 서브‑DiT들을 내재시킨다. Sandwich Rule과 Progressive Shrinking을 결합해 각 서브모델이 독립적으로 수렴하도록 하면서도 가중치를 공유함으로써 전체 파라미터는 0.4 B에 머문다. 추론 시에는 디바이스 사양에 맞춰 가장 적합한 서브‑DiT를 선택해 연산량·메모리 사용을 최적화한다.

세 번째 핵심인 K‑DMD는 기존 DMD(Distillation Matching Distribution)의 노이즈 분포 정합 목표에 ‘Knowledge Guidance’를 추가한다. 구체적으로, 4‑step 교사 모델의 중간 노이즈 예측을 KL‑Divergence로 정렬하고, CLIP‑Text‑Image 일치도와 Perceptual Loss를 결합해 학생 모델이 적은 스텝에서도 교사의 풍부한 표현력을 흡수하도록 설계했다. 결과적으로 4‑step 학생 모델이 1‑step 교사 대비 FID·Human Preference에서 10 % 이상 개선되었다.

실험에서는 iPhone 16 Pro Max에서 0.4 B 모델이 1.8 s(≈ 550 ms/step) 내에 1024 × 1024 이미지를 생성했으며, 동일 하드웨어에서 기존 U‑Net 기반 SnapGen(274 ms)보다 약 6배 빠른 속도와 비슷한 시각적 품질을 달성했다. 또한 1.6 B ‘full’ 버전은 서버‑사이드 환경에서 Flux·SD3 수준의 FID를 기록하면서도 메모리 사용량을 30 % 이하로 감소시켰다.

한계점으로는 현재 1K 해상도에 최적화돼 있어 2K·4K 확장은 추가적인 메모리 압축 기법이 필요하고, 초대형 교사 모델(>10 B 파라미터) 의존도가 높은 점이 있다. 향후 연구에서는 멀티‑모달 프롬프트와 동적 샘플링 전략을 결합해 더욱 가벼운 모델로도 다양한 생성 태스크를 지원하는 방향이 기대된다.