擴(kuò)散模型「讀腦術(shù)」，自動化所MindDiffuser清晰重建人腦視覺畫面

從腦信號中重建相應(yīng)的視覺刺激都是一項(xiàng)有意義且充滿挑戰(zhàn)性的工作，先前已經(jīng)有研究成功得到了與原圖結(jié)構(gòu)相近的重建結(jié)果，如：一些自然圖像的輪廓、大小，但是缺乏明確的語義信息，難以辨別。

近年來，借助多模態(tài)預(yù)訓(xùn)練大模型與生成能力更強(qiáng)的 AI ，不少工作都得到了在語義上與原圖十分接近的重建圖像，但這些圖像在位置、朝向等結(jié)構(gòu)信息上不可控。

為了同時解決以上兩個問題，中科院自動化所的盧一卓、杜長德等人借助 Stable Diffusion 和 CLIP 新提出了一種基于擴(kuò)散模型的兩階段圖像重建模型 MindDiffuser，相關(guān)論文已被國際多媒體大會 ACM MM 2023 接受。結(jié)果表明，本文提出的模型的重建結(jié)果已超過目前最優(yōu)的模型?？梢暬Y(jié)果也證明了本文的模型設(shè)計(jì)在神經(jīng)生物學(xué)上的可解釋性。

• 論文地址：https://arxiv.org/pdf/2308.04249.pdf

• 代碼地址：https://github.com/ReedOnePeck/MindDiffuser

第一行為原圖，第二行為從大腦 fMRI 信號中重建圖像

研究背景

人類視覺神經(jīng)系統(tǒng)能夠高效地感知和理解現(xiàn)實(shí)世界中的復(fù)雜視覺刺激，這種突出的能力是當(dāng)下的人工智能系統(tǒng)無法比擬的。通過神經(jīng)編碼和解碼模型來研究不同腦區(qū)的功能可以使我們更加深入地理解人類視覺感知系統(tǒng)。

視覺神經(jīng)編碼指的是將外部視覺刺激擬合成神經(jīng)活動信號的過程，視覺神經(jīng)解碼則是要找到一個從神經(jīng)信號到對應(yīng)視覺刺激的映射，根據(jù)解碼難度和目標(biāo)的不同，可以把它們劃分為刺激分類、刺激識別與刺激重建。在刺激分類中，大腦活動被用來預(yù)測所呈現(xiàn)刺激的離散客體類別。刺激識別的目的是從一組已知的刺激圖像中識別出與給定的大腦活動模式相對應(yīng)的特定刺激。刺激重建則是要直接通過給定的 fMRI 信號生成圖像，并且要求其形狀、位置、朝向等細(xì)節(jié)與對應(yīng)刺激圖像對齊。本研究專注于研究刺激重建。

在詳細(xì)調(diào)研圖像重建領(lǐng)域已有的研究后，根據(jù)其模型結(jié)構(gòu)，作者把目前的圖像重建范式歸納為：生成式模型和優(yōu)化式模型。 

生成式模型在訓(xùn)練階段，提取刺激圖像和相應(yīng)文本描述的特征，利用相應(yīng)的腦響應(yīng)去擬合這些特征，在測試階段，用腦響應(yīng)預(yù)測出對應(yīng)的特征，將其輸入訓(xùn)練好的生成模型，即可重建出對應(yīng)的圖像刺激。

優(yōu)化式模型在訓(xùn)練階段提取刺激圖像的層次特征并用相應(yīng)的腦響應(yīng)去擬合，在測試階段，首先隨機(jī)初始化生成模型的隱空間，并且提取出隨機(jī)生成的初始圖像的層次特征，接著把腦響應(yīng)預(yù)測出的對應(yīng)的層次特征作為約束，通過反向傳播優(yōu)化隱空間，最終得到收斂的重建圖像。

目前這兩種重建范式都各自存在一些問題，針對以上問題本研究提出了一個新的框架來解決他們。

可控圖像重建模型 ——MindDiffuser

如圖 (a) 所示，首先提取出圖像的 CLIP 視覺特征，VQVAE 特征，以及對應(yīng)的文本描述特征，接著使用相應(yīng)的腦信號去擬合這三種特征。然后通過兩個階段來實(shí)現(xiàn)圖像重建任務(wù)。接下來將詳細(xì)地介紹這兩個階段。

在階段一，首先利用腦信號解碼出 VQVAE 特征，將其通過一個前向擴(kuò)散過程隱式的為擴(kuò)散模型隱空間注入圖像信息，接著在反向去噪過程中通過交叉注意力機(jī)制引入腦信號解碼出的 CLIP 文本特征，至此完成了對重建圖像的語義信息融合。

在階段二，為了進(jìn)一步對齊重建圖像的結(jié)構(gòu)信息，研究人員以 CLIP 圖像特征為約束，通過反向傳播梯度不斷調(diào)整擴(kuò)散模型的隱空間，得到語義和結(jié)構(gòu)都與原圖對齊的重建結(jié)果。

實(shí)驗(yàn)結(jié)果

作者在目前規(guī)模最大的神經(jīng)影像數(shù)據(jù)集 NSD 上進(jìn)行了詳細(xì)的實(shí)驗(yàn)，在四種評價指標(biāo)上的結(jié)果表明，提出的 MindDiffuser 重建出的圖像與原圖無論在語義相似度，還是結(jié)構(gòu)相似度上都超過了目前的 sota 模型。

由于不同個體的大腦在解剖結(jié)構(gòu)和功能連接上有著差異 [61]，因此即使給不同的被試觀看相同的圖像刺激，采集到的 fMRI 信號也不同，為了驗(yàn)證本文提出的圖像重建模型可以適應(yīng)不同被試間的差異，研究團(tuán)隊(duì)在不加任何其他額外調(diào)整的情況下，對被試 1、2、5、7 的測試集圖像進(jìn)行了重建，結(jié)果如下圖所示： 

從圖 4-13 中可以看出對于相同的刺激圖像，由于不同的被試在采集 fMRI 信號時的大腦主觀反應(yīng)以及特征解碼階段的準(zhǔn)確率有差異導(dǎo)致一些圖像的重建結(jié)果并不令人滿意，比如：被試 7 把 “桌子上的花” 錯誤的重建為 “桌子和椅子”，被試 5 無法重建出 “黃昏下的飛機(jī)”。但是大部分重建出來的圖像對于研究人員選取的每一個被試在語義與結(jié)構(gòu)上都表現(xiàn)出了與原圖較好的對齊，這表明他們提出的圖像重建模型可以很好的適應(yīng)不同被試間的差異。

在特征解碼階段，本文使用 L2 正則化的線性回歸模型自動挑選體素來擬合三種特征：語義特征 c，細(xì)節(jié)特征 z，結(jié)構(gòu)特征 zCLIP，為了驗(yàn)證提出的模型在神經(jīng)科學(xué)上的可解釋性，本文將 fMRI 每個體素在預(yù)測不同特征時的權(quán)重投影到展開的大腦皮層上。

從圖 (a) 中可以看出在解碼語義特征 c 的時候，被篩選出來的體素絕大多數(shù)集中在 IPS、LO、MT、MST、PHC、VO 等處理高級語義的腦區(qū)，且分布在高級視皮層區(qū)域的體素權(quán)重要大于分布在低級視皮層區(qū)域的體素。從圖 (b) 中可以看出，用于監(jiān)督重建圖像結(jié)構(gòu)特征的 CLIP 底層線性層特征主要 V1、V2、V3、V3ab、hV4 等處理局部形狀、紋理信息的低級腦區(qū)擬合。這表明，研究團(tuán)隊(duì)的模型在神經(jīng)解碼階段，語義特征主要由高級腦區(qū)解釋，結(jié)構(gòu)特征主要由低級腦區(qū)解釋，這一發(fā)現(xiàn)與神經(jīng)科學(xué)先前的研究結(jié)果相一致。從圖 (c) 中可以看出，高級視皮層與低級視皮層區(qū)域的體素都參與了解碼細(xì)節(jié)特征 z，且權(quán)重大小相當(dāng)，這說明通過解碼 z，使得交叉注意力機(jī)制為重建圖像融入了更加細(xì)粒度的語義與結(jié)構(gòu)信息。以上可視化結(jié)果說明，MindDiffuser 在重建圖像時使用的多模態(tài)信息可以很好被解釋。

綜上，本文提出了一個兩階段圖像重建模型 MindDiffuser，該模型將重建圖像的語義和結(jié)構(gòu)信息與圖像刺激進(jìn)行對齊。在 NSD 數(shù)據(jù)集上，MindDiffuser 在定性和定量對比結(jié)果上都優(yōu)于目前的 sota 模型。同時，實(shí)驗(yàn)表明，MindDiffuser 對個體間的差異具有良好的適應(yīng)性，在沒有任何額外調(diào)整的情況下，可以為被試 1、2、5 和 7 的刺激實(shí)現(xiàn)很好的重建結(jié)果。此外，可視化結(jié)果還證明，本文模型中利用的多模態(tài)信息可以通過神經(jīng)科學(xué)中相應(yīng)的腦響應(yīng)來解釋，從而驗(yàn)證了模型設(shè)計(jì)的合理性與可解釋性。研究團(tuán)隊(duì)認(rèn)為，MindDiffuser 在促進(jìn)腦機(jī)接口的精確和可控刺激重建方面將發(fā)揮重要作用。 

本論文由中科院自動化所直博生（尚未入學(xué)）盧一卓和副研究員杜長德作為共同第一作者，何暉光研究員為通訊作者，周瓊怡博士是共同作者。研究工作得到了科技部科技創(chuàng)新 2030—“新一代人工智能” 重大項(xiàng)目、基金委項(xiàng)目、以及 CAAI - 華為 MindSpore 學(xué)術(shù)獎勵基金及智能基座等項(xiàng)目的支持。