효율적인 대규모 Diffusion Transformer 압축: Amber‑Image 시리즈

초록

Amber‑Image는 60계층 Dual‑Stream MMDiT 기반 Qwen‑Image를 단계적 깊이 프루닝과 하이브리드 스트림 변환으로 압축해 10B·6B 파라미터 규모의 경량 모델을 만든다. 레이어 중요도 평가·가중 평균 초기화·계층별 지식 증류·전체 파인튜닝을 결합해 70% 파라미터를 줄이고, 전체 파이프라인을 2,000 GPU‑hour 이하로 수행한다. DPG‑Bench·LongText‑Bench 등에서 대형 모델에 근접한 이미지 품질과 텍스트 재현성을 보인다.

상세 분석

본 논문은 현재 텍스트‑투‑이미지(T2I) 분야에서 가장 성능이 뛰어난 Diffusion Transformer(DiT) 계열 모델이 갖는 “대규모·고비용”이라는 근본적인 한계를 압축 기법으로 해결하고자 한다. 핵심 아이디어는 기존 60계층, 20 B 파라미터 규모의 Dual‑Stream MMDiT(Qwen‑Image)를 두 단계에 걸쳐 구조적·학습적 변환을 가함으로써, 10 B와 6 B 두 가지 경량 변형을 얻는 것이다.

1. Depth Pruning with Timestep‑Sensitive Importance

   • 레이어 중요도는 전역 Ablation 기반 스코어링으로 측정한다. 각 레이어를 일시적으로 비활성화하고, 다양한 프롬프트·시간 단계(t)에서 발생하는 노이즈 예측 차이(δ)와 시간 가중치(ωₜ)를 곱해 누적 손실을 구한다.
   • 이 방식은 기존 Gradient‑based Taylor 근사보다 시계열 의존성을 반영해, 초기 고노이즈 단계에서의 오류가 전체 의미 구조에 미치는 영향을 크게 가중한다.
   • 중요도 점수가 낮은 30개 레이어를 제거해 50% 깊이 감소, 파라미터는 약 10 B로 축소된다.

2. Local Weight Averaging (LWA) 초기화

   • 프루닝된 레이어 주변의 연속된 원본 레이어 가중치를 산술 평균해 보존 레이어를 초기화한다. 이는 인접 레이어가 수행하는 점진적 변환을 하나의 평균 변환으로 근사함으로써, 초기 학습 단계에서 발생하는 급격한 성능 저하를 방지한다.
   • LWA는 간단하지만 효과적인 Warm‑Start 전략으로, 이후 증류·파인튜닝 단계에서 빠른 수렴을 가능하게 한다.

3. Two‑Stage Recovery: Layer‑wise Distillation → Global Fine‑tuning

   • 첫 단계에서는 LWA로 재초기화된 레이어만을 학습 대상으로 삼아, 원본 60‑layer 교사 모델의 해당 레이어 출력(히든 스테이트)을 목표로 하는 Layer‑wise Knowledge Distillation을 수행한다.
   • 두 번째 단계에서는 전체 파라미터를 해제하고 표준 Diffusion 손실로 전역 Fine‑tuning을 진행한다. 이 두 단계는 각각 표현 복구와 전체 일관성 정합을 담당한다.

4. Hybrid‑Stream Architecture for Amber‑Image‑6B

   • 10 B 모델을 교사로 삼아, 30‑layer 백본의 뒤쪽 20계층을 Single‑Stream 구조로 전환한다. 초기 10계층은 기존 Dual‑Stream(텍스트·이미지 별) 형태를 유지해 modality‑specific 특성을 보존하고, 이후 레이어는 이미지 스트림 가중치를 공유한다.
   • 이 설계는 Cross‑modal redundancy를 활용해 파라미터를 추가 40% 절감하면서도, 이미지 스트림이 제공하는 공간·픽셀 수준의 priors를 그대로 이어받는다.
   • 전환된 레이어에 대해서는 Local Distillation을 수행해, 교사 모델의 텍스트·이미지 히든 스테이트를 concat 형태로 목표 삼는다. 이후 가벼운 전체 파인튜닝을 통해 최종 모델을 완성한다.

5. 효율성 및 성능

   • 전체 파이프라인(10 B → 6 B)은 8×A100 GPU 10일(≈2,000 GPU‑hour) 내에 완료된다. 이는 동일 규모 모델을 처음부터 학습하는 비용의 10배 이상 절감한다.
   • DPG‑Bench, LongText‑Bench 등에서 시각적 충실도와 텍스트 재현성이 20‑30 B 파라미터 모델에 근접하거나 일부 경우 능가한다는 실험 결과를 제시한다.
   • 데이터 측면에서도 대규모 크롤링·클리닝 파이프라인 없이, 고품질 소규모 데이터셋(수천~수만 이미지)만으로 충분히 성능을 회복한다는 점이 주목할 만하다.

핵심 기여는 (1) 시간 가중치를 포함한 전역 Ablation 기반 레이어 중요도 평가, (2) 인접 레이어 가중치 평균을 통한 손쉬운 초기화, (3) 단계적 증류·파인튜닝을 결합한 압축‑재학습 프레임워크, (4) 깊은 레이어의 Dual‑Stream → Single‑Stream 전환이라는 구조적 단순화이다. 이 네 가지 요소가 서로 보완하면서, 대규모 Diffusion Transformer를 실용적인 수준으로 “압축·경량화”하는 새로운 패러다임을 제시한다.