EdgeFusion: 온디바이스 텍스트 기반 이미지 생성

Tairen Piao
Research Engineer, Nota AI

소개

텍스트 기반 이미지 생성을 위한 Stable Diffusion(SD)는 요구되는 연산 수준이 높아 실제 배포가 어렵습니다. 최근 연구에서는 이를 해결하기 위해 Latent Consistency Model(LCM)에서처럼 샘플링 단계를 줄이거나 가지치기(Pruning)나 지식 증류(knowledge distillation)와 같은 아키텍처 최적화를 적용하려는 경향이 있습니다. 저희는 이와는 다른 방식으로, 더 컴팩트한 BK-SDM SD 변형을 시작점으로 하는 EdgeFusion을 제안했습니다. 일반적으로 사용되는 웹 크롤링 데이터셋을 이용하여 BK-SDM에 LCM을 직접 적용한 결과 만족스러운 성과를 얻지 못했기 때문에, 이를 극복하기 위해 두 가지 핵심 전략을 마련했습니다: (1) 주요 생성 모델에서 만든 고품질 합성 이미지-텍스트 쌍을 활용하고, (2) LCM에 맞춘 고급 증류 프로세스를 도입하는 것입니다. 또한, EdgeFusion 모델에 모델 수준 분할(MLT)과 양자화를 적용했을 때, 신경망 처리 장치(NPU)와 같이 자원이 제한된 엣지 디바이스에서도 단 두 번의 디노이징 단계만으로 포토리얼리스틱한 텍스트 기반이미지를 빠르게 생성할 수 있습니다. 그림 1은 EdgeFusion의 전체 성능을 보여줍니다.

이 논문은 CVPR 2024 Workshop on Efficient and On-Device Generation (EDGE)에 채택되었습니다.

그림 1. 텍스트 기반 이미지 생성 결과입니다. 개선된 데이터로 훈련된 EdgeFusion은 몇 단계의 디노이징만으로도 높은 품질의 이미지를 생성할 수 있습니다.

논문의 주요 메시지

• EdgeFusion은 주요 생성 모델에서 생성된 고품질 이미지-텍스트 쌍을 활용해 컴팩트 SD 모델의 성능을 향상시킵니다.

• EdgeFusion은 LCM에 맞춘 고급 지식 증류 프로세스를 통해 적은 단계로도 고품질 추론을 할 수 있도록 합니다.

• 모델 수준 분할과 양자화를 통해, EdgeFusion은 삼성 엑시노스 NPU에서 1초 만에 실제적인 이미지를 생성할 수 있습니다.

연구의 중요성

이 연구는 리소스가 제한된 환경(NPU 등)에서 SD 모델을 실제로 배포하기 위한 포괄적인 최적화 전략을 제안합니다. NPU의 계산 성능 한계와 메모리 제약을 해결하면서 엣지 디바이스에서 효율적으로 텍스트 기반 이미지를 생성할 수 있는 새로운 통찰력을 제공합니다.

방법론 요약

• LCM을 위한 고급 증류(distillation)

1. EdgeFusion은 BK-SDM-Tiny의 압축된 U-Net 디자인으로 시작하며, Realistic Vision V5.1을 고급 교사 모델로 사용하여 지식 증류를 통해 BK-SDM-Adv-Tiny라는 향상된 학생 모델을 얻습니다.

2. 원래의 Stable Diffusion 대신 더 발전된 Realistic Vision V5.1 모델을 교사로 사용하여 BK-SDM-Adv-Tiny라는 더 나은 초기화된 학생 모델을 얻습니다.

3. BK-SDM-Adv-Tiny는 LCM 스케줄러를 사용해 미세 조정되지만, 학생 모델과 같은 작은 모델 대신 원래의 큰 교사 모델로부터 증류가 수행됩니다.

그림 2. 단계가 감소된 컴팩트한 SD. (a) 기본 LCM 적용: BK-LCM-Tiny를 BK-SDM-Tiny의 가중치로 초기화한 후, 샘플링 단계를 줄이기 위해 증류 방식으로 훈련합니다. (b) 우리의 접근법: 더 나은 교사 모델을 통해 LCM의 학생 모델 초기화를 효과적으로 개선합니다. 또한, LCM 훈련 단계에서 원래의 교사 모델을 사용하면 성능이 향상됩니다. 두 단계 모두 고품질 데이터를 활용하는 것이 중요합니다.

• 데이터 품질 향상

LAION 데이터셋에서 저품질 이미지와 부적절한 캡션과 같은 한계가 있었고, 이를 해결하기 위해 다음의 방법들을 도입했습니다:

1. 데이터 전처리: 데이터 중복 제거, 작은 이미지 제거, 이미지 자르기 최적화. 이로써 데이터셋 크기를 44% 감소할 수 있었습니다.

2. 합성 캡션 생성: 대형 언어 모델(Sphinx)을 사용해 전처리된 이미지에 대해 상세하고 관련성 높은 캡션을 생성함으로써텍스트-이미지 일치도를 개선합니다.

3. 완전 합성 데이터 생성: GPT-4를 프롬프트 생성에 사용하고 Stable Diffusion XL(SDXL)을 이미지 합성에 사용하여 고품질의 합성 이미지-텍스트 쌍을 생성합니다. 이를 통해 인구 통계학적 속성 등 여러 속성을 조정하여 더 다양한 구성의 편향 없는 데이터셋을 확보할 수 있습니다.

4. 수동 데이터 선별 작업: 프롬프트 수정과 아티팩트 이미지 제거를 통해 29.9k개의 합성 샘플을 추가로 선별하여 품질을 높였습니다.

• NPU에 배포

1. 우리는 모델을 더 작은 세그먼트로 나누어 신경망 처리 장치의 메모리 제한에 맞추고 DRAM 접근을 줄이기 위한 모델 수준 분할을 적용했습니다. 또한, 삼성 엑시노스 NPU에서 효율적인 실행을 위해 커널 퓨전 및 기타 그래프 최적화 기술도 적용했습니다.

2. U-Net 모듈에서는 W8A16 혼합 정밀 양자화, 다른 곳에서는 16비트 부동소수점 (FP16)을 적용하여 모델 크기를 줄이고 추론 속도를 가속화했습니다.

실험 결과

• 주요 결과

  표 1에 따르면 BK-SDM-Adv-Tiny는 29.9k 샘플(합성 데이터로 큐레이션됨)으로 미세 조정했음에도 불구하고 BK-SDM-Tiny보다 IS 점수가 5.28(+17.4%) 높고, Contrastive Language–Image Pretraining (CLIP) 점수는 0.021 더 우수합니다. 프레셰 인셉션 거리 점수가 낮은 것은 합성 데이터의 분포가 COCO와 다르기 때문이며, Realistic Vision v5.1이 SD-v1.4보다 FID가 더 나쁘기 때문일 가능성이 큽니다. 그러나 인간 평가 결과(표2)에서 BK-SDM-Adv-Tiny(60.2% 승률)에 대한 강한 선호를 받은 것을 보면, COCO를 이상적인 데이터 분포라고 가정하는 것이 인간의 선호와 반드시 일치하지 않을 수 있음을 시사합니다. 우리의 최종 모델인 EdgeFusion은 182.2k 샘플(L-pps 및 합성 데이터)에서 BK-SDM-Adv-Tiny로부터 LCM-증류되어 BK-SDM-Tiny에 근접한 성능을 보여주고 CLIP 점수에서는 이를 능가합니다(표 1 참조).

표 1. 다양한 데이터셋에서 훈련된 BK-SDM-Tiny 아키텍처의 결과입니다. COCO 데이터셋에서 25단계로 평가되었습니다.

• 학생 모델을 먼저 미세 조정하는 것의 영향

  고급 증류 파이프라인에서 BK-SDM-Tiny의 미세 조정 과정을 생략하면 성능이 저하되어 인셉션 점수(IS)는 2.67 감소하고 프레셰 인셉션 거리(FID)는 2.72 증가하며, CLIP 점수는 0.01 감소합니다. 인간 선호도 평가(표 2)에서 이 결과를 확인할 수 있으며, EdgeFusion이 61.7%의 승률을 보였습니다.

표 2. 인간 선호도 평가 결과입니다. 개선된 데이터와 학생 모델 미세 조정이 없는 동일한 아키텍처와 비교했을 때의 우리 모델의 승률이 보고되었습니다 (1500회 비교, 21명 참가)

• 양자화

32비트 부동소수점 (FP32)과 W8A16 양자화된 EdgeFusion 간의 성능을 비교했습니다. 그림 3은 3가지 이미지 생성 요청에 대한 결과를 보여줍니다. W8A16 양자화를 적용한 후에도 성능 저하가 없었습니다.

그림 3. 32비트 부동소수점(FP32)과 W8A16 양자화된 EdgeFusion 모델(2단계) 간의 성능을 비교한 그림입니다.

• NPU 벤치마크

  표 3은 모델 수준 분할을 사용하지 않은 경우와 사용한 경우의 크로스 어텐션 블록에 대한 상대적인 시간 비율을 보여줍니다. 결과적으로 모델 수준 분할을 사용할 때 지연 시간이 약 73% 개선되었으며, 직접 메모리 접근(Direct Memory Access, 이하 DMA)이 효과적으로 감소되었습니다. 표 4는 EdgeFusion의 구성 요소별 추론 효율성을 보여줍니다.

표 3. 모델 수준 분할을 사용하면 추론 시간이 단축된다는 것을 확인할 수 있습니다. 크로스 어텐션 블록에서는 모델 수준 분할을 적용하지 않은 작업과 적용한 작업을 비교하여 상대적인 시간 비율을 계산했습니다.

표 4. Exynos 2400에서 모델 수준 분할을 적용한 경우와 적용하지 않은 경우의 벤치마크입니다.

결론

본 연구에서는 자원이 제한된 디바이스에서 SD 모델을 효율적으로 실행할 수 있도록 최적화된 EdgeFusion을 제안합니다. 연구 결과에 따르면, 유익한 합성 이미지-텍스트 쌍을 활용하는 것이 컴팩트한 SD 모델의 성능을 향상시키는 데 매우 중요하다는 것을 알 수 있었습니다. 또한, 고품질의 간단한 단계별 추론을 가능하게 하기 위해 LCM을 위한 고급 증류 과정을 설계했습니다. 배포 기술을 세밀하게 탐구함으로써 EdgeFusion은 삼성 엑시노스 NPU에서 프롬프트로부터 1초 만에 좋은 품질의 이미지를 생성할 수 있습니다.

이 연구에 대해 추가로 궁금한 사항이 있으시면, 📧 contact@nota.ai.로 문의해 주세요.

또한, AI 최적화 기술에 관심이 있으시면 🔗 netspresso.ai를 방문해 보세요.