arXiv:2603.22283cs.CVcs.AI2026-03被引 1

统一训练图像与分子生成,无需分步训练或预训练模型。

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

  • 共享编码器同时完成图像分词与生成,单阶段联合优化。
  • 在ImageNet上达FID 2.12(基础版)和1.73(大模型),接近顶尖水平。
  • 适合想从零训练扩散模型的研究者,尤其关注高效架构设计。

潜在扩散模型(LDMs)通过在学习到的潜在空间中操作实现高保真合成。然而,训练最先进的LDM需要复杂分阶段流程:先训练分词器,再在冻结的潜在空间中训练扩散模型。我们提出UNITE——一种用于统一分词与潜在去噪的自编码器架构。UNITE包含一个生成编码器,通过参数共享同时充当图像分词器和潜在生成器。核心洞察是:分词与生成可视为在不同条件下的同一潜在推理问题——分词从完整图像推断潜在变量,而生成则从噪声结合文本或类别条件推断。受此启发,我们设计了单阶段训练流程,通过两次前向传播共同优化两项任务。共享参数使梯度协同塑造潜在空间,促进形成‘统一潜在语言’。在图像与分子模态上,UNITE无需对抗损失或预训练编码器(如DINO),即达到接近最先进性能,在ImageNet 256×256上,基础版与大模型分别取得FID 2.12与1.73。我们进一步通过表示对齐与压缩角度分析生成编码器。结果表明,从零开始单阶段联合训练分词与生成是可行的。

原文摘要 · Abstract (English)

Latent diffusion models (LDMs) enable high-fidelity synthesis by operating in learned latent spaces. However, training state-of-the-art LDMs requires complex staging: a tokenizer must be trained first, before the diffusion model can be trained in the frozen latent space. We propose UNITE - an autoencoder architecture for unified tokenization and latent diffusion. UNITE consists of a Generative Encoder that serves as both image tokenizer and latent generator via weight sharing. Our key insight is that tokenization and generation can be viewed as the same latent inference problem under different conditioning regimes: tokenization infers latents from fully observed images, whereas generation infers them from noise together with text or class conditioning. Motivated by this, we introduce a single-stage training procedure that jointly optimizes both tasks via two forward passes through the same Generative Encoder. The shared parameters enable gradients to jointly shape the latent space, encouraging a "common latent language". Across image and molecule modalities, UNITE achieves near state of the art performance without adversarial losses or pretrained encoders (e.g., DINO), reaching FID 2.12 and 1.73 for Base and Large models on ImageNet 256 x 256. We further analyze the Generative Encoder through the lenses of representation alignment and compression. These results show that single stage joint training of tokenization & generation from scratch is feasible.

扩散模型自编码器统一训练潜在空间

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。