ReVQ (Quantize-then-Rectify,量化后修正)
扩散模型:一位“文物修复大师”
扩散模型通过一个两阶段过程来生成数据。第一阶段是固定的前向加噪过程 ,它在多个步骤中,逐步地向原始数据(如图片)添加少量的高斯噪声,直到数据最终变成完全无规律的纯噪声 。这个过程的每一步加多少噪声是预先定义好的,因此我们可以精确计算出从原始图片直接得到任意一个噪声版本的数学公式 。
模型学习的核心在于第二阶段,即一个可学习的逆向去噪过程。它的目标是学习前向过程的逆操作,从纯噪声图逐步还原出清晰的图片 。为了实现这一点,模型的训练目标被巧妙地设定为:在任意一个噪声步骤t,给定当前的噪声图片 x t x_t xt,网络需要预测出当初为了得到 x t x_t xt 而被添加进去的那部分噪声 ϵ \\epsilon ϵ 。这个目标使得损失函数变得非常简单,通常就是预测噪声和真实噪声之间的均方误差 。
当模型训练完成后,生成新数据的过程便开始了。我们先凭空创造一张完全随机的噪声图,作为去噪过程的起点 x T x_T xT 。然后,模型开始迭代工作:在每一步,它都会预测并移除当前噪声图中的一部分噪声,从而得到一个稍微清晰一点的图 x t − 1 x_{t-1} xt−1。这个步骤被重复执行,从 t=T t=T t=T 一直回溯到 t=0 t=0 t=0,最终将纯粹的噪声还原成一张全新的、清晰的、符合训练数据分布的图片。
一个 “先破坏,再修复”的游戏。
第一步:破坏文物 (前向过程)
- 你有一幅价值连城的古画(这就是你的原始数据 x 0 x_0 x0) 。
- 你故意要破坏它,但方式很特别:你每分钟都往画上均匀地喷一层薄薄的、几乎看不见的灰尘(高斯噪声) 。
- 第一分钟后,画作变得模糊了一点点(得到 x 1 x_1 x1)。第二分钟后,在 x 1 x_1 x1 的基础上又多了一层灰,变得更模糊了(得到 x 2 x_2 x2)。
- 你重复这个过程几百上千次。最后,这幅画完全被灰尘覆盖,变成了一片混沌的、毫无规律的噪点图(得到 x T x_T xT) 。
这个“破坏”的过程是固定的、可预测的。我们完全知道每一步加了多少“灰尘”,因为这是我们自己设定的 。
第二步:训练修复大师 (学习逆向过程)
现在,我们的目标是训练一个AI,让它成为一位顶级的文物修复大师。这位大师要学会的,就是如何“撤销”我们刚才的破坏过程。
- 大师如何学习?
我们不能直接让它看着一团噪声就变出古画,这太难了。所以我们用一种“作弊”的方式来教它:- 我们随机从破坏过程中拿出一个中间状态,比如被破坏了250次的画作 x 250 x_{250} x250。
- 我们对大师说:“看,这是 x 250 x_{250} x250。你猜猜,从第249次到第250次,我们刚刚往上喷的那层‘灰尘’是什么样的?”
- 大师会尽力去猜一个“预测的灰尘” ϵ θ \\epsilon_\\theta ϵθ。
- 因为我们知道标准答案(那层“真实的灰尘” ϵ \\epsilon ϵ),所以我们可以直接评判大师猜得准不准。
- 训练的核心目标:调整大师的大脑(神经网络),让它“预测的灰尘”和“真实的灰尘”之间的差距(均方误差)越小越好 。
我们重复这个训练过程数百万次,用不同的古画,在不同的破坏阶段(比如第10次、第500次、第800次)反复提问。久而久之,大师就练就了火眼金睛,无论看到哪个阶段的、被破坏的画,都能极其精准地猜出最后那一层加上去的灰尘是什么。
第三步:大师施展身手 (生成新画)
训练完成后,这位大师就出师了。我们如何让他创作一幅全新的画呢?
- 我们凭空制造一幅纯粹的、随机的噪声图(一幅全是“灰尘”的画),交给大师,告诉他这是“被破坏了1000次”的成品 x T x_T xT 。
- 大师看了一眼 x T x_T xT,凭借经验猜出了第1000层灰尘的样子,然后小心翼翼地把它“抹掉”,得到了稍微清晰一点的 x 999 x_{999} x999 。
- 接着,他看着 x 999 x_{999} x999,猜出并抹掉第999层灰尘,得到 x 998 x_{998} x998。
- ……他不断重复这个“猜灰尘、抹灰尘”的过程。
- 当他从 x 1 x_1 x1 抹掉最后一层灰尘后,一幅清晰的、前所未见的、但风格与我们训练数据一致的“新古画” x 0 x_0 x0 就诞生了!
这就是扩散模型的精髓:通过学习如何“预测噪声”,来一步步地“去噪”,最终从纯噪声中生成全新的数据。
变分自编码器 (VAE):一位“图书管理员”
VAE是一种包含编码器和解码器的生成模型,其关键在于它的编码过程是概率性的 。当输入一个数据样本时,编码器并不会输出一个确定的潜在向量,而是输出一个概率分布的参数(具体来说是高斯分布的均值 μ \\mu μ 和方差 σ 2 \\sigma^2 σ2) 。然后,系统会从这个由参数定义的分布中随机采样一个点 z z z ,这个点 z z z 才是真正传递给解码器用于重建原始数据的潜在向量。
VAE的训练目标由两部分损失函数共同决定,以达到精妙的平衡 。第一部分是重建损失 ,它要求从潜在向量 z z z 解码出的数据必须与原始输入尽可能相似,确保模型能够有效编码和解码信息 。第二部分是KL散度损失 ,它作为一个正则化项,强制编码器产生的所有概率分布都必须向一个标准的、简单的先验分布(通常是标准正态分布 N(0,I) \\mathcal{N}(0, I) N(0,I))看齐 。
这个双重目标使得训练好的VAE具备了生成能力。KL散度损失将潜在空间规整成一个连续、有结构的“地图” 。因此,要生成一个全新的数据,我们不再需要原始输入,只需从那个作为“地图”基准的标准正态分布中随机采样一个新的潜在向量 z n e w z_{new} znew,然后将其直接送入解码器 。解码器便能将这个随机点“翻译”成一个与训练数据风格一致的、全新的、高品质的输出样本。
假设你要管理一个巨大的图书馆,里面只存放各种“猫”的照片。
一个糟糕的管理员 (标准自编码器 AE)
- 他给每张猫的照片编一个死板的号码(比如“001号”、“002号”)。当你给他“001号”,他能找到对应的猫照片。
- 问题:这个编码系统是离散的、没有规律的。如果你问他“001.5号”是什么,他会不知所措,给你一堆乱码 。所以,他只能存取,不能创造新的猫照片。
一位聪明的管理员 (VAE)
这位管理员不给照片编死板的号码,他用一种更高级的方式管理。
-
编码:用一个“模糊区域”来描述
当你给他一张新的猫的照片时,他不会说“这张放在5号架子上”,而是说:“这张照片的特征大概在5号架子附近,可能偏左或偏右一点点。”
他为每张照片定义了一个“概率区域”(也就是高斯分布),这个区域由两个参数描述:中心位置(均值 μ \\mu μ)和模糊范围(方差 σ 2\\sigma^2 σ2) 。 -
解码:从区域中绘制图像
管理员的另一项能力是,你只要在他管理的“地图”上随便指一个点,他就能根据这个点的位置,画出对应的猫的照片 。
管理员的两条黄金准则 (损失函数)
为了让这个系统有效运作,管理员在接受培训时必须遵守两条铁律 :
-
第一条:重建必须准确 (Reconstruction Loss)
“我给你一张猫的照片,你把它定位到地图上的一个‘模糊区域’,再从这个区域里随机选个点把它画出来,画出来的必须和原图一模一样!” 。这保证了信息没有丢失。 -
第二条:地图必须整齐 (KL Loss)
“你为所有猫的照片定义的那些‘模糊区域’,整体来看,必须紧密地、有规律地聚集在地图的中心,形成一个漂亮的圆形云团(标准正态分布)。” 。
这条规则是VAE的灵魂!它强迫管理员把所有猫的特征都整理在一个连续、平滑的空间里。比如,“波斯猫”的区域和“橘猫”的区域会被安排得很近,而它们之间的空间,会被“波斯橘猫”的特征平滑地填满 。
最终的魔法:创造新猫
当管理员训练好后,他的“猫咪地图”就变得完美了。现在,要创造一只世界上独一无二的猫,你只需要:
- 闭上眼睛,在那片漂亮的“中心云团”里随便戳一个点 。
- 把这个点的坐标告诉管理员 。
- 管理员就会根据这个坐标,为你画出一只全新的、栩栩如生的猫!
这就是VAE的精髓:通过强迫潜在空间变得“规整且连续”,从而拥有了从这个空间中采样并生成新数据的能力。
强化学习 (RL):“训练一只小狗”
强化学习的核心是训练一个智能体 (Agent) 在一个环境 (Environment) 中,通过不断试错来学习一个最优的策略 (Policy),以最大化其能获得的长期累积奖励 (Reward) 。整个过程基于马尔可夫决策过程框架,即智能体在某个状态 s s s下,根据其策略执行一个动作 a a a,环境会反馈一个新的状态 s ′ s\' s′ 和一个即时奖励 r r r 。智能体的最终目标不是最大化眼前单步的奖励,而是最大化从长远来看的折扣奖励总和,即回报 (Return) 。
为了判断决策的好坏,RL引入了价值函数的概念来评估不同状态或动作的优劣。其中,动作价值函数 Q(s,a) Q(s, a) Q(s,a) 尤为重要,它表示在状态 s s s下执行动作 a a a后,遵循当前策略到底,预期能获得的总回报是多少 。通过学习和比较不同动作的Q值,智能体就能知道在某个状态下哪个动作是更好的选择。几乎所有RL算法都围绕着如何通过与环境的交互,更准确地估计这些价值函数而展开。
强化学习的算法主要分为两大类。基于价值 (Value-based) 的方法,如Q-learning,专注于学习精确的Q值函数,其策略是隐式的(例如总是选择Q值最高的动作) 。而基于策略 (Policy-based) 的方法,如Policy Gradient,则不直接学习价值函数,而是直接对策略本身进行参数化和优化 。它会通过试验,直接增大那些能带来高回报的动作序列的发生概率,反之则减小。此外,还有结合两者的演员-评论家 (Actor-Critic) 方法,能实现更高效的学习 。
你想训练你的小狗(智能体 Agent)学会“捡飞盘”这个游戏。
- 环境 (Environment):公园的草地。
- 状态 (State):小狗的位置、飞盘的位置、风速等等 。
- 动作 (Action):小狗能做的,比如跑、跳、叫、摇尾巴。
- 奖励 (Reward):你给的零食。捡到飞盘给一大块,跑近了给一小块 。
核心目标: 教会小狗一套“策略” (Policy) ,让它知道在任何状态下,做什么动作,才能让自己未来能得到的“零食总和”最大化。
两种主流训练哲学
-
基于价值 (Value-Based):教小狗“思考”
- 核心思想:不直接教小狗怎么做,而是教它评估“做某件事有多好”。这个“好”的程度,就是Q值 (Q-Value)。
- Q值是什么? Q ( s , a ) Q(s, a) Q(s,a) 代表在状态 s s s (比如“我离飞盘很近”),做动作 a a a (比如“向前冲”),预计未来能拿到的零食总和是多少 。
- 如何学习?(时序差分 TD):小狗在草地上乱跑。它在位置A(状态s),选择“向前冲”(动作a),到达了位置B(状态s’),你马上给了它一小块零食R。小狗会这样更新自己的大脑笔记:
“哦,我刚刚在A点向前冲,立刻拿到了零食R,并且到达了B点。我知道B点本身就很有价值(因为离飞盘更近了,预估价值是 Q ( s ′ , a ′ ) Q(s\', a\') Q(s′,a′))。那么,我现在觉得,当初在A点向前冲这个决策的价值,应该等于
即时拿到的零食R
+到达B点后的长期价值
。” - 它会用这个“新估算出的价值”去微调脑子里对“在A点向前冲”的旧价值 。这个过程不断重复,小狗的Q值笔记会越来越准。
- 如何决策?:当Q值学得差不多了,小狗的决策就简单了:在任何状态下,看看笔记,做那个Q值最高的动作就行了!
-
基于策略 (Policy-Based):直接教小狗“行动”
- 核心思想:我们不教小狗思考每个动作的价值,而是直接调整它的“行动直觉”(策略 π \\pi π)。
- 如何学习?(策略梯度):我们让小狗完整地玩一轮捡飞盘的游戏。
- 如果成功了:游戏结束,小狗拿到了很多零食。我们会回顾它刚才做的所有动作,然后想:“太棒了!刚才这一套连招(比如跑、调整方向、跳)都很不错,我们让小狗以后做出这些连招的概率都高一点点。”
- 如果失败了:小狗什么也没拿到。我们就回顾它刚才的所有动作,然后让它以后做出这些动作的概率都低一点点。
- 通过成千上万次的完整尝试,小狗的“直觉”会被不断地强化,好的动作组合概率越来越高,坏的越来越低。
演员-评论家 (Actor-Critic):最强训练法
这是上面两种方法的结合,就像一个团队在训练小狗:
- 演员 (Actor):就是小狗的大脑,负责做动作(基于策略) 。
- 评论家 (Critic):像一个拿着笔记本跟在旁边的教练,负责打分(基于价值) 。
工作流程:演员(小狗)做出一个动作,评论家(教练)马上在旁边喊:“好球!这个动作让你离飞盘更近了,价值+5分!”或者“臭棋!这个动作让你跑远了,价值-3分!”。演员听到教练的即时反馈,立刻调整自己的策略,而不需要等到游戏结束。这使得学习效率大大提高。
Quantize-then-Rectify: Efficient VQ-VAE Training
这篇论文提出了一种名为ReVQ (Quantize-then-Rectify,量化后修正) 的新框架,旨在解决当前视觉基础模型中一个核心的痛点:向量量化变分自编码器 (VQ-VAE) 的训练成本极高。VQ-VAE作为连接连续视觉信号(如图像)和离散语言模型(LLM)的桥梁(即视觉分词器),其重要性不言而喻,但其高昂的训练开销(动辄需要数千GPU小时)限制了其在更广泛研究和应用中的普及 。
论文的核心思想是:我们不必从零开始训练一个VQ-VAE,而是可以高效地将一个已经预训练好的、强大的连续特征自编码器(VAE)转化为一个高性能的VQ-VAE 。通过这种方式,论文实现了一个在训练效率、压缩率和重建质量之间取得极佳平衡的模型。
1. 研究背景、现状与挑战
研究背景
随着大型语言模型(LLM)的成功,多模态学习,特别是视觉-语言模型的融合,成为了前沿热点 。在这个过程中,视觉分词器(Visual Tokenizer) 扮演着至关重要的角色 。它的任务是将高维、连续的图像信息,转化为离散的、类似文本的Token序列,以便LLM能够理解和处理 。VQ-VAE就是实现这一目标的主流技术之一 。
研究现状
当前的VQ-VAE研究主要分为两大阵营 :
- 高压缩率、高成本方法:这类方法(如MaskBit)可以把图像压缩到非常短的Token序列(例如小于256个Token),重建质量也高,但训练代价极为昂贵,通常需要在A100集群上花费超过3000个GPU小时 。这使得只有少数大型机构能够负担得起。
- 高效率、低压缩率方法:这类方法(如TokenBridge, CODA)巧妙地利用预训练的VAE来加速训练,但代价是压缩率很低,生成的Token序列过长(例如超过2048个Token),这对于后续的生成任务并不理想 。
发现的问题与挑战
显然,当前领域面临一个核心的权衡困境(Trade-off):训练效率和压缩率似乎鱼与熊掌不可兼得 。研究者们迫切需要一种既能实现高压缩率(短Token序列),又能快速、低成本训练的VQ-VAE框架 。这就是本文要解决的核心挑战。
2. 核心研究动机与目标
核心研究动机
论文的作者们通过分析发现,VQ-VAE训练的计算瓶颈主要在于其庞大的编码器和解码器部分,特别是处理高分辨率输入的浅层网络 。同时,他们通过实验观察到一个有趣的现象:一个预训练好的VAE对一定程度的噪声具有很强的容忍度 。只要量化过程引入的误差(可以看作一种噪声)在VAE的容忍阈值之内,VAE的解码器依然能够重建出高质量的图像 。
这个发现是整篇论文的基石和核心动机:既然VAE能容忍噪声,那么我们就可以将量化误差控制在一定范围内,直接在一个预训练好的、冻结的VAE上插入一个量化器,从而将VAE“改造”成VQ-VAE,极大地节约训练成本 。
核心研究目标
基于以上动机,论文的核心目标是:
- 解决效率与压缩率的矛盾:开发一个全新的框架(ReVQ),使其能够在保持高压缩率(例如512或256个Token)和高重建质量的同时,将训练时间从数千GPU小时降低到几十个GPU小时 。
- 降低技术门槛:使得研究者在单张消费级显卡(如NVIDIA 4090)上,仅用一天左右的时间就能完成整个训练过程 ,极大地推动了该领域技术的可及性。
这个研究具有非常重要的实际意义。一个高效、低成本的视觉分词器能够加速多模态大模型的迭代,并为图像/视频生成、理解等下游应用提供强大的基础能力,使其更容易落地。
3. 研究内容、技术路线与创新点
为了实现上述目标,论文提出了两个关键的研究内容,共同构成了ReVQ框架。
研究内容一:设计高效且强大的量化器
- 动机:直接在VAE的特征上进行量化会遇到两大难题:1) 码本容量有限,导致量化误差大;2) 优化不稳定,容易出现“码本崩溃”(codebook collapse),即大量码本条目从未被激活和使用 。
- 核心内容与技术路线:
- 通道多组量化 (Channel Multi-Group Quantization):传统方法通常在空间维度上共享同一个码本,这隐含了一个很强但并不合理的假设,即所有空间位置的特征分布都相同 。作者通过实验发现,沿通道维度切分特征,其各组特征的分布独立性更高 。因此,他们提出沿通道维度将特征分组,并为每一组分配独立的码本 。
- 非激活重置 (Non-Activation Reset):为了解决“码本崩溃”问题,作者提出一个极其简单但有效的策略。在每个训练周期(epoch)结束时,检查并统计哪些码本条目没有被激活(使用次数为0)。然后,将这些“死掉”的码本条目重置到那些最常被激活的码本条目附近(加上一个微小的随机扰动)。这个操作相当于“复活”了无效的码本,并让它们去分担那些热门码本的压力,从而促进码本的充分利用 。
- 创新点:
- 通道多组量化:颠覆了传统沿空间维度切分的做法,找到了一个更适合视觉特征的量化策略,有效扩大了码本的等效自由度 。
- 非激活重置:提供了一种“即插即用”的、几乎无额外计算开销的码本优化策略 ,相比于复杂的优化算法或损失函数,更加简洁高效。
研究内容二:构建“量化后修正” (Quantize-then-Rectify) 框架
- 动机:尽管量化器设计得到了优化,但仅靠量化器在高压缩率下(如压缩到256个Token)产生的量化误差仍然可能超出预训练VAE的容忍极限,导致重建质量下降 。
- 核心内容与技术路线:
- 冻结预训练VAE:选择一个性能优越的预训练VAE模型(本文使用DC-AE ),并将其编码器和解码器的参数完全冻结,不参与训练 。
- 插入量化器和修正器:在冻结的VAE编码器和解码器之间,插入前述设计的量化器 (Quantizer) 和一个轻量级的、可学习的修正器 (Rectifier) 。
- 训练过程:图像首先通过冻结的VAE编码器得到潜层特征 Z e Z_e Ze。然后, Z e Z_e Ze 经过量化器得到离散的 Z q Z_q Zq。 Z q Z_q Zq 并不直接送入解码器,而是先经过修正器 g ( ⋅ ) g(\\cdot) g(⋅) 进行一次“纠错”,得到修正后的特征 Z q ′ Z_q\' Zq′ 。最后,训练的目标是让修正后的特征 g ( q ( Z e ) ) g(q(Z_e)) g(q(Ze)) 与原始特征 Z e Z_e Ze 尽可能接近(通过最小化 L 2 L_2 L2损失),而无需计算庞大的VAE梯度 。
- 创新点:
- ReVQ框架:这是一个全新的范式。它将繁重的图像编解码任务交给了强大的预训练VAE,而自身只专注于 “如何用最小的误差完成量化” 这一核心矛盾。通过引入修正器,进一步弥补了量化过程中的信息损失,使得在高压缩率下依然能保持低误差 。
- 极高的训练效率:由于只训练轻量级的量化器和修正器,整个框架的计算开销极低,实现了超过两个数量级的训练加速 。
4. 实验设计与验证
实验设置 (Setting)
- 数据集:ImageNet (128万张训练图像) 。
- 基础模型 (Backbone):使用预训练好的DC-AE作为冻结的VAE编解码器 。
- 硬件:在单张NVIDIA RTX 4090上进行训练 。
- 评价指标:
- 重建质量:PSNR, SSIM, LPIPS (感知损失), rFID (重建FID) 。
- 训练效率:GPU训练总小时数 。
基线模型 (Baseline)
论文与当前主流的SOTA(State-of-the-Art)方法进行了全面对比,涵盖了:
- 从零训练 (From Scratch) 的方法:VQGAN, MaskBit, ViT-VQGAN等 。
- 基于微调 (Fine Tuning) 的方法:TiTok, CODA等 。
- ReVQ(本文方法) 则属于冻结VAE (Frozen VAE) 的类别 。
核心想法验证(对比与消融实验)
-
如何验证“量化器设计”的有效性?
- 通道vs空间多组量化:在相同条件下,分别使用空间切分和通道切分的量化策略进行训练,对比它们的rFID指标。实验结果表明,通道切分在512和256两种Token长度下都显著优于空间切分 。
- 非激活重置策略:对比“使用”和“不使用”重置策略时,码本的利用率和最终的损失。实验图表清晰地显示,不使用该策略时,随着码本增大,利用率急剧下降至65.3%;而使用后,利用率始终保持在97%以上 。
-
如何验证“量化后修正框架”的有效性?
- 修正器的作用:在多个不同的Token长度设置下,对比“有修正器”和“没有修正器”的重建损失。结果显示,修正器显著降低了所有配置下的重建损失,尤其是在压缩率更高(模型基线更弱)的情况下,性能提升更明显(例如在64 Token长度下,损失降低了23.3%)。
- 修正器的架构选择:实验对比了使用ViT、CNN、MLP三种不同架构作为修正器的性能。结果表明,ViT架构在所有设置下都表现最好 。
- 非对称设计的合理性:作者还进行了一个有趣的探索,即在量化器前增加一个与修正器对称的编码器。实验发现这反而会急剧增加训练难度并导致性能下降 ,从而证明了当前非对称、轻量化的修正器设计是最优的。
5. 总结与评价
核心贡献总结
这篇论文的核心贡献在于:
- 提出了ReVQ框架:一个新颖、高效的范式,通过**“量化后修正”**的思想,成功地将一个预训练的连续VAE高效地转换为一个离散的VQ-VAE。
- 实现了效率与性能的统一:在保持与SOTA方法相媲美的重建质量和高压缩率的同时,将训练成本降低了超过两个数量级 ,使得在单张消费级显卡上进行快速训练成为可能。
- 提出了有效的技术组件:通道多组量化和非激活重置策略,为解决VQ-VAE训练中的码本容量和稳定性问题提供了简洁而强大的解决方案。
“利用预训练VAE的噪声容忍度”——非常深刻,并以此构建了一个逻辑清晰、执行简单的解决方案。
评价:
-
优点:
- 创新性强:它没有在现有VQ-VAE上“卷”网络结构或优化算法,而是提出了一种全新的、釜底抽薪式的“改造”范式,思路非常巧妙。
- 实践价值巨大:极大地降低了高性能视觉分词器的训练门槛,使得更多的研究者和开发者能够参与到多模态大模型的研究中,有力地推动了社区的发展。
- 实验扎实:消融实验设计得非常充分,有力地验证了每一个技术组件的有效性和设计选择的合理性。
-
潜在局限与未来展望:
- 论文也坦诚,在极高的压缩率下(如32个Token),ReVQ目前还无法匹敌像TiTok这样的顶尖模型 。作者认为这与修正器的架构设计有关,未来可以探索更强大的修正器模型来突破这一极限 。
- 该框架的性能上限在一定程度上取决于所使用的预训练VAE的性能。但反过来看,这也是一个优点:随着未来出现更强大的VAE模型,ReVQ框架的性能也可以“水涨船高”,轻松地享受到基础模型进步带来的红利。
从一张猫的图片到重建
我们有一张高清的猫的图片(比如,256x256
像素),我们的目标是把它“翻译”成一个很短的、由512个“单词”组成的句子(即512个Token),然后再用这个句子把猫的图片大致画出来。
整个流程分为训练和使用(推理)两个阶段。
第一阶段:训练(耗时约22小时)
这个阶段的目标是训练出两个“专家”:一个 “分组翻译家” (量化器 Quantizer) 和一个 “润色画家” (修正器 Rectifier) 。
我们还有一个已经训练好的、非常强大的 “艺术大师VAE”,它由一个 “高级素描师” (VAE Encoder) 和一个 “上色大师” (VAE Decoder)组成。这位艺术大师非常厉害,但我们不准他再学习了,让他当评委和工具人(即参数冻结)。
具体步骤如下:
第1步:高级素描师出马 (Frozen VAE Encoder)
- 输入: 一张猫的图片,数据形态是
[1, 3, 256, 256]
(1张图片, 3个颜色通道, 256x256尺寸)。 - 操作: “高级素描师”把这张彩色的、复杂的图片,压缩成一幅信息量极大的“素描草稿”。这幅草稿不是一张图,而是一个数字矩阵(潜层特征)。根据论文,这个草稿的总信息量是2048个数字 。为了方便后续处理,我们把它看成一个
8x8
的网格,每个格子里有32个数字。 - 输出: 一个潜层特征图
z_e
,数据形态是[1, 8, 8, 32]
(1张图, 8x8网格, 每个格点32维特征)。这幅“草稿”就是我们后续所有操作的黄金标准。
第2步:分组翻译家上场 (Channel Multi-Group Quantizer)
- 输入: “高级素描师”画的草稿
z_e
,形态[1, 8, 8, 32]
。 - 目标: 把这幅草稿“翻译”成512个“单词”(Token)。
- 核心操作 - 通道多组量化:
- 特征拆分: “分组翻译家”不把
8x8
网格上的64个位置分开翻译,而是把总共8*8*32 = 2048
个特征数字,重新分成512组,每组4个数字 (2048 / 512 = 4
)。现在我们有了512个小向量,每个向量都是4维的。 - 查字典: “分组翻译家”有512本不同的字典(512个独立的码本)。每本字典里有16384个“标准词汇”(码本大小为16384),每个词汇也是一个4维向量。
- 翻译: 对于第1个4维向量,他就在第1本字典里找一个最像的“标准词汇”来代替它。对于第2个向量,就在第2本字典里找…以此类推,直到512个向量全部被字典里的“标准词汇”替换掉。
- 这个过程产生了512个“单词”,也就是512个离散的码本索引,这就是我们想要的Token。
- 特征拆分: “分组翻译家”不把
- 输出: 一幅用“标准词汇”拼凑回来的“生硬版草稿”
z_q
。它的数据形态和输入一样,还是[1, 8, 8, 32]
,但里面的数字已经被字典里的标准值替换了,所以信息有损失,比较粗糙。
第3步:润色画家进行修正 (Rectifier)
- 输入: “分组翻译家”翻译的“生硬版草稿”
z_q
,形态[1, 8, 8, 32]
。 - 操作: “润色画家”是一位轻量级专家(比如一个小的ViT网络),他的任务是看着这幅生硬的草稿,进行“润色和修正”,让它看起来更接近“高级素描师”最开始画的原始草稿。他会调整线条,弥补细节,但不会改变整体布局。
- 输出: 一幅“精修版草稿”
z_q_rectified
,数据形态依然是[1, 8, 8, 32]
。
第4步:评判与学习 (计算损失并反向传播)
- 操作: 我们拿出“高级素描师”的原始草稿
z_e
和“润色画家”的精修版草稿z_q_rectified
进行对比。计算它们之间的差距( L _ 2 L\\_2 L_2损失)。 - 学习: 如果差距很大,就告诉“润色画家”和“分组翻译家”:“你们俩配合得不好,翻译和润色的结果离原作差太远了!”。然后,他们会根据这个反馈调整自己的技能(更新网络参数和码本)。注意,“艺术大师VAE”全程不参与学习,只当评委。
第5步:周期性复盘 (Non-Activation Reset)
- 操作: 在一轮训练(一个epoch)结束后,我们会检查“分组翻译家”的512本字典。如果发现某本字典里有些“标准词汇”一次都没被用过(死词汇),就把这些死词汇扔掉,然后从这本字典里最热门的词汇旁边抄几个过来,稍微改一改,作为新的词汇放进去。这能保证字典里的每个词都有机会被用到,避免资源浪费。
这个过程重复进行,直到“分组翻译家”和“润色画家”配合默契,能将任何“素描草稿”都处理得和原作非常接近为止。
第二阶段:使用(推理)
训练完成后,我们有了一个成熟的“翻译-润色”团队。
- 拿来一张新的猫图片,交给 “高级素描师” (VAE Encoder) ,得到原始草稿
z_e
。 - 将
z_e
交给 “分组翻译家” (Quantizer),得到512个“单词”(Tokens)和 “生硬版草稿”z_q
。 - 将
z_q
交给 “润色画家” (Rectifier),得到 “精修版草稿”z_q_rectified
。 - 最后,将这幅 “精修版草稿”
z_q_rectified
交给一直没出手的 “上色大师” (VAE Decoder) 。 - “上色大师”根据这幅高质量的草稿,就能轻松地将原始的猫图片重建出来。
PyTorch伪代码
import torchimport torch.nn as nnfrom torch.optim import AdamW# ==============================================================================# 1. 定义模型组件# ==============================================================================# 假设我们有预训练好的、冻结的VAE模型 (工具人)# 论文中使用DC-AEclass FrozenVAEEncoder(nn.Module): def __init__(self): super().__init__() # ... 这里是复杂的网络结构,例如ResNet或ViT ... self.requires_grad_(False) # 关键:冻结所有参数,不参与训练 def forward(self, image_batch): # image_batch: [B, 3, 256, 256], B是批量大小 # 输出潜层特征图 z_e # z_e: [B, 8, 8, 32], 这是假设的维度,总特征数为 B * 8 * 8 * 32 z_e = self.internal_model(image_batch) # z_e.shape -> [B, 8, 8, 32] return z_eclass FrozenVAEDecoder(nn.Module): def __init__(self): super().__init__() # ... 这里是与Encoder对称的复杂网络结构 ... self.requires_grad_(False) # 关键:冻结所有参数 def forward(self, z_q_rectified): # z_q_rectified: [B, 8, 8, 32], 润色后的特征 # reconstructed_image: [B, 3, 256, 256], 重建的图像 reconstructed_image = self.internal_model(z_q_rectified) return reconstructed_image# 我们真正要训练的核心模型 ReVQclass ReVQ(nn.Module): def __init__(self, num_tokens=512, codebook_size=16384, feature_dim_total=2048): super().__init__() self.num_tokens = num_tokens # B_tokens in paper, e.g., 512. 这是我们要生成的Token数量 self.codebook_size = codebook_size # N in paper, e.g., 16384. 每个码本的大小 self.feature_dim_total = feature_dim_total # 2048, VAE输出的总特征数 # 计算每个Token代表的特征维度 # 2048 / 512 = 4 self.group_dim = feature_dim_total // num_tokens # --- 2.1 \"分组翻译家\" (Quantizer) --- # 使用ModuleList来存储512个独立的码本 # nn.Embedding 是一个高效的查找表,完美契合码本的需求 self.codebooks = nn.ModuleList([ nn.Embedding(self.codebook_size, self.group_dim) for _ in range(self.num_tokens) ]) # 用于记录码本激活次数,为 \"非激活重置\" 作准备 for i in range(self.num_tokens): self.register_buffer(f\'activation_counts_{i}\', torch.zeros(self.codebook_size)) # --- 2.2 \"润色画家\" (Rectifier) --- # 论文中使用轻量级的EfficientViT Block # 这里用一个简单的MLP代替,保持输入输出维度一致 self.rectifier = nn.Sequential( nn.Linear(feature_dim_total, feature_dim_total), nn.ReLU(), nn.Linear(feature_dim_total, feature_dim_total) ) def quantize(self, z_e): # z_e: [B, 8, 8, 32] -> VAE Encoder的输出 B, H, W, D_grid = z_e.shape # (B, 64, 32) -> (B, 2048) z_e_flat = z_e.view(B, H * W * D_grid) # 关键的 \"通道多组\" 拆分 # (B, 2048) -> (B, 512, 4) z_groups = z_e_flat.view(B, self.num_tokens, self.group_dim) quantized_groups = [] token_indices = [] # 遍历512个组,在各自的码本中查找 for i in range(self.num_tokens): group_features = z_groups[:, i, :] # 当前组的特征, shape: [B, 4] codebook_i = self.codebooks[i].weight # 当前组的码本, shape: [16384, 4] # 计算距离 (欧式距离的平方) # dist_shape: [B, 16384] dist = torch.sum(group_features**2, dim=1, keepdim=True) + \\ torch.sum(codebook_i**2, dim=1) - \\ 2 * torch.matmul(group_features, codebook_i.t()) # 找到最近的码本向量的索引 # min_indices_shape: [B] min_indices = torch.argmin(dist, dim=1) token_indices.append(min_indices.unsqueeze(1)) # 根据索引从码本中取出对应的向量 # quantized_group_shape: [B, 4] quantized_group = self.codebooks[i](min_indices) quantized_groups.append(quantized_group.unsqueeze(1)) # 更新激活计数 (用于非激活重置) self.activation_counts_[i].index_add_(0, min_indices, torch.ones_like(min_indices, dtype=torch.float)) # 合并所有量化后的组 # z_q_flat_shape: [B, 512, 4] z_q_flat = torch.cat(quantized_groups, dim=1) # 转换回原始草稿的形状 # z_q_shape: [B, 8, 8, 32] z_q = z_q_flat.view(B, H, W, D_grid) # 最终的离散tokens # tokens_shape: [B, 512] tokens = torch.cat(token_indices, dim=1) return z_q, tokens def forward(self, z_e): # z_e: [B, 8, 8, 32], VAE Encoder的输出 # 2. 量化 z_q, _ = self.quantize(z_e) # z_q: [B, 8, 8, 32] # 3. 修正 # 为了送入MLP,先展平 z_q_flat = z_q.view(z_q.size(0), -1) # -> [B, 2048] z_q_rectified_flat = self.rectifier(z_q_flat) # 转换回原始草稿的形状,以便计算损失 z_q_rectified = z_q_rectified_flat.view_as(z_e) # -> [B, 8, 8, 32] return z_q_rectified def non_activation_reset(self): print(\"执行非激活重置策略...\") for i in range(self.num_tokens): counts = getattr(self, f\'activation_counts_{i}\') # 找到未激活 (count=0) 和已激活的码 unactivated_indices = torch.where(counts == 0)[0] activated_indices = torch.where(counts > 0)[0] if len(unactivated_indices) == 0: continue # 找到最常被激活的码 most_activated_indices = torch.topk(counts, k=len(unactivated_indices), largest=True).indices # 重置操作: 将未激活的码重置到最激活的码附近 with torch.no_grad(): new_codes = self.codebooks[i].weight[most_activated_indices] # 添加微小随机扰动 new_codes += torch.randn_like(new_codes) * 1e-4 self.codebooks[i].weight[unactivated_indices] = new_codes # 重置计数器 counts.zero_()# ==============================================================================# 3. 训练和推理流程# ==============================================================================# 实例化所有组件frozen_vae_encoder = FrozenVAEEncoder()frozen_vae_decoder = FrozenVAEDecoder()revq_model = ReVQ() # 这是我们要训练的模型# 优化器只关注ReVQ的参数optimizer = AdamW(revq_model.parameters(), lr=1e-3)loss_fn = nn.MSELoss() # 论文中使用L2损失# --- 训练循环 ---def train_one_epoch(dataloader): revq_model.train() for image_batch, _ in dataloader: optimizer.zero_grad() # 1. 高级素描师出马,得到原始草稿 z_e = frozen_vae_encoder(image_batch) # Shape: [B, 8, 8, 32] # 2. ReVQ模型处理,得到精修版草稿 z_q_rectified = revq_model(z_e) # Shape: [B, 8, 8, 32] # 4. 计算损失:在潜层空间中对比原始草稿和精修草稿 loss = loss_fn(z_q_rectified, z_e.detach()) # 用z_e.detach()确保梯度不流向VAE Encoder # 5. 反向传播和优化 loss.backward() optimizer.step() # 6. Epoch结束后,执行非激活重置 revq_model.non_activation_reset()# --- 推理/使用函数 ---def inference(image): revq_model.eval() with torch.no_grad(): # 1. 编码 z_e = frozen_vae_encoder(image) # 2. 量化并获得Tokens z_q, tokens = revq_model.quantize(z_e) # 3. 修正 z_q_rectified_flat = revq_model.rectifier(z_q.view(z_q.size(0), -1)) z_q_rectified = z_q_rectified_flat.view_as(z_e) # 4. 解码/重建 reconstructed_image = frozen_vae_decoder(z_q_rectified) return reconstructed_image, tokens # 返回重建图像和离散Tokens
这篇论文提出了一种极其巧妙且高效的 “寄生式” 训练范式。它不再从零开始构建和训练一个庞大而昂贵的VQ-VAE,而是聪明地利用了一个已经存在的、强大的预训练VAE模型作为基础。其核心流程是:首先,将预训练VAE的编码器和解码器冻结,仅作为特征提取和图像重建的工具;然后,在这两者之间插入一个轻量级的、可学习的 “量化-修正”模块 (ReVQ)。该模块通过通道多组量化策略来提升码本的表达能力,并通过非激活重置机制来保证码本的有效利用,最后再由一个修正器网络来弥补量化过程中的信息损失。整个训练过程只优化这个轻量级模块,并且直接在VAE的潜层空间计算损失,从而完全绕开了对图像像素的直接操作和庞大网络的梯度计算,最终以降低超过两个数量级的计算成本,在单张消费级显卡上实现了与SOTA方法相媲美的高质量图像分词和重建效果。
扩散模型 (Diffusion Model)
扩散模型是一种强大的生成模型,它通过两个过程来学习生成数据:一个固定的前向过程(不断加噪)和一个学习的 逆向过程(不断去噪)。
核心思想
我们有一张清晰的图片。
- 前向过程 (Forward Process):我们逐步、反复地向这张图片添加少量的高斯噪声。经过足够多的步骤后,原始图片会变得面目全非,最终看起来就像一张纯粹的、无意义的高斯噪声图 。这个过程是固定的、无需学习的。
- 逆向过程 (Reverse Process):模型的核心任务是学习如何“撤销”这个加噪的过程。它从一张纯噪声图开始,逐步、反复地去除噪声,最终恢复出一张清晰、有意义的图片 。通过学习这个“去噪”的逆向过程,模型就具备了从无到有生成新数据的能力。
前向过程 (Forward Process) 的数学表达
前向过程定义了如何从时刻 t−1 t-1 t−1 的图片 x t − 1 x_{t-1} xt−1 得到时刻 t t t 的图片 x t x_t xt。这通过向 x t − 1 x_{t-1} xt−1 添加高斯噪声来实现。
每一步的加噪过程可以表示为:
x t = α t x t − 1+ 1 − α t ϵ t − 1 其中 ϵ t − 1∼ N ( 0 , I ) x_{t} = \\sqrt{\\alpha_{t}}x_{t-1} + \\sqrt{1-\\alpha_{t}}\\epsilon_{t-1} \\quad \\text{其中} \\quad \\epsilon_{t-1} \\sim \\mathcal{N}(0, I) xt=αtxt−1+1−αtϵt−1其中ϵt−1∼N(0,I)
这里:
- x t − 1 x_{t-1} xt−1 是上一步的图片(信号)。
- ϵ t − 1 \\epsilon_{t-1} ϵt−1 是一个标准高斯噪声 。
- α t \\alpha_t αt 是一个预先设定的超参数(称为 “schedule”),它控制着每一步中信号和噪声的比例 。通常,随着 t t t 增大, α t \\alpha_t αt 会变小,意味着信号的比例越来越小,噪声的比例越来越高 。
这个公式意味着 x t x_t xt 的分布是一个以 α t x t − 1 \\sqrt{\\alpha_{t}}x_{t-1} αtxt−1 为均值, (1− α t )I (1-\\alpha_{t})I (1−αt)I 为方差的高斯分布 :
q ( x t ∣ x t − 1) = N ( x t ; α t x t − 1, ( 1 − α t ) I ) q(x_{t}|x_{t-1}) = \\mathcal{N}(x_{t}; \\sqrt{\\alpha_{t}}x_{t-1}, (1-\\alpha_{t})I) q(xt∣xt−1)=N(xt;αtxt−1,(1−αt)I)
一个非常重要的特性是,我们可以推导出从原始图片 x 0 x_0 x0 直接得到任意时刻 t t t 的噪声图片 x t x_t xt 的 “封闭形式”(closed-form) 表达式。通过对上述公式进行递归展开和利用高斯分布的可加性 ,可以得到 :
x t = α ‾ t x 0 + 1 − α ‾ t ϵ 其中 α ‾ t = ∏ i = 1 tα i x_{t} = \\sqrt{\\overline{\\alpha}_{t}}x_{0} + \\sqrt{1-\\overline{\\alpha}_{t}}\\epsilon \\quad \\text{其中} \\quad \\overline{\\alpha}_{t} = \\prod_{i=1}^{t} \\alpha_i xt=αtx0+1−αtϵ其中αt=i=1∏tαi
这个公式极为关键,因为它允许我们在训练时,可以随机采样一个时间步 t t t,然后一步就计算出对应的噪声图片 x t x_t xt,而不需要进行 t t t 次迭代,极大地提高了训练效率。
逆向过程 (Reverse Process) 与训练目标
逆向过程的目标是学习分布 p( x t − 1 ∣ x t ) p(x_{t-1}|x_t) p(xt−1∣xt),即给定 t t t 时刻的噪声图,如何预测出 t−1 t-1 t−1 时刻相对更清晰的图 。
挑战:直接计算真实的后验分布 q( x t − 1 ∣ x t ) q(x_{t-1}|x_t) q(xt−1∣xt) 是极其困难的,因为它需要知道整个数据集的分布 。
解决方案:利用贝叶斯定理,我们可以证明,如果我们 额外知道原始图片 x 0 x_0 x0,那么这个后验分布 q( x t − 1 ∣ x t , x 0 ) q(x_{t-1}|x_t, x_0) q(xt−1∣xt,x0) 是可以被精确计算的,并且它也是一个高斯分布 。它的均值和方差都可以用 x t x_t xt, x 0 x_0 x0 和预设的 α t \\alpha_t αt 参数来表示 。
于是,模型的学习目标就变成了:用一个神经网络 p θ ( x t − 1 ∣ x t ) p_{\\theta}(x_{t-1}|x_t) pθ(xt−1∣xt) 去拟合这个理论上可解的分布 q( x t − 1 ∣ x t , x 0 ) q(x_{t-1}|x_t, x_0) q(xt−1∣xt,x0) 。我们同样假设神经网络的输出是一个高斯分布 :
p θ ( x t − 1∣ x t ) = N ( x t − 1; μ θ ( x t , t ) , Σ θ ( x t , t ) ) p_{\\theta}(x_{t-1}|x_{t}) = \\mathcal{N}(x_{t-1}; \\mu_{\\theta}(x_{t},t), \\Sigma_{\\theta}(x_{t},t)) pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),Σθ(xt,t))
训练目标的简化:预测噪声
为了让两个高斯分布(理论值 q q q 和网络预测值 p θ p_\\theta pθ)接近,我们需要让它们的均值和方差尽可能一致 。
- 方差:研究发现,逆向过程的方差可以被设定为一个与 α t \\alpha_t αt 相关的固定值,而不需要神经网络去学习 。这大大简化了问题。
- 均值:现在,我们只需要让神经网络预测的均值 μ θ ( x t , t ) \\mu_{\\theta}(x_t, t) μθ(xt,t) 尽可能接近理论上的均值 μ q ( x t , x 0 ) \\mu_q(x_t, x_0) μq(xt,x0)。
然而,理论均值 μ q \\mu_q μq 依赖于我们不知道的原始图片 x 0 x_0 x0 。这里的神来之笔是,我们再次使用前向过程的封闭形式公式:
x 0 = x t − 1 − α ‾ t ϵ α ‾ t x_{0}=\\frac{x_{t}-\\sqrt{1-\\overline{\\alpha}_{t}}\\epsilon}{\\sqrt{\\overline{\\alpha}_{t}}} x0=αtxt−1−αtϵ
将这个 x 0 x_0 x0 的表达式代入理论均值 μ q \\mu_q μq 的公式中,经过一系列复杂的代数化简后 ,可以发现预测均值 μ θ \\mu_{\\theta} μθ 本质上等价于预测噪声 ϵ \\epsilon ϵ 。
因此,扩散模型的训练目标被极大地简化了:我们不再直接预测去噪后的图片,而是训练一个神经网络(通常是U-Net结构),记为 ϵ θ ( x t ,t) \\epsilon_{\\theta}(x_t, t) ϵθ(xt,t),让它根据当前的噪声图片 x t x_t xt 和时间步 t t t,去预测当时被添加进去的噪声 ϵ \\epsilon ϵ 。
损失函数也因此变得非常简单,就是一个均方误差(MSE Loss):
L = ∣ ∣ ϵ − ϵ θ ( x t , t ) ∣ ∣ 2 2 \\mathcal{L} = ||\\epsilon - \\epsilon_{\\theta}(x_t, t)||_2^2 L=∣∣ϵ−ϵθ(xt,t)∣∣22
其中 ϵ \\epsilon ϵ 是真实的噪声, ϵ θ \\epsilon_{\\theta} ϵθ 是网络预测的噪声。
训练与采样
-
训练 (Training) :
- 从数据集中随机选择一张原始图片 x 0 x_0 x0。
- 随机选择一个时间步 t t t。
- 生成一个标准高斯噪声 ϵ \\epsilon ϵ。
- 使用封闭形式公式直接计算出噪声图片: x t = α ‾ tx 0 + 1 −α ‾ t ϵ x_{t} = \\sqrt{\\overline{\\alpha}_{t}}x_{0} + \\sqrt{1-\\overline{\\alpha}_{t}}\\epsilon xt=αtx0+1−αtϵ。
- 将 x t x_t xt 和 t t t 输入到神经网络 ϵ θ \\epsilon_{\\theta} ϵθ,得到预测噪声 ϵ θ ( x t , t ) \\epsilon_{\\theta}(x_t, t) ϵθ(xt,t)。
- 计算损失 ∣ ∣ ϵ − ϵ θ ( x t , t ) ∣ ∣ 2 2 ||\\epsilon - \\epsilon_{\\theta}(x_t, t)||_2^2 ∣∣ϵ−ϵθ(xt,t)∣∣22 并通过反向传播更新网络参数 θ \\theta θ。
-
采样 (Sampling) :
- 从一个标准高斯分布中采样得到一张纯噪声图片 x T x_T xT。
- 从 t = T t=T t=T 开始,循环直到 t = 1 t=1 t=1:
a. 将当前的 x t x_t xt 和时间步 t t t 输入网络,得到预测的噪声 ϵ θ ( x t , t ) \\epsilon_{\\theta}(x_t, t) ϵθ(xt,t)。
b. 使用推导出的去噪公式,结合预测的噪声,计算出 x t − 1 x_{t-1} xt−1 。
x t − 1 = 1α t ( x t − 1 − α t 1 − α ‾ t ϵ θ ( x t , t ) ) + σ t z x_{t-1} = \\frac{1}{\\sqrt{\\alpha_{t}}}(x_{t}-\\frac{1-\\alpha_{t}}{\\sqrt{1-\\overline{\\alpha}_{t}}}\\epsilon_{\\theta}(x_{t},t)) + \\sigma_{t}z xt−1=αt1(xt−1−αt1−αtϵθ(xt,t))+σtz
(其中 z z z 是一个随机噪声,在某些采样方法中会加入以增加随机性) - 当循环结束时,得到的 x 0 x_0 x0 就是一张由模型生成的新图片。
变分自编码器 (Variational Autoencoder, VAE)
VAE是一种具有优秀生成能力的自编码器模型。它通过强制让潜在空间(latent space)服从一个特定的概率分布来实现这一点。
从自编码器 (AE) 到 VAE
-
标准自编码器 (AE):由一个编码器(Encoder)和一个解码器(Decoder)组成。编码器将输入数据压缩成一个低维的潜在向量(latent vector),解码器再从这个向量重建出原始数据 。AE的目标是让重建结果和原始输入尽可能相似。但是,AE的潜在空间是“无组织的”,在潜在空间中两个点之间进行插值,解码后的结果通常是无意义的,因此AE不具备好的生成能力 。
-
变分自编码器 (VAE):为了解决这个问题,VAE对潜在空间施加了一个强约束。编码器不再输出一个确定的向量,而是输出一个概率分布的参数(通常是高斯分布的均值 μ \\mu μ 和方差 σ 2\\sigma^2 σ2)。然后,我们从这个分布中采样一个点 z z z 作为潜在向量,再送入解码器 。这种设计强制让潜在空间变得连续、规整,从而拥有了强大的生成能力 。VAE可以被看作是在AE的基础上增加了正则化,以确保潜在空间具有良好的生成属性 。
VAE的数学原理:证据下界 (ELBO)
VAE的最终目标是最大化观测数据 x x x 的对数似然 log p(x) log~p(x) log p(x)。但这个值通常难以直接计算,因为涉及到对所有潜在变量 z z z 的积分 p(x)=∫p(x∣z)p(z)dz p(x) = \\int p(x|z)p(z)dz p(x)=∫p(x∣z)p(z)dz,这个积分没有解析解 。
因此,VAE引入了变分推断(Variational Inference)的思想。我们使用一个由神经网络构成的编码器 q ϕ (z∣x) q_{\\phi}(z|x) qϕ(z∣x) 来近似真实但不可解的后验分布 p(z∣x) p(z|x) p(z∣x) 。通过数学推导,我们可以得到一个非常重要的关系 :
l o g p ( x ) = D K L( q ϕ ( z ∣ x ) ∣ ∣ p ( z ∣ x ) ) + L E L B O log~p(x) = D_{KL}(q_{\\phi}(z|x) || p(z|x)) + \\mathcal{L}_{ELBO} log p(x)=DKL(qϕ(z∣x)∣∣p(z∣x))+LELBO
这里:
- D K L D_{KL} DKL 是KL散度,衡量两个分布的相似性,它永远大于等于0。
- L E L B O \\mathcal{L}_{ELBO} LELBO 被称为证据下界 (Evidence Lower Bound, ELBO)。
因为 D K L ≥0 D_{KL} \\geq 0 DKL≥0,所以 log p(x)≥ L E L B O log~p(x) \\geq \\mathcal{L}_{ELBO} log p(x)≥LELBO。这意味着ELBO是 log p(x) log~p(x) log p(x) 的一个下界。因此,最大化ELBO就等价于间接地最大化我们真正关心的数据似然 。同时,这也等价于最小化近似分布 q ϕ (z∣x) q_{\\phi}(z|x) qϕ(z∣x) 和真实后验 p(z∣x) p(z|x) p(z∣x) 之间的KL散度。
VAE的损失函数
通过对ELBO项进行展开,我们可以得到VAE最终的、可操作的损失函数。最大化ELBO等价于最小化以下损失函数:
L V A E= − L E L B O= − E q ϕ ( z ∣ x )[ l o g p θ ( x ∣ z ) ] + D K L( q ϕ ( z ∣ x ) ∣ ∣ p ( z ) ) \\mathcal{L}_{VAE} = - \\mathcal{L}_{ELBO} = - \\mathbb{E}_{q_{\\phi}(z|x)}[log~p_{\\theta}(x|z)] + D_{KL}(q_{\\phi}(z|x) || p(z)) LVAE=−LELBO=−Eqϕ(z∣x)[log pθ(x∣z)]+DKL(qϕ(z∣x)∣∣p(z))
这个损失函数由两部分组成 :
-
重建损失 (Reconstruction Loss): − E q ϕ ( z ∣ x ) [ l o g p θ ( x ∣ z ) ] -\\mathbb{E}_{q_{\\phi}(z|x)}[log~p_{\\theta}(x|z)] −Eqϕ(z∣x)[log pθ(x∣z)] 。
- 作用:这一项鼓励解码器 p θ ( x ∣ z ) p_{\\theta}(x|z) pθ(x∣z) 能够根据潜在向量 z z z 完美地重建出原始数据 x x x。
- 实现:通常实现为输入 x x x 和重建输出 x ^ \\hat{x} x^ 之间的均方误差(L2 Loss)或交叉熵损失 。
-
KL散度/正则化项 (KL Divergence / Regularization Term): D K L ( q ϕ ( z ∣ x ) ∣ ∣ p ( z ) ) D_{KL}(q_{\\phi}(z|x) || p(z)) DKL(qϕ(z∣x)∣∣p(z)) 。
- 作用:这一项是一个正则化器。它要求编码器 q ϕ ( z ∣ x ) q_{\\phi}(z|x) qϕ(z∣x) 产生的分布,必须与一个预设的先验分布 p ( z ) p(z) p(z) 相似。我们通常选择 p ( z ) p(z) p(z) 为标准正态分布 N ( 0 , I ) \\mathcal{N}(0, I) N(0,I) 。这会强制编码器将所有输入数据都映射到这个以原点为中心的高斯分布“云”中,从而使得潜在空间变得规整和连续。
- 实现:由于我们假设 q ϕ ( z ∣ x ) q_{\\phi}(z|x) qϕ(z∣x) 也是一个高斯分布,这个KL散度有解析解: 1 2 ∑ ( σ 2 + μ 2 − l o g ( σ 2 ) − 1 ) \\frac{1}{2}\\sum(\\sigma^2 + \\mu^2 - log(\\sigma^2) - 1) 21∑(σ2+μ2−log(σ2)−1) 。
这两项损失的作用 :
- 只用重建损失,潜在空间中的点虽然能重建,但分布杂乱无章,没有区分度 。
- 只用KL损失,所有点都挤在了一起,无法区分,失去了重建信息 。
- 两者都用,潜在空间既规整紧凑,不同类别又能很好地分开,兼顾了重建和生成能力 。
训练与生成
-
训练 (Training) :
- 输入数据 x x x 送入编码器,得到均值 μ \\mu μ 和对数方差 l o g σ 2 log\\sigma^2 logσ2。
- 使用重参数化技巧 (Reparameterization Trick) 进行采样: z = μ + σ ⊙ ϵ z = \\mu + \\sigma \\odot \\epsilon z=μ+σ⊙ϵ,其中 ϵ ∼ N ( 0 , I ) \\epsilon \\sim \\mathcal{N}(0, I) ϵ∼N(0,I)。这个技巧让采样过程变得可导,从而使梯度能够从解码器流回编码器。
- 将 z z z 送入解码器,得到重建数据 x ^ \\hat{x} x^。
- 计算重建损失和KL散度损失,相加后进行反向传播。
-
生成 (Sampling) :
- 从先验分布(即标准正态分布 N ( 0 , I ) \\mathcal{N}(0, I) N(0,I))中随机采样一个向量 z n e w z_{new} znew 。
- 将 z n e w z_{new} znew 送入解码器,得到输出 x n e w x_{new} xnew 。这就是一个全新的、由模型生成的数据样本。
强化学习 (Reinforcement Learning, RL)
RL研究的是智能体(Agent)如何在一个环境中通过与环境交互来学习一个最优策略,以最大化其获得的累积奖励。
核心概念
-
马尔可夫决策过程 (Markov Decision Process, MDP):RL问题的数学框架 。它由一系列状态、动作、奖励组成 ( s t , a t , r t + 1 , s t + 1 , . . . s_t, a_t, r_{t+1}, s_{t+1}, ... st,at,rt+1,st+1,...) 。核心是马尔可夫性质:未来只与当前状态有关,与过去无关。
-
回报 (Return, U t U_t Ut):从 t t t 时刻开始,未来所有奖励的折扣总和。 γ ∈ [ 0 , 1 ) \\gamma \\in [0, 1) γ∈[0,1) 是折扣因子,表示未来的奖励没有当前的奖励重要 。
U t = R t + γ R t + 1 + γ 2 R t + 2 + ⋯ = ∑ k = t ∞ γ k − tR k U_t = R_t + \\gamma R_{t+1} + \\gamma^2 R_{t+2} + \\dots = \\sum_{k=t}^{\\infty} \\gamma^{k-t} R_k Ut=Rt+γRt+1+γ2Rt+2+⋯=k=t∑∞γk−tRk
这个公式也可以递归地表示为: U t = R t + γ U t + 1 U_t = R_t + \\gamma U_{t+1} Ut=Rt+γUt+1 。RL的目标就是最大化期望回报。 -
策略 (Policy, π ( a ∣ s ) \\pi(a|s) π(a∣s)):智能体的“大脑”,定义了在给定状态 s s s 下,选择动作 a a a 的概率 。
-
价值函数 (Value Function):用来评估一个状态或一个状态-动作对有多“好”。
- 状态价值函数 (State-Value Function, V π ( s ) V_{\\pi}(s) Vπ(s)):从状态 s s s 出发,遵循策略 π \\pi π,能够获得的期望回报 。它回答了“处于状态 s s s 有多好?”
- 动作价值函数 (Action-Value Function, Q π ( s , a ) Q_{\\pi}(s, a) Qπ(s,a)):在状态 s s s 下,执行动作 a a a,然后遵循策略 π \\pi π,能够获得的期望回报 。它回答了“在状态 s s s 下执行动作 a a a 有多好?”
- 关系: V π ( s ) = ∑ a ∈ A π ( a ∣ s ) Q π ( s , a ) V_{\\pi}(s) = \\sum_{a \\in A} \\pi(a|s) Q_{\\pi}(s,a) Vπ(s)=∑a∈Aπ(a∣s)Qπ(s,a) 。
贝尔曼方程 (Bellman Equation)
贝尔曼方程是RL中最重要的公式之一,它给出了价值函数之间的递归关系,是几乎所有RL算法的基础。
- V函数贝尔曼方程:
V π ( s ) = E π [ R t + γ V π ( S t + 1) ∣ S t = s ] V_{\\pi}(s) = \\mathbb{E}_{\\pi}[R_t + \\gamma V_{\\pi}(S_{t+1}) | S_t=s] Vπ(s)=Eπ[Rt+γVπ(St+1)∣St=s] - Q函数贝尔曼方程:
Q π ( s , a ) = E π [ R t + γ Q π ( S t + 1, A t + 1) ∣ S t = s , A t = a ] Q_{\\pi}(s, a) = \\mathbb{E}_{\\pi}[R_t + \\gamma Q_{\\pi}(S_{t+1}, A_{t+1}) | S_t=s, A_t=a] Qπ(s,a)=Eπ[Rt+γQπ(St+1,At+1)∣St=s,At=a]
RL算法的核心思想
几类主流算法的思想:
-
基于价值的方法 (Value-Based)
- 核心思想:学习一个精确的价值函数(通常是Q函数),然后根据这个价值函数来隐式地导出一个最优策略(例如,总是选择Q值最大的那个动作)。
- 时序差分学习 (Temporal Difference, TD) :这是这类方法的核心更新机制。它不像蒙特卡洛(MC)方法那样需要等待一个完整的轨迹结束才更新 ,而是在执行一步之后,就用“当前的奖励”加上“对下一状态价值的估计”来更新“当前状态的价值估计”。这个过程叫做自举(bootstrapping)。
- TD更新规则:
Q ( s , a ) ← Q ( s , a ) + α [ R + γ Q ( s ′ , a ′ ) − Q ( s , a ) ] Q(s,a) \\leftarrow Q(s,a) + \\alpha [R + \\gamma Q(s\',a\') - Q(s,a)] Q(s,a)←Q(s,a)+α[R+γQ(s′,a′)−Q(s,a)]
这里的 R + γ Q ( s ′ , a ′ ) R + \\gamma Q(s\',a\') R+γQ(s′,a′) 被称为 TD目标 ,它是一个比旧的 Q ( s , a ) Q(s,a) Q(s,a) 更准确的估计。
-
基于策略的方法 (Policy-Based)
-
核心思想:不学习价值函数,而是直接学习一个参数化的策略 π θ ( a ∣ s ) \\pi_{\\theta}(a|s) πθ(a∣s) 。目标是找到最优的参数 θ \\theta θ 来最大化期望回报 J ( θ ) J(\\theta) J(θ) 。
-
策略梯度 (Policy Gradient):这类方法使用梯度上升来优化策略参数 。其核心是策略梯度定理,它给出了目标函数 J ( θ ) J(\\theta) J(θ) 相对于参数 θ \\theta θ 的梯度:
∇ θ J ( θ ) = E π θ [ ∇ θ l o g π θ ( A ∣ S ) ⋅ Q π ( S , A ) ] \\nabla_{\\theta}J(\\theta) = \\mathbb{E}_{\\pi_{\\theta}}[\\nabla_{\\theta}log~\\pi_{\\theta}(A|S) \\cdot Q_{\\pi}(S,A)] ∇θJ(θ)=Eπθ[∇θlog πθ(A∣S)⋅Qπ(S,A)] -
直观理解 :这个梯度的意思是,如果一个动作 A A A 最终导向了一个好的结果(即 Q π ( S , A ) Q_{\\pi}(S,A) Qπ(S,A) 值很高),我们就应该增大这个动作被选中的概率(即增大 l o g π θ ( A ∣ S ) log~\\pi_{\\theta}(A|S) log πθ(A∣S))。反之,如果动作导向了坏的结果,就减小其概率。
-
-
演员-评论家方法 (Actor-Critic)
- 核心思想:结合了上述两种方法的优点。它同时学习一个策略(演员 Actor)和一个价值函数(评论家 Critic)。
- 演员 (Actor):负责根据当前状态选择动作,即策略 π θ ( a ∣ s ) \\pi_{\\theta}(a|s) πθ(a∣s)。
- 评论家 (Critic):负责评估演员选择的动作有多好,即价值函数 Q w ( s , a ) Q_w(s,a) Qw(s,a) 或 V w ( s ) V_w(s) Vw(s)。
- 工作流程:演员做出动作,环境给出反馈,评论家根据反馈和贝尔曼方程来评估动作的好坏,并指导演员更新其策略。演员则使用策略梯度进行更新,但其中用于评估动作好坏的 Q π ( S , A ) Q_{\\pi}(S,A) Qπ(S,A) 由评论家来提供。