CoLoGen: Progressive Learning of Concept-Localization Duality for Unified Image Generation¶
会议: CVPR 2026
arXiv: 2602.22150
代码: 暂无
领域: 统一图像生成 / 扩散模型
关键词: 统一生成框架, 概念-定位对偶性, 渐进式训练, 专家路由, FLUX
一句话总结¶
提出 CoLoGen,一个基于"概念-定位对偶性"(Concept-Localization Duality)的统一图像生成框架,通过渐进式分阶段训练和 Progressive Representation Weaving(PRW)动态专家路由架构,在指令编辑、可控生成和个性化生成三大任务上同时达到或超越专用模型水平。
研究背景与动机¶
统一多模态图像生成(涵盖 mask inpainting、visual grounding、可控生成、个性化生成、指令编辑)面临的核心困境是表征冲突:
- 概念表征 \(\mathcal{R}_c\):编码语义一致性和物体级理解,可控生成(如 canny/depth/seg 条件)主要依赖此能力
- 定位表征 \(\mathcal{R}_l\):编码空间对齐、几何和结构一致性,个性化生成需精确定位参考图中的身份特征
现有统一框架将这两种异质表征强行共享,导致概念理解和空间精度互相干扰(联合优化 \(f_c\) 可能损害 \(f_l\))。这解释了为何现有通用模型往往在部分任务上表现良好而在其他任务上退化。
方法详解¶
整体框架¶
基于 FLUX.1 架构构建,核心由两部分组成:
- PRW(Progressive Representation Weaving):嵌入每个 MMDiT block 的动态专家路由模块
- 渐进式分阶段训练:5 步训练策略,从基础能力到复杂任务逐步递进
关键设计¶
-
Progressive Representation Weaving (PRW):在每个多模态注意力块中,为 source latent 的 KV 投影引入专家池 \(\{E_k\}_{k=1}^N\) 和动态路由器 \(G\)。路由器通过带噪声的 top-1 softmax 选择最相关专家:
\(\mathbf{w} = hW_r + \epsilon \odot \text{softplus}(hW_n), \quad \epsilon \sim \mathcal{N}(0, \mathbf{I})\) \((K_{\hat{h}}, V_{\hat{h}}) = \text{KV\_proj}_{\text{base}}(h) + \sum_{k \in \mathcal{S}} \text{softmax}(\mathbf{w})_k E_k(h)\)
注意力分两步:先让 source latent 自注意力融入专家信息,再让 noisy/text latent 与之交互。设计动机:让不同任务自动激活不同专家,避免表征混淆。
-
渐进式分阶段训练策略:5 步从易到难:
- Step 0-1(内生预训练):Mask Inpainting(3M 合成数据)学概念 + Visual Grounding(1M 数据)学定位
- Step 2(条件注入):Controllable Generation(20M 数据)适配 Canny/Depth/HED/Lineart/Seg
- Step 3-4(指令-图像对齐):Customized Generation(200K)+ Instruction Editing(1.6M)
每步只解锁新专家 \(E_{N-1}\),历史专家冻结保留已学知识,类似终身学习。
-
Veteran Gate Routing Supervision:为平衡新旧专家利用率,引入辅助损失约束新专家的路由密度:
\(\mathcal{L}_{\text{veteran}} = \alpha \cdot |U_t - \rho|, \quad U_t = \frac{1}{L_n} \sum_{i=1}^{L_n} \mathbb{I}(e_i = N-1)\)
其中 \(\rho = 0.8\) 表示期望新专家被激活 80%,剩余 20% 保留给历史专家。总损失 \(\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \mathcal{L}_{\text{veteran}}\)。
损失函数 / 训练策略¶
- 主损失:标准扩散生成损失 \(\mathcal{L}_{\text{task}}\)(Flow Matching)
- 辅助损失:Veteran Gate Routing Supervision \(\mathcal{L}_{\text{veteran}}\),权重 \(\alpha = 0.5\)
- PRW 专家使用 LoRA(rank=128)实现参数高效
- 各阶段训练 200K-400K iterations,global batch size 128-256
- 总训练数据:~25M 样本
实验关键数据¶
主实验¶
| 任务 / 数据集 | 指标 | CoLoGen | 之前 SOTA | 对比 |
|---|---|---|---|---|
| 指令编辑 / Emu Edit | DINO ↑ | 0.843 | 0.831 (UniReal) | +0.012 |
| 指令编辑 / MagicBrush | DINO ↑ | 0.932 | 0.879 (Emu Edit) | +0.053 |
| 指令编辑 / MagicBrush | CLIP_out ↑ | 0.301 | 0.308 (UniReal) | -0.007 |
| 可控生成 / MultiGen-20M | Canny CLIP-S ↑ | 33.31 | 32.15 (ControlNet) | +1.16 |
| 可控生成 / MultiGen-20M | Depth RMSE ↓ | 31.79 | 33.83 (PixWizard) | -2.04 |
| 个性化生成 / DreamBench | DINO ↑ | 0.714 | 0.702 (UniReal) | +0.012 |
| 个性化生成 / DreamBench | CLIP-T ↑ | 0.315 | 0.326 (UniReal) | -0.011 |
消融实验¶
| 配置 | CLIP-T ↑ | CLIP-I ↑ | DINO ↑ | 说明 |
|---|---|---|---|---|
| Baseline (w/o \(\mathcal{R}_l\) & \(\mathcal{R}_c\)) | 0.260 | 0.889 | 0.901 | MagicBrush 上无专家 |
| w \(\mathcal{R}_l\) only | 0.279 | 0.922 | 0.927 | 定位表征提升结构保持 |
| w \(\mathcal{R}_c\) only | 0.302 | 0.881 | 0.905 | 概念表征提升指令跟随 |
| Co-training (\(\mathcal{R}_c\) & \(\mathcal{R}_l\)) | 0.269 | 0.918 | 0.922 | 联合训练反而不如单独 |
| CoLoGen (渐进式) | 0.301 | 0.931 | 0.932 | 渐进训练解决冲突 |
关键发现¶
- Co-training 策略在个性化生成上 DINO 和 CLIP-I 甚至低于 baseline,验证了"表征冲突"假设
- 渐进式训练在各指标上全面优于 co-training,证明分阶段学习能有效缓解概念-定位冲突
- Veteran Gate Routing \(\rho = 0.8\) 最优;\(\alpha\) 过大反而限制灵活性
- LoRA rank=128 为最佳设置
亮点与洞察¶
- 概念-定位对偶性是对统一图像生成困境的深刻理论洞察——将"不同任务需要不同表征"形式化为两个竞争子空间
- PRW 架构巧妙复用 MoE 思想,但将其限制在 KV 投影层且仅 top-1 路由,保持轻量
- 渐进训练 + 专家冻结是终身学习在生成模型中的有效应用,有效缓解灾难性遗忘
- 数据工程细致:mask inpainting 的 3 种 mask 类型(random/object-shaped/Bessel curve 不规则)、20:40:40 采样比例
局限与展望¶
- 随着任务和专家数量增加,PRW 的内存占用持续增长,可扩展性受限
- 当前仅验证了 5 个任务,对更多条件类型(如 pose、sketch)的泛化能力未知
- 部分指标略低于 UniReal 等专用模型,统一模型在极致性能上仍有差距
- LoRA rank=128 参数量已不算"轻量",与 full fine-tuning 的性能差距未报告
相关工作与启发¶
- 与 OmniGen(统一多模态生成但无显式表征管理)相比,CoLoGen 通过 PRW 实现了更好的编辑/定制平衡
- 与 PixWizard 和 UniReal 等通用编辑模型相比,CoLoGen 的核心优势在于可控生成方面的兼顾
- 启发:统一生成不应追求"一个表征适配所有任务",而应通过动态路由让模型学会任务自适应的表征切换
评分¶
- 新颖性: ⭐⭐⭐⭐ 概念-定位对偶性视角新颖,PRW + 渐进训练的组合设计系统性强
- 实验充分度: ⭐⭐⭐⭐ 覆盖编辑/可控/个性化三大任务线,6 个基准,消融深入
- 写作质量: ⭐⭐⭐⭐ 问题定义清晰,图示质量高,训练策略描述详尽
- 价值: ⭐⭐⭐⭐ 为统一图像生成提供了有理论基础的实用方案,PRW 可迁移到其他多任务场景