EdgeDiT: 모바일 기기용 하드웨어 인식 확산 트랜스포머 (Hardware-Aware Diffusion Transformers)

1. 문제 정의 (Problem Statement)

• 배경: 확산 모델 (Diffusion Models), 특히 확산 트랜스포머 (DiT) 는 고화질 이미지 생성 분야에서 새로운 최첨단 (SOTA) 성능을 달성했습니다.
• 도전 과제: 기존 DiT 아키텍처는 방대한 계산 복잡도와 메모리 요구 사항으로 인해 리소스가 제한된 엣지 장치 (모바일, 임베디드) 에 직접 배포하기 어렵습니다.
• 현황: 클라우드 기반 추론은 이러한 한계를 우회하지만, 개인정보 보호 문제, 네트워크 의존성, 높은 에너지 소비 등의 단점이 있습니다.
• 목표: 고화질 생성 능력을 유지하면서 모바일 신경 처리 장치 (NPU, 예: 퀄컴 Hexagon, 애플 ANE) 에 최적화된 경량화된 확산 트랜스포머를 개발하여 온디바이스 (On-Device) 생성 AI 를 실현하는 것.

2. 방법론 (Methodology)

저자들은 DiT-XL/2 를 기반으로 하드웨어 효율성을 극대화하는 EdgeDiT 패밀리를 제안합니다. 핵심 접근 방식은 다음과 같습니다.

• 하드웨어 인식 아키텍처 탐색 (Hardware-Aware Architecture Search):

  • DiT 의 구조적 중복성을 식별하고 제거하기 위해 구조적 단순화 및 가지치기 전략을 적용합니다.
  • 서로게이트 (Surrogate) 블록 설계:
    1. 블록 제거 (Block Removal): 연속된 두 DiT 레이어를 하나의 레이어로 축소.
    2. MLP 비율 수정: FFN(Feed-Forward Network) 의 확장 비율 (Expansion Ratio) 을 변경하여 파라미터 및 FLOPs 감소.
    3. 잠재 차원 축소 (Hidden Dimension Reduction): 트랜스포머 블록의 내부 차원을 축소하여 저랭크 표현 생성.
  • 이러한 수정 사항들을 조합하여 하드웨어 친화적인 탐색 공간 (Search Space) 을 구성합니다.

• 특징 기반 지식 증류 (Feature-wise Knowledge Distillation, FwKD):

  • 모든 후보 아키텍처를 처음부터 학습하는 것은 계산 비용이 너무 큽니다. 이를 해결하기 위해 분할 정복 (Divide-and-Conquer) 전략을 사용합니다.
  • 사전 학습된 큰 모델 (Teacher, DiT-XL/2) 의 레이어 출력을 작은 서로게이트 블록 (Student) 이 모방하도록 특징 기반 지식 증류를 수행합니다.
  • 이 과정은 병렬화 가능하여 학습 비용을 크게 절감합니다.

• 다목적 베이지안 최적화 (Multi-Objective Bayesian Optimization, MOBO):

  • 생성 품질 (FID 점수) 과 하드웨어 효율성 (엣지 장치 지연 시간) 간의 파레토 최적 (Pareto-optimal) 해를 찾기 위해 MOBO 를 적용합니다.
  • 탐색 공간 내에서 가장 균형 잡힌 아키텍처를 식별하여 EdgeDiT 의 기본 구조로 선정합니다.

• 종단 간 학습 (End-to-End Training):

  • 최적화된 아키텍처가 선정된 후, 표준 확산 학습 목적 함수를 사용하여 모델 파라미터를 완전히 적응 (Fine-tuning) 시킵니다.

3. 주요 기여 (Key Contributions)

1. EdgeDiT 제안: 모바일 NPU 에 최적화된 하드웨어 인식 확산 트랜스포머 아키텍처 패밀리 개발.
2. 구조적 최적화 전략: 모바일 하드웨어의 데이터 흐름에 맞춰 설계된 구조적 단순화 및 가지치기 전략 도입.
3. 성능 - 효율성 트레이드오프 개선: 기존 DiT 변형 및 최적화된 모바일 U-Net 대비 생성 품질 (FID) 과 추론 효율성 간의 파레토 최적 관계를 입증.
4. 실제 배포 가능성 증명: 대규모 기반 모델을 고성능 GPU 에서 사용자의 손끝 (모바일 기기) 으로 확장할 수 있는 확장 가능한 청사진 제공.

4. 실험 결과 (Experimental Results)

• 성능 향상:
  • 파라미터: 20~30% 감소.
  • FLOPs: 36~46% 감소.
  • 지연 시간 (Latency): 온디바이스 추론 속도가 1.65 배 향상 (삼성 갤럭시 S25 Ultra 기준).
• 생성 품질:
  • ImageNet 256x256 기준, EdgeDiT 모델은 파라미터가 20~30% 적음에도 불구하고 Baseline 인 DiT-XL/2 보다 더 낮은 FID 점수 (더 좋은 품질) 를 기록했습니다.
  • 시각적 비교에서도 EdgeDiT 가 생성한 이미지가 원본 DiT-XL/2 와 동등하거나 일부 경우 더 우수한 품질을 보였습니다.
• 하드웨어 벤치마크:
  • Apple A18 Pro (iPhone 16 Pro Max): Baseline 대비 1.45 배 빠른 추론 속도.
  • Qualcomm 8850 (Samsung Galaxy S25 Ultra): Baseline 대비 1.65 배 빠른 추론 속도.
  • GMACs 는 28~40% 감소했습니다.

5. 의의 및 결론 (Significance & Conclusion)

• 기술적 의의: EdgeDiT 는 확산 트랜스포머의 확장성과 엣지 배포가 상충되지 않음을 입증했습니다. 레이어 단위 증류와 베이지안 최적화된 아키텍처 탐색을 통해 SOTA 모델의 30% 를 압축하면서도 성능을 유지하거나 향상시켰습니다.
• 실용적 가치: 이 연구는 대규모 생성 모델을 클라우드에 의존하지 않고도 개인정보 보호, 낮은 지연 시간, 에너지 효율성을 보장하며 모바일 기기에서 직접 실행할 수 있는 길을 열었습니다.
• 미래 전망: 현재는 DiT-XL/2 를 기반으로 하지만, 이 프레임워크는 SiT 나 MDT 와 같은 다른 확산 트랜스포머 계열이나 AR/VR/XR 장치 등 다양한 하드웨어 가속기로 확장될 수 있는 잠재력을 가지고 있습니다.

요약하자면, EdgeDiT 는 하드웨어의 제약을 고려하여 확산 모델의 구조를 체계적으로 재설계하고, 지능적인 탐색 알고리즘을 통해 모바일 기기에서도 고품질 이미지 생성이 가능하도록 만든 획기적인 연구입니다.