Transformer 及其核心的注意力機(jī)制在自然語(yǔ)言處理和計(jì)算機(jī)視覺(jué)等領(lǐng)域帶來(lái)了革命性進(jìn)展，展現(xiàn)出強(qiáng)大的深度上下文建模和數(shù)據(jù)間復(fù)雜依賴關(guān)系捕捉能力。

然而，其在處理視覺(jué)數(shù)據(jù)時(shí)面臨兩大核心挑戰(zhàn)：

1. 二次計(jì)算復(fù)雜度使其難以高效處理高分辨率圖像等長(zhǎng)上下文數(shù)據(jù)；
2. 忽略空間結(jié)構(gòu)，將多維圖像視為無(wú)結(jié)構(gòu)的一維標(biāo)記序列，破壞了圖像固有的空間連貫性，而這種信息對(duì)于依賴空間關(guān)系的視覺(jué)任務(wù)至關(guān)重要。

為克服效率瓶頸，近期研究如線性注意力和狀態(tài)空間模型（如 Mamba） 致力于將復(fù)雜度降低至線性。

然而，這些方法在提升效率的同時(shí)，依然未能有效保留和利用圖像的關(guān)鍵二維空間結(jié)構(gòu)信息，本質(zhì)上仍是序列化處理。

嘗試將一維光柵掃描（raster scan）擴(kuò)展至二維的線掃描方法（line scan）是增強(qiáng)空間連貫性的一種思路。

但二維線性傳播面臨嚴(yán)峻挑戰(zhàn)：標(biāo)量權(quán)重變?yōu)檫B接像素與前序鄰居的矩陣權(quán)重。在傳播過(guò)程中累積的矩陣乘法極易導(dǎo)致穩(wěn)定性問(wèn)題——矩陣特征值過(guò)大引發(fā)指數(shù)增長(zhǎng)（不穩(wěn)定），過(guò)小則導(dǎo)致信號(hào)迅速衰減（信息消失）。

因此，在二維空間中同時(shí)實(shí)現(xiàn)穩(wěn)定性和維持長(zhǎng)距離上下文的有效傳播，是一個(gè)亟待解決的難題。

針對(duì)上述挑戰(zhàn)，來(lái)自英偉達(dá)、香港大學(xué)和UCSD的研究人員提出廣義空間傳播網(wǎng)絡(luò)（GSPN），一種專為視覺(jué)任務(wù)優(yōu)化的新型注意力機(jī)制，其核心優(yōu)勢(shì)在于直接操作空間連貫的圖像數(shù)據(jù)，通過(guò)高效的線掃描方法建立密集的像素間連接。

論文地址：https://arxiv.org/abs/2501.12381

項(xiàng)目主頁(yè)：https://whj363636.github.io/GSPN/

代碼：https://github.com/NVlabs/GSPN

GSPN成功的關(guān)鍵在于其提出的穩(wěn)定性-上下文條件（Stability-Context Condition），該條件確保了跨二維序列的穩(wěn)定長(zhǎng)上下文傳播，并將具有N個(gè)元素的圖像的復(fù)雜度顯著降低至√N(yùn) 量級(jí)。

因此，GSPN能夠在保持卓越空間保真度的同時(shí)，實(shí)現(xiàn)極高的計(jì)算效率，并在ImageNet分類、類引導(dǎo)圖像生成及文本到圖像生成等任務(wù)中達(dá)到先進(jìn)性能。例如，在生成16K圖像時(shí)，GSPN相比基于softmax注意力的SD-XL加速超過(guò)84倍。

論文第一作者為王弘焌，香港大學(xué)統(tǒng)計(jì)系博士三年級(jí)學(xué)生，目前為NVIDIA research intern，研究方向包括高效基礎(chǔ)模型、開(kāi)放世界理解。

GSPN方法

二維線性傳播

二維線性傳播通過(guò)逐行或逐列的順序處理進(jìn)行。對(duì)于二維圖像，其遵循線性循環(huán)過(guò)程，隱藏層通過(guò)前一行的隱藏狀態(tài)和當(dāng)前輸入計(jì)算得出。

將隱藏狀態(tài)和輸入的行向量連接成序列后，可表示為輸入與一個(gè)下三角矩陣的乘積，輸出則為輸入的加權(quán)和，該公式可類比為帶因果掩碼的非歸一化線性注意力機(jī)制，其中額外的傳播矩陣調(diào)制注意力強(qiáng)度。

穩(wěn)定性-上下文條件

在傳播過(guò)程中上述累積的矩陣乘法極易導(dǎo)致穩(wěn)定性問(wèn)題。

為實(shí)現(xiàn)穩(wěn)定且有效的長(zhǎng)距離傳播，研究人員引入定理1和定理2（統(tǒng)稱為穩(wěn)定性-上下文條件）。

定理1指出，若所有矩陣均為行隨機(jī)矩陣，則滿足各元素加權(quán)和為1

定理2表明，行隨機(jī)矩陣可確保傳播過(guò)程的穩(wěn)定性。行隨機(jī)矩陣的定義為元素非負(fù)且每行元素之和為1，乘積仍為行隨機(jī)矩陣，這為穩(wěn)定傳播提供了數(shù)學(xué)基礎(chǔ)。

傳播層的關(guān)鍵實(shí)現(xiàn)

對(duì)于二維線性循環(huán)過(guò)程，研究人員對(duì)前序狀態(tài)的三鄰居連接來(lái)計(jì)算當(dāng)前時(shí)刻的隱藏層（每個(gè)像素連接前一行的三個(gè)相鄰像素）以提高參數(shù)效率。

文中同時(shí)還提出GSPN的兩種變種，全局GSPN和局部GSPN：

全局GSPN捕捉整個(gè)序列的長(zhǎng)距離依賴，局部GSPN通過(guò)將空間維度劃分為非重疊組來(lái)限制傳播序列長(zhǎng)度，提高效率。

最后，通過(guò)四方向集成確保全像素連接，形成密集成對(duì)連接。

對(duì)每個(gè)傳播方向的矩陣元素應(yīng)用 sigmoid 函數(shù)并歸一化，以保證行隨機(jī)性。

通過(guò)定制的CUDA內(nèi)核實(shí)現(xiàn)線性傳播層，采用并行化結(jié)構(gòu)，在批量、通道和與傳播方向正交的行/列上實(shí)現(xiàn)全并行化，有效減少內(nèi)核循環(huán)長(zhǎng)度，實(shí)現(xiàn)高效可擴(kuò)展的線性傳播。

GSPN架構(gòu)

GSPN是一個(gè)通用序列傳播模塊，可無(wú)縫集成到各種視覺(jué)任務(wù)的神經(jīng)網(wǎng)絡(luò)中。針對(duì)判別任務(wù)和生成任務(wù)設(shè)計(jì)了不同的GSPN塊，均基于核心GSPN模塊構(gòu)建：

• GSPN模塊：通過(guò)共享1×1卷積進(jìn)行降維，再通過(guò)三個(gè)獨(dú)立的1×1卷積生成依賴于輸入的參數(shù)，用于二維線性傳播，這些投影和傳播封裝在模塊化的GSPN單元中。
• 圖像分類架構(gòu)：采用Swin-Transformer的四級(jí)分層架構(gòu)，通過(guò)堆疊設(shè)計(jì)良好的GSPN塊，在相鄰層級(jí)間進(jìn)行下采樣操作，平衡計(jì)算效率和表示能力。
• 類條件圖像生成架構(gòu)：重新設(shè)計(jì)生成架構(gòu)，通過(guò)向量嵌入加法集成時(shí)間步和條件信息，包含跳躍連接和線性投影，去除位置嵌入并引入FFN進(jìn)行通道混合。
• 文本到圖像生成架構(gòu)：將GSPN模塊直接集成到Stable Diffusion架構(gòu)中，替換所有自注意力層，利用預(yù)訓(xùn)練權(quán)重初始化參數(shù)，加速訓(xùn)練。

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

圖像分類

在ImageNet-1K分類任務(wù)中，GSPN在參數(shù)數(shù)量相當(dāng)?shù)那闆r下優(yōu)于現(xiàn)有序列模型，GSPN在從小型到基礎(chǔ)配置的模型規(guī)模上表現(xiàn)出一致的性能提升，證明了其可擴(kuò)展性。

類條件圖像生成

與多種基線方法相比，GSPN-XL/2在ImageNet 256×256類條件生成任務(wù)中建立了新的最先進(jìn)性能，GSPN-L/2僅使用先前模型65.6%的參數(shù)就獲得了更優(yōu)的FID和IS分?jǐn)?shù)，GSPN-B/2在收斂時(shí)僅使用DiT-XL/2 20.3%的參數(shù)就實(shí)現(xiàn)了有競(jìng)爭(zhēng)力的性能，驗(yàn)證了GSPN的效率和可擴(kuò)展性。

文本到圖像生成

GSPN由于其歸一化權(quán)重滿足穩(wěn)定性-上下文條件，無(wú)需額外歸一化即可適應(yīng)任意分辨率，在不使用任何預(yù)訓(xùn)練權(quán)重且在相同訓(xùn)練輪數(shù)內(nèi)達(dá)到了與baseline相當(dāng)?shù)男阅堋?/span>

此外，GSPN在單塊A100 GPU上生成16K×8K分辨率圖像可實(shí)現(xiàn)約84倍的加速。

總結(jié)

研究人員提出了廣義空間傳播網(wǎng)絡(luò)（GSPN），這是一種用于視覺(jué)任務(wù)中并行序列建模的新型注意力機(jī)制。

通過(guò)穩(wěn)定性-上下文條件確保穩(wěn)定且上下文感知的傳播，GSPN在保持效率的同時(shí)將序列復(fù)雜度減少到√N(yùn)

實(shí)驗(yàn)表明，GSPN在多個(gè)視覺(jué)任務(wù)中實(shí)現(xiàn)了最先進(jìn)的結(jié)果和顯著的加速，展示了其在視覺(jué)任務(wù)中的效率和潛力。

未來(lái)，GSPN有望在更多視覺(jué)領(lǐng)域及視覺(jué)多模態(tài)模型中發(fā)揮重要作用，推動(dòng)下一代視覺(jué)理解和生成基礎(chǔ)結(jié)構(gòu)的發(fā)展。