Chimera:设计与 Chinchilla 扩展的混合视觉扩散 Transformer Chimera: Designing and Chinchilla-Scaling Hybrid Visual Diffusion Transformers
混合线性/全局注意力的视觉扩散 Transformer,以模块级扩展律实现 7.3 倍算力效率
前置知识
Diffusion Transformer (DiT) 与整流流 (Rectified Flow)
DiT 用 Transformer 替代 U-Net 作为扩散模型的去噪骨干,把图像/视频经 VAE 编码为 latent token 序列后逐步去噪。整流流是流匹配的一种,用线性插值 $z_\tau = \tau z_1 + (1-\tau) z_0$ 连接噪声 $z_0$ 与干净数据 $z_1$,预测速度 $v_\tau = z_1 - z_0$,路径比传统 DDPM 更直、采样更高效。
Chimera 的整个骨干、损失函数和 timestep 调制都建立在 DiT + 整流流之上,这是理解其训练目标以及 KDA 如何接入去噪过程的基础。
Kimi Delta Attention (KDA) 与线性/门控注意力
KDA 源自 Kimi Linear / Gated Delta-Net,是一种线性复杂度 $O(N)$ 的注意力变体。它维护与序列长度无关的固定大小递归状态 $S_t$,按因果顺序扫描:用通道级遗忘门 $\alpha_t$ 对状态逐通道衰减,再用标量 $\beta_t$ 控制写入,效果像自适应窗口的滑动注意力,但不像 softmax 那样产生 $N \times N$ 矩阵,长序列显存恒定。
KDA 是 Chimera 长上下文效率的核心,3/4 的注意力层都用它;理解其状态更新机制才能看懂 NoPE 和零样本长度外推从何而来。
Multi-head Latent Attention (MLA)
MLA 来自 DeepSeek-V2,对标准多头注意力做键值压缩:先把 hidden state 压成低维 latent $U$,再分别解压出 key/value,并额外保留一个跨头共享的 $K_{direct}$。这样 KV cache 体积大幅缩小,但仍保留完整的双向 softmax 注意力和精确的全局召回能力。
Chimera 每 4 层插入 1 层 MLA 来恢复 KDA 受限状态无法精确召回的全局交互,二者的互补性正是混合架构的理论依据。
Mixture-of-Experts (MoE) 稀疏激活
MoE 把单个 FFN 换成多个小专家 FFN 的集合,路由器为每个 token 选 top-K 个专家计算并加权求和,其余专家不激活。这样总参数增大而单 token 激活计算量保持不变,常用于在控住算力下扩容。Chimera 用 56 个路由专家、每 token 激活 8 个(比例 1/7),并采用无辅助损失的偏置均衡策略。
MoE 是 Chimera 把总参数推到 11B 而激活参数控制在 2B 的关键;论文还发现 MoE 在扩散里的收益远低于语言模型,是个值得记住的重要洞察。
Chinchilla 计算最优扩展律与超参迁移 (µP)
Chinchilla 用参数化损失 $\hat{L}(N,D) = E + A N^{-a} + B D^{-b}$ 拟合给定算力 $C$ 下最优的模型规模 $N$ 与训练 token 数 $D$ 的分配。µP(maximal-update parameterization)则通过重参数化让在小 proxy 模型上调好的学习率/初始化能迁移到不同宽度深度的大模型,避免每个规模重新调参。
本文的两大方法论贡献 HeteroP 和 Chinchilla 扩展律都建立在这两者之上;不懂它们就读不懂为什么要做模块级 ratio 和三种独立估计器。
研究动机
视觉生成正进入 token 密集区:一张 $4096 \times 4096$ 的图按 $16 \times 16$ patchify 就有约 65K 个 token,再叠加长视频、密集条件和多轮交互,序列长度进一步膨胀,成为显存、延迟和训练成本的主导因素。然而主流 Diffusion Transformer 仍依赖全自注意力,其成对交互随序列长度二次增长 $O(N^2)$,在长序列下很快成为瓶颈。语言模型早已走过这一关,发展出稀疏注意力、KV 压缩、线性/循环算子和稀疏激活等方案,也积累了超参迁移与计算最优扩展律。但这些方案不能直接照搬到视觉扩散:视觉数据必须保持二维/三维网格上的局部结构、支持跨模态的双向交互、并在去噪目标下稳定训练;而混合骨干里不同张量的 fan-in 增长速率不一致,单一全局宽度比会破坏 µP 的迁移规则。同时视觉训练算力不仅取决于 token 数,还取决于图像/视频的数据配比。这两个缺口相互耦合,构成了本文要解决的难题。
本文的目标是本文的双目标是耦合的:一方面设计一个面向 token 密集生成的高效上下文建模架构,把文本、图像、视频表示为单一 token 流并用线性复杂度算子处理长序列,同时保留全局建模与局部时空结构能力;另一方面为这种异构架构建立可预测的扩展框架,能在固定算力预算下给出最优的激活模型规模、训练 token 数与图像/视频配比。具体到可量化的指标:训练一个 11B 总参数、2B 激活参数的 Chimera,在约 600 H100 Days 的预算内,相对同算力的全注意力 Wan-2.1 (2B) 基线在预训练扩散损失上实现数倍的计算效率提升,并在无长度微调的前提下支持从 5 秒训练片段到 30 秒视频的零样本外推。
与已有工作不同的是,本文的独特切入在于把架构设计与扩展方法学当作一个整体联合考虑,而非孤立优化。架构上,它没有像以往视觉线性注意力工作那样堆砌复杂的多方向扫描来凑出二维/三维局部性,而是抓住一个被忽视的点:RoPE 的三种位置功能(位置选择、近邻衰减、多维布局)其实可以拆给专门的模块——递归扫描提供顺序、短卷积提供位置选择、KDA 的通道衰减提供近邻性——从而彻底去掉位置编码并解除对训练长度的依赖。方法学上,它识别到异构骨干不能用单一全局 ratio 迁移超参,于是从每个张量的功能 fan-in 推导模块级 ratio;并把 Chinchilla 扩展律第一次扩展到图像/视频数据配比这一维度。这种『架构即受控模型族、扩展即预测前沿』的配对视角是此前视觉扩散工作所缺乏的。
核心方法
直觉上,Chimera 之于视觉扩散,类似于现代混合线性/全局语言模型之于早期 dense Transformer:把大部分全注意力层换成线性状态跟踪,只周期性保留少量全局注意力层做精确召回,再为视觉数据的局部结构补上专门机制。技术路线是单流骨干:文本 token(umT5-XXL 编码、≤512 token)与去噪视觉 token(冻结 3D VAE + (1,2,2) patchify)拼成一条序列,按时间主序的栅格顺序展平。每 4 层中有 3 层是 KDA(线性复杂度 $O(N)$)、1 层是 MLA(全双向、压缩 KV),比例为 3:1。KDA 更新前先过一个模态感知短卷积:文本走因果 1D、视觉走 3D,把局部邻域信息注入 Q/K/V。整个骨干不用任何位置编码(NoPE),位置信息来自扫描顺序、卷积相对偏移和 KDA 的通道衰减三者。容量上用稀疏 MoE(中间层)、稳定性上用 4 流 iHC 和三明治 RMSNorm。配套的 HeteroP 让超参从小 proxy 迁移,Chinchilla 律给出算力分配。
全文最核心的创新是把 RoPE 的三种位置归纳偏置解耦到三个专用模块,从而既腾出注意力通道专注内容匹配、又解除模型对训练长度的耦合。作者先用 Qwen3-4B 的 previous-token head 做频谱分解,证明 RoPE 的位置选择几乎只在高频、短距离有效,而大振幅低频对实际只贡献与偏移无关的背景——说明一个完整注意力头常被浪费在固定的 n-token 平移上。Chimera 因此用短卷积(少量参数即可覆盖短距离选择)、KDA 的通道遗忘门 $\alpha_t$(内容自适应近邻衰减)、以及卷积结构本身(1D/3D 天然编码布局)来分别承担这三项功能。第二个关键创新是 HeteroP:因为 KDA head、MLA 压缩秩、MoE 专家宽、router 等张量的 fan-in 增长速率不同,套用单一全局宽度比会错配,HeteroP 从每个张量的功能 fan-in 单独算 $m_W$,配合深度因子 $m_L$,使 proxy 上调好的 $\eta_{base}=10^{-3}$ 在 56 倍宽度、8 倍深度上保持近最优。这两点共同支撑了零样本长度外推与可靠的扩展律。
方法步骤详情
完整流程分七步。(1) 编码:图像/视频经冻结因果 3D VAE(16 通道、(4,8,8) 压缩)+ (1,2,2) patchify 成视觉 token;文本经冻结 umT5-XXL 编码(≤512 token)。(2) 拼流:序列 $x = [P_{text}(c), P_{vis}(z_\tau)]$,视觉按时间主序栅格展平。(3) 整流流训练:$z_\tau = \tau z_1 + (1-\tau) z_0$,目标 $v_\tau = z_1 - z_0$,损失 $\mathcal{L}=\mathbb{E}\|\text{Chimera}(c,z_\tau,\tau;\theta)-v_\tau\|_2^2$ 仅在视觉位置计算;timestep 从 logit-normal 采样并按 $\tau' = \tau/(\tau+s(1-\tau))$、$s=\sqrt{L_{vis}/256}$ 做噪声偏移。(4) 每 block:AdaLN 把 timestep 映射为调制与门控;iHC-Read 把 4 流聚成 1 输入,pre/post RMSNorm 包住注意力(KDA/MLA 按 3:1)与 FFN(首尾 dense、中间 MoE),再 iHC-Write 回 4 流。(5) KDA 状态更新 $S_t = (I-\beta_t k_t k_t^\top)\mathrm{Diag}(\alpha_t) S_{t-1} + \beta_t k_t v_t^\top$,卷积在前注入局部结构。(6) HeteroP:proxy(22M 激活)上调好基超参,按张量角色算 $m_W$、$m_L$ 实例化各尺度。(7) 拟合 $\hat{L}(N,D)=E+A N^{-a}+B D^{-b}$ 推出 $N_{opt}(C), D_{opt}(C)$ 选最终配置。
技术新颖性
和以往技术相比,新颖性集中在三点。其一,相对视觉线性注意力(多方向扫描 + 1D 卷积),Chimera 用模态感知短卷积 + 单一栅格扫描就恢复了二维/三维局部性,并用融合 Triton kernel 取得 2.2–2.3× 前向加速、最多 4× 显存下降,无需复杂扫描编排。其二,相对 µP/CompleteP 这类单全局 ratio 的超参迁移,HeteroP 首次为异构视觉扩散骨干给出模块级 ratio,并能证明标准参数化会让扩展律指数偏移 0.08–0.10、误开『重参数、轻数据』的药方。其三,iHC 把 mHC 的 Sinkhorn 双随机投影简化为恒等迁移,既保留多流自适应读写又省去逐 token 的 $M \times M$ 矩阵运算;MoE 用无辅助损失偏置均衡,MaxVio 稳定在约 0.5。这三点共同让 11B/2B 模型在小预算下稳定收敛。
实验结果
三组核心发现。计算效率:在 $5 \times 10^{20}$ FLOPs 匹配预算下,达到预训练损失 0.149 时 Wan-2.1 需 $4.29 \times 10^{20}$、Z-Image 需 $3.75 \times 10^{20}$,而 Chimera-dense(不含 MoE/HeteroP)仅需 $2.55 \times 10^{20}$(1.7×)、完整 Chimera 仅 $6.27 \times 10^{19}$(6.8–7.3×)。组件消融 dense→MoE 约 1.5×→+iHC 1.7×→+HeteroP 4.1×,HeteroP 贡献最大。扩展律:图像预训练 $N_{opt} \propto C^{0.48-0.52}$,参数与 token 近乎均衡;视频 $N_{opt} \propto C^{0.53-0.56}$,略偏模型容量;参数化拟合 $\hat{L}_{image}=0.126+5.28 N^{-0.315}+33.8 D^{-0.336}$、$\hat{L}_{video}=0.124+8.07 N^{-0.330}+145.2 D^{-0.394}$,图像 $R^2=0.993$。零样本外推:从 5 秒训练片段到 30 秒视频、仅看末 5 秒,FID 从 77.1 升到 82.1(+6.5%),对比 Wan2.1 +50.5%、HunyuanVideo-1.5 +53.6%;FVD +20.9% 对比 +33.9%/+33.7%。系统效率上 KDA/MLA 骨干在单卡 A100-80GB 支持 255k token(MHA/MLA 在 152k OOM,1.68×),255k 时快 2.14×;MoE 无辅助损失均衡使 MaxVio 稳定约 0.5。基准上 GenEval 0.82、DPG 85.12,且仅训练约 600 H100 Days。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| GenEval 文生图组合准确率 | Overall Score | 0.82 | FLUX.1-dev 0.82 / Z-Image-Turbo 0.82 / Seedream 3.0 0.84 | 持平 FLUX/Z-Image,仅落后 Seedream 3.0 |
| DPG-Bench 提示对齐 | Overall Score | 85.12 | Z-Image-Turbo 84.86 / FLUX.1-dev 84.00 | +0.26 / +1.12,超越二者 |
| 预训练计算效率(损失达 0.149) | 所需 FLOPs | 6.27e19 | Wan-2.1 需 4.29e20 | 6.8–7.3× 算力效率 |
| 零样本视频长度外推(5→30s 末5秒) | FID 相对退化 | +6.5% | Wan2.1 +50.5% / HunyuanVideo-1.5 +53.6% | 退化幅度约为基线的 1/8 |
| 单卡长序列支持(A100-80GB) | 最大 token 数 / 前向速度 | 255k / 基准 | MHA/MLA 在 152k OOM | 1.68× 序列长度,255k 时快 2.14× |
局限与改进
作者承认的局限:MoE 在扩散里的计算效率收益只有约 1.5×,远低于语言模型常见的 $\sqrt{sparsity}$ 启发式预测,专家路由专门化很弱;Muon 优化器系统性地输给 AdamW(扩散优化可能依赖 Adam 的隐式低秩偏置);扩展族仍把若干结构轴固定(iHC 流数、MoE 配置、MLA 压缩比),视频拟合指数的 bootstrap 区间较宽(约 [0.47, 0.80]);目前只评测文生图/文生视频,未做多模态理解。我额外观察到:模型只在 5 秒片段和 1K 图上训练,更长的真实长视频和高分辨率生成仍需微调验证;GenEval 上 Seedream 3.0 以 0.84 领先,DPG 优势也只是个位数;绝对 FID(77–82)在常见基准里偏高,部分源于『只评末 5 秒』的特殊协议;7.3× 的效率是对 Wan 特定基线、在选定损失参考点 0.149 下测得。
独立分析的弱点
独立来看有几个弱点。其一,MoE 在扩散中收益疲软:视觉 token 连续、空间冗余、高度相关,且各 timestep 共享大量公共去噪成分,专家难专门化——改进方向是引入 timestep 条件路由或模态路由专家,或干脆放弃 MoE 改用条件化稠密计算。其二,长度外推只在低分辨率(448×288)、30 秒内验证,真实部署的长时长 + 高分辨率场景未测;与自回归 rollout 的衔接也未验证,可结合其已因果的卷积做世界模型化。其三,NoPE 放弃了绝对位置,文字渲染虽有定性样例但缺严格基准,对需要精确定位的排版任务可能不稳。其四,单一时间主序栅格扫描未必对所有空间任务最优,可探索可学习扫描。其五,扩展律只在 256²/180p 拟合,高分辨率下分配规律可能不同。其六,视频配比最优落在 1:1 视频重边界,无法断言真实最优。
未来方向
作者明确指出三方向:把 HeteroP 扩展到结构轴随尺度变化(专家数、路由稀疏、残差流数、MLA 压缩比都设为尺度相关而非预设);把 Chimera 的设计迁移到 Qwen3.5/K3 这类单流混合线性-全局架构,迈向『理解 + 生成统一』的单模型;用 timestep 直接条件化 iHC 的读写映射,从而吸收掉独立的 AdaLN 调制,让混合扩散骨干与混合语言骨干更对齐。基于成果可延伸的:由于卷积已是时间因果,骨干天然适合改造成自回归世界模型;可探索 timestep 感知的 MoE 路由以挽救扩散里 MoE 的疲软收益;在视频数据更充足时重做长时长扩展律拟合;并补充人类偏好评测,因为 FID/FVD 难以反映真实生成质量。
复现评估
复现门槛较高。代码与权重未公开(Adobe Research,闭源),训练数据是专有精选语料。完整复现需巨大算力:扩展律研究本身要 36 次图像从零训练 + 30 次视频从零训练,模型族横跨 22M 到 19.3B 总参数;最终 11B 总/2B 激活模型约用 600 H100 Days。不过 HeteroP 的模块级迁移规则在 Table 1 中列得很完整,KDA/MLA 来自公开的 Kimi Linear 与 DeepSeek-V2,整流流目标也是标准做法,理论上可在小尺度(如 proxy 22M)验证超参迁移与扩展律趋势。难点集中在自定义 Triton kernel(融合多模态短卷积、分块 KDA、MoE token 排序)和大规模分布式训练(HSDP + 负载均衡)。综合难度:全量复现高,小尺度原理验证中等。
论文图表
左上展示混合线性-全局扩散架构(单流输入 → 稀疏 MoE + MLA/KDA + 模态感知短卷积),右上展示 Chinchilla 风格扩展方法学(Proxy 调参 → HeteroP 迁移 → Sweep → 拟合扩展律)。下方 A1/A2/S1/S2 四块分别给出预训练计算效率、长序列显存/延迟、跨尺度超参迁移、可预测的损失-算力前沿的实证结果。
这张图是全文的导航地图,一图概括『架构即受控模型族、扩展即预测前沿』这一核心论点,是理解论文双贡献结构的必读图。
在标准参数化下重做 256² 图像扩展研究:训练轨迹更噪、大模型后期有损失尖峰;拟合出的 $N_{opt}(C) \propto C^{0.588/0.581}$、$D_{opt}(C) \propto C^{0.392/0.410}$,模型规模指数被抬高约 0.08–0.10,会误开『重参数、轻数据』的药方,外推三档算力后模型约偏大 2×。
这是 HeteroP 必要性的反证,说明尺度相关的学习率错配会系统性地污染扩展律,是方法学贡献的关键论据。