立方离散扩散:基于高维表示token的离散视觉生成 Cubic Discrete Diffusion: Discrete Visual Generation on High-Dimensional Representation Tokens
首次在768维高维表示token上实现离散扩散生成,达到离散生成SOTA
前置知识
离散扩散模型(Discrete Diffusion)
离散扩散是一类生成模型,通过迭代去噪过程从噪声生成数据。与连续扩散不同,它在离散token空间操作,使用掩码预测而非连续噪声去除。MaskGIT是典型的离散扩散方法,通过掩码预测实现并行生成。核心思想是将生成过程建模为从完全掩码状态到完全可见状态的迭代精化过程,每步预测被掩码位置的token。
本文的核心创新在于将离散扩散范式扩展到高维token空间,理解离散扩散的基本原理是理解本文贡献的基础
向量量化(Vector Quantization, VQ)
向量量化是将连续向量映射到有限码本中最近邻的技术。传统VQ方法如VQVAE使用固定大小的码本来压缩高维特征到低维离散token(通常8-32维)。这种方法在高维空间面临码本覆盖不足的问题,因为需要指数级增长的码本大小才能覆盖高维空间。
本文采用的dimension-wise quantization是对传统VQ的替代方案,理解VQ的局限性有助于理解为什么需要新的量化方法
Dimension-wise Quantization(维度级量化)
维度级量化是一种独立处理每个维度的量化方法。对于768维的特征向量,它将每个维度独立映射到L个离散级别。这避免了VQ的码本覆盖问题,因为每个维度只需要L个值,总共L的768次方种可能组合。这种方法保持了特征的语义质量,同时实现可处理的离散化。
这是本文使用的核心量化技术,使得在高维表示空间进行离散生成成为可能
表示编码器(Representation Encoder)
表示编码器是预训练的视觉特征提取器,如DINOv2或SigLIP2。它们产生语义丰富的高维特征(768-1024维),捕获局部细节和全局语义结构。与传统VAE编码器不同,这些编码器专注于理解任务而非重建,产生更具语义的特征表示。
本文的核心观点是利用这些高维表示token进行生成,理解这些编码器的特性是理解本文动机的关键
Classifier-Free Guidance (CFG)
Classifier-Free Guidance是一种条件生成技术,通过同时训练有条件和无条件模型,在推理时进行线性组合来增强条件遵循。通常对生成质量有显著提升,但会增加计算成本。在本文中发现高维表示token对CFG的依赖性降低,这是一个重要的发现。
本文发现高维表示token对CFG的依赖性降低,无guidance时仍能达到优秀性能
研究动机
当前离散视觉生成方法主要依赖低维潜在token(通常8-32维),这牺牲了理解任务所需的语义丰富性。例如,MaskGIT使用16维token,LlamaGen使用8维token,VAR使用32维token。这些低维压缩空间虽然使生成更易处理,但损失了原始高维表示的语义信息。与此同时,高维预训练表示(如DINOv2的768维特征)在理解任务中表现出色,但直接在其上进行离散生成面临巨大挑战:对于16x16x768的配置,需要生成196,608个离散token,直接自回归生成需要O(hwd)步骤,而朴素的并行方法无法捕获结构化张量内的复杂依赖关系。
本文的目标是本文的目标是首次实现直接在原生高维表示token上进行离散视觉生成,无需对表示空间进行压缩或重组。具体目标包括:证明离散扩散可以处理768维的高维token空间;保持表示token的理解能力,使同一token可用于生成和理解;在ImageNet 256x256类条件生成任务上达到竞争性性能。
与已有工作不同的是,本文的独特切入角度是提出立方离散扩散(CubiD),将h x w x d的表示张量视为统一的建模空间,在空间和维度两个轴上进行细粒度掩码。与传统MaskGIT仅在空间位置上掩码不同,CubiD在元素级别独立掩码,使模型能够观察沿两个轴的部分信息。这种方法将不可处理的顺序建模问题转化为可管理的迭代精化,同时捕获空间内和维度间的依赖关系。关键洞察是:高维特征具有内在冗余性,激进掩码(高掩码比例)在训练中是有效的。
核心方法
CubiD的方法可以概括为三个阶段:首先使用冻结的预训练视觉编码器(如DINOv2或SigLIP2)提取高维连续特征;然后通过维度级量化将连续特征离散化为h x w x d的离散token;最后使用离散扩散模型在空间和维度两个轴上进行掩码预测训练。整体思路是将高维表示张量视为一个立方体,在训练时随机掩码这个立方体中的元素,让模型从可见上下文预测被掩码的元素。在推理时,从完全掩码状态开始,迭代精化直到生成完整图像。这种方法的关键优势是并行生成能力:与自回归方法需要196,608顺序步骤不同,CubiD只需数百次迭代。
CubiD的核心创新是维度级掩码策略。传统离散扩散方法(如MaskGIT)仅在空间维度上掩码,即掩码整个空间位置的所有特征维度。而CubiD在元素级别独立掩码,对于h x w x d的张量,每个元素(共h x w x d个)都被独立随机掩码。这种细粒度掩码策略有两个关键优势:允许模型观察沿两个轴的部分信息,利用双向注意力捕获跨张量的依赖关系;避免了per-spatial掩码(掩码整个空间位置)导致的局部不一致性和per-dim掩码(掩码整个维度)导致的纹理伪影。实验表明,per-element掩码的gFID为5.33,而per-spatial为22.22,per-dim完全失败(120.03)。
方法步骤详情
CubiD的训练和推理过程包含以下步骤:特征提取阶段,给定输入图像,预训练编码器E产生特征图z,其中h=H/p,w=W/p,d为特征维度(768-1024)。维度级量化阶段,独立量化每个连续值为L个离散级别,其中L为量化级别数(DINOv2用L=8,SigLIP2用L=16)。训练掩码阶段,采样掩码比例r从截断高斯分布,均值偏向激进掩码,随机选择相应数量的位置进行掩码。掩码token设计阶段,使用可学习的mask token,在训练中演化为与内容token最大可区分。预测阶段,Transformer从可见token预测被掩码token,使用交叉熵损失。推理阶段,从完全掩码开始,迭代精化,512步达到饱和(gFID=5.25)。
技术新颖性
CubiD的技术新颖性体现在多个方面:首先,这是首个在原生高维表示token(768维)上进行离散生成的方法,所有现有方法都在低维空间(8-128维,多数低于32维)操作。其次,维度级掩码策略是全新的,与传统空间级掩码有本质区别。第三,可学习的mask token设计优于固定值和随机采样,因为随机采样导致模型无法区分内容token和掩码token(gFID从5.33退化到56.38)。第四,发现高维特征对激进掩码敏感,最优标准差为0.10,太小(0.05)导致过拟合到高掩码模式。第五,表示token对classifier-free guidance的依赖性降低,无guidance时CubiD-XXL达到2.02 gFID,优于大多数有guidance的VAE方法。
实验结果
本文在ImageNet 256x256类条件生成任务上进行了全面实验。主要发现包括:CubiD-XXL达到1.88 gFID(有guidance)和2.02 gFID(无guidance),这是离散生成方法的state-of-the-art。与之前的SOTA VFMTok-3B(2.04 gFID)相比,CubiD-XXL在无guidance时就已超越。模型scaling行为良好:从946M参数(gFID=5.25)到3.7B参数(gFID=4.68),性能持续提升,表明方法能有效利用增加的模型容量。表示token对guidance的依赖性降低:无guidance时CubiD-XXL的2.02 gFID优于MaskGIT(6.18)和LlamaGen-XXL(14.6)的无guidance结果。推理效率方面,尽管需要生成196,608个token,但只需512步迭代,远少于自回归方法的顺序步骤。量化保真度方面,dimension-wise量化保持了原始连续表示的语义质量,SigLIP2-DQ在GQA上达到63.1 vs 连续特征的63.2,在TextVQA上达到59.8 vs 59.6。跨编码器泛化方面,DINOv2和SigLIP2都有效,分别达到5.25和5.87 gFID。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet 256x256 类条件生成 | gFID分数(越低越好) | 1.88 (CubiD-XXL w/ cfg) | VFMTok-3B: 2.04 | 7.8%相对提升 |
| ImageNet 256x256 类条件生成(无guidance) | gFID分数(越低越好) | 2.02 (CubiD-XXL) | MaskGIT: 6.18, LlamaGen-XXL: 14.6 | 显著优于传统离散方法 |
| ImageNet 512x512 低维token生成 | gFID分数(越低越好) | 1.58 (CubiD) | USiT-2B: 1.72 | 8.1%相对提升 |
| 理解任务GQA | 准确率 | 63.1 (SigLIP2-DQ) | 63.2 (连续SigLIP2) | 几乎无损失 |
| 理解任务TextVQA | 准确率 | 59.8 (SigLIP2-DQ) | 59.6 (连续SigLIP2) | 几乎无损失 |
局限与改进
作者承认的局限性包括:对表示编码器的依赖,由于CubiD操作在冻结的预训练编码器特征上,重建质量(约18dB PSNR)设定了生成质量的上界,限制了生成图像的细粒度细节。与连续生成的差距,尽管离散生成在统一多模态建模中有优势(共享交叉熵目标),但与连续扩散方法(如RAE)相比仍有差距。推理效率方面,CubiD需要数百到上千步生成,多于连续扩散模型,作者提出可通过离散语言模型的加速技术来改进。此外还有以下观察:实验仅在ImageNet 256x256和512x512上验证,未在更高分辨率或更复杂数据集上测试;未与最新的连续扩散方法(如SiT、DiT变体)进行直接比较;768维token的巨大空间(196,608个token)虽然通过并行迭代处理,但计算成本仍然很高。
独立分析的弱点
基于独立分析,CubiD存在以下弱点及改进方向:量化级别固定方面,DINOv2用L=8,SigLIP2用L=16,不同编码器需要不同设置,可探索自适应量化或学习最优量化级别。掩码比例分布方面,使用截断高斯分布,均值1.0偏向激进掩码,可探索课程学习策略,从易到难逐步增加掩码比例。mask token设计方面,虽然可学习mask token有效,但未探索更复杂的mask token策略,如多尺度mask token或条件mask token。推理效率方面,512步迭代仍较多,可借鉴离散语言模型的加速技术,如并行解码、投机解码等。应用范围方面,仅验证类条件生成,未探索文本条件生成或其他条件模态。模型架构方面,使用标准Transformer,未探索针对高维张量的专门架构设计。
未来方向
作者提出和基于成果可延伸的未来方向包括:改进表示编码器的重建能力,当前18dB PSNR限制了生成质量上限,可探索专门为生成优化的表示学习目标。缩小与连续生成的差距,通过改进离散生成建模技术。加速离散扩散推理,借鉴离散语言模型的技术。扩展到更高分辨率和更复杂数据集。统一多模态架构方面,由于表示token保持了理解能力,可探索同一token用于生成和理解的统一架构。探索文本到图像生成,当前仅验证类条件生成,可扩展到文本条件。视频生成方面,将立方离散扩散扩展到时空张量。与连续扩散方法的混合,探索离散-连续混合生成范式。
复现评估
复现评估:开源情况方面,论文未明确提到代码开源计划,但提供了详细的实现细节,包括训练超参数(batch size 2048、fp16、EMA动量0.9999、学习率warmup 100 epochs)。数据方面,使用标准ImageNet-1K数据集(1,281,167图像,1000类),公开可用。算力方面,使用分布式GPU训练,具体GPU数量未明确,但提到使用fp16减少内存消耗,暗示大规模计算资源,3.7B参数模型需要显著计算资源。难度方面,方法相对复杂,涉及预训练编码器、维度级量化、离散扩散训练、可学习mask token等多个组件,但核心思想清晰,实现细节充分,有经验的研究者应该能够复现。依赖方面,需要预训练的DINOv2或SigLIP2编码器,这些模型公开可用。
论文图表
对比展示两种生成范式:左侧展示低维token生成,需要h x w顺序步骤(如16x16=256步);右侧展示高维token生成,只需T步(T远小于h x w),通过并行迭代精化。图中展示了token空间的维度差异,低维通常8-32维,高维768-1024维。
这张图直观展示了本文的核心动机和创新点:高维token虽然数量巨大(196,608个),但通过并行迭代只需数百步,而低维token虽然数量少但需要顺序生成。
比较CubiD与现有离散生成方法。CubiD-XXL达到1.88 gFID(有guidance),2.02(无guidance),为离散生成SOTA。比较了低维离散扩散、低维自回归和高维离散三类方法。
这是论文的核心结果表,展示CubiD达到离散生成SOTA。
展示CubiD在32维token(DC-AE-f32c32)上ImageNet 512x512的结果。CubiD达到1.58 gFID,优于USiT-2B(1.72 gFID)。
这个表格展示CubiD不仅适用于高维token,在低维token上也有效。