On-the-fly Repulsion in the Contextual Space for Rich Diversity in Diffusion Transformers

一句话总结

本文提出在扩散 Transformer（DiT）多模态注意力块内部的”上下文空间”（被图像信息富化后的文本 token）上施加即时排斥力，在几乎不增加计算开销的前提下显著提升文生图的样本多样性，同时不破坏画质与语义对齐。

研究背景

• 领域现状：现代文生图扩散模型（尤其是 DiT 架构）语义对齐能力很强，但为了追求画质和人类偏好，往往对同一 prompt 收敛到很窄的一批”典型”视觉解，即所谓 typicality bias。Classifier-Free Guidance（CFG）会把概率分布锐化到单一模式，进一步压缩了多样性。
• 核心痛点：现有多样性方法都卡在”干预时机与位置”的根本权衡上。上游方法（改初始噪声或 prompt 嵌入）缺乏来自成形图像的结构反馈，要么语义漂移，要么需要昂贵的优化搜索；下游方法（在图像 latent 空间做排斥）作用得太晚——此时视觉结构已经固定，硬推样本会把它们推出数据流形，产生伪影。在少步”Turbo”蒸馏模型上这个问题尤其严重，因为生成路径几乎瞬间就定了。
• 本文 idea：找到一个既”结构上已被告知”又”概念上仍可塑”的中间表示——DiT 多模态注意力块中被图像特征反复富化的文本 token（上下文空间）。在这里施加排斥，可以在模型确定具体视觉模式之前重新引导其”生成意图”。

方法

整体框架：把一个 batch 的 \(B\) 个样本当作相互作用的粒子，沿用 Particle Guidance 的排斥思想，但不作用在图像 latent 上，而是作用在多模态注意力块之间流动的上下文文本 token \(\hat{f}_T\) 上。排斥在 Transformer 前向过程中”即时”完成，不需要对模型层反向传播，因此计算代价很小；且只在去噪早期若干时间步内介入。

flowchart LR
  A["同一 prompt 编码 p"] --> B["MM-Attention 块 l"]
  B --> C["富化文本 token f_T-hat 与图像 token f_I-hat"]
  C --> D["取上下文 token f_T-hat 计算多样性损失 L_div"]
  D --> E["按梯度做 M 次块内排斥更新"]
  E --> F["更新后的 token 送入下一块"]
  F --> G["仅在早期时间步介入，得到多样化图像"]

关键设计：

1. 上下文空间的定义与优势：在每个 MM-Attention 块中，文本特征 \(f_T\) 与图像特征 \(f_I\) 双向交互，文本 token 吸收当前正在形成的图像的空间构图信息，变成”既懂 prompt 又与具体图像同步”的富化表示 \(\hat{f}_T\)。关键性质是它有固定的 token 顺序：不像图像 latent 空间那样同一语义在不同样本里会落在不同空间位置，上下文空间里每个 token 索引在整个 batch 中大体对应同一概念，因此跨样本可直接比较、排斥不会撕裂几何结构。

2. 即时排斥机制：对样本 \(i\) 的上下文 token 做梯度式更新 \[\hat{f}_{T,i}^{(l)\prime} = \hat{f}_{T,i}^{(l)} + \frac{\eta}{M}\, \nabla_{\hat{f}_{T,i}^{(l)}} \mathcal{L}_{div}\big(\{\hat{f}_{T,j}^{(l)}\}_{j=1}^{B}\big)\] 其中 \(\eta\) 是总排斥强度，每个块内做 \(M\) 次迭代精修。由于每个样本的条件都从同一份未改动的 prompt 编码初始化，排斥只改变”同一 prompt 如何被视觉化实现”，避免了永久性语义漂移。

3. 基于 Vendi Score 的多样性目标：把每个样本在第 \(l\) 块的 \(N\) 个上下文 token 展平成向量 \(c_i^{(l)}\)，用余弦相似度构造核矩阵 \(K\)，其中 \[K_{ij} = \frac{\langle c_i^{(l)}, c_j^{(l)}\rangle}{\lVert c_i^{(l)}\rVert\, \lVert c_j^{(l)}\rVert}\] 对归一化核 \(\tilde{K} = \frac{1}{B}K\) 求特征值 \(\{\lambda_k\}\)，损失取其冯·诺依曼熵的负值 \[\mathcal{L}_{div} = -\sum_{k=1}^{B} \lambda_k \log \lambda_k\] 这相当于把上下文 token 推向更高维的流形，抵消 CFG 造成的语义坍缩。

4. 只在早期介入：去噪早期时间步对最终语义与全局构图最关键，也是 CFG 偏置最强的地方，因此排斥被限制在前若干个时间步的一个区间内，兼顾多样性与后期结构稳定。

作者还通过插值/外推实验佐证上下文空间的性质：在 VAE latent 空间做线性插值会因空间未对齐导致鬼影与结构模糊，外推更会迅速跌出流形；而在上下文空间做同样操作能得到平滑、语义连贯的高保真过渡，说明该空间把”生成意图”与”固定空间结构”解耦了。

实验结果

在 MS-COCO 2017 验证集上采样 1000 条 prompt、每条生成 4 张图（共 4000 张），用 Vendi Inception Score 衡量语义多样性，用 ImageReward（人类偏好）、VQAScore（细粒度 prompt 对齐）、KID（分布保真度）衡量质量，绘制 Pareto 前沿。方法在 Flux-dev、SD3.5-Turbo、SD3.5-Large 三种架构上均取得优于所有基线的多样性-质量折中；一项 45 人 450 次的用户研究也显示本方法在多样性、画质、对齐和总体偏好上普遍胜过基线，唯一与其接近的是 SGI。

下表为生成一组 4 张图的运行时对比（主实验之一），凸显本方法相对基模型仅增加约 20%–30% 开销，而基于优化搜索的 SGI 随候选数增加代价急剧上升：

方法	SD3.5-Large	SD3.5-Turbo	Flux-dev
基模型	13.83s	4.18s	10.34s
本文（Contextual）	18.12s	5.52s	12.80s
SGI（8 候选）	66.79s	13.15s	47.47s
SGI（16 候选）	76.79s	23.73s	56.32s
SGI（64 候选）	145.14s	91.30s	113.99s

消融实验表明：随排斥强度 \(\eta\) 增大，结构性变化（加塔楼、换季节、改地形）逐步增多且画质保持；把同样的排斥改施加在图像注意力 token \(\hat{f}_I\) 上时，因空间刚性导致布局停滞、只能靠改局部纹理制造”多样性”并产生伪影，Pareto 前沿明显劣于上下文空间。此外方法可零改动地嵌入图像编辑模型 Flux-Kontext，产出多样但保持编辑语义的编辑结果。

亮点与局限

• 亮点：
  • 识别出 DiT 特有的”上下文空间”这一兼具语义灵活性与结构信息的干预层级，理念清晰且有插值/外推实验支撑。
  • 即时前向干预、无需反传、无额外显存，开销仅增 20%–30%，尤其能在传统轨迹式方法失效的少步 Turbo/蒸馏模型上生效。
  • 通用性强，跨三种 DiT 架构有效，还能直接迁移到图像编辑（Flux-Kontext）。
• 局限：
  • 无法直接控制”哪些属性”发生变化，有时偏向粗粒度语义改动而非用户想要的细节变化。
  • 干预集中在生成早中期，如何与后期阶段协同、或与其他控制机制结合仍是开放问题。

延伸思考

• 作者提出的未来方向很值得追：用一个文本线索（如 “color”/”size”）把排斥力偏置到上下文空间中的某个语义方向，实现”可解释的定向多样性”，让变化集中在选定属性上而其余保持稳定。
• 该工作与一系列基于注意力/富化文本表示做 training-free 编辑与控制的 DiT 研究一脉相承（同组的 Stable Flow 等），说明”在正确的表示层级介入”正在成为操控 DiT 的通用范式。
• 论文也点出一个评测困境：标准指标偏爱”典型”样本，可能系统性低估真正有创意的多样性，这提示多样性任务需要更贴合人类判断的评价体系。