DreamLite：面向端侧设备的轻量级图文生成与编辑统一模型 DreamLite: A Lightweight On-Device Unified Model for Image Generation and Editing
0.39B参数统一扩散模型，在手机端实现<1秒图文生成与编辑
前置知识
扩散模型（Diffusion Models）
扩散模型是一类基于概率的生成模型，通过逐步向数据添加高斯噪声将数据分布转化为噪声分布，然后学习逆向去噪过程来生成新样本。在图像生成领域，扩散模型如Stable Diffusion、FLUX等已经成为主流范式。本文基于Flow Matching框架，通过线性插值构建噪声潜变量 $z_t = t \cdot z + (1-t) \cdot \epsilon$，模型学习预测从噪声到数据的速度场 $v_t = z - \epsilon$。理解扩散模型的基本原理对于把握本文的技术路线至关重要。
本文的核心是在扩散模型框架下实现高效的统一生成与编辑，所有技术设计都围绕扩散模型展开。
U-Net架构
U-Net是一种编码器-解码器架构，最初用于医学图像分割，在扩散模型中被广泛用作去噪网络。它包含下采样路径（提取特征）和上采样路径（恢复分辨率），并通过跳跃连接将编码器特征传递给解码器。本文基于SnapGen的压缩版SDXL U-Net，通过一系列优化（剪枝、移除高分辨率层自注意力、使用分离卷积、多查询注意力等）将2.5B参数基线压缩至389M参数。
本文的骨干网络是经过极致压缩的U-Net，理解U-Net架构有助于理解本文的模型设计决策。
潜变量扩散模型（Latent Diffusion Models）
潜变量扩散模型不是在像素空间直接进行扩散，而是先用变分自编码器（VAE）将图像编码到低维潜空间，在潜空间中执行扩散过程，最后再解码回像素空间。这种方法大幅降低了计算成本。本文采用TinyVAE（仅2.5M参数，8x8下采样因子）将图像映射为4通道潜变量，实现高效的训练和推理。
本文的整个架构建立在潜变量扩散框架之上，VAE的选择直接影响模型效率和生成质量。
In-context Conditioning（上下文条件机制）
In-context conditioning是本文提出的核心创新，通过在空间维度（宽度）上拼接目标图像和条件图像的潜变量来统一生成和编辑任务。具体地，构造双面板潜变量 $z_{pair} = \text{Concat}(z_{tgt}, z_{cond})$，其中左侧为目标输出，右侧为视觉条件。对于文生图任务，条件面板设为全黑空白图像；对于编辑任务，条件面板为源图像。这种设计使模型无需额外模块即可从生成扩展到编辑。
这是本文的核心技术贡献，理解这一机制是理解本文如何统一生成和编辑的关键。
任务渐进式联合预训练（Task-progressive Joint Pretraining）
这是本文提出的训练策略，将预训练分为三个渐进阶段：(1) T2I预训练：建立强大的生成先验；(2) 编辑训练：激活in-context条件机制，使视觉条件与生成潜空间对齐；(3) 统一联合训练：在T2I和编辑数据的混合上进行训练。这一策略的关键是在联合优化前引入中间编辑预训练阶段，缓解任务干扰，确保紧凑模型的稳定收敛。
统一训练生成和编辑任务面临优化目标冲突的问题，这一训练策略是确保模型稳定工作的关键。
强化学习对齐（RLHF / ReFL）
强化学习从人类反馈（RLHF）是一种对齐技术，利用奖励模型的信号来优化生成模型的行为。本文采用ReFL（Reward Feedback Learning）框架，通过奖励模型对去噪轨迹进行梯度反向传播。对于文生图任务使用HPSv3作为奖励模型（阈值b=11），对于编辑任务使用EditReward（阈值b=2.5），采用ReLU截断奖励公式 $\mathcal{L}_{rl} = -\max(0, r(c, \hat{x}) - b)$ 来确保训练稳定性。
RLHF对齐显著提升了生成质量和指令遵循能力，是模型性能超越基线的重要因素。
分布匹配蒸馏（DMD2）
分布匹配蒸馏是一种将多步扩散模型压缩为少步生成器的技术。DMD2通过最小化真实图像分布 $p_{real,t}$ 和生成器输出分布 $p_{fake,t}$ 之间的近似KL散度来实现蒸馏。本文使用DMD2将采样过程压缩到仅4步去噪，同时配合GAN损失 $\mathcal{L}_{GAN}$ 保持生成多样性和真实感，使得在移动端实现近实时推理成为可能。
端侧部署的关键瓶颈是推理速度，DMD2蒸馏是实现<1秒推理的核心技术。
研究动机
当前大规模扩散模型在文生图（T2I）和文本引导图像编辑（I2I）方面取得了显著进展，但存在严重的部署障碍。以FLUX系列为例，其DiT骨干网络扩展到12B参数，带来极高的内存需求和推理延迟，无法在消费级设备上高效部署。虽然近期研究探索了轻量级架构如SANA（约2B参数），但在移动硬件上实现实时性能仍是重大挑战。更重要的是，现有的端侧扩散模型（如SnapFusion、Mobile Diffusion、SnapGen等）主要聚焦于T2I生成，缺乏对图像编辑的支持。而在实际应用中，创作者需要在单一应用中无缝集成生成和编辑功能。部署两个独立模型会显著增加系统复杂性和资源消耗，特别是在内存受限的移动设备上。这种生成能力与编辑能力的割裂，以及大模型与端侧部署之间的矛盾，构成了亟待解决的核心问题。
本文的目标是本文的具体目标是设计一个紧凑的统一扩散模型（compact unified diffusion model），能够在单一网络中同时支持文生图生成和文本引导图像编辑，并且能够在移动设备上实现实时部署。量化目标包括：在标准基准测试上达到与大型模型相当的性能水平，在智能手机上实现1024x1024分辨率图像的生成或编辑在1秒内完成。具体而言，团队希望模型参数量控制在0.5B以下，推理步数压缩到4步，并在GenEval、DPG、ImgEdit和GEdit等主流基准上取得有竞争力的结果。
与已有工作不同的是，本文的独特切入角度在于同时解决统一性和高效性两个维度的挑战。现有工作要么专注于大参数量的统一模型（如FLUX 2、OmniGen2等，参数量在4B-12B），要么专注于轻量级但仅支持单一任务的模型（如SnapGen仅支持生成、EditMGT仅支持编辑）。本文的创新在于：(1) 提出in-context空间拼接机制，在输入层面统一生成和编辑任务，避免了传统Pix2Pix式通道拼接对生成先验的退化；(2) 设计任务渐进式联合预训练策略（T2I到编辑到统一联合），解决小容量模型多任务训练的稳定性问题；(3) 通过系统化的架构压缩（从2.5B到389M）和蒸馏（4步采样），实现真正的端侧实时推理。这种小而统一的技术路线填补了现有研究的空白。
核心方法
DreamLite的整体技术路线可以分为三个层面理解。从直觉上看，本文的核心思想是用空间换通道。传统方法（如InstructPix2Pix）将条件图像与噪声潜变量在通道维度拼接，这会破坏预训练T2I模型的生成先验；而DreamLite采用空间维度拼接，将目标图像和条件图像并排放置，生成任务用（目标|空白）配置，编辑任务用（目标|源图像）配置。技术路线方面，首先基于SnapGen的压缩SDXL U-Net构建0.39B参数的骨干网络，配合TinyVAE（2.5M）和Qwen3-VL-2B文本编码器。然后通过三阶段渐进式预训练建立统一能力：T2I预训练建立生成先验、编辑训练对齐视觉条件、统一联合训练巩固多任务能力。最后通过SFT、RLHF对齐和DMD2蒸馏实现高质量的4步推理。这套方案使得在Xiaomi 14手机上，1024x1024图像的生成或编辑可在不到1秒内完成。
本文的核心创新是in-context条件机制及其配套的任务渐进式联合预训练策略，这两个设计点与已有方法有本质区别。首先，区别于InstructPix2Pix范式的通道维度拼接，DreamLite在空间（宽度）维度进行拼接：$z_{pair} = \text{Concat}(z_{tgt}, z_{cond})$，其中左侧为目标输出，右侧为视觉条件。这种设计的关键优势在于：生成被视为in-context引用的特殊情况（条件为空白图像），使得模型可以利用统一的空间先验处理不同行为。其次，为解决紧凑模型统一训练的稳定性问题，本文提出任务渐进式联合预训练：不是直接从T2I跳到联合训练，而是引入中间编辑预训练阶段，在联合优化前将视觉条件表示与生成潜空间对齐，从而缓解任务干扰。此外，为处理编辑任务中编辑区域通常很小的问题，本文设计了前景强调掩码策略，通过像素差异、膨胀、连通分量过滤、最大池化下采样的四步流程估计编辑区域，并使用对数加权函数 $w(x) = \log_2(x) + 1$ 平衡训练稳定性与小编辑的敏感性。
方法步骤详情
DreamLite的完整方法包含以下步骤：(1) 模型架构搭建：采用压缩的U-Net骨干（389M参数），通过移除高分辨率层自注意力、使用扩展分离卷积、设置FFN扩展比为3、采用多查询注意力（MQA）单KV头、阶段对齐（下采样和上采样从3层减至2层）、添加QK-RMSNorm和轻量文本投影器等优化，将SDXL 2.5B基线压缩至0.39B。配合TinyVAE（2.5M参数，4通道潜变量，8倍下采样）和Qwen3-VL-2B文本编码器。(2) T2I预训练：使用Flow Matching目标 $\mathcal{L}_{fm} = \mathbb{E}_{t,z,\epsilon,y}[\|v_\theta(z_t, t, y) - (z - \epsilon)\|^2]$，通过渐进分辨率课程（256到512到1024）训练，使用logit-normal噪声采样器和动态时间偏移。(3) 编辑预训练：激活in-context条件机制，使用前景强调掩码加权的Flow Matching损失处理局部编辑任务。(4) 统一联合训练：在T2I和编辑数据的1:1混合比例上训练，使用[Generate]和[Edit]任务标记。(5) SFT：在精选的0.5M高质量样本上微调。(6) RLHF：使用ReFL框架，T2I用HPSv3（b=11），编辑用EditReward（b=2.5）。(7) DMD2蒸馏：压缩至4步采样，配合GAN损失。
技术新颖性
DreamLite的技术新颖性体现在多个层面。首先，in-context空间拼接机制是对传统InstructPix2Pix通道拼接范式的根本性创新。实验表明，在相同编辑训练阶段（T2I到编辑），in-context机制将ImgEdit从3.67提升至3.88；在统一训练阶段（T2I到编辑到统一），GenEval从0.61提升至0.71。这证明空间拼接更适合多任务统一框架。其次，任务渐进式联合预训练是解决小容量模型多任务训练稳定性的新方案。实验显示，直接联合训练（T2I到统一）导致GenEval仅0.65、ImgEdit仅3.14，而渐进式方案达到0.71和3.94，甚至超越单独任务的性能。第三，前景强调掩码策略是对编辑任务中编辑区域小导致梯度信号被淹没问题的创新解决方案，通过自适应加权使模型更关注实际编辑区域。第四，系统化的U-Net压缩路径（8步优化，从2.5B到389M，FLOPs减少80%以上）为端侧统一模型提供了可复用的架构设计范式。
实验结果
DreamLite在多个主流基准上取得了令人瞩目的结果，验证了紧凑统一模型的可行性。在文生图任务上，GenEval基准中DreamLite（0.39B）以0.72的总分超越了SANA-0.6B（0.64）、SnapGen（0.70）和SnapGen++（0.66），甚至接近OmniGen2（4B，0.80）和BAGEL（7B，0.82）。在DPG基准上，DreamLite达到85.8分，超越SnapGen（81.1）和SANA-0.6B（83.6），接近LongCat-Image（6B，86.6）。在图像编辑任务上，ImgEdit基准中DreamLite以4.11分大幅超越EditMGT（0.96B，2.89）和VIBE（1.6B，3.85），成为所有轻量级模型中的SOTA。在GEdit-EN基准上，DreamLite的总体分为6.88，感知质量分达7.54，甚至超越了DeepGen1.0（2B，7.47）和LongCat-Image（6B，7.49）。消融研究进一步验证了各组件的贡献：in-context机制优于Pix2Pix式拼接，渐进式预训练优于直接联合训练，RLHF对齐同时提升生成和编辑性能（ImgEdit从3.94提升至4.11，GenEval从0.71提升至0.72），DMD2蒸馏将采样压缩至4步仅带来轻微性能下降。在端侧部署方面，Snapdragon 8 Gen3上U-Net单步推理仅需103.84ms，4步总推理时间约415.36ms（不含VAE），加上VAE解码（约22ms）和系统开销，端到端用户体验接近1秒阈值。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 文生图（Text-to-Image Generation） | GenEval Overall | 0.72 | SnapGen 0.70, SANA-0.6B 0.64, Nitro-E-GRPO 0.70 | 超越所有小于1B参数基线，接近OmniGen2（4B，0.80） |
| 文生图（Text-to-Image Generation） | DPG Overall | 85.8 | SnapGen 81.1, SANA-0.6B 83.6, Nitro-E-GRPO 81.1 | 超越所有小于1B参数基线，接近LongCat-Image（6B，86.6） |
| 图像编辑（Image Editing） | ImgEdit Overall | 4.11 | EditMGT 2.89, VIBE 3.85, DeepGen1.0 4.03 | 超越所有轻量级模型，接近DeepGen1.0（2B，4.03） |
| 图像编辑（Image Editing） | GEdit-EN Overall | 6.88 | EditMGT 6.33, VIBE 7.28, DeepGen1.0 7.54 | 在0.39B规模下实现有竞争力的编辑能力 |
| 端侧推理（On-device Inference） | 1024x1024生成/编辑时间 | <1秒（Xiaomi 14） | 传统多步扩散模型需数秒至数十秒 | 4步采样，U-Net单步103.84ms（Snapdragon 8 Gen3） |
局限与改进
尽管DreamLite展现了紧凑统一模型的潜力，但仍存在若干局限性。首先，文本编码器规模问题：虽然U-Net骨干仅0.39B参数，但当前管线仍依赖2B参数的Qwen3-VL-2B文本编码器，在端侧部署时引入了不可忽视的延迟。目前的解决方案是预计算常用提示词的嵌入，但这限制了用户交互的灵活性。其次，VAE重建保真度不足：TinyVAE仅1.2M参数，在处理复杂结构细节（如文本生成、文本编辑和肖像编辑中的身份保持）时可能出现信息丢失或重建模糊。这反映在GEdit基准上相对较低的分数（6.88 vs LongCat-Image的7.55）。第三，虽然DMD2蒸馏将采样压缩至4步，但在复杂语义对齐任务上仍存在性能下降（GenEval从0.72降至0.70，ImgEdit从4.11降至3.8）。第四，从作者自己的观察来看，模型在极端精细纹理的人像渲染上与SANA-1.6B等较大模型存在边际差距。此外，当前的评估主要集中在英文提示词上，对多语言场景的支持尚不明确。
独立分析的弱点
从独立分析的角度来看，DreamLite存在以下几个值得关注的弱点：(1) 文本编码器瓶颈：Qwen3-VL-2B（2B参数）相对于U-Net骨干（0.39B）过大，形成了头重脚轻的架构。预计算嵌入的方案虽然绕过了推理延迟，但牺牲了用户交互的实时性和灵活性。改进方向是开发专用的轻量级文本编码器（小于1B），或者探索文本编码器的量化和蒸馏技术。(2) VAE容量不足：TinyVAE的1.2M参数和8倍下采样因子可能导致信息损失，特别是在需要精确重建的场景（如文字、人脸）。可以考虑训练稍大的高保真VAE，或采用多尺度VAE策略。(3) 编辑能力的泛化性：虽然在ImgEdit基准上表现优异，但模型在面对分布外的复杂编辑指令时的鲁棒性尚不明确。可以引入更多样化的编辑数据和更复杂的评估场景。(4) 多模态对齐：当前的RLHF仅使用单一奖励模型，可能无法全面捕捉人类偏好。可以探索多奖励模型的组合或更先进的对齐技术。
未来方向
基于本文的成果，未来研究方向可以从多个维度展开。作者明确提出的方向包括：(1) 开发轻量级文本编码器（小于1B）以实现全管线端侧推理；(2) 训练稍大的高保真VAE以改善重建质量；(3) 针对文本生成、人脸编辑等特定任务进行专门微调；(4) 探索更先进的步蒸馏方案，在后训练阶段结合奖励模型以更好地对齐人类审美偏好。基于本文成果可延伸的方向包括：(5) 将in-context条件机制扩展到更多视觉任务（如超分辨率、修复、视频编辑等）；(6) 探索更高效的架构压缩技术，进一步降低参数量和FLOPs；(7) 研究多语言支持，使模型能够处理中文等非英文提示词；(8) 将任务渐进式预训练策略应用于其他多任务统一模型的设计中；(9) 探索联邦学习或持续学习方案，使端侧模型能够根据用户偏好进行个性化调优。
复现评估
从复现评估的角度来看，本文的复现存在一定挑战但也提供了较好的基础。有利因素包括：(1) 论文详细描述了架构演进过程（图3展示了8步压缩路径）和各阶段的训练细节（学习率、批量大小、数据配比等）；(2) 使用了开源的基线架构（SnapGen基于SDXL）和公开的评估基准（GenEval、DPG、ImgEdit、GEdit）；(3) 奖励模型（HPSv3、EditReward）和蒸馏方法（DMD2）均有公开实现。挑战因素包括：(1) 训练数据规模庞大（20M T2I样本+1.74M编辑样本），完整复现需要大量算力（批量大小576，多阶段训练）；(2) 部分数据来源为内部数据集或商业数据，可能难以完全获取；(3) 论文未明确说明是否开源模型权重或训练代码；(4) 端侧部署涉及模型量化（W8A8）和平台优化（Snapdragon 8 Gen3、Dimensity 9300），复现这些结果需要特定硬件和优化经验。总体而言，学术研究者可能难以完全复现端侧部署结果，但核心算法和技术路线应该可以基于公开信息复现。
论文图表