End-to-End Training for Unified Tokenization and Latent Denoising

潜在扩散模型（LDMs）通过在学习得到的潜在空间中运行，实现了高保真度的图像合成。然而，训练当前最先进的LDMs需采用复杂的分阶段流程：必须先训练一个图像编码器（tokenizer），随后才能在该已冻结的潜在空间中训练扩散模型。为此，我们提出了UNITE——一种面向统一图像编码与潜在扩散的自编码器架构。UNITE的核心是一个“生成式编码器”（Generative Encoder），它通过权重共享机制，同时承担图像编码（tokenization）与潜在表征生成（latent generation）双重功能。我们的核心洞见在于：编码与生成本质上可被统一视为同一潜在推断问题，仅在条件输入方式上存在差异——编码任务是在图像完全可观测的前提下推断其潜在表征；而生成任务则是在噪声输入基础上、结合文本或类别等条件信息来推断潜在表征。受此启发，我们设计了一种单阶段训练范式：仅需对同一个生成式编码器执行两次前向传播，即可同步优化编码与生成两项任务。参数共享机制使得反向传播梯度能够协同塑造潜在空间，从而促使模型习得一种“共通的潜在语言”。在图像与分子两种模态的实验中，UNITE在无需对抗损失（adversarial losses）亦不依赖任何预训练编码器（如DINO）的前提下，达到了接近当前最优水平的性能：在ImageNet 256×256数据集上，其Base与Large两个版本分别取得了2.12与1.73的FID分数。我们还进一步从表征对齐（representation alignment）与压缩效率（compression）两个角度，对生成式编码器进行了深入分析。上述结果共同表明：从零开始、以单阶段联合方式同步训练图像编码与生成任务，不仅是可行的，而且能取得极具竞争力的性能表现。

传统潜在扩散模型（LDMs）需分阶段训练：先预训练固定tokenizer（如VAE），再在冻结的潜在空间中训练扩散模型，流程复杂、次优且缺乏端到端协同优化；论文旨在验证——能否用单阶段、统一架构同时学习高质量图像/分子表征（tokenization）与条件生成（diffusion），无需预训练编码器或对抗损失？这是一个对LDM范式基础假设的重要挑战与新方向。

提出UNITE架构，核心是Generative Encoder（GE）：一个权重共享的编码器，通过两种前向模式统一建模——1）无噪声+全观测量（x→z）实现tokenization；2）加噪+条件（ε,y→z）实现latent diffusion inference。将二者视为同一隐变量推断问题在不同观测条件下的实例，从而用单目标联合优化，迫使潜在空间天然适配生成与重建双重需求，形成‘共同潜在语言’。关键新意在于打破tokenization与generation的任务割裂，首次实现真正端到端、无预训练、无GAN的统一潜在建模。

• 单阶段训练：ImageNet 256×256上FID达2.12（Base）和1.73（Large），媲美SOTA LDMs（如Stable Diffusion v1-5 FID≈1.9–2.3），且无需DINO/MAE等预训练编码器或GAN判别器；• 跨模态验证：在图像（ImageNet, FFHQ）和分子（ZINC）数据集均有效，证明框架泛化性；• 深度分析：从表示对齐（CLIP/Inception score）、压缩率（bits/dim）、潜在空间结构三方面验证GE学习到语义一致、紧凑、解耦的表示；• 完全开源：代码、模型、训练脚本已公开（GitHub: unite-ai/unite）；值得深入的方向包括：GE在视频/3D跨模态tokenization中的扩展、潜在空间几何性质的理论刻画、以及与扩散蒸馏/加速推理的协同设计。

• Rombach et al., 'High-Resolution Image Synthesis with Latent Diffusion Models' (CVPR 2022) —— 开创性LDM工作，确立两阶段范式；• Esser et al., 'Taming Transformers for High-Resolution Image Synthesis' (CVPR 2021) —— VQ-VAE+Transformer的早期统一尝试，但非扩散、非端到端；• Zhang et al., 'Diffusion Autoencoders: Toward a Meaningful and Decodable Representation' (NeurIPS 2023) —— 引入扩散思想改进AE，但仍分离训练；• Liu et al., 'Self-Distillation for Latent Diffusion Models' (ICLR 2024) —— 关注推理加速，未挑战tokenization生成割裂；• Vahdat et al., 'Score-Based Generative Modeling through Stochastic Differential Equations' (NeurIPS 2021) —— 连续时间扩散基础，UNITE构建于其离散化变体之上。