UDT: Reconciling U-Nets and Diffusion Transformers with Data-Adaptive Token Reduction¶

会议: NeurIPS 2026
arXiv: 2608.01298
论文: https://arxiv.org/abs/2608.01298
代码: https://github.com/JN-Yun/UDT
领域: 图像生成
关键词: 扩散模型, 扩散Transformer, Token合并, U型网络架构, 图像生成

一句话总结¶

针对传统各向同性 Diffusion Transformer 编码器-解码器不平衡及既有 U 型 DiT 固定局部下采样破坏自注意力与特征维度的缺陷，UDT 提出基于无参数数据自适应 Token 合并（ToMe）与精确还原的 U 型架构，在保持特征通道维度不变的同时大幅降低计算量并加速训练收敛。

研究背景与动机¶

扩散 Transformer（Diffusion Transformers, DiTs）凭借卓越的参数可扩展性与统一表征能力，已成为当前二维图像生成乃至文本到图像（T2I）多模态合成的基石主干网络。DiT 沿用了 ViT 的各向同性（isotropic）级联结构，其深层网络在去噪目标驱动下逐步学习语义特征。然而去噪扩散的本质要求最终阶段重建高频细粒度信息，导致模型表征质量往往在网络中后层达到顶峰后迅速下降；这使得各向同性 DiT 在实际功能上呈现出“极深的编码器与极短的解码器”这种不对称失衡状态。既有表征对齐工作（如 REPA）通过引入外部冻结自监督视觉编码器（如 DINOv2）进行多层特征正则化，促使表征提早形成，但并未从模型内在骨干架构层面彻底解决这一结构失衡问题。

受经典 CNN 架构 U-Net 编码器-解码器对称多尺度设计与跳跃连接（skip connections）优势的启发，近期探索（如 U-DiT、SiT↓）尝试将 U 型拓扑移植到扩散 Transformer 中以加速训练收敛。然而，这些现有 U 型 DiT 普遍机械照搬 CNN 下采样的局部网格范式，采用可学习的线性投影或卷积压缩通道维度，并将局部 $2\times 2$ 空间邻域的 4 个相邻 token 简单拼接为 1 个新 token。这种硬性局部聚合策略不仅割裂了全局跨区域语义相似度、损害了 Transformer 全局自注意力机制的建模多样性，更由于空间分辨率突变（如 256 token 骤降至 64 token）以及通道尺寸变化，给外部表征对齐（REPA）带来了维度失配与复杂的流形损失开销。

为了化解 U 型分层设计与 Transformer 统一表征之间的内在张力，本文提出根本性变革：放弃固定局部卷积降采样，改用基于特征相似度的数据自适应 Token 归约机制。核心 idea：在完整保留 DiT 原始特征隐藏维度的前提下，构建对称 U 型骨干网络 UDT，利用无参数的 Token Merging（ToMe）与逆向精确 Unmerging 实现自适应渐进下/上采样，既保留细粒度关键区域细节又消除冗余背景计算，原生兼容自注意力机制与外部表征正则化。

方法详解¶

整体框架¶

UDT（U-Net Diffusion Transformer）在网络宏观拓扑上划分为对称的编码器（Encoder）、瓶颈层（Bottleneck）与解码器（Decoder），严格复用标准 DiT 的层数、注意力头数与隐藏通道维度 $D$。整个网络输入带有时间步 $t$ 与类别标签条件 $c$ 的潜空间噪声图 $z_t \in \mathbb{R}^{H \times W \times C}$，经 Patchify 切块后得到长度为 $T_0$ 的全分辨率 token 序列。在第一层编码器保持全分辨率以确保全局上下文感知，随后各层逐步应用二分图软匹配（Bipartite Soft Matching）自适应合并语义相似的 token，直至在瓶颈层收敛至预设的压缩目标数量 $N_{\text{Merge}}$；解码器侧则根据编码阶段记录的动态合并索引（merge indices）逐层无损解合并，最终通过长程残差跳跃连接融合全分辨率浅层特征，输出每个原始 patch 的去噪预测值。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入噪声潜变量<br/>Patchify 全分辨率序列 T0=256"] --> B["第一层编码器<br/>全分辨率自注意力建模"]
    B --> C["渐进自适应 Token 合并<br/>基于 Key 相似度二分匹配 r 缩减"]
    C --> D["瓶颈层两层低分辨率块<br/>到达目标 N_Merge 紧凑表征"]
    D --> E["渐进自适应 Token 解合并<br/>依据 Recorded Merge Indices 精确还原"]
    B -.->|长程残差跳跃连接<br/>Full-Resolution Skip Connection| F["最终全分辨率解码器<br/>恢复 T0 长度预测原始噪声"]
    E --> F
    F --> G["输出去噪预测<br/>速度场 v 或噪声 epsilon"]

关键设计¶

1. 保持通道维度的无参数数据自适应 Token 归约：解决固定网格下采样破坏全局交互的缺陷

以往的 U 型 DiT 普遍通过可学习卷积或线性映射压缩 token 并在空间上硬性拼接相邻 patch，导致 token 数量骤降（如 $256 \to 64$），抹杀了注意力机制捕捉全图长距离依赖的能力。UDT 保持每个 token 的特征隐藏维度 $D$ 恒定不变，完全在序列长度维度 $T$ 上执行动态调度。在编码器各层自注意力模块计算中，提取各 token 的投影键向量（Key），将 token 集合划分为大小相等的两组，构建稀疏二分图计算跨组余弦相似度，并筛选出相似度最高的前 $r$ 对 token 进行加权均值合并： $$ \mathbf{x}_{1,2}^{\text{m}} = \frac{s_1 \mathbf{x}_1 + s_2 \mathbf{x}_2}{s_1 + s_2} $$ 其中 $s_1, s_2$ 分别记录合并前 token 所代表的原始 patch 数量（初始为 1），合并后的 token 尺度更新为 $s = s_1 + s_2$。为补偿 token 聚合带来的注意力权重偏差，模型引入比例注意力（Proportional Attention）： $$ \text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d}} + \log \mathbf{s}\right) \mathbf{V} $$ 这一设计使得平坦单一的背景区域 token 被优先归约，而高频纹理、前景边缘等信息关键区域得以完整保留，既大幅节省了计算量（GFLOPs 减少约 25%），又规避了信息混叠失真。

2. 基于索引反查的精确解合并与对称跳跃连接：保障密集像素级去噪预测保真度

扩散模型与分类/识别任务的关键差异在于其属于密集去噪预测，必须为每个原始像素/patch 输出精确的去噪向量。如果在解码端使用插值或转置卷积反采样，极易引入空间伪影。UDT 在编码器阶段执行每次合并时，轻量缓存所选 token 对的配对索引（merge indices）。在解码器对应的对称层，直接根据所记录的索引将合并后的 token 特征复制恢复至原始两个位置，同时还原其初始尺度因子 $s$。此外，UDT 在网络首尾两端分别建立全分辨率跳跃连接（首层编码块输出直接通过残差连接与最终解码块相加），使深层高阶去噪表征与浅层高频几何结构无缝互补，彻底杜绝了分辨率升降带来的信息不可逆丢失。

3. 原生表征对齐与无缝扩展架构：无需额外流形适配的 REPA 即插即用支持

针对以往 U 型网络由于下采样特征尺寸（如 64）与预训练外部视觉大模型（如 DINOv2 的 256 tokens）空间失配而必须引入额外升采样 MLP 及流形正则化损失的工程弊病，UDT 提供了极简的表征对齐通路。由于 UDT 具备精确的解合并算子，在将瓶颈层特征与外部 ViT 对齐时，只需利用已存合并索引将 $N_{\text{Merge}}$ 长度的瓶颈特征一键无损投影回原始 $T_0$ 分辨率，即可通过简单的线性 MLP 投影头直接与 DINOv2 patch 特征计算余弦相似度对齐损失，省去了所有定制化的流形约束辅助损失。此外，由于保持隐藏维度 $D$ 不变，UDT 在引入 RoPE（仅施加于全分辨率的首末块）、SwiGLU 前馈网络以及对长序列输入（如 patch size 1、高分辨率 $512\times 512$）执行早期快速合并时均展现出卓越的适应性。

损失函数 / 训练策略¶

UDT 默认遵循 SiT 的线性插值流匹配（Rectified Flow）框架与速度场预测（$v$-prediction）目标： $$ \mathcal{L}{\text{flow}}(\theta) = \mathbb{E} \left| v_\theta(x_t, t) - (x_1 - x_0) \right|2^2 $$ 其中 $x_t = (1-t)x_0 + t x_1$ 为噪声插值轨迹。当启用 REPA 辅助表征对齐时，总训练损失为流匹配损失与解码还原后的瓶颈层表征对齐损失的加权和： $$ \mathcal{L}}} = \mathcal{L{\text{flow}} + \lambda}} \mathcal{L{\text{repa}}, \quad \lambda = 0.5 $$ 采用 AdamW 优化器，学习率设为 }$10^{-4}$，在采样推理阶段使用 SDE Euler-Maruyama 求解器（250 步），并结合无分类器引导（CFG）与时间区间截断策略进一步提升生成样本的分布贴合度。

实验关键数据¶

主实验¶

在 ImageNet $256\times 256$ 条件图像生成基准下，对比各向同性基线模型与各类先进扩散架构（包含/不包含 CFG 评估）：

模型	参数量	GFLOPs	训练轮数	CFG	FID ↓	IS ↑
SiT-B/2	130M	23.0	80	✗	33.0	-
UDT-B/2 (本文基线)	135M	17.7	80	✗	24.0	-
UDT-B/2+ (本文增强)	135M	17.7	80	✗	20.7	-
SiT↓-B/2† (UREPA)	192M	24.7	80	✗	20.9	-
SiT-XL/2	675M	118.6	80	✗	17.2	-
SiT↓-XL/2† (UREPA)	946M	112.0	80	✗	9.2	-
UDT-XL/2 (本文基线)	707M	91.6	80	✗	7.7	-
UDT-XL/2+ (本文增强)	707M	91.6	80	✗	6.1	-
UDT-XL/2+ + REPA	715M	91.6	80	✗	5.7	-
SiT-XL/2 (1400 ep)	675M	118.6	1400	✓	2.06	270.3
SiT-XL/2 + REPA	683M	118.7	800	✓	1.42	305.7
UDT-XL/2+ (SD-VAE)	707M	91.6	350	✓	1.42	308.4
UDT-XL/2+ + REPA (SD-VAE)	715M	91.6	320	✓	1.38	306.3
UDT-XL/1+ (VA-VAE)	708M	158.1	500	✓	1.35	322.0

消融实验¶

在 UDT-B/2 模型（ImageNet $256\times 256$, 80 轮, 无 CFG）上测试瓶颈层保留 Token 数量、下采样算子以及 Token Merging 各核心组件的影响：

实验组	配置项 / 变体	序列长度变化	FID ↓	说明
下采样机制对比	SiT-B/2 (各向同性基线)	256 (不变)	33.0	无空间分辨率压缩
	可学习线性下采样投影算子	$256 \to 112$	45.5	强制 Token 线性混合破坏结构
	数据自适应 Token 合并 (UDT)	$256 \to 112$	24.0	动态保留语义与细粒度结构
瓶颈 Token 数量 $N_{\text{Merge}}$	$N_{\text{Merge}} = 96$	$256 \to 96$	24.3	过度压缩略有损失
	$N_{\text{Merge}} = 112$ (默认最佳)	$256 \to 112$	24.0	压缩率与表征保持的最佳平衡点
	$N_{\text{Merge}} = 128$	$256 \to 128$	24.2	计算收益略有降低
	$N_{\text{Merge}} = 160$	$256 \to 160$	25.1	欠压缩导致 U 型瓶颈效益弱化
Token Merging 组件分解	朴素随机/平均合并	$256 \to 112$	25.9	无自适应引导
	+ 基于 Key 的余弦相似度	$256 \to 112$	24.5	精准匹配相似区域
	+ 尺度加权平均更新	$256 \to 112$	24.2	规避反复合并后的权重偏差
	+ 比例注意力（完整 UDT）	$256 \to 112$	24.0	消除注意力偏置，达到最优

关键发现¶

收敛速度呈代际级飞跃：在 XL 模型尺寸下，开箱即用的 UDT 仅需 80 个 epoch（无 CFG 模式下 FID 7.7）即可超越原始 SiT-XL 训练 1400 个 epoch 的生成质量（FID 7.9），实现约 $20\times$ 的训练收敛加速；若结合 REPA 辅助对齐，仅需 40 个 epoch 即可达到 FID 7.6（加速约 $40\times$）。
计算开销显著节省：相比同规格各向同性 Transformer，UDT 通过在网络内部逐层减少活跃 token 数量，减少了约 23%~25% 的 GFLOPs（如 XL 规模从 118.6 GFLOPs 降至 91.6 GFLOPs），且在小 patch（如 patch size 1）时计算复杂度增幅由原本的 $4.5\times$ 抑制至 $1.3\times$。
数据匮乏场景泛化出色：在仅使用 10% ImageNet 训练集（128K 样本）的少数据设置下，UDT-L/2 在 300 轮即可达到 FID 10.3，显著优于在全量数据上训练 80 轮的 SiT-L/2（FID 18.8），训练墙钟时间降低约 $2.67\times$。

亮点与洞察¶

无参数、自适应的下/上采样架构创新：跳出了既有扩散 U-Net 机械照搬 CNN 步长卷积的传统范式，利用基于 Key 相似度的二分软匹配将无参 Token Merging 转化为 Transformer 原生的下采样算子，并通过索引记录实现解码器的无损精确还原，兼具轻量性与全局表征保真度。
天然的外部表征对齐相容性：传统 U 型 DiT 的分辨率硬压缩导致无法与标准 ViT 编码器直接对齐，UDT 借助反向索引一键复原全分辨率特征，省去了繁杂的跨尺度投影层和流形正则化损失，以最纯粹的架构实现了与 REPA 的高效即插即用协同。
极具通用性的通用主干替换能力：不仅适用于标准连续连续时间 flow 匹配模型，还可平滑迁移至 $\epsilon$-预测 DiT、纯像素级无潜编码器扩散模型（JiT）、增强型 VAE（VA-VAE）以及文本到图像模型（MMDiT），均带来稳定显著的加速与质量跃升。

局限与展望¶

动态索引追踪的工程开销：在编码器各层记录合并索引并在解码器反查虽然在理论上无额外训练参数，但在跨多 GPU 分布式通信或特定硬件编译加速（如 TensorRT / TorchDynamo 静态图优化）时，动态变长计算与非连续访存可能略微削弱实际推理加速比。
位置编码与长程合并约束：由于 Token 动态合并打破了规整的网格空间拓扑，RoPE 等旋转位置编码只能受限施加于未发生合并的首尾全分辨率块，未能贯穿整个深层瓶颈网络。
未来方向：探索更贴合硬件静态图编译的固定块结构自适应规约策略，并将 UDT 推广至大尺度的自回归多模态生成模型以及长视频扩散生成任务中。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙跳出 CNN 式卷积降采样框架，首次将数据自适应 Token 归约确立为 Transformer 原生 U 型分层算子。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 256/512 分辨率、多模型尺度、消融实验、表征线性探测、少样本场景以及 4 类主流 DiT 变体的即插即用验证。
写作质量: ⭐⭐⭐⭐⭐ 逻辑严密流畅，图表翔实精美，对架构机理与表征瓶颈的阐述鞭辟入里。
价值: ⭐⭐⭐⭐⭐ 为扩散生成模型提供了兼具高收敛效率与高性能的新一代架构基准，应用落地与工程推广价值巨大。