当前位置: 首页 > news >正文

MAR:无需量化的掩码自回归图像生成模型

MAR:无需量化的掩码自回归图像生成模型

TL; DR:用扩散损失(diffusion loss)将自回归条件分布建模为(更一般的)连续分布,并将掩码生成模型(如 MaskGIT)统一到自回归生成的视角下,提出了无需量化的掩码自回归图像生成模型(Masked Autoregressive Models, MAR)。

方法

整体来看,本文提出的生成模型也是先将像素图片进行 tokenize,映射到隐层空间,然后在隐层空间进行下一个 token 预测的序列生成模型。但是之前 AR image gen 的工作都是用类似 VQVAE 的离散 tokenizer,将连续的图片像素转化成离散的 token,本文则是采用连续的 tokenizer,提出用 diffusion loss 实现的连续 token 的序列生成模型建模。并将掩码生成模型统一到自回归生成的视角下。

连续型数据的自回归建模

自回归式生成,实际上就是根据已知 token,预测下一个 token,即建模条件分布: p ( x i ∣ x < i ) p(x^i|x^{<i}) p(xix<i)。在常规 AR text/image gen 中,一般是将这个分布建模为一个离散的类别分布(Categorical Distribution),从而能够简单方便地用 cross entropy loss 进行训练,用 softmax 进行采样。本文作者指出,自回归的生成范式其实并没有要求这个分布一定要是离散的,连续型分布同样可以用于自回归生成。本节首先回顾常见的离散式自回归,分析出生成模型 训练时损失函数 和 推理时采样器 两个基本的要求,然后提出对连续型数据进行建模的 diffusion loss 和对应采样器。

离散自回归

首先回顾一下常规的离散 token 的自回归模型。我们用 x x x 表示下一个位置 token 的 ground truth。通过离散 tokenizer, x x x 可以用一个整数 0 ≤ x < K 0\le x< K 0x<K 来表示,其中 K K K 是词表(或码本)的尺寸。离散自回归模型一般是输出一个连续值的特征向量 z ∈ R D z\in\mathbb{R}^D zRD,然后通过一个 K K K 类的分类矩阵 W ∈ R K × D W\in\mathbb{R}^{K\times D} WRK×D ,并进行 softmax,最终可以得到一个离散的类别分布(categorical probability distribution)。

可以发现,从一般的角度来说,生成式建模的概率分布需要有两点关键属性:

  1. 一是在训练时,需要有一个损失函数,可以用来衡量估计的分布与目标真实分布之间的差异。在离散的类别分布中,直接用 cross entropy 就好了;
  2. 二是在推理生成时,需要有一个采样器(sampler),可以从分布 x ∼ p ( x ∣ z ) x\sim p(x|z) xp(xz) 中采样。在离散的类别分布中,可以直接通过 softmax 进行采样 p ( x ∣ z ) = softmax ( W z / τ ) p(x|z)=\text{softmax}(Wz/\tau) p(xz)=softmax(Wz/τ)

通过以上分析可以看到,在自回归生成式建模中,离散的 token 形式并不是必须的,关键是要对分布进行建模。离散型可以用类别分布,其损失函数和采样器很容易定义。而如果要使用连续型,我们需要设计对应的用于连续分布建模的损失函数和采样器。

在这里插入图片描述

Diffusion Loss

如何建模连续分布呢?其实我们有一个现成的工具,那就是扩散模型。理论上来说,扩散模型可以建模任意的分布。在图像生成领域,一般用扩散模型建模的是所有像素或者所有 token 的联合分布。实际上我们当然也可以用它来建模这里的连续型自回归分布,本文中,就是用扩散模型来建模单个 token 的分布

用连续型向量 x ∈ R d x\in\mathbb{R}^d xRd 表示下一个位置 token 的 ground truth。自回归模型 f f f 在当前位置输出一个连续型向量 z ∈ R D z\in\mathbb{R}^D zRD,我们的目标是建模 x x x 在条件 z z z (即模型输出的特征向量)下的连续型分布,即条件分布 p ( x ∣ z ) p(x|z) p(xz)。本文提出用扩散模型来定义连续自回归模型的损失函数和采样器。

损失函数

要建模条件分布 p ( x ∣ z ) p(x|z) p(xz) ,其扩散模型损失函数可以写成如下去噪的形式:
L ( z , x ) = E ϵ , t [ ∣ ∣ ϵ − ϵ θ ( x t ∣ x , z ) ∣ ∣ ] where  x t = α ˉ t x + 1 − α ˉ t ϵ \mathcal{L}(z,x)=\mathbb{E}_{\epsilon,t}[||\epsilon-\epsilon_\theta(x_t|x,z)||] \\ \text{where}\ x_t=\sqrt{\bar{\alpha}_t}x+\sqrt{1-\bar{\alpha}_t}\epsilon \notag \\ L(z,x)=Eϵ,t[∣∣ϵϵθ(xtx,z)∣∣]where xt=αˉt x+1αˉt ϵ
这里 ϵ ∈ R d \epsilon\in\mathbb{R}^d ϵRd 是采样自标准高斯分布 N ( 0 , I ) \mathcal{N}(0,\mathbf{I}) N(0,I) 的噪声向量, x t x_t xt 是对 x x x 进行加噪得到, α ˉ t \bar\alpha_t αˉt 是超参数,它定义了 noise schedule, t t t 是 noise schedule 的时间步。 ϵ θ \epsilon_\theta ϵθ 是噪声估计网络,在本文中,就是一个简单的 MLP 网络。从得分模型的角度来看,这个式子是一种得分匹配(score matching),学习的是 p ( x ∣ z ) p(x|z) p(xz) 的得分函数 ∇ x log ⁡ p ( x ∣ z ) \nabla_x\log p(x|z) xlogp(xz)。另外,扩散损失也是一个含参的损失函数,与对抗损失或感知损失类似。训练时,计算该损失函数的梯度,同时更新自回归模型 f f f 和噪声估计网络 ϵ θ \epsilon_\theta ϵθ。由于噪声估计网络是一个小型 MLP,因此实际训练时对于同一个 z z z,可以采样多次(4 次) t t t 进行训练,提高训练效率。

采样器

通过引入 diffusion loss,我们已经建模并训练拟合了条件分布 p ( x ∣ z ) p(x|z) p(xz)。在推理生成时,需要从这个条件分布中采样。扩散模型的采样,是通过反向扩散过程进行:
x t − 1 = 1 α t ( x t − 1 − α t 1 − α ˉ t ϵ θ ( x t ∣ t , z ) ) + σ t δ x_{t-1}=\frac{1}{\sqrt{\alpha_t}}\left(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\epsilon_\theta(x_t|t,z)\right)+\sigma_t\delta \notag \\ xt1=αt 1(xt1αˉt 1αtϵθ(xtt,z))+σtδ
其中 δ \delta δ 是采样自标准高斯分布的噪声, σ t \sigma_t σt 是时间步 t t t 的噪声层级。这样从 x T ∼ N ( 0 , I ) x_T\sim\mathcal{N}(0,\mathbf{I}) xTN(0,I) 开始,一步步到去噪,最终就可以得到符合分布的 x 0 ∼ p ( x ∣ z ) x_0\sim p(x|z) x0p(xz)

在采用离散的类别分布时,自回归采样过程可以通过控制超参数 τ \tau τ 来调节采样过程的多样性。实际中文本和图像生成很多工作也都表明了这个温度系数对于自回归生成很关键。因此这里也想让 diffusion sampler 有类似的调节能力。我们看采样时加入温度超参数,实际上就是想改为从重新归一化(renormalized)的分布 p ( x ∣ z ) 1 τ p(x|z)^{\frac{1}{\tau}} p(xz)τ1 中采样,对应的,其得分函数为 1 τ ∇ x log ⁡ p ( x ∣ z ) \frac{1}{\tau}\nabla_x\log p(x|z) τ1xlogp(xz)。参考 guided diffusion 的做法,要么把 ϵ θ \epsilon_\theta ϵθ 除一个 τ \tau τ,要么把噪声 σ t δ \sigma_t\delta σtδ 乘一个 τ \tau τ,本文选择了后者。实际上,扩散模型中的 τ \tau τ 就是通过控制噪声的方差,来控制采样的多样性。

使用diffusion loss进行连续型自回归建模

上面我们介绍了如何利用 diffusion loss,对条件分布 p ( x ∣ z ) p(x|z) p(xz) 进行建模。

在自回归中,我们是有一个 token 序列 { x 1 , x 2 , … , x n } \{x^1,x^2,\dots,x^n\} {x1,x2,,xn},其中上标 1 ≤ i ≤ n 1\le i\le n 1in 表示序列的某种顺序。自回归模型实际上就是根据已有的 token,预测下一个 token:
p ( x 1 , x 2 , … , x n ) = ∏ i = 1 n p ( x i ∣ x 1 , … , x i − 1 ) p(x^1,x^2,\dots,x^n)=\prod_{i=1}^np(x^i|x^1,\dots,x^{i-1}) p(x1,x2,,xn)=i=1np(xix1,,xi1)
一般用一整个模型来拟合这个条件分布 p ( x i ∣ x < i ) p(x^i|x^{<i}) p(xix<i)。而在本文中,实际上是将这个过程拆成了两个部分,首先使用一个神经网络(AR transformer)来根据已有的 token 产生条件向量 z i = f ( x < i ) z^i=f(x^{<i}) zi=f(x<i),然后用 diffusion loss 来建模条件分布 p ( x i ∣ z i ) p(x^i|z^i) p(xizi),根据 z i z^i zi 作为条件,预测下一个(连续) token x i x^i xi。公式(1)中的 diffusion 就可以应用于用来拟合 $ p(xi|zi)$。梯度反向传播到 z i z^i zi,以更新 f ( ⋅ ) f(\cdot) f() 的参数。

在这里插入图片描述

统一自回归模型与掩码生成模型

本节讨论了掩码生成模型(比如 MaskGIT、MAGE 等),同样可以看作是一种广义的自回归模型,作者称为 MAR,其特点是 1. 可用双向注意力;2. 随机顺序;3.每次可预测多个 tokens。

双向注意力的自回归

主流的自回归一般采用 causal attention Transformer 的模型结构。实际上,自回归的建模方式是指 “根据已有的 token 预测下一个 token”,并没有限定之前 token 与下一个 token 之间的实际位置关系,所以其实 bidirectional attention 也可以做自回归。

作者实际采用了类似 MAE 中的双向注意力机制。参见下图 (b)。具体来说,首先用 MAE encoder 对已知 tokens(带有位置嵌入)进行编码。然后,将编码序列与掩码 tokens 拼接起来(再次添加位置嵌入),将该序列输入到 MAE decoder 中。掩码 tokens 上的位置嵌入可以让 decoder 知道哪些位置需要预测。只对掩码 token 计算损失。

这样,所有已知 token 相互可见,也所有掩码 token 也可见所有已知 token,是一种 full attention 的形式。相较于 causal attention,这种 full attention 能进行更好的上下文交互。在推理时,每步可以生成一个或多个 token,这是一种广义的自回归形式。虽然这样就不能使用 causal attention 中的 kv cache 来加速推理了,但由于我们可以同时生成多个 token,因此效率还是很高。full attention + 多 token 生成。可以显著提高质量,并在速度/质量之间进行更好的权衡。

在这里插入图片描述

随机顺序的自回归

掩码生成模型,可以视作为随机顺序的一种自回归。每个样本序列进行不同的随机重排 (见图 3b),预测未知 token。在这种情况下,模型需要能够知道下一个待预测 token 的位置。这里也是用 MAE 的策略,在 decoder 中添加位置嵌入(原位置顺序),从而确定要预测的位置。

在这里插入图片描述

预测多个 tokens 的自回归

在掩码生成模型中,模型基于已有的 tokens,预测一组 “随机” tokens。这看作是将 token 序列进行随机重排,然后基于已知 tokens 预测多个掩码 tokens。从这个角度来看,掩码自回归模型就是一种 ”预测下一组 tokens“ 的自回归模型,即建模条件分布 p ( { x i , x i + 1 , … , x j } ∣ x 1 , … , x i − 1 ) p(\{x^i,x^{i+1},\dots,x^j\}|x^1,\dots,x^{i-1}) p({xi,xi+1,,xj}x1,,xi1),其中 { x i , x i + 1 , … , x j } , i ≤ j \{x^i,x^{i+1},\dots,x^j\},\ i\le j {xi,xi+1,,xj}, ij 是本次要预测的 tokens。

这种自回归模型可以写为:
p ( x 1 , … , x n ) = p ( X 1 , … , X K ) = ∏ k K p ( X k ∣ X 1 , … , X k − 1 ) p(x^1,\dots,x^n)=p(X^1,\dots,X^K)=\prod_k^K p(X^k|X^1,\dots,X^{k-1}) p(x1,,xn)=p(X1,,XK)=kKp(XkX1,,Xk1)
这样将一组 tokens X k = { x i , x i + 1 , … , x j } X^k=\{x^i,x^{i+1},\dots,x^j\} Xk={xi,xi+1,,xj} 视为一个预测单元,就是一种广义的自回归模型了,本文将这种模型称为 MAR(Masked Autoregressive Models),总结来说,MAR 是一种每次同时预测多个 tokens 的随机顺序的自回归模型。

MAR 与 MaskGIT、MAGE 类似。然而,MAR 通过在每个 token 的概率分布上用 τ \tau τ 来对 token采样进行重归一化(这是 GPT 等生成式语言模型的标准做法)。而 MaskGIT/MAGE 则是在待预测 token 的位置采样时应用了一个温度,这样不是完全随机的顺序,在训练时和推理时的行为之间造成了差距。

实验

下表展示了图像生成中 AR/MAR 用本文提出的 diffusion loss 与常规的 cross entropy loss 的对比结果。可以看到,无论是在哪种自回归形式下,diff loss 都优于 cross entropy loss,在常规的 causal (raster/rand order) 的自回归形式下优势不是很显著,但是在 bidirect 下,表现大幅提升,而且即使时 cross entropy loss 下也比 causal 有很大的提升。这个实验充分说明了本文两个创新点:diffusion loss 和 MAR 在图像生成领域的有效性。

在这里插入图片描述

其他实验就不一一贴了。

总结

本文指出自回归生成式建模不一定要用离散的 token 和分布,只要能够建模自回归过程中生成单个 token 的条件概率 p ( x ∣ z ) p(x|z) p(xz) 即可,还是很有启发性的。之前一直以为想要用 AR,就必须要构造成 vocab/codebook 这样离散的 categorical distribution,没有去深入理解 AR。

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.dtcms.com/a/214160.html

相关文章:

  • RuoYi前后端分离框架集成UEditorPlus富文本编辑器
  • 嵌入式学习(基本操作)day1
  • 在麒麟系统(Kylin OS)上安装`geckodriver`
  • 认识微服务
  • DNS 详情 新增 DNS 自适应服务器 ip
  • 【部署】在离线服务器的docker容器下升级dify-import程序
  • leetcode 3559. Number of Ways to Assign Edge Weights II
  • 【Hive基础】01.数据模型、存储格式、排序方式
  • 内网映射有什么作用,如何实现内网的网络地址映射到公网连接?
  • 【图像处理基石】什么是色彩模式?
  • 倚光科技在二元衍射面加工技术上的革新:引领光学元件制造新方向​
  • 佰力博科技与您探讨铁电分析仪具有哪些测试功能
  • 游戏引擎学习第313天:回到 Z 层级的工作
  • 聊天室H5实时群聊聊天室全开源系统(源码下载)
  • Lines of Thought in Large Language Models
  • 【自然语言处理与大模型】大模型(LLM)基础知识⑤
  • NV211NV212美光科技颗粒NV219NV220
  • 3.python操作mysql数据库
  • 01 NLP的发展历程和挑战
  • 大语言模型 19 - MCP FastAPI-MCP 实现自己的MCP服务 快速接入API
  • Spring AI 系列之一个很棒的 Spring AI 功能——Advisors
  • 临床试验中的独立数据监查委员会
  • 动态规划-918.环形子数组的最大和-力扣(LeetCode)
  • 【bug排查记录】由Redission配置引发的Satoken血案
  • 树莓派超全系列教程文档--(47)如何使用内核补丁
  • leetcode 525. 连续数组
  • 如何给自研MCP加上安全验证
  • OpenSSL 文件验签与字符串验签原理及 C 语言实现详解
  • 行为型:状态模式
  • AI时代新词-大模型(Large Language Model)