High-resolution image reconstruction with latent diffusion models from human brain activity

论文地址：https://doi.org/10.1101/2022.11.18.517004

项目地址：https://sites.google.com/view/stablediffusion-with-brain/

Abstract

从人类大脑活动中重建视觉体验，为理解大脑如何代表世界，以及解释计算机视觉模型和我们的视觉系统之间的联系提供了独特的方法。虽然深度生成模型最近被用于这一任务，但重建具有高语义保真度的真实图像仍然是一个具有挑战性的问题。在这里，我们提出了一种基于扩散模型(DM)的新方法来重建通过功能磁共振成像(fMRI)获得的人脑活动图像。更具体地说，我们依赖于被称为稳定扩散的潜在扩散模型(LDM)。该模型降低了DM的计算成本，同时保留了它们的高生成性能。我们还通过研究LDM的不同组成部分（如图像$z$的潜在向量、条件输入$c$和去噪U-Net的不同元素）如何与不同的大脑功能相关，来描述LDM的内部机制。我们表明，我们提出的方法可以以直接的方式以高保真度重建高分辨率图像，不需要对复杂的深度学习模型进行任何额外的训练和微调。我们还从神经科学的角度提供了不同LDM成分的定量解释。总的来说，我们的研究提出了一种很有前途的从人脑活动重建图像的方法，并为理解DM提供了一个新的框架。

1. Introduction

计算机视觉的一个基本目标是构建能够像人类视觉系统一样看待和识别世界的人工系统。群体大脑活动测量的最新进展，结合深度神经网络模型实施和设计的进展，已经允许直接比较生物大脑中的潜在表征和人工网络的架构特征，为这些系统如何运作提供了重要的见解[3,8 - 10,13,18,19,21,42,43,54,55]。这些努力包括从大脑活动中重建视觉体验（感知或图像），并检查与生物和人工系统相关的计算过程之间的潜在对应[2，5，7，24，25，27，36，44–46]。

从大脑活动重建视觉图像，例如用功能性磁共振成像(fMRI)测量的图像，是一个有趣但具有挑战性的问题，因为大脑中的潜在表征在很大程度上是未知的，而且通常与大脑数据相关的样本量相对较小[17,26,30,32]。近年来，研究人员已经开始使用深度学习模型和算法来解决这一任务，包括生成对抗网络(GANs)和自监督学习[2,5,7,24,25,27,36,44 - 46]。此外，最近的研究通过显式地使用图像的语义内容作为重建的辅助输入来提高语义保真度[5,25]。然而，这些研究需要从头开始用fMRI数据训练新的生成模型，或者对fMRI实验中使用的特定刺激进行微调。这些努力在像素和语义保真度方面取得了令人印象深刻但有限的成功，部分原因是神经科学中的样本数量很少，部分原因是学习复杂的生成模型带来了许多挑战。

扩散模型(Diffusion models, DMs)[11,47,48,53]是近年来备受关注的深度生成模型。DM在条件图像生成[4,39,49]、图像超分辨率[40]、图像彩色[38]和其他相关任务[6,16,33,41]等多项任务中都取得了最先进的性能。此外，最近提出的潜扩散模型(LDM)[37]通过利用其自动编码组件产生的潜空间进一步降低了计算成本，使训练和推理阶段的计算更有效。LDM的另一个优势是它们能够生成具有高语义保真度的高分辨率图像。然而，由于LDM是最近才引入的，我们仍然对其内部机制缺乏满意的了解。具体来说，我们仍然需要发现它们是如何在DM的每一层中表示潜在信号的，在去噪过程中潜在表示是如何变化的，以及添加噪声是如何影响条件图像生成的。

在这里，我们试图通过使用一种名为稳定扩散的LDM从fMRI信号重建视觉图像来解决上述挑战。这种架构是在大型数据集上训练的，具有很高的文本到图像生成性能。我们表明，我们的简单框架可以重建具有高语义保真度的高分辨率图像，而无需对复杂的深度学习模型进行任何训练或微调。我们还提供了LDM每个组成部分的生物学解释，包括正向/反向扩散过程，U-Net和不同噪声水平的潜在表示。

我们的贡献如下：(i) 我们证明了我们的简单框架可以以高语义保真度从大脑活动重建高分辨率（512×512）图像，而不需要训练或微调复杂的深度生成模型（图1）；(ii) 我们从神经科学的角度定量解释LDM的每个组成部分，通过将特定的组成部分映射到不同的大脑区域；(iii) 我们对LDM实现的文本到图像转换过程如何合并由条件文本表示的语义信息，同时保持原始图像的外观给出了客观的解释。

2. Related Work

2.1. Reconstructing visual image from fMRI

从功能磁共振成像(fMRI)活动中解码视觉体验已经以不同的方式进行了研究。例如明确呈现的视觉刺激[17,26,30,32]、呈现的刺激的语义内容[15,31,52]、想象的内容[13,29]、感知的情绪[12,20,51]以及许多其他相关应用[14,28]。一般来说，这些解码任务由于低信噪比和与fMRI数据相关的相对较小的样本量而变得困难。

虽然早期的尝试使用手工制作的特征来重建fMRI的视觉图像[17,26,30,32]，但最近的研究已经开始使用在大量自然主义图像上训练的深度生成模型[2,5,7,24,25,27,36,44 - 46]。此外，一些研究使用与图像相关的语义信息，包括分类信息或文本信息，来提高重建图像的语义保真度[5,25]。为了产生高分辨率的重建，这些研究需要训练并可能微调生成模型，如GANs，使用与fMRI实验相同的数据集。这些要求带来了严重的限制，因为训练复杂的生成模型通常具有挑战性，而且神经科学中的样本数量相对较少。因此，即使是现代的实现也很难产生分辨率最高为256×256的图像，并且具有较高的语义保真度，除非使用大量的工具和技术来增强它们。DM和LDM是最近的图像生成算法，可以潜在地解决这些限制，这要归功于它们能够生成各种高分辨率图像，具有高文本条件反射的语义保真度和高计算效率。然而，据我们所知，没有先前的研究使用DM进行视觉重建。

2.2. Encoding Models

为了从生物学的角度理解深度学习模型，神经科学家采用了编码模型:从深度学习模型的不同组成部分中提取特征，构建大脑活动的预测模型，然后检查模型表示和相应的大脑过程之间的潜在联系[3,8 - 10,13,18,19,21,42,43,54,55]。由于大脑和深度学习模型具有相似的目标（例如，对世界的识别），因此可以实现类似的功能，在这两个结构之间建立联系的能力为我们提供了对深度学习模型基础架构的生物学解释，否则被视为黑箱。例如，在CNN的早期和晚期层中观察到的激活模式与从视觉皮层的早期和晚期层中测量到的神经活动模式相对应，这表明CNN的潜在表征与大脑中存在的表征之间存在等级对应关系[9,10,13,19,54,55]。这种方法主要应用于视觉科学，但最近已扩展到其他感官模式和更高的功能[3,8,18,21,42,43]。

与CNN等生物学启发的架构相比，DM和大脑之间的对应关系不那么明显。通过检查DM的每个组成部分和过程与相应的大脑活动之间的关系，我们能够获得DM的生物学解释，例如潜在向量、去噪过程、条件操作和U-net组件如何对应于我们的视觉流。据我们所知，之前没有研究调查DM和大脑之间的关系。

总之，我们的首要目标是使用DM进行高分辨率视觉重建，并使用大脑编码框架更好地理解DM的潜在机制及其与大脑的对应关系。

3. Methods

图2展示了我们的方法的概述。

3.1. Dataset

我们使用自然场景数据集(NSD)进行这个项目[1]。详情请访问国家SD网站。简单地说，NSD提供了从7-特斯拉功能磁共振扫描仪获得的数据，在30-40次实验中，每个受试者重复观看三次10,000张图像。我们分析了8名受试者中4名完成所有成像疗程的数据(subj01, subj02, subj05和subj07)。NSD实验中使用的图像从MS COCO中检索，并裁剪到425×425（如果需要）。我们使用每个受试者来自NSD的27750个试验(NSD未公开发布的3万个试验中有2250个试验)。在这些试验的一个子集（N= 2770个试验）中，所有四名受试者都观看了982张图像。这些试验被用作测试数据集，而剩余的试验（N=24,980）被用作训练数据集。

对于功能性数据，我们使用NSD提供的预处理扫描（分辨率为1.8 mm）。预处理协议详见附录A。我们使用了从广义线性模型和感兴趣区域估计的单次试验$\beta$权重(ROIs)。对于测试数据集，我们使用与每张图像相关的三次试验的平均值。对于训练数据集，我们使用了三个独立的试验，没有平均。

3.2. Latent Diffusion Models

DMs是概率生成模型，通过迭代去噪将采样变量从高斯噪声恢复到学习数据分布的样本。对于给定的训练数据，扩散过程通过逐渐添加高斯噪声来破坏数据的结构。每个时间点的样本定义为
x
t
=
α
t
x
+
1
−
α
t
ϵ
t
\mathbf{x}_t=\sqrt{\alpha_t} \mathbf{x}_0+\sqrt{1-\alpha_t} \epsilon_t
xt​=αt​​x0​+1−αt​​ϵt​，其中
x
t
x_t
xt​是输入
x
x_0
x0​的噪声版本，$t\in \{1,\dots ,T\}$，$\alpha$为超参数，$ϵ$为高斯函数。 应用神经网络
f
θ
(
x
t
,
t
)
f_\theta\left(\mathrm{x}_t, t\right)
fθ​(xt​,t)对逆扩散过程进行了建模，以恢复原始输入。学习目标为
f
θ
(
x
,
t
)
≈
ϵ
t
f_\theta(\mathbf{x}, t) \approx \epsilon_t
fθ​(x,t)≈ϵt​[11,47]。U-Net通常用于神经网络
f
θ
f_\theta
fθ​。

该方法可以推广到通过在神经网络中插入辅助输入$c$来学习条件分布。如果我们将文本序列的潜在表示设置为$c$，它可以实现文本到图像模型。最近的研究表明，通过使用大型语言和图像模型，DM可以从文本输入中创建逼真的高分辨率图像。此外，当我们从有输入文本的源图像开始时，我们可以通过编辑图像生成新的文本条件图像。在这种图像到图像的转换中，原始图像的退化程度由一个参数控制，该参数可以调整以保留原始图像的语义内容或外观。

在像素空间中操作的DM计算成本很高。LDM通过使用自动编码器压缩输入来克服这一限制（图2，左上）。具体来说，自编码器首先使用图像数据进行训练，扩散模型使用U-Net架构生成其潜在表示$z$。在此过程中，它通过交叉注意引用条件输入。与基于像素的DM相比，这允许进行轻量级推理，并支持非常高质量的文本到图像和图像到图像实现。在这项研究中，我们使用了一种称为稳定扩散的LDM，它建立在LDM之上，并在一个非常大的数据集上进行训练。该模型可以根据文本输入生成和修改图像。文本输入由预训练的文本编码器(CLIP)[34]投影到固定的潜在表示。我们使用模型的1.4版本。训练协议详见附录A。

我们定义$z$为被自编码器压缩的原始图像的潜在表示，$c$为文本的潜在表示(与每张MS COCO图像相关的平均五个文本注释)，
z
c
z_c
zc​为由$c$修改的模型生成的$z$的潜在表示。我们在下面描述的解码/编码模型中使用了这些表示。

3.3. Decoding: reconstructing images from fMRI

我们使用LDM对fMRI信号进行了视觉重建，分为以下三个简单步骤（图2，左下）。在我们的方法中，唯一需要的训练是构建将fMRI信号映射到每个LDM组件的线性模型，并且不需要对深度学习模型进行训练或微调。我们使用LDM ²作者提供的图像到图像和文本到图像代码的默认参数，包括用于DDIM采样器的参数。详见附录A。

(i) 首先，我们从早期视觉皮层的fMRI信号中预测了所呈现图像X的潜在表征z。然后用自编码器的解码器对z进行处理，得到粗解码图像
X
z
X_z
Xz​，大小为320×320，再调整大小为512×512。

(ii) 通过自编码器的编码器对
X
z
X_z
Xz​进行处理，然后通过扩散过程添加噪声。

(iii) 我们从高级（腹侧）视觉皮层的fMRI信号中解码潜在文本表征$c$。将粗图像的加噪潜在表征
z
T
z_T
zT​和解码后的$c$作为去噪U-Net的输入，生成
z
c
z_c
zc​。最后将
z
c
z_c
zc​作为自编码器解码模块的输入，得到最终的重建图像
X
z
c
X_zc
Xz​c，大小为512×512。

为了构建从fMRI到LDM组件的模型，我们使用L2正则化线性回归，并且所有模型都是在每个受试者的基础上构建的。从训练数据中估计权重，并在训练过程中使用5倍交叉验证探索正则化参数。我们将原始图像的大小从425×425调整为320×320，但确认将它们调整为更大的大小（448×448）并不影响重建的质量。作为对照分析，我们也只使用$z$或$c$生成图像。为了生成这些对照图像，我们简单地分别从上面的步骤(iii)中省略了$c$或$z$。

通过评估原始测试图像（N=982张图像）是否能从生成的图像中识别出来，客观地(perceptual similarity metrics，PSMs)和主观地（人类评分者，N=6）评估图像重建的准确性。作为PSM的相似度度量，我们使用了CLIP的早/中/后期层和CNN (AlexNet)[22]。简而言之，我们进行了双向识别实验：检查fMRI重建图像是否比随机选取的重建图像更接近相应的原始图像。详见附录B和其他结果。

3.4. Encoding: Whole-brain Voxel-wise Modeling

接下来，我们试图通过将LDM映射到LDM的内部操作。为此，我们为以下四种设置构建了全脑体素编码模型（实现细节见图2右和附录A）：

(i) 我们首先建立线性模型，从以下三个LDM的潜在表征中独立预测体素活动:$z$, $c$和
z
c
z_c
zc​。

(ii) 虽然
z
c
z_c
zc​和z产生的图像不同，但它们在皮层上产生的预测图相似（见4.2.1）。因此，我们将它们合并到一个单一的模型中，并通过将每个特征解释的唯一方差映射到皮层[23]上，进一步研究它们之间的差异。为了控制原始图像的外观和条件文本的语义保真度之间的平衡，我们改变了添加到$z$的噪声水平。这种分析使图像到图像过程的定量解释成为可能。

(iii) 虽然LDM被描述为一个迭代的去噪过程，但去噪过程的内部动态尚不清楚。为了深入了解这个过程，我们研究了
z
c
z_c
zc​在去噪过程中的变化。为此，我们从去噪的前、中、后期步骤中提取
z
c
z_c
zc​。然后，我们构建了上述分析(ii)中含有$z$的组合模型，并将其唯一方差映射到皮层上。

(iv) 最后，为了检查与LDM相关的最后一个黑箱，我们从U-Net的不同层中提取特征。针对去噪的不同步骤，采用不同的U-Net层分别构建编码模型：第一阶段2个，bottleneck阶段1个，第二阶段2个。然后，我们为每个体素和每个步骤确定精度最高的层。

使用L2正则化线性回归从训练数据中估计模型权重，然后应用于测试数据（详见附录A）。为了进行评估，我们使用了预测和测量fMRI信号之间的Pearson相关系数。我们通过比较估计的相关性与两个相同长度的独立高斯随机向量(N=982)之间的相关性的零分布来计算统计显著性（单侧）。统计阈值设置为$P<0.05$，并使用FDR程序进行多次比较校正。我们展示了来自单个随机种子的结果，但我们验证了不同的随机种子产生了几乎相同的结果（见附录C）。我们通过应用主成分分析，通过估计训练数据中的成分，将所有特征维数降至6400。

4. Results

4.1. Decoding

图3显示了一个受试者(subj01)的视觉重建结果。我们为每个测试图像生成5张图像，并选择生成的PSM最高的图像。一方面，仅使用$z$重构的图像在视觉上与原始图像一致，但未能捕捉到原始图像的语义内容；另一方面，仅使用$c$生成的图像重构图像具有较高的语义保真度，但视觉不一致。最后，使用
z
c
z_c
zc​重建的图像可以生成具有高语义保真度的高分辨率图像（更多示例参见附录B）。

图4显示了同一图像的所有受试者的重建图像（所有图像都是使用
z
c
z_c
zc​生成的。其他例子见附录B）。总体而言，重建质量稳定且准确。我们注意到，重建图像的具体细节缺乏一致可能是受试者感知经验的差异，而不是重建的失败。或者，它可能只是反映了不同受试者之间数据质量的差异。事实上，fMRI解码精度高(subj01)和低(subj07)的受试者分别是数据质量指标高和低的受试者（见附录B）。

图5绘制了定量评价结果。在客观评价中，使用
z
c
z_c
zc​重建的图像在不同指标上通常比仅使用z或c重建的图像具有更高的精度值。当只使用$z$时，来自CLIP和CNN早期层的PSM的精度值特别高。另一方面，当只使用$c$时，来自后期层的PSM的精度值更高。在主观评价中，$c$方法得到的图像精度值高于$z$方法，而
z
c
z_c
zc​方法与其他两种方法相比精度最高（所有比较$P<0.01$，双面符号秩检验，FWE校正）。总之，这些结果表明，我们的方法不仅捕捉了低水平的视觉外观，而且还捕捉了原始刺激的高水平语义内容。

很难将我们的结果与以前大多数研究报告的结果进行比较，因为他们使用了不同的数据集。在以前的研究中使用的数据集包含的图像要少得多，图像复杂度也低得多（通常是位于图像中心的单个对象），并且缺乏NSD提供的那种全文注释。迄今为止，只有一项研究[25]使用NSD进行视觉重建，他们报告了一名受试者(subj01)使用基于Inception V3的PSM的精度值为78±4.5%。很难与这项研究进行直接比较，因为它与我们的研究在几个方面有所不同（例如，它使用了不同的训练和测试样本量，以及不同的图像分辨率）。尽管存在这些差异，但他们报告的值与我们的相同主题在相似的范围内(77%使用CLIP, 83%使用AlexNet, 76%使用Inception V3)。然而，这项先前的研究依赖于广泛的模型训练和特征工程，具有比我们研究中采用的更多的超参数，包括训练复杂的生成模型的必要，对MS COCO进行微调，数据增强和特征的任意阈值。我们没有使用上述任何一种技术——相反，我们的简单管道只需要构建两个从fMRI活动到LDM潜在表征的线性回归模型。

此外，当我们使用与图像相关的分类信息而不是全文注释时，我们观察到语义保真度的降低。我们还发现，当我们使用语义地图而不是原始图像时，$z$的语义保真度有所增加，尽管在这种情况下视觉相似性有所降低（见附录B）。

4.2. Encoding Model

4.2.1 Comparison among Latent Representations

图6显示了与LDM相关的三种潜在表征的编码模型的预测精度：$z$，原始图像的潜在表征；$c$，潜在的图像文本注释表示；
z
c
z_c
zc​是交叉注意到$c$的反向扩散过程后$z$的噪声附加潜伏表示。

尽管这三个组成部分在大脑的视觉皮层产生了很高的预测性能，但它们显示出鲜明的对比。其中$z$在视觉皮层的后部，即早期视觉皮层具有较高的预测能力。在视皮层前部也有显著的预测价值，即较高的视皮层，但在其他区域的预测价值较小。另一方面，$c$在高级视觉皮层的预测性能最高。该模型在大范围的皮层上也表现出较高的预测性能。
z
c
z_c
zc​所携带的表示法与$z$非常相似，对早期视觉皮层表现出较高的预测性能。尽管考虑到它们内在的相似性，这在一定程度上是可以预测的，但这仍然很有趣，因为这些表示对应于视觉上不同的生成图像。我们还观察到，正如预期的，使用注入$z$的噪声水平降低的
z
c
z_c
zc​产生与从z获得的预测图更相似的预测图（参见附录C）。这种相似性促使我们进行额外的分析，以比较这两个模型解释的唯一方差，详细内容将在下一节中介绍。各科目成绩见附录C。

4.2.2 Comparison across different noise levels

虽然先前的结果表明$z$和
z
c
z_c
zc​的预测精度图呈现类似的轮廓，但它们并没有告诉我们作为不同噪声水平的函数，每个特征解释了多少独特的方差。为了增强我们对上述问题的理解，我们接下来构建了编码模型，同时将$z$和
z
c
z_c
zc​合并到一个模型中，并研究了每个特征的独特贡献。我们还改变了为生成
z
c
z_c
zc​而添加到$z$的噪声级别。

图7显示，当添加少量噪声时，$z$比
z
c
z_c
zc​更好地预测整个皮层的体素活性。有趣的是，当我们增加噪声水平时，
z
c
z_c
zc​比$z$更好地预测了高级视觉皮层内的体素活动，表明图像的语义内容逐渐被强调。这个结果很有趣，因为如果没有这样的分析，我们只能观察随机生成的图像，我们无法检查以文本为条件的图像到图像的过程是如何在语义内容和原始视觉外观之间取得平衡的。

4.2.3 Comparison across different diffusion stages

接下来，我们询问在迭代去噪过程中噪声添加的潜在表示如何变化。

从图8可以看出，在去噪过程的早期，$z$信号主导了fMRI信号的预测。在去噪过程的中间阶段，
z
c
z_c
zc​比$z$更能预测高级视觉皮层的活动，表明大部分语义内容出现在这一阶段。这些结果展示了LDM如何从噪声中提炼和生成图像。

4.2.4 Comparison across different U-Net Layers

最后，我们询问U-Net的每一层正在处理什么信息。

图9显示了去噪过程的不同步骤（早期、中期、晚期）以及U-Net的不同层的编码模型的结果。在去噪过程的早期阶段，U-Net的bottleneck层（橙色）产生了最高的跨皮层预测性能。然而，随着去噪的进行，U-Net的早期层（蓝色）预测早期视觉皮层的活动，bottleneck层转向高级视觉皮层的超强预测能力。

这些结果表明，在反向扩散过程开始时，图像信息在bottleneck层内被压缩。随着去噪的进行，U-Net层之间的功能分离在视觉皮层内出现：即，第一层倾向于表示早期视觉区域的精细尺度细节，而bottleneck层对应于更腹侧的语义区域的高阶信息。

5. Conclusions

我们提出了一种新的基于LDM的视觉重建方法。我们表明，我们的方法可以从人脑活动中重建具有高语义保真度的高分辨率图像。与之前的图像重建研究不同，我们的方法不需要对复杂的深度学习模型进行训练或微调：它只需要从fMRI到LDM内潜在表征的简单线性映射。

我们还通过构建编码模型为LDM的内部组件提供了定量解释。例如，我们演示了整个逆扩散过程中语义内容的出现，我们执行了U-Net的分层表征，并且我们提供了具有不同噪声级别的图像到图像转换的定量解释。尽管DM正在迅速发展，但对其内部过程仍然知之甚少。这项研究首次从生物学角度提供了定量解释。

Supplementary Material（附录）

https://www.biorxiv.org/content/10.1101/2022.11.18.517004v2.supplementary-material