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

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

기존 확산 모델은 노이즈에서 이미지를 생성하기 위해 수십억 개의 파라미터가 필요하며, 특히 이미지 편집을 위해서는 원본 이미지를 입력받는 별도의 구조가 추가되어야 하므로 모바일 기기의 제한된 메모리에서 실행하기 어렵다는 한계가 있다. 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의 정보 손실로 인한 재구성 흐릿함이 존재한다.