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

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

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

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

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

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

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

研究背景

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

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

在詳細調研圖像重建領域已有的研究后，根據其模型結構，作者把目前的圖像重建范式歸納為：生成式模型和優(yōu)化式模型。 

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

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

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

可控圖像重建模型 ——MindDiffuser

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

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

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

實驗結果

作者在目前規(guī)模最大的神經影像數據集 NSD 上進行了詳細的實驗，在四種評價指標上的結果表明，提出的 MindDiffuser 重建出的圖像與原圖無論在語義相似度，還是結構相似度上都超過了目前的 sota 模型。

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

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

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

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

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

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