为自监督学习重构去噪扩散模型

在这项研究中,作者检验了最初用于图像生成的去噪扩散模型(DDM)的表示学习能力。其理念是解构DDM,逐渐将其转化为经典的去噪自动编码器(DAE)。这一解构过程让大家能够探索现代DDM的各个组成部分如何影响自监督的表征。观察到,只有极少数现代组件对学习好的表征至关重要,而其他许多组件则不重要。研究最终得出了一种高度简化的方法,在很大程度上类似于经典的DAE。

来自:Deconstructing Denoising Diffusion Models for Self-Supervised Learning

目录

背景概述

去噪是当前计算机视觉和其他领域生成模型的核心。如今,这些方法通常被称为去噪扩散模型(DDM),学习了一种去噪自编码器(DAE),它可以去除由扩散过程驱动的多个级别的噪声。这些方法实现了令人印象深刻的图像生成质量,尤其是高分辨率、逼真的图像。这不禁让人想到,这些生成模型非常好,似乎对理解视觉内容有很强的表示。

虽然DAE是当今生成模型的强大力量,但它最初是为了以自监督的方式从数据中学习表示而提出的。在当今的表示学习社区中,DAE最成功的变体可以说是基于"masking noise",例如预测语言中的缺失文本(例如,BERT)或图像中的缺失patch(例如,MAE)。然而,在概念上,这些基于mask的变体与去除加性噪声(例如,高斯噪声)存在显著不同:mask的token明确指定未知与已知内容,但在分离加性噪声的任务中没有明确的信号可用。然而,今天的DDM主要基于加性噪声,这意味着它们可以在不显式标记未知和已知内容的情况下学习表示。

最近,人们对DDM的表征学习能力越来越感兴趣。特别是,这些研究直接来自最初用于生成的预训练DDM,并评估其表示质量以进行识别。他们报告说,使用这些模型取得了令人鼓舞的结果。然而,这些开创性的研究显然留下了悬而未决的问题:这些现成的模型是为生成而设计的,而不是识别;目前还不清楚表示能力是通过去噪过程还是扩散过程获得的。

在这项工作中,作者对这些最近的探索相关初始化的方向进行了更深入的研究。不像之前使用面向生成的现成(off-the-shelf)DDM,作者训练面向识别的模型。核心是解构DDM,逐步将其转变为经典的DAE。通过这个解构性的研究过程,作者考察了现代DDM的每一个方面,目的是学习表征。这一研究过程使我们对DAE学习良好表示的关键组件有了新的理解。

令人惊讶的是,作者发现主要的关键组件是标记器(tokenizer),它创建了一个低维的潜在空间。有趣的是,这一观察结果在很大程度上独立于tokenizer的具体情况:作者探索了标准VAE、patch-wise VAE、patch-wise AE和patch-wise PCA编码器。发现使DAE能够实现良好的表示的是低维潜在空间,而不是tokenizer的细节。

由于PCA的有效性,解构轨迹最终达到了一个与经典DAE高度相似的简单架构(图1)。使用patch-wise PCA将图像投影到潜在空间上,添加噪声,然后通过逆PCA将其投影回来。然后训练一个自动编码器来预测去噪图像。作者将这种架构称为"潜在去噪自动编码器"(l-DAE,latent Denoising Autoencoder)。

  • 图1:这种简单的架构在很大程度上类似于经典的DAE(主要区别是将噪声添加到潜在的DAE中),并实现了有竞争力的自监督学习性能。

解构轨迹还揭示了DDM和经典DAE之间的许多其他有趣的特性。例如,即使使用单个噪声水平(即,不使用DDM的噪声scheduling),也可以让l-DAE获得不错的结果。使用多个级别的噪声的作用类似于一种形式的数据增强,这可能是有益的,但不是一个促成因素。根据这一点和其他观察结果,作者认为DDM的表示能力主要是通过去噪驱动的过程获得的,而不是扩散驱动的过程。

最后,作者将结果与以前的基线进行比较。一方面,我们的结果比现有的结果要好得多:这正如预期的那样,因为这是解构的起点。另一方面,结果没有达到基线对比学习方法和基于mask的方法,但差距缩小了。研究表明,在DAE和DDM的方向上还有更多的研究空间。

相关工作

在机器学习和计算机视觉的历史上,图像(或其他内容)的生成与无监督或自监督学习的发展密切相关。生成方法在概念上是非监督或自监督学习的形式,在没有标记数据的情况下训练模型,学习输入数据的基本分布。人们普遍认为,模型生成高保真度数据的能力表明了其学习良好表示的潜力。例如,生成对抗性网络(GAN)激发了人们对对抗性表示学习的广泛兴趣。变分自动编码器(VAE)最初被概念化为近似数据分布的生成模型,现已发展成为学习局部表示("tokens")的标准,例如VQV AE和变体。图像修复本质上是条件图像生成的一种形式,它产生了一系列现代表示学习方法,包括上下文编码器Context Encoder和masked自动编码器(MAE)。

类似地,去噪扩散模型出色的生成性能因其在表示学习中的潜力而引起关注。目前已经开始通过评估现有的预训练DDM来研究这一方向。然而,注意到,虽然模型的生成能力表明了一定程度的理解能力,但它并不一定转化为对下游任务有用的表示。因此该研究深入探讨了这些问题。

另一方面,尽管去噪自动编码器(DAE)为基于自动编码的表示学习奠定了基础,但它们的成功主要局限于涉及基于mask破坏的场景。最近很少或根本没有研究报告具有加性高斯噪声的经典DAE的结果,作者认为根本原因是简单的DAE基线(图2a)表现不佳。

  • 图2:一个经典DAE和一个现代DDM。a:在图像空间中添加并预测噪声的经典DAE,b:在潜在空间上操作的最先进的DDM(例如DIT),在latent space中添加并预测噪声。

去噪扩散模型

解构研究始于去噪扩散模型。简要描述了DDM如下。

扩散过程从干净的数据点 z 0 z_{0} z0开始,按顺序添加噪声。在指定的时间step t t t,加噪后的数据 z t z_{t} zt为: z t = γ t z 0 + σ t ϵ z_{t}=\gamma_{t}z_{0}+\sigma_{t}\epsilon zt=γtz0+σtϵ其中, ϵ ∼ N ( 0 , I ) \epsilon\sim N(0,I) ϵ∼N(0,I)是一个从高斯分布采样的noise map, γ t \gamma_{t} γt和 σ t \sigma_{t} σt分别定义了信号和噪声的scaling因子。此外, γ t 2 + σ t 2 = 1 \gamma_{t}^{2}+\sigma_{t}^{2}=1 γt2+σt2=1。

以时间step t t t为条件,学习去噪扩散模型来去除噪声。不同于DAE预测干净的输入,DDM预测噪声 ϵ \epsilon ϵ。损失为最小化: ∣ ∣ ϵ − n e t ( z t ) ∣ ∣ 2 ||\epsilon-net(z_{t})||^{2} ∣∣ϵ−net(zt)∣∣2其中, n e t ( z t ) net(z_{t}) net(zt)是网络输出。在给定以时间步长 t t t为条件的噪声调度下,针对多个噪声水平来训练网络。在生成过程中,迭代地应用经过训练的模型,直到它达到干净信号 z 0 z_{0} z0。

DDM可以在两种类型的输入空间上操作。一个是原始像素空间,其中原始图像 x 0 x_{0} x0直接用作 z 0 z_{0} z0。另一种选择是在tokenizer产生的潜在空间上构建DDM,见图2b。在这种情况下,预训练的tokenizer f f f(通常是一个自编码器),比如VQVAE被用于 z 0 = f ( x 0 ) z_{0}=f(x_{0}) z0=f(x0)。

DiT:Diffusion Transformer

该研究从DiT开始,选择这种基于Transformer的DDM有几个原因:i.与其他基于UNet的DDM不同,基于Transformer的架构可以与Transformer驱动的其他自监督学习基线进行更公平的比较;ii.DiT在编码器和解码器之间有更清晰的区别,而UNet的编码器和解码器通过skip连接,并且在评估编码器时可能需要在网络改动上付出额外的努力;iii.DiT的训练速度比其他基于UNet的DDM快得多,同时实现了更好的生成质量。

作者使用DiT-Large(DiT-L)变体作为DDM基线。在DiT-L中,编码器和解码器加在一起的大小为ViT-L(24个blocks)。作者评估编码器的表示质量(linear probe accuracy),编码器有12个blocks,称为 1 2 L \frac{1}{2}L 21L(half large)。

Tokenizer

DiT是潜在扩散模型(LDM)的一种形式,它使用VQGAN tokenizer。VQGAN tokenizer将(256,256,3)的输入图像转换到 ( 32 , 32 , 4 ) (32,32,4) (32,32,4)的latent map。

开始的baseline

默认情况下,作者在ImageNet上以256×256像素的分辨率训练400个epochs的模型。DiT baseline的结果见表1(第一行)。使用DiT-L,作者报告使用其 1 2 L \frac{1}{2}L 21L编码器的线性探头精度为57.5%。该DiT-L模型的生成质量为11.6。

重构扩散模型

解构轨迹分为三个阶段。作者首先调整了DiT中以生成为中心的设置,使其更倾向于自监督学习。接下来,逐步解构和简化tokenizer。最后,作者试图扭转尽可能多的DDM驱动设计,将模型推向经典的DAE。

用于自监督学习的DDM重定向

虽然DDM在概念上是DAE的一种形式,但它最初是为了生成图像而开发的。DDM中的许多设计都面向生成任务。有些设计对于自监督学习是不合法的(例如,涉及类别标签);如果不考虑视觉质量,则其他一些是不必要的。在本节中,作者为了自监督学习的目的重新调整了DDM基线,总结见表1。

  • 表1:从DiT baseline开始,并在ImageNet上评估其线性探针精度(acc)。每一行都基于对前一行的修改。使用类标签的灰色条目不是自监督学习的结果。

移除类别-条件

高质量的DDM通常在类标签上进行条件训练,这可以在很大程度上提高生成质量。但是,在自监督学习中,标签的使用是不行的。作为第一步,作者在baseline中删除类条件。令人惊讶的是,尽管生成质量如预期的那样受到极大的损害(FID从11.6提高到34.2),但去除类别条件显著提高了线性探针的精度,从57.5%提高到62.1%(表1)。作者假设,直接将模型条件化为类标签会减少模型对编码的信息需求,删除类条件可以强制模型学习更多语义。

重构VQGAN

在baseline中,LDM的VQGAN tokenizer使用多个损失项进行训练:i.自动编码重建损失;ii.KL散度正则化损失;iii.基于针对ImageNet分类训练的VGGNet的感知损失;iv.具有判别器的对抗性损失。

由于感知损失涉及有监督的预训练网络,使用用这种损失训练的VQGAN是不合法的。相反,作者在其中消除了感知损失。使用这种tokenizer将线性探针的准确率从62.5%显著降低到58.4%(表1)。这种比较表明,用感知损失(带有类标签)训练的tokenizer本身提供了语义表示。

作者训练下一个VQGAN tokenizer,它可以进一步消除对抗性损失。它将线性探头的精度从58.4%略微提高到59.0%(表1)。tokenizer在这一点上本质上是一个VAE。还注意到,消除这两种损失都会损害生成质量。

替换噪声调度

在生成任务中,目标是逐步将噪声图转换为图像。因此,原始噪声调度在噪声很大的图像上花费了许多时间步长(图3)。如果模型不是面向生成的,那么这是没有必要的。

为了进行自监督学习,作者研究了一种更简单的噪声调度。具体来说,让 γ t 2 γ^{2}_{t} γt2在[0,1]的范围内线性衰减(图3)。这一变化大大提高了线性探头的精度,从59.0%提高到63.4%(表1),表明原始调度过于关注噪声较大的状态。另一方面,正如预期的那样,这样做会进一步损害生成能力,导致FID为93.2。

  • 图3:原始调度器设置 γ t 2 = ∏ s = 1 t ( 1 − β s ) \gamma_{t}^{2}=\prod_{s=1}^{t}(1-\beta_{s}) γt2=∏s=1t(1−βs),如果使用线性调度器,可以减少step。让模型一开始就关注比较干净的数据,而不是纯纯的噪声。

总体而言,表1中的结果表明,自监督学习表现与生成质量无关。DDM的表示能力不一定是其生成能力的结果

重构Tokenizer

接下来,通过进行实质性的简化来进一步解构VAE tokenizer。作者将以下四种自动编码器变体作为tokenizer进行比较,每种变体都是前一种的简化版本。

卷积VAE

到目前为止,解构被引导到VAE tokenizer,现在将VAE的编码器和解码器设置为卷积网络 f ( ⋅ ) f(\cdot) f(⋅)和 g ( ⋅ ) g(\cdot) g(⋅)。卷积VAE被如下损失最小化: ∣ ∣ x − g ( f ( x ) ) ∣ ∣ 2 + K L [ f ( x ) ∣ N ( 0 , I ) ] ||x-g(f(x))||^{2}+KL[f(x)|N(0,I)] ∣∣x−g(f(x))∣∣2+KL[f(x)∣N(0,I)]其中, x x x是VAE的输入图像。

Patch-wise VAE

接下来考虑简化情况,VAE的编码器和解码器都是线性投影,并且VAE的输入 x x x是一个patch,最小化损失为: ∣ ∣ x − U T V x ∣ ∣ 2 + K L [ V x ∣ N ( 0 , I ) ] ||x-U^{T}Vx||^{2}+KL[Vx|N(0,I)] ∣∣x−UTVx∣∣2+KL[Vx∣N(0,I)]这里, x x x表示flatten为 D D D维向量的patch。 U U U和 V V V都是 d × D d×D d×D矩阵,其中 d d d是潜在空间的维数。patch大小设置为16×16像素。

Patch-wise AE

作者通过去除正则化项对VAE进行进一步简化: ∣ ∣ x − U T V x ∣ ∣ 2 ||x-U^{T}Vx||^{2} ∣∣x−UTVx∣∣2因此,这个tokenizer本质上是patch上的自动编码器(AE),编码器和解码器都是线性投影。

Patch-wise PCA

最后,考虑一个更简单的变体,它在patch空间上执行主成分分析(PCA)。很容易证明PCA等效于AE的一个特殊情况: ∣ ∣ x − V T V x ∣ ∣ 2 ||x-V^{T}Vx||^{2} ∣∣x−VTVx∣∣2其中, V V V符合 V V T = I ∈ R d × d VV^{T}=I\in R^{d\times d} VVT=I∈Rd×d,PCA可以简单地通过在一大组随机采样的patch上进行特征分解来计算,不需要基于梯度的训练。

由于使用patch的简单性,对于三个patch tokenizer,我们可以在patch空间中可视化它们的滤波器(图4)。

  • 图4:patch tokenizer的可视化。每个滤波器对应于线性投影矩阵 V : d × D V:d×D V:d×D的一行,为了可视化,将其整形为16×16×3。

表2总结了使用这四种tokenizer变体的DiT的线性探针准确性。作者展示了关于潜在维度"per token"的结果。PCA tokenizer的有效性在很大程度上将现代DDM推向经典DAE。

  • 表2:线性探针精度与latent维度的权衡。利用DiT模型,作者研究了用于计算潜在空间的tokenizer的四种变体。作者改变潜在空间的维度 d d d(per token)。通过上面的图显示出来。尽管在体系结构和损失函数方面存在差异,但tokenizer的所有四种变体都表现出相似的趋势。

高分辨率、基于像素的DDM在自监督学习方面较差

在继续之前,报告了一个与上述观察结果一致的额外的消融实验。具体来说,考虑一种naive tokenizer,它对从resized image中提取的patch执行identity mapping。在这种情况下,token是由patch的所有像素组成的展平向量。在图5中,作者展示了这种"pixel-based" tokenizer的结果,该tokenizer分别对256、128、64和32的图像大小进行操作,patch大小为16、8、4、2。这些token化空间的"潜在"维度分别为768、192、48和12。在所有情况下,Transformer的序列长度都保持不变(256)。

  • 图5:pixel-based的tokenizer的线性探测结果,分别在256、128、64和32的图像大小上操作,补丁大小为16、8、4、2。这些token化空间的"潜在"维度分别为768、192、48和12。与之前研究的其他tokenizer类似,这种基于像素的标记器呈现出类似的趋势:相对较小的潜在空间维度是最优的。

这些比较表明,tokenizer和由此产生的潜在空间对于DDM或DAE在自监督学习场景中至关重要。特别是,在像素空间上应用具有加性高斯噪声的经典DAE会导致较差的结果。

走向经典的DAE

接下来,继续解构轨迹,目标是尽可能接近经典的DAE。作者试图消除当前基于PCA的DDM和经典DAE之间仍然存在的每一个方面。通过这个解构过程,作者更好地理解了每一个现代设计如何影响经典DAE。表3给出了下面讨论的结果。

  • 表3:从Patch-wise PCA标记器开始,迈向经典DAE。每一行都基于对前一行的修改。

预测干净的数据(而不是噪声)

虽然现代DDM通常预测噪声 ε ε ε,但经典DAE预测的是干净数据。作者通过最小化以下损失函数来检验这种差异: λ t ∣ ∣ z 0 − n e t ( z t ) ∣ ∣ 2 \lambda_{t}||z_{0}-net(z_{t})||^{2} λt∣∣z0−net(zt)∣∣2其中, z 0 z_{0} z0是干净数据(latent空间中), λ t = γ t 2 \lambda_{t}=\gamma_{t}^{2} λt=γt2是基于 t t t的loss权重,为不同等级的噪声引入去噪贡献。直观地说,它只是对更干净的数据的损失项给予了更多的重视(large γ t 2 \gamma_{t}^{2} γt2)。

随着预测干净数据(而不是噪声)的修改,线性探头的精度从65.1%下降到62.4%(表3)。这表明预测目标的选择会影响表示质量。尽管在这一步中受到了下降影响,但从现在起,作者依然坚持这种修改,因为目标是向经典DAE迈进。

移除输入scaling

在现代DDM中,输入按 γ t γ_t γt因子缩放。这在传统DAE中并不常见。接下来,作者研究去除输入缩放因子,即设置 γ t = 1 γ_t=1 γt=1。由于 γ t γ_t γt是固定的,我们需要直接在 σ t σ_t σt上定义一个噪声调度。作者简单地将 σ t σ_t σt设置为从0到 2 \sqrt{2} 2 的线性调度。此外,根据经验将等式中的权重设置为 λ t = 1 / ( 1 + σ t 2 ) λ_t=1/(1+σ^{2}_{t}) λt=1/(1+σt2)。此后获得了63.6%的准确率(表3),这有缩放的62.4%相比是有利的。这表明缩放数据是不必要的。

使用逆PCA对图像空间进行运算

到目前为止,对于探索过的所有条目,该模型在tokenizer产生的潜在空间上运行。理想情况下,作者希望可以直接在图像空间上工作,同时仍然具有良好的准确性。利用主成分分析,可以通过逆主成分分析来实现这一目标。这个想法如图1所示。特别地,通过PCA基(即 V V V)将输入图像投影到潜在空间中,在潜在空间中添加噪声,并通过逆PCA基( V T V^{T} VT)将带噪的潜在空间投影回图像空间。图1(底部)显示了在潜在空间中添加了噪声的示例图像。有了这个有噪声的图像作为网络的输入,我们可以应用直接对图像进行操作的标准ViT网络,就好像没有tokenizer一样。

进一步将其应用于输出侧(即,用逆PCA预测图像空间上的输出)具有63.9%的准确度。结果表明,用逆PCA在图像空间上操作可以获得与在潜在空间上操作类似的结果。

预测原始图像

虽然逆PCA可以在图像空间中产生预测目标,但该目标不是原始图像。这是因为PCA对于任何降维都是有损编码器。相比之下,直接预测原始图像是一种更自然的解决方案。当让网络预测原始图像时,引入的"噪声"包括两部分:i.加性高斯噪声,其固有维数为 d d d,ii.PCA重建误差,其固有维数为 D − d D−d D−d( D D D为768)。作者对这两部分的损失进行了不同的加权。

形式上,使用干净的原始图像 x 0 x_{0} x0和网络预测 n e t ( x t ) net(x_{t}) net(xt),可以计算投影到全PCA空间上的残差 r r r: r = V ( x 0 − n e t ( x t ) ) r=V(x_{0}-net(x_{t})) r=V(x0−net(xt)),这里 V V V是PCA基 D × D D\times D D×D。然后最小化: λ t ∑ i = 1 D w i r i 2 \lambda_{t}\sum_{i=1}^{D}w_{i}r_{i}^{2} λti=1∑Dwiri2其中, i i i表示向量 r r r的第 i i i维。 w i w_{i} wi对PCA重建误差的损失进行加权。采用该公式,预测原始图像的线性探头精度达到64.5%(表3)。

这个变体在概念上非常简单:它的输入是一个噪声图像,其噪声来自PCA潜在空间中,它的预测是原始的干净图像(图1)。

单级噪声

最后,出于好奇,作者进一步研究了具有单级噪声的变体。并注意到,由噪声调度给出的多级噪声是由DDM中的扩散过程驱动的;在传统DAE中,这在概念上是不必要的。

作者将噪声水平 σ σ σ固定为常数( 1 / 3 \sqrt{1/3} 1/3 )。使用这种单级噪声实现了61.5%的良好精度,与多级噪声对应(64.5%)相比,降低了3%。使用多级别的噪声类似于DAE中的一种数据增强形式:这是有益的,但不是一个促成因素。这也意味着DDM的表示能力主要是通过去噪驱动过程获得的,而不是扩散驱动过程。

总结

作者解构了现代DDM,并将其推向经典DAE(图6)。其中撤销了许多现代设计,在概念上只保留了从现代DDM继承的两种设计:i.添加了噪声的低维潜在空间;ii.多级噪声。作者使用表3末尾的条目作为最后的DAE实例化(如图1所示),并将这种方法称为"潜在去噪自动编码器",简称l-DAE。

  • 图6:从现代DDM到l-DAE的整体解构轨迹。每一行都基于对前一行的修改。
相关推荐
Halo咯咯3 分钟前
无限免费 OpenClaw:接入本地模型后,你的 AI Agent 就可以 24 小时自动干活(Mac Mini 可用)
人工智能
NAGNIP12 小时前
一文搞懂深度学习中的通用逼近定理!
人工智能·算法·面试
冬奇Lab13 小时前
一天一个开源项目(第36篇):EverMemOS - 跨 LLM 与平台的长时记忆 OS,让 Agent 会记忆更会推理
人工智能·开源·资讯
冬奇Lab13 小时前
OpenClaw 源码深度解析(一):Gateway——为什么需要一个"中枢"
人工智能·开源·源码阅读
AngelPP17 小时前
OpenClaw 架构深度解析:如何把 AI 助手搬到你的个人设备上
人工智能
宅小年17 小时前
Claude Code 换成了Kimi K2.5后,我再也回不去了
人工智能·ai编程·claude
九狼17 小时前
Flutter URL Scheme 跨平台跳转
人工智能·flutter·github
ZFSS17 小时前
Kimi Chat Completion API 申请及使用
前端·人工智能
天翼云开发者社区19 小时前
春节复工福利就位!天翼云息壤2500万Tokens免费送,全品类大模型一键畅玩!
人工智能·算力服务·息壤
知识浅谈19 小时前
教你如何用 Gemini 将课本图片一键转为精美 PPT
人工智能