NanoFLUX: Distillation-Driven Compression of Large Text-to-Image Generation Models for Mobile Devices

기술 요약: NanoFLUX

문제 정의

FLUX 제품군(최대 17B 파라미터) 및 Qwen-Image(20B 파라미터)와 같은 대규모 텍스트-이미지(T2I) 확산 모델은 최첨단 시각적 품질을 달면성했습니다. 그러나 이들의 거대한 규모는 온디바이스 배포에 상당한 장벽을 형성하며, 비용이 많이 드는 GPU 인프라나 클라우드 기반 추론을 요구합니다. 기존의 온디바이스 솔루션은 대등한 품질을 갖춘 더 작은 모델에 의존하거나, 대규모 확산 트랜스포머(DiT) 및 관련 텍스트 인코더의 구조적 중복성을 완전히 해결하지 못하는 압축 기술(프루닝, 양자화, 단계 축소)을 채택하는 경우가 많습니다. 과제는 생성 품질을 유지하면서 모바일 기기에서 실행 가능하도록 이 거대한 모델들을 압축하고 추론 지연 시간을 최소화하는 것입니다.

방법론: NanoFLUX

저자들은 17B FLUX.1-Schnell 교사 모델로부터 증류된 2.4B 파라미터 규모의 T2I 모델인 NanoFLUX를 소개합니다. 압축은 가장 무거운 두 구성 요소인 12B 확산 트랜스포머(DiT)와 5B T5-XXL 텍스트 인코더를 대상으로 하는 점진적인 증류 기반 파이프라인을 통해 이루어집니다.

1. 트랜스포머 압축 (12B $\to$ 2B)

저자들은 DiT 크기를 줄이기 위해 4단계의 점진적 압축 전략을 사용합니다:

• 단계 C1: 어텐션 헤드 프루닝 (12B $\to$ 5B):

  • 분석: 어텐션 헤드 출력에 대한 특이값 분해(SVD) 결과, 모든 헤드가 독립적이지 않으며 많은 헤드가 중복된 정보를 포함하고 있음이 밝혀졌습니다.
  • 조치: 어텐션 헤드의 수($H$)를 24개에서 16개로 균일하게 감소시켰습니다.
  • 증류: 남은 헤드들이 제거된 헤드의 정보를 집합적으로 보존할 수 있도록 특징 수준의 손실(feature-level loss)을 도입했습니다. 학생 모델은 교사 모델의 첫 16개 헤드로 초기화되었습니다.

• 단계 C2: 헤드 차원 축소 (5B $\to$ 3B):

  • 분석: 5B 모델의 특징에 대해 다시 SVD를 적용했습니다.
  • 조치: 헤드 차원($d_H$)을 128에서 96으로 축소하여 상위 특이 성분을 유지했습니다. 이를 통해 3B 파라미터 모델을 얻었습니다.

• 단계 C3: 블록 병합을 통한 깊이 프루닝 (3B $\to$ 2.5B):

  • 분석: 트랜스포머 블록에 대한 입력-출력 코사인 유사도를 계산했습니다. 싱글 스트림(SS) 블록은 높은 유사도(>0.85)를 보였으며, 이는 기능적 중복성을 나타냅니다.
  • 조치: 이러한 블록들을 단순히 삭제하는 대신, 저자들은 15개의 중복된 SS 블록 체인을 매개변수 평균화를 통해 하나의 블록으로 병합했습니다. 이를 통해 트랜스포머 깊이를 38개 블록에서 24개 블록으로 줄였습니다.

• 단계 C4: 정적 레이어 정규화 (2.5B $\to$ 1.8B):

  • 분석: 모델이 예측한 적응형 레이어 정규화(AdaLN) 계수들이 샘플 간에 낮은 분산을 보였는데, 이는 CLIP 프로젝션으로부터의 동적 컨디셔닝이 중복적임을 시사합니다.
  • 조치: 동적 AdaLN 계수를 사전 계산된 평균값(Static-LN)으로 대체했습니다. 이를 통해 약 0.7B 파라미터를 제거했습니다. 품질 회복을 위해 고정된 계수를 사용하여 짧게 미세 조정(fine-tuning)을 수행했습니다.

2. 점진적 토큰 다운샘플링

상세도를 희생하지 않으면서 셀프 어텐션의 이차 복잡도 문제를 해결하기 위해:

• 메커니즘: ResNet 기반의 다운샘플링 모듈($D$)과 업샘플링 모듈($U$)을 트랜스포머에 삽입했습니다.
• 하이브리드 전략: 모델은 하이브리드 해상도 체계로 작동합니다. 전역 구조를 포착하는 초기 확산 타임스텝은 다운샘플링된 토큰(토큰 수를 4배 감소)을 처리하고, 세부 사항을 정교화하는 후기 타임스텝은 전체 해상도의 토큰을 처리합니다.
• 학습: 학습을 안정화하기 위해 "점진적 토큰 다운샘플링(PTD)" 전략을 사용했습니다. 모듈을 점진적으로 삽입 및 학습하여, 반복(iteration)에 따라 다운샘플링 지점을 네트워크의 더 앞쪽으로 이동시켰습니다.
• 결과: 43개의 트랜스포머 블록 중 23개가 다운샘플링된 토큰을 사용하여 작동함으로써, 필요한 곳에서는 고해상도 처리를 유지하면서도 추론 지연 시간을 크게 줄였습니다.

3. 텍스트 인코더 증류 (5B $\to$ 330M)

• 대상: 5B T5-XXL 인코더를 330M T5-Large 모델로 증류합니다.
• 방법: 2단계 증류 프로세스를 사용합니다:
  1. 임베딩 매칭: 학생 모델에 MLP를 추가하여 교사의 출력 차원과 일치시키고, 프롬프트 임베딩 간의 MSE를 최소화합니다.
  2. 블록 단위 증류: 학생 모델이 특정 트랜스포머 레이어에서 교사의 프롬프트 은닉 상태(hidden states)를 일치시키도록 학습됩니다. 결정적으로, 이 방법은 오차 누적을 완화하기 위해 조인트 어텐션(joint attention)을 통해 디노이저의 초기 레이어로부터 오는 시각적 신호를 활용하여 텍스트 인코더를 가이드합니다.
  3. 안정성: 불안정한 감독(supervision)을 피하기 위해 고노이즈 타임스텝($t < \hat{t}$)에서는 그래디언트를 분리(detach)하며, 후기 타임스텝과 처음 세 개의 트랜스포머 레이어에 집중하여 증류를 수행합니다.

주요 결과

• 모델 크기: 최종 NanoFLUX 모델은 2.4B 파라미터(2B DiT + 330M 텍스트 인코더)로, 교사 모델인 17B FLUX.1-Schnell 대비 7배 감소했습니다.
• 성능:
  • 모바일 기기에서 NanoFLUX는 512 $\times$ 512 이미지를 약 2.5초 만에 생성합니다.
  • 정량적 벤치마크(One-IG, DPG, GenEval, HPDv3)에서 NanoFLUX는 교사 모델과 대등한 성능을 유지하며, 특정 지표(예: HPSv3가 11.45에서 10.41로 하락)에서만 미미한 저하를 보였습니다.
  • 정성적 결과는 모델이 원래의 17B 모델과 비교할 만한 이미지 품질과 프롬프트 준수 능력을 보존함을 나타냅니다.
• 효율성: 점진적 토큰 다운샘플링 전략은 추론 지연 시간과 생성 품질 사이의 균형을 성공적으로 맞추어, 자원이 제한된 하드웨어에서도 효율적으로 실행될 수 있게 합니다.

의의 및 주장

본 논문은 NanoFLUX가 FLUX와 같은 모델을 구체적으로 겨냥하여 대규모 모델 압축을 종합적으로 연구한 최초의 프레임워크라고 주장합니다. 이 연구의 의의는 다음과 같습니다:

1. 온디바이스 생성의 실현 가능성: 클라우드 인프라에 의존하지 않고도 모바일 기기에서 고품질 T2I 생성이 가능하다는 것을 입증했습니다.
2. 구조적 통찰: 단순한 범용 양자화나 단계 축소에 의존하는 대신, 특정 구조적 중복성(어텐션 헤드, 블록 깊이, 정규화 역학)을 식별하고 활용함으로써 대규모 생성 모델을 축소하는 실질적인 경로를 제공했습니다.
3. 일반화 가능성: 향후 등장할 대규모 생성 모델에도 적용 가능한 점진적 증류 파이프라인과 통찰력을 제공합니다.

저자들은 이 연구를 컴퓨팅 자원이 제한된 사용자들에게 고품질 생성형 AI에 대한 접근성을 넓히는 단계로 규정합니다.