让每个时间步、每个图像区域只支付它需要的计算

DiT block 与 TDW 路由

\[\mathbf{X}\leftarrow\mathbf{X}+\alpha\operatorname{MHSA}(\gamma\mathbf{X}+\beta),\qquad \mathbf{X}\leftarrow\mathbf{X}+\alpha'\operatorname{MLP}(\gamma'\mathbf{X}+\beta')\]\[\mathbf{S}_{\mathrm{head}}=R_{\mathrm{head}}(\mathbf{E}_t),\qquad \mathbf{S}_{\mathrm{channel}}=R_{\mathrm{channel}}(\mathbf{E}_t)\]\[\mathbf{M}=\mathbb{1}[\mathbf{S}\gt0.5]\]

被激活的子网络

\[\operatorname{MHSA}(\mathbf{X})=\sum_{h:M^{h}_{\mathrm{head}}=1}\mathbf{X}^{h}_{\mathrm{attn}}\mathbf{W}^{h}_{O},\qquad \operatorname{MLP}(\mathbf{X})=\sum_{h:M^{h}_{\mathrm{channel}}=1}\sigma(\mathbf{X}\mathbf{W}^{h}_{1})\mathbf{W}^{h}_{2}\]

若激活 head 数为 \(\widetilde H_{\mathrm{head}}\)，TDW 让 attention 复杂度按 \(\widetilde H_{\mathrm{head}}/H\) 缩减；MLP 同理按 \(\widetilde H_{\mathrm{channel}}/H\) 缩减。每层强制保留按 magnitude 排名最重要的一组，避免全部关闭。

完整复杂度变化

\[\mathcal{O}\!\left(H(4NCC_H+2N^2C_H)\right)\;\longrightarrow\;\mathcal{O}\!\left(\widetilde H_{\mathrm{head}}(4NCC_H+2N^2C_H)\right)\]\[\mathcal{O}\!\left(H\,2NCD_H\right)\;\longrightarrow\;\mathcal{O}\!\left(\widetilde H_{\mathrm{channel}}\,2NCD_H\right)\]

\(C_H=C/H\) 是每个 attention head 的特征宽度；若 MLP 完整中间宽度为 \(D\)，则每个 channel group 的宽度为 \(D_H=D/H\)。SDT 再将第二式中的 \(N\) 替换为实际进入 MLP 的 \(\widetilde N\)。第一式覆盖 Q/K/V、attention 乘法和输出投影；第二式覆盖 MLP 的两次线性变换。TDW 的 head 与 channel router 都是“单个线性层 + Sigmoid”，输入为时间步嵌入 \(\mathbf{E}_t\)，输出维度均为 \(H\)。SDT 的公开 DiT 实现使用两层小型 MLP 产生逐 token 概率；论文正文只把它抽象为 \(R_{\mathrm{token}}\)。

SDT token 路由

\[\mathbf{S}_{\mathrm{token}}=R_{\mathrm{token}}(\mathbf{X})\in[0,1]^N,\qquad M^{i}_{\mathrm{token}}=\mathbb{1}[S^{i}_{\mathrm{token}}\gt0.5]\]\[\widetilde N=\sum_{i=1}^{N}M^{i}_{\mathrm{token}}\]

实际推理应 gather \(\widetilde N\) 个 token 进入 MLP，再 scatter 回 \(N\) 个位置。MHSA 对全部 token 运行。训练时用 Gumbel-Sigmoid 与 straight-through estimator 让二值选择可以反向传播。

FLOPs-aware 训练

\[\mathcal{L}_{\mathrm{FLOPs}}=\left(\frac{1}{B}\sum_{b=1}^{B}\frac{F^{t_b}_{\mathrm{dynamic}}}{F_{\mathrm{static}}}-\lambda\right)^2\]\[\mathcal{L}=\mathcal{L}_{\mathrm{DiT}}+w\mathcal{L}_{\mathrm{FLOPs}},\qquad w=1.0\]

\(\lambda\) 约束 batch 内均匀采样时间步的平均 FLOPs 比例。它没有为每个样本设置硬上限，因此实际 FLOPs 会围绕目标波动。论文的 warm-up 阶段另保留完整 DiT 分支监督，随后移除；DiT-XL λ=.5 使用 30k warm-up 和总计 200k 追加迭代。

FLUX SingleBlock

\[\mathbf{X}\leftarrow\mathbf{X}+\operatorname{FC}_2\!\left(\operatorname{MHSA}'(\mathbf{X})\,\Vert\,\operatorname{FC}_1(\mathbf{X})\right)\]

DoubleBlock 为图像与文本保留并行 MLP；SDT 只跳过图像 MLP。SingleBlock 将 attention 和 MLP 中间特征拼接，\(\operatorname{FC}_2\) 的有效输入宽度同时受 head 与 channel router 控制。token router 在图像 token、池化文本 token 与时间步 AdaLN 融合后作判断。

DyFLUX 两阶段训练

\[\mathcal{L}_{\mathrm{distill}}=u\sum_{\ell\in\mathcal{S}}\operatorname{MSE}(\mathbf{X}^{\ell}_{d},\mathbf{X}^{\ell}_{s})+v\operatorname{MSE}(\mathbf{Y}_{d},\mathbf{Y}_{s})\]\[u=10^{-4},\qquad v=0.1\]

第一阶段每隔四个 block 蒸馏静态 FLUX-Lite 的中间表示与输出，并使用 \(w=5.0\) 的预算损失；第二阶段进行 CFG distillation，把 conditional 与 unconditional 两次前向蒸馏到带 guidance embedding 的单次前向。v1 使用 λ=.7、内部数据集、500k 迭代、1024² 图像；数据规模、batch、optimizer、硬件和训练时长未公开。

TD-LoRA

\[\mathbf{X}\mathbf{W}'=\mathbf{X}\mathbf{W}+(\mathbf{X}\mathbf{A})\left(\sum_{i=1}^{M}\mu_i\widehat{\mathbf{B}}_i\right)\]\[\boldsymbol{\mu}=\operatorname{Softmax}(R_{\mathrm{expert}}(\mathbf{E}_t))\]

默认 \(M=3\)、TD-LoRA rank \(=4\)，AdaLN 普通 LoRA rank \(=8\)。推理时先按时间步融合专家，再合并到原权重；论文报告额外融合延迟相对批处理前向可以忽略。训练显存从 34GB 降到 25GB，FID 从 full fine-tune 的 2.07 变为 2.23。