DreamLite: 이미지 생성 및 편집을 위한 경량 온디바이스 통합 모델

왜 중요한가

기존의 고성능 이미지 생성 모델은 수십억 개의 파라미터를 가져 모바일 기기에서 실행하기 무거웠고, 생성과 편집을 별도의 모델로 처리해야 하는 불편함이 있었다. DreamLite는 단일 네트워크로 두 기능을 통합하면서도 크기를 획기적으로 줄여, 최신 스마트폰에서 고해상도 이미지를 실시간으로 처리할 수 있는 환경을 제공한다.

핵심 기여

온디바이스용 통합 확산 모델 DreamLite 개발

0.39B 파라미터 규모의 단일 네트워크에서 텍스트 기반 이미지 생성(T2I)과 이미지 편집(I2I)을 모두 지원하는 최초의 온디바이스 통합 모델이다.

인컨텍스트 컨디셔닝 메커니즘 도입

잠재 공간에서 타겟 이미지와 조건 이미지를 가로로 결합하여 입력하는 방식을 통해, 추가적인 파라미터 없이 생성과 편집 작업을 단일 구조로 통합했다.

단계적 작업 진행형 공동 사전학습 전략 수립

T2I 학습에서 시작해 편집 학습을 거쳐 통합 공동 학습으로 이어지는 3단계 전략을 통해, 용량이 작은 모델에서도 두 작업 간의 간섭을 최소화하고 안정적인 수렴을 달성했다.

DMD2 기반 4단계 추론 구현

수십 번의 노이즈 제거 과정을 4단계로 압축하여 모바일 기기에서 1024x1024 해상도의 이미지를 1초 미만의 속도로 처리할 수 있게 했다.

핵심 아이디어 이해하기

기존 확산 모델은 노이즈에서 이미지를 생성하기 위해 수십억 개의 파라미터가 필요하며, 특히 이미지 편집을 위해서는 원본 이미지를 입력받는 별도의 구조가 추가되어야 하므로 모바일 기기의 제한된 메모리에서 실행하기 어렵다는 한계가 있다. DreamLite는 이를 해결하기 위해 '인컨텍스트 공간 결합' 방식을 사용한다. 이는 모델의 입력 단계에서 생성할 영역과 참고할 영역을 가로로 붙여서 넣어주는 방식이다. 생성 작업 시에는 참고 영역을 빈 이미지로 채우고, 편집 작업 시에는 원본 이미지를 넣어줌으로써 모델이 동일한 연산 경로를 통해 두 작업을 자연스럽게 구분하고 수행하도록 유도한다. 또한, 모델의 크기가 작아지면 여러 작업을 동시에 배울 때 성능이 급격히 떨어지는 문제가 발생한다. 연구진은 이를 '단계적 학습'으로 해결했다. 먼저 이미지 생성 능력을 충분히 확보한 뒤, 이미지 편집 능력을 정렬하고, 마지막으로 두 능력을 통합하는 과정을 거쳐 작은 모델임에도 불구하고 대형 서버급 모델에 근접한 정교한 제어력을 갖추게 되었다.

방법론

아키텍처는 SDXL을 기반으로 가지치기된 모바일 최적화 U-Net을 사용한다. 트랜스포머 블록 수를 줄이고 채널 차원을 축소했으며, Multi-Query Attention(MQA)과 QK-RMSNorm을 도입해 연산 효율을 극대화했다. 인컨텍스트 컨디셔닝은 z_pair = Concat(z_tgt, z_cond) 수식을 통해 구현된다. [타겟 잠재 벡터 z_tgt와 조건 잠재 벡터 z_cond를 가로 방향으로 이어 붙여] → [U-Net의 입력으로 전달하여] → [단일 네트워크 연산을 수행하고] → [최종 결과물을 생성한다]. 학습 전략은 T2I 사전학습, 편집 사전학습, 통합 공동 학습의 3단계로 구성된다. 편집 학습 시에는 편집 영역이 작을 경우 손실 함수 값이 묻히는 것을 방지하기 위해 w(x) = log2(x) + 1 형태의 가중치를 적용한다. [전체 면적 대비 편집 면적의 비율 x를 입력으로] → [로그 연산을 거쳐 가중치 w를 계산하고] → [편집 영역의 손실값에 곱해줌으로써] → [작은 변화도 모델이 민감하게 학습하도록 만든다]. 사후 학습 단계에서는 0.5M 규모의 고품질 데이터셋으로 Supervised Fine-Tuning(SFT)을 진행한 후, HPSv3와 EditReward 모델을 보상 모델로 사용하는 Reinforcement Learning(RL)을 통해 인간의 선호도에 맞게 정렬한다. 마지막으로 DMD2 증류 기법을 통해 4단계 추론을 구현한다.

주요 결과

DreamLite는 GenEval 벤치마크에서 0.72점을 기록하여 SnapGen(0.70)이나 SANA-0.6B(0.64) 같은 기존 온디바이스 모델을 능가했다. 특히 객체 수 세기와 색상 속성 부문에서 높은 성능을 보였다. 이미지 편집 성능을 측정하는 ImgEdit 벤치마크에서는 4.11점을 획득하여, 12B 파라미터의 Kontext-Dev(3.76)나 7B의 BAGEL(3.42) 등 훨씬 거대한 서버급 모델들보다 우수한 편집 정확도를 입증했다. 실제 모바일 기기(Xiaomi 14) 배포 환경에서 1024x1024 해상도의 이미지 생성 및 편집을 1초 미만(약 0.42s, VAE 제외)에 완료하는 실시간 성능을 확인했다.

기술 상세

모델은 TinyVAE, Qwen3-VL-2B 텍스트 엔코더, 그리고 389M 파라미터의 경량화된 U-Net 백본으로 구성된다. U-Net은 고해상도 단계에서 Self-Attention 층을 제거하여 연산 복잡도를 줄였으며, 표준 컨볼루션을 Depthwise Separable Convolution으로 대체했다. 작업 구분을 위해 텍스트 프롬프트 앞에 [Generate] 또는 [Edit] 토큰을 추가하는 명시적 작업 라우팅을 사용한다. 이는 추가적인 파라미터 없이 공유된 파라미터 공간 내에서 모델이 작업 모드를 동적으로 전환하게 돕는다. 편집 학습 시 Foreground-emphasis Mask 기법을 사용하여 배경 보존과 편집 영역의 변화 사이의 균형을 맞춘다. 픽셀 단위 차이 계산, 팽창, 연결 요소 필터링을 거쳐 편집 마스크를 생성하고 이를 학습 가중치로 활용한다. 추론 가속을 위해 DMD2를 적용하여 Classifier-Free Guidance(CFG) 없이도 높은 품질을 유지하며 4단계 샘플링을 수행한다.

한계점

현재 텍스트 엔코더로 사용되는 Qwen3-VL-2B가 모델 전체 크기 대비 비중이 커서 온디바이스 배포 시 지연 시간을 유발한다. 또한, 매우 미세한 질감 표현이나 복잡한 인물 편집 시 대형 모델 대비 품질 저하가 발생할 수 있으며, TinyVAE의 정보 손실로 인한 재구성 흐릿함이 존재한다.

실무 활용

모바일 앱 내에서 별도의 서버 통신 없이 실시간으로 고품질 이미지 생성 및 편집 기능을 구현하는 데 최적화되어 있다.

모바일 사진 앱의 AI 지우개 및 개체 추가 기능
스마트폰용 실시간 텍스트-이미지 생성 도구
온디바이스 개인 정보 보호가 강조되는 이미지 편집 서비스

코드 공개 여부: 공개

코드 저장소 보기

키워드

Diffusion Model(확산 모델)On-Device AI(온디바이스 AI)Image Editing(이미지 편집)Model Compression(모델 압축)U-Net(U-네트)

용어 해설

인컨텍스트 공간 결합(In-context Spatial Concatenation): — 잠재 공간에서 타겟 이미지와 조건 이미지를 가로로 이어 붙여 모델의 입력으로 사용하는 방식이다. 추가적인 파라미터나 전용 브랜치 없이도 단일 네트워크가 생성과 편집 작업을 동시에 학습하고 수행할 수 있게 하며, 작업 간의 공간적 일관성을 유지하는 데 중요한 역할을 한다.
단계적 증류(Step Distillation): — 수십 단계의 반복적인 노이즈 제거 과정을 거치는 확산 모델의 추론 과정을 단 몇 단계로 압축하는 기법이다. 교사 모델의 출력 분포를 학생 모델이 학습하도록 유도하여, 시각적 품질 저하를 최소화하면서도 모바일 기기에서 실시간에 가까운 추론 속도를 달성하게 한다.
플로우 매칭(Flow Matching): — 노이즈와 데이터 분포 사이의 직선 경로를 학습하여 이미지 생성 과정을 최적화하는 훈련 목적 함수이다. 기존의 확산 확률 모델보다 학습이 안정적이고 수렴 속도가 빠르며, 특히 고해상도 이미지 생성에서 더 정교한 궤적 예측을 가능하게 하여 생성 품질을 높인다.
멀티 쿼리 어텐션(Multi-Query Attention): — 여러 개의 Attention Head가 Key와 Value 값을 공유하도록 설계하여 메모리 사용량과 연산량을 줄이는 기법이다. 모델의 파라미터 효율성을 높여주며, 특히 메모리가 제한된 모바일 환경에서 대규모 언어 모델이나 확산 모델을 효율적으로 실행하기 위해 필수적으로 사용된다.
인간 피드백 기반 강화학습(RLHF): — 모델의 출력 결과에 대해 인간의 선호도를 반영한 보상 모델을 학습시키고, 이를 통해 모델을 미세 조정하는 강화학습 기법이다. 이미지 생성에서는 구도, 색감, 지시어 준수 여부 등 주관적인 품질 요소를 개선하여 사용자가 원하는 결과물에 더 가깝게 정렬하는 데 사용된다.

DreamLite: 이미지 생성 및 편집을 위한 경량 온디바이스 통합 모델

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 피드