<aside>

🧠 Paper Idea

Task-Aware Expert Routing: 给 MoE 路由加入任务语义信息，让专家按任务分工，从而缓解统一图像生成中的任务干扰。 本质是：把 global intent 注入 local decision.

</aside>

一 问题

<aside> 💡

❗ 多任务干扰（Task Interference）

统一模型（生成 + 编辑）中：

• 局部编辑 → 需要 强保留
• 主体生成 → 需要 强变化

👉 冲突：同一套参数难以同时满足 → 性能下降

不同图像生成任务的目标之间存在内在冲突。 编辑任务（例如风格迁移、物体移除）要求在修改局部时精确保留区域信息，而定制化任务（例如主体驱动的生成）则要求在新的上下文中保持强一致的身份特征。 由于缺乏对这些截然不同且往往相互竞争的需求的显式建模，现有方法难以自适应地满足用户多样化的意图，限制了鲁棒性和泛化能力。

</aside>

<aside> 💡

❗ MoE 的不足

MoE 是：每个词（token）完全独立地、单独地选择自己的专家，独立计算，最后独立加权。 (a) 词 A 选 专家 2、5 (b) 词 B 可以选 专家 0、3 (c)互相完全不干扰。

虽然 MoE 可以提升容量，但：

• 路由是 task-agnostic
• 只基于 token feature
• 不感知任务目标

👉 结果：

• 专家分工混乱
• 无法真正解耦任务 </aside>

二 方法简述

<aside>

🧩 (1) 任务语义建模（Task Descriptor）

将任务结构化为：

• Scope：局部编辑 / 全局编辑
• Type：属性编辑 / 动作编辑
• Preserve：主体保持 / 风格保持 / 结构保持

得到统一任务表示

</aside>

<aside>

🔀 (2) Task-Aware MoE Routing

传统：router(token)

现在：router(token, task_semantics)

👉 路由同时依赖：

• 局部信息（token）
• 全局信息（task） </aside>

<aside>

🎯 (3) 预测对齐正则（核心）

核心约束： routing pattern → 可以预测 task semantics

等价于：

• 相似任务
• 使用相似专家分布

👉 本质：把“任务信息”压进路由空间

</aside>

三 Key Insight

<aside>

✅ Insight 1：MoE 的瓶颈在“路由”，不是“专家数”

• 不是专家不够，是不会用

✅ Insight 2：任务级语义是缺失变量

• token-level 信息不足
• 必须引入 global condition

✅ Insight 3：中间层监督很关键

• 不是只监督 output，而是监督 routing behavior </aside>

四 方法详情

<aside>

🧩 模型架构与训练（MM-DiT + MoE）

<aside>

TAG-MoE 包括：(1) 带有 MoE 层的 MM-DiT；(2) 一种分层任务语义标注，使用原子任务描述符对训练数据进行标注；(3) 一种新颖的语义对齐路由器，通过预测性对齐正则化将 MoE 路由行为与任务语义显式对齐

</aside>

🧠 输入建模（Unified Token Framework）

模型基于 MM-DiT，将多模态输入统一为一个 Token 序列：

• 文本输入：使用预训练 MLLM 编码为文本 embedding 序列：C
• 图像输入：条件图像：I_c → z_c（通过 VAE 编码器 ℰ） 目标图像：I_0 → z_0
• 扩散过程：对 z_0 加噪得到：z_t
• Token 化： z_c 和 z_t → 切片为视觉 tokens

🔗 最终输入序列

将以下内容拼接为一个统一序列输入 MM-DiT：

文本 tokens：C | 条件图像 tokens：z_c| 带噪图像 tokens：z_t

结合 时间步 embedding：t

⚙️ MoE 架构设计

• 将 Transformer 中图像流的 后几层FFN 层替换为 MoE（高层语义更受益）

🧠 MoE 组成

• 专家网络：${E₁, ..., E_N}$
• 门控网络：𝒢(x)（输出每个专家的权重）

🔀 路由机制

对每个 token：

1. 门控网络计算专家分布 𝒢(x)
2. 选择 top-k 专家 𝒯
3. 输出为加权和：

$$ \text{MoE}(x) = \sum_{E_i \in \mathcal{T}(x)} \mathcal{G}(x)_i \cdot E_i(x) $$

👉 即：稀疏激活 + 专家加权融合

🚀 优势

• 在计算预算不变下：
• 更适合：

🎯 训练目标

• 使用 Flow Matching 进行端到端训练

📌 Key Takeaway

• 将多模态输入统一为单序列 Transformer 建模
• 用 MoE 替换 FFN → 提升容量但保持计算可控
• MoE 主要放在深层，专注高层语义

</aside>

<aside>

</aside>

<aside>

</aside>

🧩 3.4 数据集构建