엣지디트 하드웨어 인식 확산 트랜스포머

본 논문은 최신 이미지 생성 모델인 확산 트랜스포머(DiT)의 뛰어난 성능을 모바일 환경에 적용하기 위한 연구로, EdgeDiT라는 하드웨어‑인식 경량 모델군을 제시한다. 기존 DiT‑XL/2는 28개의 트랜스포머 블록과 높은 차원(1152)·MLP 확장비(4)를 갖추어 높은 연산량과 메모리 요구사항을 발생시킨다. 이러한 구조는 GPU와 같은 고성능 장비에서는 문제되지 않지만, Qualcomm Hexagon, Apple ANE와 같은 모바일 NPU에서는 실시간 추론이 불가능하다. **문제 정의 및 접근법** 저자는 “연산량 감소 ≠ 지연시간 감소”라는 현상을 지적하고, NPU가 최적화된 연산(GEMM, 행렬 곱)과 메모리 접근 패턴을 고려한 구조적 압축이 필요하다고 주장한다. 이를 위해 두 단계의 하드웨어‑인식 서브스페이스를 설계한다. 1. **블록 제거(Block Removal)** – 인접한 두 DiT 레이어를 하나로 합쳐 레이어 깊이를 절반으로 줄인다. 이는 NPU에서 레이어 호출 오버헤드와 메모리 캐시 사용을 감소시킨다. 2. **구조 단순화(Structural Simplification)** – 각 레이어에 대해 (a) MLP 확장비를 4→2 로 낮추고, (b) 내부 차원을 1152→512 로 축소한다. 두 옵션을 조합해 4⁽²⁸⁾개의 후보를 만든다. **특징‑와이즈 지식 증류** 전체 후보를 처음부터 학습하는 것은 계산 비용이 과다하므로, 원본 DiT의 각 블록 출력을 교사 신호로 삼아 서브 블록이 동일한 특성을 모방하도록 L2 손실을 최소화한다. 블록 제거 단계에서는 두 교사 레이어를 하나의 학생 레이어가 학습하고, 구조 단순화 단계에서는 각각의 교사 레이어와 대응하는 학생 레이어를 독립적으로 학습한다. 이 과정은 병렬화가 가능해 전체 학습 시간을 크게 단축한다. **다목적 베이지안 최적화** 사전 학습된 서브 블록을 조합해 후보 아키텍처를 즉시 구성한 뒤, 다목적 베이지안 최적화(MOBO)를 통해 파레토 전선을 탐색한다. 목표 함수는 (i) 생성 품질(FID)와 (ii) 하드웨어 비용(지연시간·피크 메모리)이며, 기대 하이퍼볼륨 개선(EHVI) 획득 함수를 사용한다. 탐색 과정에서 이미지 품질이 급격히 떨어지는 조합(예: 3개 블록 동시 제거, MLP 비율 1)은 자동으로 배제된다. **실험 및 결과** - **모델 선택**: EdgeDiT‑1(파라미터 471 M, FLOPs 36 % 감소)과 EdgeDiT‑6(파라미터 530 M, FLOPs 30 % 감소)를 최종 선택. - **성능**: ImageNet 256×256 조건부 생성에서 DiT‑XL/2(FID 16.23) 대비 EdgeDiT‑1(FID 12.3)과 EdgeDiT‑6(FID 12.4)로 품질 향상. IS, Precision, Recall에서도 큰 격차가 없으며, 시각적으로도 원본과 거의 구분이 어려움. - **하드웨어**: Qualcomm Hexagon 및 Apple ANE에서 평균 1.65배 지연시간 감소를 기록, 메모리 사용량도 현저히 낮아짐. **기여 및 의의** 1. NPU 특성을 반영한 구조적 서브스페이스 정의. 2. 블록‑단위 특징‑와이즈 지식 증류를 통한 효율적 사전 학습. 3. 다목적 베이지안 탐색을 통한 파레토 최적 모델 자동 선정. 4. 실제 모바일 NPUs에서 고품질 이미지 생성이 가능함을 실증. **한계 및 향후 연구** - 현재는 ImageNet 256×256에 초점을 맞추었으며, 더 높은 해상도나 다른 도메인(예: 비디오, 3D)으로 확장할 필요가 있다. - 증류 손실이 L2 기반이라 텍스처 디테일 보존에 한계가 있을 수 있어, perceptual loss나 adversarial loss를 결합한 증류 방안이 고려될 수 있다. - 현재 탐색은 블록 수준에 국한되었지만, 토큰 샘플링, 스케줄링, 양자화 등 하드웨어‑친화적 기법과 결합하면 추가적인 효율성을 얻을 수 있다. 종합하면, EdgeDiT는 하드웨어‑인식 설계와 효율적 증류·탐색 파이프라인을 통해 대규모 확산 트랜스포머를 모바일 NPUs에 실시간으로 적용할 수 있는 실용적인 청사진을 제공한다. 이는 프라이버시 보호, 네트워크 의존도 감소, 에너지 효율성 측면에서 중요한 진전이며, 향후 다양한 생성 AI 응용에 확장될 가능성이 크다.