大数跨境
0
0

图像生成别只知道扩散模型,还有基于梯度去噪的分数模型:NCSN

图像生成别只知道扩散模型,还有基于梯度去噪的分数模型:NCSN 极市平台
2023-01-31
1
导读:“回到过去”启发未来
↑ 点击蓝字 关注极市平台
作者丨CW不要無聊的風格@知乎(已授权)
来源丨https://zhuanlan.zhihu.com/p/597490389
编辑丨极市平台

极市导读

 

本文带大家一起来“回看” Diffusion 诞生之前的一种生成范式—基于数据分布相关的梯度,它被称为分数模型。特别地,本文主要介绍其中的一种类型—NCSN(Noise Conditional Score Networks) 。 >>加入极市CV技术交流群,走在计算机视觉的最前沿

前言

2022 年可谓是 AIGC 元年,CV 有 Diffusion Models,NLP 有 ChatGPT,它们都大大超越了之前视觉和语言生成模型的效果,十分强烈地让大家感觉到 AI 离人类更近了不止一步。

于是,近来只要谈到图像生成,大家脱口而出的就是 diffusion,我理解这是主流效应,很正常。但如果仅把目光局限于此而不去探索其它更多的可能性,那么这波主流难免在一定程度上阻碍了 AIGC 的进步与发展。

把目光放宽看远很重要,探索是面向未知领域的,但目光可以往回看。“回到过去”往往能够启发未来,科学研究与工业发展都需要有历史观,站在宏观的视角下才能把路铺得更远,以至于一直走下去,从而实现永恒的发展与进步。

CW 想通过本文与大家一起来“回看” Diffusion (这里主要指 DDPM 及之后的扩散模型)诞生之前的一种生成范式 —— 基于数据分布相关的梯度,它被称为 分数模型(Score Networks) 。特别地,本文主要介绍其中的一种类型 —— NCSN(Noise Conditional Score Networks) 。它与 DDPM 类似:也是对数据进行加噪,并且其生成过程的实质也是在去噪。往大的层面来看,DDPM 可看作是分数模型的一种特例

聊聊生成模型

你是否想过,当我们在玩生成模型时,实际上在玩什么?

玩多了之后,不难发现,其实我们玩生成模型,就是想让它“表示”概率分布,从而能够通过其实现采样生成。

至于如何表示,不同的模型有不同的方法,但它们大致都可被归纳为两种范式,分别是:对数据的采样过程建模对概率密度进行建模。前者不纠结于数据分布的概率密度,而是通过其它方法达到表示概率分布的目的,因此也被称为 隐式(implicit)生成 模型;而后者则直接让模型去估计概率密度,不绕圈子,于是被称为 显式(explicit)生成 模型。

隐式(Implicit)生成

这种范式下最出圈的当属大名鼎鼎的生成对抗网络 —— GAN(Generative Adversarial Networks)。它是近几年来最主流的玩法,长期霸占着 SOTA 榜单,在炼丹界拥有众多粉丝,直到最近被扩散模型 Diffusion Models 抢了风头,掉粉十分严重..

由于我们几乎无法知道数据的分布,因此很难精确估计概率密度,于是 GAN 就绕开了这一点,采用“博弈”的思想,让生成器和判别器不断 PK,直到判别器无法判断到底是真实的数据样本还是由生成器生成的样本时,就可认为生成效果很棒了。

这种方法的优点是对网络架构没有什么限制,而且采样生成的质量也通常不错;但缺点也不少,它们就是广为炼丹师们所诟病的训练艰难和不稳定、无法计算似然(likelihood)、无法基于统一客观的评估指标去比较不同模型的优劣(no principled model comparisons)

显式(Explicit)生成

这类模型也叫基于似然(likelihood) 的模型,它们通过极大似然(maximum likelihood)估计的方式来进行学习。这个家族的成员也不少,如上图所示。

这类方法的首要优点就是能够明确计算出似然函数,因此可以用统一的指标来比较分析不同模型的好坏。

但由于概率必须位于 [0, 1] 之间,因此这类模型需要指定的网络架构来构建归一化(normalized)的概率模型(如自回归模型、流模型 flow models 等),这限制了网络架构的灵活性,从而也制约了模型的表示能力;另外,由于真实的概率密度无法获知,因此通常会使用代理的损失函数来进行训练(如 VAE 和 基于能量的模型)。

从本质上来说,最近大火的扩散模型可被收在这个家族们下。扩散模型可以看作是级联的 VAE,其使用的训练方法也同样为极大似然法,并且它对似然函数的处理方式也类似于 VAE 用了变分推断(variational inference)来近似。只不过,扩散模型在扩散过程中的每一步都走得非常小(也天然地,因此需要的步数比较多,导致采样生成的过程慢),“破坏性”不大,因此模型也比较容易学习,学习效果也比较好。

OK,稍微聊了下一些主流的生成模型玩法,现在把目光聚焦到我们接下来将要介绍的分数模型,它的做法与以上提到的家伙们都不太一样。

分数模型简介与动机

首先简单讲讲分数模型(score-based models) 是怎么回事,以一句话总结来说就是:

它估计 数据分布相关的梯度 并基于 朗之万动力学(Langevin Dynamics) 的思想来生成样本。

What is Score?

讲了这么多,到底何谓“分数”?

其实,上面说的“数据分布相关的梯度”实质上是 对数概率密度函数对于输入数据的梯度( ),全称应该是Stein Score(https://arxiv.org/pdf/1602.03253.pdf) ,这就是分数。而我们玩分数模型,就是要训练它让它学会估计(预测)分数

What does Score Really Mean?

那么,分数到底意味着什么呢?

我们知道,数据往往是多维的。由分数的定义以及从数学的角度出发来看,它应当是一个“矢(向)量场”(vector field) 。既然是向量,那么就有方向,这个方向就是:对于输入数据(样本)来说,其对数概率密度增长最快的方向

箭头代表分数,圆弧描绘了概率密度的情况

Why Estimate the Score?

So.. 为何要训练网络来估计分数呢?首先要想明白最终生成的样本应该达到怎样的效果。

和原来的数据一模一样?那岂不就是 copy-paste(不由地想到 Billie 的 copy cat,有段时间超爱这首), 谈何生成!

但是随意生成也不行,比如我拿猫的图片去训练,我希望最终生成的是猫而非鱼,并且这张猫的图片应该是训练集中不存在的,而是通过模型“创造”出来的。

综上可知,其实最终的效果应该是在保证“创新性”的同时也要符合(靠近)原数据分布

结合上一节的解析,我们可以启发式地想到:

如果在采样过程中沿着分数的方向走,就能够走到数据分布的高概率密度区域,最终生成的样本就会符合原数据分布。

沿着分数的方向进行采样生成

What is the Langevin Dynamics?

Langevin dynamics 本指朗之万动力学方程,它是描述物理学中布朗运动(悬浮在液体/气体中的微小颗粒所做的无规则运动)的微分方程,借鉴到这里作为一种生成样本的方法。

概括地来说,该方法首先从先验分布随机采样一个初始样本,然后利用模型估计出来的分数逐渐将样本向数据分布的高概率密度区域靠近。为保证生成结果的多样性,我们需要采样过程带有随机性。 正好!布朗运动就是带有随机性的,朗之万动力学方程因此也带有随机项。

极简版 NCSN 概要

对分数模型有了基本的理解后,现在来简要概述下本文的主角 NCSN 的做法。

训练目标

使用不同强度(不同方差) 的高斯噪声对数据进行扰动(加噪),训练网络去联合估计加噪后的数据对应的分数。

“联合”代表对于不同噪声级别的分数,都使用同一个网络去估计

另外,要注意是加噪后的数据对应的分数,也就是加噪后的数据的对数概率密度对加噪的数据求导的结果,而非对原数据求导

采样生成

使用一种称为 “退火朗之万动力学”(Annealed Langevin Dynamics) 的方法,这种方法实质上是在噪声强度递减的情况下使用朗之万动力学采样,在每个噪声级别下都有一个朗之万动力学采样生成的过程

首先从最高强度的噪声级别开始使用朗之万动力学进行一定步数的采样生成,这个阶段结束后生成的样本会作为下一个噪声级别的初始样本,然后继续进行朗之万动力学的采样生成过程。重复这种做法,直至在最小的噪声级别下也完成了朗之万动力学的采样生成过程,就得到了最终生成的样本。

这种方式之所以叫“退火”的朗之万动力学,是因为噪声级别不断地减小,而每个噪声级别都在进行朗之万动力学采样,这期间采样生成的样本也不断靠近原数据分布的高概率密度区域,还蛮形象的吧~

分数模型的玩法

现在,CW 向大家介绍下分数模型是怎么玩的。

Score Matching(分数匹配)

分数模型需要估计分数,它的训练目标可表示为:

其中, 代表网络参数, 是网络估计的分数, 其减去的右边那项代表真实的分数, 表示在整体数据中求期望。

然而, 这就要求我们知道概率密度 , 但通常我们是无法知道训练集服从的分布的。并且, 如果提前知道了数据分布, 那么直接从其中采样生成样本就好了, 还搞那么多飞机来炼丹干嘛..

这么看来,玩不下去了?

不! 有一种叫作 score matching(分数匹配) 的方法, 可以让我们在不了解概率密度的情况下也可 以成功训练出模型来估计分数, 这种方法利用了分部积分, 从而将公式 的训练目标等价转化为:

其中, 代表 的 雅克比(Jocabian) 矩阵, 表示矩阵的迹(trace), 即矩阵主对角线元素的总和。

嗯, 看似有玩下去的可行性, 但或许你也看 式不顺眼吧, 这坨东西实在有点难受.. 是的, 当遇到高维数据(图像特征通的通道数常都成百上千..)时, 这东西在计算上十分"难受"一需要经历多次(维数多少就需要多少次)反向传播来求偏微分:

模型通过前向过程估计分数

由于是偏导数,因此需要多次反向传播来分别对每个分量计算

OMG! 刚看到些希望,现在又破灭了..

但是,超人迪加告诉人们:要相信光!

下面,迪迦让 CW 替他向大家传达:有两种方法解决以上这个痛点,它们分别是:sliced score matching & denoising score matching

  • Sliced Score Matching(切片分数匹配)

这种方法的思想根源是: 既然 对于高维雅克比矩阵这么难算, 那么就不直接算它, 先 把雅克比矩阵降到一维(从而变成雅克比向量, matrix->vector)后再去算。

具体做法就是使用简单分布(如多元标准高斯分布)的随机向量 来对 实施随机映射(random projections), 从而得到:

这东西好算!首先网络通过前向过程得到 , 然后令其与随机向量做内积得到 , 接着通过一次反向传播计算出 , 最后再次与随机向量做内积即可。这么一搞, 训练目标就 从 式变身为:

虽说比起原先来好算, 但其实它比起正常的网络训练过程来说, 是多了一次反向传播的。因为上述的反向传播是为计算训练目标而进行的, ( 式计算出来后还要进行一次反向传抙得到的才是模型基于损失而获得的梯度。

  • Denoising Score Matching(去噪分数匹配)

感觉 sliced score matching 还是太麻烦?那么来看看这里的 denoising score matching 吧!

这种方法是作者对 NCSN 采取的默认方法, 它也是不算 , 也不做降维, 而是直接回归 到 式的做法。

可是, 你会说:我们不是不知道概率密度 吗?

没关系, 我们“作弊"一下下—— 自己搞一个数据分布, 从而明确知道它的概率密度函数。

做法就是对原始数据加噪, 使得其结果满足我们预先定义好的分布, 比如高斯分布。这样, 我们就知道了现在的概率密度了, 于是就可以使用 式那种做法来进行训练。

记训练样本为 , 加噪后的数据为 , 预定义的分布为 , 方差 预先定义好。加噪后的数据分布表示为 , 我们在这个数据分布上计算的分数表示为:

由于 不相干, 于是导数为 0 , 因此最右边那项没有了它的身影。

这时, 训练目标就表示为:

在这种方法下, 我们只要从标准高斯分布 中采样随机噪声 , 然后乘上我们预定义的 , 接着加到样本 中, 从而就能得到加噪后的样本 。根据高斯分布的性质, 这样加噪后 就会满足分布 估计分数时, 记得要将加噪后的样本 喂给模型, 而非原数据样本

另外, 要注意: 网络估计出来的分数是对应噪声数据分布 的, 而非原始数据分布 ! 因此我们需要让 尽可能与 接近, 这就要求预定义的 要足够小。 否则, 如果对原始数据扰动过大, 模型训练完后生成的样本就会严重偏离原数据分布。

Sampling with Langevin Dynamics(朗之万动力学采样)

如果我们知道整个分布中所有样本对应的分数 , 那么就可以 "跟随"分数的方向, 一步步地从随机初始样本来进行采样生成:

其中 代表前一步生成的样本, 代表当前生成的样本, 而初始样本 可以从先验分布(比如均匀分布)中采样。

但这样不行, 因为你会发现每次采样生成都会得到相同的结果(假设步数非常大, 大至无穷, 则最 终都会走到概率密度最高的那点), 即生成的结果缺乏"多样性"。

所以, 我们要在生成过程中加入随机性。这时候, 我们可以将朗之万动力学采样的思想用上, 它是根据以下方程来采样生成的:

其中 是随机项, 是预定义的步长(step size)。

由于加入了随机项, 因此生成结果也就丰富多样了。一旦我们训练好网络, 就可以用它估计出来的 分数 来替换掉 式对应的部分:

理论上, 在一定约束条件下, 并且当 时, 最终生成的样本 将会服从原数据分布

分数模型的“困局”

分数模型在实际应用时会面临一些困难。概括起来,主要是以下三点:

  • loss 不收敛
  • 模型估计的分数不准
  • 生成结果与原数据分布偏差大

接下来,CW 就以上几方点逐一进行解析,来看看造成以上“困局”背后的原因是什么。

loss 不收敛

如果你根据前文介绍的 score matching 去训练分数模型,会发现模型被“玩坏”了:loss 不断震荡,可谓“一秒天堂,一秒地狱”。

基于 sliced score matching 方式下训练的 loss

靠,为何会这样呢?明明前面分析得头头是道..

年轻人切忌心浮气躁,大家不妨先来了解下一种叫 “流形假设”(mainfold hypotheis) 的理论。

流行假设认为,生活中的真实数据大部分都倾向于分布在低维空间中,这也是大名鼎鼎的 “流形学习”(mainfold learning) 的大前提。也就是说,我们的编码空间(通过神经网络等一系列方法)的维度可能很广泛,但是数据经过编码后在许多维度上都存在冗余,实际上用一部分维度就足以表示它,说明实质上它仅分布在整个编码空间中的一部分(低维流形),并没有“占满”整个编码空间。

打个可能不太严谨的比方:假设你的数据都分布在单位矩形中,同时通过二维直角坐标系来表示它,于是编码空间就包含了坐标系中的所有点,但你的数据仅分布在其中的单位矩形区域。

OK,基于流形假设,就可以来解释 loss 不收敛的问题了:

首先,分数这个东西: ,是针对整个编码空间定义的,然而当我们的数据仅分布在编码空间的低维流形时,分数就会出现“没有定义” (undefined) 的情况。

另外,score matching 的训练目标(见公式 ( ) ) 是基于数据“占满”了整个编码空间这个前提下而推出来的。当数据仅分布在低维流形中时,score matching 的方法就不适用了。

综合以上两点,在基于 score matching 的方式下,训练分数模型就会出现 loss 震荡而不收敛的情况。

ps: 关于流形学习的解释,可以看看这个回答,CW 感觉蛮不错~


分数估计不准

除了 loss 不收敛,作者还发现有时会出现模型估计分数不准的问题。这是因为,属于低概率密度区域的数据,会由于没有足够的样本来供模型训练学习,从而导致这部分的分数估计不准,实质上就是训练不充分。

上图对比了真实的分数(左)与模型估计的分数(右),其中颜色(橘红色)深浅代表数据概率密度的大小(越深越大,对应的样本越可能出现),箭头代表分数

可以看到,右图的箭头(也就是模型估计的分数)在颜色深的地方基本和左图的箭头(也就是真实的分数)一致,如矩形框出来的部分所示。而在颜色比较浅的部分,特别是左图和右图对角线那部分,两者的箭头差异就比较大,代表模型估计的分数不准确。

生成结果偏差大

分数估计不准,带来的一个结果就是生成效果不佳,采样生成出来的样本不太符合原数据分布。

为何?因为分数模型的生成是根据模型估计的分数并且基于朗之万动力学采样来玩的啊!(见公式 ( ) )

其实也不能全怪模型预估的分数不准,因为作者发现,当数据分布是由多个分布按一定比例混合而成时,在朗之万动力学采样的玩法下,即使用真实的分数(对应公式 ( ) ),采样生成出来的结果也可能不能反应出各个分布之间的比例关系。

举个例子,就拿上一节的 Figure 2 这种情况来说吧,数据分布是由两个分布(右上和左下部分其实是两个不同的高斯分布的高概率密度区域)按不同比例混合而成并且由低概率密度区域“隔开”时,生成结果就不能反映出两个分布之间的相对关系。

如上图,在实际的数据分布中(上图左部分),右上部分的分布对应的样本比较多,左下部分的比较少;而生成的结果(上图右部分)确是两种分布对应的样本一样多,并没有反映出原始各分布之间的关系。

尽管能力有限,我们不能严格地论证出导致这现象的具体原因,但是,我们还是可以基于直觉,使用不太严谨的数学,来小装逼一下~

记两种分布为 , 分别以比例 混合而成, 于是整体的数据分布表示为

对于整体数据分布对应的分数,如果我们分别来看其中两个分布对应的分数,就会发现其实和它们的混合比例不相关

于是, 当我们使用整体数据分布对应的分数 来生成样本时, 就丟失了两个分布之间的比例关系了, 最终造成“以类似等比例的关系从两个分布中去采样生成"的效果。

作者说,其实,尽管是面对这种混合分布的情况,从理论上来说(至于什么理论 CW 也不知道..),朗之万动力学采样是可以生成与原分布相似的结果的,但可能需要很小的步长(step size)以及非常大的步数(steps)才能做到。

(⊙o⊙)… 感觉危机重重哦~ 这分数模型,还玩得下去么..

Don't worry!

接下来就有请我们的主角 NCSN 登场,这些问题在它那都不是问题!

(低调低调..)

NCSN 如何破局

在前文的 NCSN 概要部分,CW 给大家透露了它的做法,那么,为何这么做能够破解上一章提到的分数模型的困局呢?

一破 loss 不收敛

NCSN 用了高斯噪声去扰动数据,而高斯噪声是分布在整个编码空间的,因此,扰动后的数据就不会仅“驻扎”在低维流形,而是能够有机会出现在整个编码空间的任意部分,从而避免了分数(score)没有定义(undefined)的情况,同时使得 score matching(见公式 (ii)(ii) )适用,这就破解了“loss 不收敛”的困局。

二破 分数估计不准

在扰动数据的时候,尺度较大的噪声会有更多机会将原来的数据“踢”到低概率密度区域,于是原来的低概率密度区域就能获得更多的监督信号,从而提升了分数估计的准确性,这,破解了“分数估计不准”的困局。

三破 生成结果偏差大

最牛逼的,对于数据是由多个分布混合而成的情况,比如上一章展示的 Figure2,是由两个高斯分布按比例混合而成,同时被低概率密度区域“隔开”。

从下图可以看到,在低概率密度区域,分数几乎是 50% 的机率指向两个分布。 那么在使用朗之万动力学采样生成时,由于是根据分数来迭代生成的,若初始点落在了这些区域,就会造成有均等机会走向两个分布的方向,因此生成的整体分布所包含的两个分布的样本数量就会很接近,而理想的是其中一个分布的样本数应该要多一些。

箭头表示分数,红色长条代表低概率密度区域,两个圆代表两个分布的中心

而在加噪扰动后(噪声尺度要稍微大些),原本“稀疏”的低概率密度区域就能够被“填满”,并且,自然地,两个分布中混合比例高的那个分布所占这部分(“填满”后的区域)的比例会更高

扰动后,分数更多地指向混合比例高的那个分布(右上部分)

于是,扰动后的分数 将更多地由混合比例高的那个分布的分数“演变”而来。这就是说,分数的方向将更多地指向混合比例高的那个分布。 最终,基于分数去进行采样生成的时候,将更多地生成混合比例高的那个分布的样本,从而生成的整体分布就能够成功反映出原数据中两个分布的相对关系了。于是,very good 地破解了生成结果偏差大的困局。

另外,这样带来的收益是,我们不需要很“苛刻”的条件 —— 很小的步长 + 很大的步数(如上一章提到,作者说在理论上,使用很小的步长 + 很大的步数,朗之万动力学采样也能 work) 就能够生成符合原数据分布的结果,相当于起到了加速采样生成的作用。

NCSN 详解

吹完牛逼后,是时候为大家详细解析下 NCSN 本身的技术细节了。这章的内容主要分为四部分,分别是:

  • 噪声设计原则
  • 网络结构设计
  • 模型训练方法
  • 采样生成过程

其中,网络结构部分的内容在本章会讲得比较简单,主要是整体架构的选取以及重点模块的介绍,更具体的实现细节会放到源码解析部分(没错,CW 不喜欢无聊的风格,当然有源码解析!)。

噪声设计原则

我们知道, NCSN 加了不同强度的高斯噪声, 每种噪声强度所对应的分布可表示为 , 那么到底应该如何选取噪声强度 呢?

在考虑具体的选取方法时, 大家不妨先从更高层面上来想想 为何要加多尺度噪声?

OK, 可能现在有人不服, 他觉得仅加固定强度的一种噪声足矣(CW: 此言差矣!), 那么, 来试试 看会怎样:

假设只加一种强度的噪声 , 为了让扰动后的分布能够更多地“填充"原来的低概率密度区域, 于是 必须足够大, 这样才能让扰动后的分布更加广泛。嗯, 到这里没毛病吧

但是, NCSN 是以加噪后的分数 来近似原分布的分数 分别代表加噪后的数据与原数据)。 太大, 扰动过大, 抺去了原数据大部分信息, 就会造成近似效果不佳, 从而基于加噪后的分数使用朗之万动力学采样生成的结果也就不太符合原数据分布。

相反地,噪声强度小能够获得与原数据分布较为近似的效果,但是却不能够很好地“填充”低概率密度区域。

于是,trade-off 一下,较大和较小的噪声我都要(小孩子才做选择),多尺度噪声的设计就顺理成章了!

具体地, 作者采取了一个噪声序列 , 并满足: , 同时 需要足够大, 以能够充分“填充"低概率密度区域; 而 得足够小, 以获得对原数据分布良好的近似, 避免过度扰动。

至于 的具体数值, 就需要根据数据集去精心调整了, 不然怎么叫炼丹捏? (阴阴一笑 )

网络结构设计:以噪声为条件的 U-Net

由于生成的图片和原图一样大小,也就是每个像素点都需要由朗之万动力学采样生成,因此模型对于每个像素点都要估计其对应的分数。也就是说,从张量的角度来看,网络的输出(分数)要和输入图像的形状(shape)一致

自然而然地,作者联想到语义分割领域的模型结构。在脑海中思寻了没有多久,他也懒得努力了,高呼:就是你了,U-Net!就这样,U-Net 开启了它在以分数模型为范式的图像生成之旅~

当然,作者也不能脸皮太厚,直接白嫖别人的网络结构。在结合自身多年炼丹功力以及多番实验的探索结果后,他还在其中加入了 空洞卷积(dilated/atrous convolution) 和 以噪声为条件的实例归一化(conditional instance normalization++)

另外,即使对于同一个像素点,在不同的噪声强度下也要对应估计出不同的分数,于是,模型还需要将噪声强度 也作为输入。

训练方法:去噪分数匹配

对于训练方法,作者采用了去噪分数匹配。当然,前面提到的 切片分数匹配(sliced score matching)也可以训练 NCSN,只不过需要的计算步骤会多一些,因此训练过程会更慢。

通过前文介绍去噪分数匹配的内容,大家已经知道其训练目标为:

并且, 作者将样本加噪后的分布建模为 , 也就是均值为原数据样本 , 方差为预定义的 的高斯分布。于是, 对应的概率密度 则为 。将其取对数并对 求导, 就可获得其分数的表示形式:

将以上代入公式 , 就得到在某个噪声级别 对应的损失函数:

其中 代表网络在特定噪声级别 下估计的分数。而 NCSN 使用了多个噪声级别, 于是, 分别对它们的损失加权求和后再求均值, 就得到了联合的损失函数, 它表示为:

其中 代表不同噪声级别的损失的权重。

现在, 重头戏来了一如何设置 呢?

首先,作者的出发点是:加权后所有噪声级别的损失都在同一量级,不受 的大小影响。 这样也保证了“公平性”,即不会有哪个噪声级别的损失非常大/小,从而使得模型过份“重视”/“忽略”这个噪声级别所要学习的内容。

其次, 作者结合其在炼丹场上的经验, 发现网络训练到收敛时, 估计出的分数的量级 (即 norm)在 这个水平, 即:

最终, 这个经验性的结论启发了作者将 设置为 。为何? 你试着将 代入 式并结合 式, 会发现:

咦, 仔细看看 是什么! 还记得 在前文讲述去噪分数匹配时提到的吗, 由于 , 因此有 , 那么其实 就是噪声 , 这是与 无关的量, 即不受后者影响。

另外, 结合刚刚的结论: 的量级在 这个水平, 于是 的量级就会是 1.

综上, 的设置下, 加权后的损失 的量级就会与 的大小无关, 即能够保证所有噪声级别的损失都在同一量级, 使得模型能够"一视同仁"。

在这种训练方式下, 只要预先设定好 , 然后在每次迭代时随机选取其中一个噪声级别 , 并采样一个标准高斯噪声 , 对原始数据样本加噪: , 接着将加噪后的样本 与 对应的噪声级别 一并喂给网络食用, 待其吐出预测的分数 后, 就可以使用 来 计算损失了。

采样生成:退火朗之万动力学

CW 在前面已经给大家展示过朗之万动力学采样的公式(见 式),并且在 NCSN 概要部分也简单阐述过退火朗之万动力学采样的过程,这里进一步放出伪代码,让大家有更清晰的认知:

退火朗之万动力学采样的伪代码

一起对照着上面的伪代码来过一遍吧~

首先, 初始样本 从某种固定的先验分布去采样就好(作者默认采用均匀分布, 注意, 是均匀分布哦! 不是高斯分布!); 然后, 从最大的噪声级别 开始使用朗之万动力学采样, 直至最小的噪声级别 (这里你可以停顿下, 想想为何要由大到小, CW 在前面可是有给答案哦 ) 。在每个噪声级别进行采样前, 会先设定步长(step size) (为何这么设, 不告诉你 );接着, 在每个噪声级别下, 基于一定的步数 去迭代进行朗之万动力学采样。该噪声级别最后一步采样生成的样本会作为下一个噪声级别的初始样本(见上图 ); 最后, 待所有噪声级别的朗之万动力学采样过程均完成时, 就得到了最终的生成结果。

看到上面这段话的加粗部分,相信你们也知道 CW 又要啰嗦一下了,之所以这么做,是因为:

由于对数据进行了扰动,因此能够避免数据仅“驻扎”在低维流形,于是模型估计出来的分数就会比较准确(结合前面 CW 向大家分析的结论)。那么沿着分数(梯度)的方向,某个噪声级别下最终的采样生成结果就会聚敛到其对应的高概率密度区域。而由于相邻噪声级别之间的差异比较小,因此它们的高概率密度区域是非常近似的,相当于某噪声级别下样本最终聚敛到的区域也是下个噪声级别的高概率密度区域,而概率密度高,分数估计通常更准。这不是捡了个便宜嘛!对于下一个噪声级别来说,以当前噪声级别最终生成的样本作初始化无疑是白嫖一波~ 每个噪声级别都白嫖“前辈”们的,层层推波助澜,这就是所谓的站在巨人的肩膀上(突然懂了一些社会道理..)

哦, 你说学会了白嫖, 现在怪 小气, 不告诉你步长 为何设置为 (其中 是个常数, 这表示 )。

乖~别闹,现在就让你白嫖!

社会人做事都带有目的性, 美其名日: motivation。作者的 motivation 就是固定 “信噪比"(signalto-noise) 的量级。什么是信噪比? 就是这个东东: ( 是随机噪声)。现在来观察下它的量级:

以上最后是将 代入。而 是常量, 不影响结果, 于是略去。

咉! 看我们发现了什么! 前面作者凭他的经验性功力告诉过我们一个结论: , 于 是这里就有 。结合式 , 我们得到

这就说明, 的设置下, 信噪比会固定在常量级, 与噪声级别 无关。

NCSN 核心源码解析

CW 仅对 NCSN 算法本身紧密相关的代码实现进行解析,包括:loss 函数、采样生成过程 以及 作者自行设计的网络模块 —— 以噪声为条件的实例归一化(CondInstanceNom++)。

Loss 函数

通过前面的万字长文,相信 loss 公式你们几乎都会背了吧..(神马!?你告诉我你不会!?赶紧滑上去偷看多几下)

和理论描述的一致,loss 的实现非常简单,源码如下:

def anneal_dsm_score_estimation(scorenet, samples, labels, sigmas, anneal_power=2.):
    # 取出每个样本对应噪声级别下的噪声分布的标准差,即公式中的sigma_i,
    # 这里的 labels 是用于标识每个样本的噪声级别的,就是 i,实际是一种索引标识
    # (bs,)->(bs,1,1,1) 扩展至与图像一致的维度数
    used_sigmas = sigmas[labels].view(samples.shape[0], *([1] * len(samples.shape[1:])))
    # 加噪:x' = x + sigma * z (z ~ N(0,1))
    perturbed_samples = samples + torch.randn_like(samples) * used_sigmas
    
    # 目标score,本质是对数条件概率密度 log(p(x'|x)) 对噪声数据 x' 的梯度
    # 由于这里建模为高斯分布,因此可计算出结果最终如下,见前文公式(vii)
    target = - 1 / (used_sigmas ** 2) * (perturbed_samples - samples)
    # 模型预测的 score
    scores = scorenet(perturbed_samples, labels)
    target = target.view(target.shape[0], -1)
    scores = scores.view(scores.shape[0], -1)

    # 先计算每个样本在所有维度下分数估计的误差总和,再对所有样本求平均
    # 见前文公式(vii)
    loss = 1 / 2. * ((scores - target) ** 2).sum(dim=-1) * used_sigmas.squeeze() ** anneal_power

    return loss.mean(dim=0)

没有太多可言,CW 想讲的几乎都包含在以上注释中了。

可以注意下的是,paper 中以及前文都写到会将 输入网络。但实际只要让模型能够区分每个噪声级别即可,完全可以用序号(1,2,3 ..., L)来表示不同的噪声级别(上面代码中的 labels 就是这个意思),将这个这个序号输进网络,以替代 。当然,这个序号在网络中应该要被进一步编码成特征向量(通常用 embedding 的方式),以和图像特征共同交互。

采样生成

采样生成就是退火朗之万动力学那个套路,伪代码已经在前文展示过(见 Algorithm 1),这里就向大家展示下它的真面目:

def anneal_Langevin_dynamics(self, x_mod, scorenet, sigmas, n_steps_each=100, step_lr=0.00002):
    images = []

    with torch.no_grad():
        # 依次在每个噪声级别下进行朗之万动力学采样生成,噪声强度递减
        for c, sigma in tqdm.tqdm(enumerate(sigmas), total=len(sigmas), desc='annealed Langevin dynamics sampling'):
            # 噪声级别
            labels = torch.ones(x_mod.shape[0], device=x_mod.device) * c
            labels = labels.long()

            # 这个步长并非 Algorithm 1 中的 alpha,而是其中第6步的 alpha/2
            # 对应朗之万动力学采样公式(见公式(vi))的 epsilon/2
            step_size = step_lr * (sigma / sigmas[-1]) ** 2
            
            # 每个噪声级别下进行一定步数的朗之万动力学采样生成
            for s in range(n_steps_each):
                images.append(torch.clamp(x_mod, 0.01.0).to('cpu'))
                # 对应公式(vi)最后一项
                noise = torch.randn_like(x_mod) * np.sqrt(step_size * 2)
                # 网络估计的分数
                grad = scorenet(x_mod, labels)
                # 朗之万动力方程
                x_mod = x_mod + step_size * grad + noise

        return images

CondInstanceNorm++

这部分可以详细讲下,这个模块实质上就是在 instance normalization 的基础上对不同噪声级别下的特征做不同的仿射(affine)变换

其中 表示归一化后的特征, 是噪声级别 的标识, 代表图像特征的第几个通道 (channel)。

就是 instance normalization 做的事情, 但其本身也是带有可学习的仿射变换参数 的, 作者将它们取消掉, 分别以 进行替代(也是可学习的), 从而区分不同的噪声级别。

你以为就这样? 非也! 作者发现, 仅仅这样的话会噪声生成的图像中有颜色偏差的问题, 因为在归一化的时候抺去了 中带有的信息。

one downside of instance normalization is that it completely removes the information of
for different feature maps. This can lead to shifted colors in the generated images._

咋办捏?很直接 —— 加回去就好了!

但直接加末免太过粗暴, 不妨利用下神经网络的学习能力, 让 先乘上一个可学习的因子 (也是对不同噪声级别做区分), 然后再加回去。另外,在实际实现时, 作者还对 也做了归一化:计算所有通道的 的均值 和标准差

【声明】内容源于网络
0
0
极市平台
为计算机视觉开发者提供全流程算法开发训练平台,以及大咖技术分享、社区交流、竞赛实践等丰富的内容与服务。
内容 8155
粉丝 0
极市平台 为计算机视觉开发者提供全流程算法开发训练平台,以及大咖技术分享、社区交流、竞赛实践等丰富的内容与服务。
总阅读3.2k
粉丝0
内容8.2k