新加坡南洋理工大學(xué)突破：讓AI看圖速度提升28倍的"分層觀察法"

這項(xiàng)由新加坡南洋理工大學(xué)S-Lab實(shí)驗(yàn)室的周一凡、肖澤琦、魏天逸、潘新鋼團(tuán)隊(duì)，以及北京大學(xué)王選計(jì)算機(jī)研究所楊帥共同完成的研究，于2025年12月18日發(fā)表在arXiv預(yù)印本平臺(tái)，論文編號(hào)為arXiv:2512.16615v1。有興趣深入了解的讀者可以通過該編號(hào)查詢完整論文。

當(dāng)我們看一幅畫時(shí)，眼睛不會(huì)同時(shí)關(guān)注畫面的每個(gè)細(xì)節(jié)。相反，我們會(huì)先整體掃視，然后聚焦在最重要的部分。現(xiàn)在，人工智能也學(xué)會(huì)了這種"聰明的觀察方式"。研究團(tuán)隊(duì)開發(fā)了一種叫做"對(duì)數(shù)線性稀疏注意力"的新技術(shù)，讓AI在處理圖像時(shí)變得更加高效。

傳統(tǒng)的AI圖像生成模型就像一個(gè)需要同時(shí)觀察每個(gè)像素點(diǎn)的"完美主義者"。當(dāng)圖像尺寸增大時(shí)，這種全面觀察的方式會(huì)讓計(jì)算量呈指數(shù)級(jí)增長(zhǎng)。舉個(gè)例子，如果原來處理一張64×64像素的圖片需要1秒，那么處理256×256像素的圖片就可能需要16倍的時(shí)間。這就像一個(gè)人試圖同時(shí)記住房間里每一個(gè)物品的位置和每?jī)蓚€(gè)物品之間的關(guān)系一樣，隨著物品數(shù)量增加，這種全面記憶的難度會(huì)急劇上升。

研究團(tuán)隊(duì)的創(chuàng)新在于模仿了人類視覺系統(tǒng)的工作原理。他們?cè)O(shè)計(jì)了一種"分層觀察"的方法，就像我們看地圖一樣——先看整個(gè)城市的輪廓，再看具體的街區(qū)，最后關(guān)注具體的建筑物。AI首先對(duì)圖像進(jìn)行粗略觀察，識(shí)別出最重要的區(qū)域，然后逐步細(xì)化到具體細(xì)節(jié)。這種方法將計(jì)算復(fù)雜度從平方級(jí)降低到了對(duì)數(shù)線性級(jí)，意味著處理時(shí)間的增長(zhǎng)速度大幅減緩。

這項(xiàng)技術(shù)的實(shí)際效果令人印象深刻。在處理256×256像素的圖像時(shí)，新方法讓注意力推理速度提升了28.27倍，整體訓(xùn)練速度提升了6.09倍，同時(shí)保持了圖像生成的質(zhì)量。這相當(dāng)于原來需要一小時(shí)完成的工作，現(xiàn)在只需要不到三分鐘。

一、傳統(tǒng)方法的困境：全面觀察的代價(jià)

要理解這項(xiàng)研究的重要性，我們首先需要了解當(dāng)前AI圖像生成面臨的核心挑戰(zhàn)。現(xiàn)代的圖像生成模型，特別是擴(kuò)散Transformer模型，已經(jīng)成為視覺生成領(lǐng)域的佼佼者。這些模型能夠生成令人驚艷的高質(zhì)量圖像，但它們有一個(gè)致命的弱點(diǎn)：對(duì)計(jì)算資源的巨大需求。

這個(gè)問題的根源在于"自注意力"機(jī)制的工作方式。自注意力就像一個(gè)極度認(rèn)真的觀察員，需要考慮圖像中每個(gè)像素點(diǎn)與其他所有像素點(diǎn)之間的關(guān)系。當(dāng)圖像尺寸為N×N像素時(shí)，需要進(jìn)行的比較次數(shù)是N的平方。這意味著圖像尺寸每增加一倍，計(jì)算量就會(huì)增加四倍。

以目前流行的圖像生成模型為例，F(xiàn)LUX模型處理的是64×64的潛在圖像（相當(dāng)于4096個(gè)標(biāo)記），而Wan 2.1處理的視頻序列包含75600個(gè)標(biāo)記。當(dāng)我們想要生成更高分辨率的圖像或更長(zhǎng)的視頻時(shí)，這種平方級(jí)的復(fù)雜度增長(zhǎng)就成了不可逾越的障礙。

現(xiàn)有的稀疏注意力方法試圖通過只關(guān)注最重要的K個(gè)區(qū)域來解決這個(gè)問題。這種方法分為三個(gè)步驟：首先將查詢和鍵標(biāo)記壓縮成粗略表示，然后計(jì)算壓縮標(biāo)記之間的相似度分?jǐn)?shù)并選擇前K個(gè)關(guān)鍵塊，最后在選定的塊上執(zhí)行稀疏注意力。然而，這種單層設(shè)計(jì)仍然存在兩個(gè)根本性問題：選擇階段的計(jì)算成本仍然是平方級(jí)的，而且為了保持全局上下文，需要隨著序列長(zhǎng)度的增加而使用更大的K值。

二、仿生靈感：學(xué)習(xí)人類的觀察智慧

研究團(tuán)隊(duì)從人類視覺系統(tǒng)中獲得了靈感。當(dāng)我們觀察復(fù)雜場(chǎng)景時(shí)，大腦不會(huì)試圖同時(shí)處理所有細(xì)節(jié)。相反，我們采用了一種層次化的處理策略：首先獲取整體印象，然后逐步聚焦到感興趣的區(qū)域，最后處理具體細(xì)節(jié)。

這種觀察方式的優(yōu)勢(shì)在于效率。通過在不同抽象層次上組織信息，我們能夠用相對(duì)較少的計(jì)算資源處理復(fù)雜的視覺信息。大腦皮層的視覺處理區(qū)域就是按照這種層次化結(jié)構(gòu)組織的，從V1區(qū)域的簡(jiǎn)單邊緣檢測(cè)到高級(jí)視覺區(qū)域的復(fù)雜對(duì)象識(shí)別。

基于這種仿生理念，研究團(tuán)隊(duì)設(shè)計(jì)了對(duì)數(shù)線性稀疏注意力機(jī)制。這種方法將單層的全面觀察擴(kuò)展為多層的層次化觀察。在最粗糙的層次上，AI獲得圖像的整體結(jié)構(gòu)信息；在中等層次上，它識(shí)別重要的區(qū)域和模式；在最精細(xì)的層次上，它處理具體的細(xì)節(jié)。

層次化的關(guān)鍵在于每個(gè)層次使用的抽象程度不同。就像地圖有不同的縮放級(jí)別一樣，最粗糙的層次可能將16×16個(gè)像素壓縮為一個(gè)代表性特征，中等層次可能是4×4，而最精細(xì)的層次則保持原始像素的分辨率。這種設(shè)計(jì)使得AI能夠用對(duì)數(shù)級(jí)的計(jì)算量獲得原本需要平方級(jí)計(jì)算才能獲得的全局信息。

三、技術(shù)核心：分層選擇與信息增強(qiáng)

對(duì)數(shù)線性稀疏注意力的核心創(chuàng)新體現(xiàn)在兩個(gè)關(guān)鍵技術(shù)上：層次化Top-K選擇和層次化鍵值增強(qiáng)機(jī)制。

層次化Top-K選擇過程就像一個(gè)逐步細(xì)化的篩選系統(tǒng)。在最粗糙的層次上，系統(tǒng)首先計(jì)算所有粗粒度標(biāo)記之間的相似度，并選擇最相關(guān)的K個(gè)區(qū)域。然后，在下一個(gè)更精細(xì)的層次上，系統(tǒng)只在這K個(gè)已選區(qū)域內(nèi)進(jìn)行進(jìn)一步的篩選，而不是重新考慮整個(gè)圖像空間。這種遞歸式的選擇過程繼續(xù)下去，直到達(dá)到最精細(xì)的層次。

這種方法的巧妙之處在于它將原本需要在整個(gè)N×N空間中進(jìn)行的搜索，轉(zhuǎn)化為在多個(gè)較小空間中的搜索。數(shù)學(xué)分析表明，這種層次化選擇將選擇階段的復(fù)雜度從O(N?)降低到了O(N)，實(shí)現(xiàn)了根本性的效率提升。

然而，單純的層次化選擇可能會(huì)丟失重要的全局信息。為了解決這個(gè)問題，研究團(tuán)隊(duì)引入了層次化鍵值增強(qiáng)機(jī)制。這種機(jī)制的工作原理類似于一個(gè)智能的信息匯總系統(tǒng)。在最終的注意力計(jì)算中，系統(tǒng)不僅使用最精細(xì)層次選擇的鍵值對(duì)，還會(huì)包含從各個(gè)粗糙層次收集的代表性信息。

鍵值重加權(quán)是這個(gè)機(jī)制中的一個(gè)重要細(xì)節(jié)。由于不同層次的標(biāo)記代表不同數(shù)量的原始像素信息，系統(tǒng)需要相應(yīng)地調(diào)整它們的重要性權(quán)重。一個(gè)代表16個(gè)像素的粗糙標(biāo)記應(yīng)該比一個(gè)代表單個(gè)像素的精細(xì)標(biāo)記具有更高的權(quán)重。具體來說，權(quán)重設(shè)置為該層次的塊大小，確保信息的重要性與其代表的內(nèi)容量成正比。

四、工程實(shí)現(xiàn)：高效的GPU計(jì)算方案

理論上的優(yōu)勢(shì)需要通過高效的實(shí)現(xiàn)才能轉(zhuǎn)化為實(shí)際的性能提升。研究團(tuán)隊(duì)在GPU實(shí)現(xiàn)方面也進(jìn)行了重要?jiǎng)?chuàng)新，特別是在稀疏索引的處理上。

傳統(tǒng)的稀疏注意力實(shí)現(xiàn)通常使用二進(jìn)制掩碼來標(biāo)記哪些位置需要計(jì)算，哪些可以跳過。然而，構(gòu)建和處理這種掩碼本身就需要二次方的內(nèi)存和計(jì)算開銷，這與稀疏化的初衷背道而馳。

研究團(tuán)隊(duì)開發(fā)了一種直接操作稀疏索引的方法。在前向傳播過程中，系統(tǒng)直接收集選中的鍵塊，而不是構(gòu)建完整的掩碼矩陣。在反向傳播過程中，他們實(shí)現(xiàn)了一個(gè)輕量級(jí)的稀疏索引轉(zhuǎn)置內(nèi)核，能夠動(dòng)態(tài)計(jì)算稀疏索引的反向查找，避免了密集掩碼的構(gòu)建。

這種稀疏索引轉(zhuǎn)置算法基于經(jīng)典的CSR到CSC稀疏矩陣轉(zhuǎn)置方法。算法的核心思想是將所有相關(guān)查詢的索引保存在一個(gè)扁平向量中，并使用累積偏移來獲取每個(gè)鍵的查詢范圍。雖然這個(gè)過程需要原子加法操作，但由于選擇的鍵數(shù)量K相對(duì)較小，原子操作在內(nèi)存中稀疏分布，沖突概率極低，因此開銷可以忽略不計(jì)。

為了適應(yīng)二維視覺數(shù)據(jù)，研究團(tuán)隊(duì)還引入了索引重排序方案。傳統(tǒng)的光柵掃描順序不能有效地在一維池化過程中聚集相似像素，而索引重排序確保空間相鄰的像素在扁平化序列中也保持鄰近關(guān)系。這種重排序類似于Z-order曲線或希爾伯特曲線的思想，保持了空間局部性。

五、驗(yàn)證與評(píng)估：從理論到實(shí)踐的飛躍

為了全面驗(yàn)證新方法的有效性，研究團(tuán)隊(duì)設(shè)計(jì)了一系列嚴(yán)格的實(shí)驗(yàn)。他們選擇了像素空間的圖像生成作為測(cè)試場(chǎng)景，這是一個(gè)特別具有挑戰(zhàn)性的任務(wù)，因?yàn)樗皇褂萌魏螆D像壓縮或編碼技術(shù)，直接在原始像素級(jí)別進(jìn)行操作。

實(shí)驗(yàn)設(shè)置覆蓋了從128×128到512×512像素的多種分辨率，使用FFHQ和ImageNet數(shù)據(jù)集。在128×128的FFHQ數(shù)據(jù)集上，與傳統(tǒng)的Top-K稀疏注意力方法相比，對(duì)數(shù)線性稀疏注意力在僅使用K=8的情況下，就超越了使用K=32的基線方法。這個(gè)結(jié)果特別令人印象深刻，因?yàn)樗砻鲗哟位I值增強(qiáng)機(jī)制確實(shí)能夠用更少的計(jì)算資源保持更好的全局上下文。

在訓(xùn)練效率方面，新方法在256×256像素序列上實(shí)現(xiàn)了6.09倍的訓(xùn)練加速，同時(shí)保持了生成質(zhì)量。這種加速不僅來自于理論復(fù)雜度的降低，也得益于高效的GPU實(shí)現(xiàn)。反向傳播內(nèi)核在不同序列長(zhǎng)度下保持了近乎恒定的吞吐量，證實(shí)了線性復(fù)雜度的實(shí)際實(shí)現(xiàn)。

消融研究揭示了各個(gè)組件的重要性。單獨(dú)使用Top-K注意力難以匹配完整注意力的性能，F(xiàn)ID分?jǐn)?shù)從24.91上升到28.21。啟用鍵值增強(qiáng)后，性能顯著改善，F(xiàn)ID降低到26.09。正確設(shè)置粗糙標(biāo)記的重要性權(quán)重進(jìn)一步提升了模型質(zhì)量，F(xiàn)ID達(dá)到24.18，甚至略優(yōu)于完整注意力基線。

與其他稀疏注意力方法的比較顯示，對(duì)數(shù)線性稀疏注意力在相同有效標(biāo)記數(shù)量下，始終獲得最佳的FID分?jǐn)?shù)和最高的訓(xùn)練吞吐量。這種優(yōu)勢(shì)在不同分辨率下都保持一致，證明了方法的魯棒性和通用性。

六、實(shí)際應(yīng)用：從實(shí)驗(yàn)室到現(xiàn)實(shí)世界

研究團(tuán)隊(duì)還將這項(xiàng)技術(shù)集成到了PixelFlow模型中，在ImageNet-256數(shù)據(jù)集上進(jìn)行了大規(guī)模驗(yàn)證。PixelFlow是一個(gè)多階段像素?cái)U(kuò)散模型，能夠?qū)⒌头直媛蕡D像逐步上采樣到高分辨率。在這個(gè)更具挑戰(zhàn)性的場(chǎng)景中，對(duì)數(shù)線性稀疏注意力再次展現(xiàn)了其優(yōu)勢(shì)，在FID和Inception Score兩個(gè)關(guān)鍵指標(biāo)上都優(yōu)于現(xiàn)有的稀疏注意力方法。

特別值得注意的是，這項(xiàng)技術(shù)在處理真正長(zhǎng)序列時(shí)的表現(xiàn)。在512×512像素的圖像生成任務(wù)中，單層設(shè)計(jì)的稀疏注意力由于二次方選擇成本而無法在合理時(shí)間內(nèi)收斂，而層次化設(shè)計(jì)顯著提升了訓(xùn)練吞吐量，并且隨著層次數(shù)量的增加，性能進(jìn)一步改善。

噪聲重縮放技術(shù)也為像素空間訓(xùn)練做出了重要貢獻(xiàn)。通過調(diào)整噪聲強(qiáng)度而不是輸入縮放，研究團(tuán)隊(duì)能夠在不同分辨率下保持一致的信噪比，大大加速了高分辨率模型的收斂。這種方法比傳統(tǒng)的時(shí)間步偏移或?qū)?shù)正態(tài)采樣器更加有效。

索引重排序的引入進(jìn)一步優(yōu)化了二維視覺數(shù)據(jù)的處理。實(shí)驗(yàn)顯示，使用空間連貫的索引排序比默認(rèn)的光柵掃描排序能獲得更好的生成質(zhì)量，證明了空間局部性在稀疏注意力中的重要作用。

七、技術(shù)影響：開啟高效AI視覺的新時(shí)代

這項(xiàng)研究的意義遠(yuǎn)超技術(shù)層面的創(chuàng)新，它為AI視覺生成領(lǐng)域帶來了范式轉(zhuǎn)變。傳統(tǒng)上，高質(zhì)量的圖像生成需要在計(jì)算效率和生成質(zhì)量之間做出妥協(xié)。研究團(tuán)隊(duì)證明了通過巧妙的算法設(shè)計(jì)，可以同時(shí)實(shí)現(xiàn)兩者的優(yōu)化。

從計(jì)算復(fù)雜度的角度看，從O(N?)到O(N log N)的降低是質(zhì)的飛躍。這意味著處理序列長(zhǎng)度翻倍時(shí)，計(jì)算時(shí)間的增長(zhǎng)從四倍降低到了略多于兩倍。對(duì)于實(shí)際應(yīng)用而言，這種改進(jìn)使得在消費(fèi)級(jí)硬件上生成高分辨率內(nèi)容成為可能。

該技術(shù)的模塊化設(shè)計(jì)也為進(jìn)一步的創(chuàng)新奠定了基礎(chǔ)。層次化注意力的思想可以擴(kuò)展到其他類型的序列建模任務(wù)，不僅限于圖像生成。視頻生成、文檔處理、音頻合成等領(lǐng)域都可能從這種方法中受益。

從工程實(shí)踐的角度看，高效的GPU實(shí)現(xiàn)證明了理論優(yōu)勢(shì)可以轉(zhuǎn)化為實(shí)際的性能提升。稀疏索引的直接操作避免了傳統(tǒng)掩碼方法的開銷，為大規(guī)模稀疏計(jì)算提供了新的實(shí)現(xiàn)范式。

說到底，這項(xiàng)研究展現(xiàn)了仿生學(xué)在人工智能領(lǐng)域的巨大潛力。通過學(xué)習(xí)人類視覺系統(tǒng)的層次化處理策略，AI系統(tǒng)不僅變得更加高效，也更加符合自然智能的工作原理。這種從生物系統(tǒng)中汲取靈感的做法，為未來的AI算法設(shè)計(jì)提供了重要的指導(dǎo)方向。

歸根結(jié)底，對(duì)數(shù)線性稀疏注意力技術(shù)的成功證明了一個(gè)重要觀點(diǎn)：在AI發(fā)展中，效率和質(zhì)量不是非此即彼的選擇，而是可以通過創(chuàng)新的算法設(shè)計(jì)同時(shí)實(shí)現(xiàn)的目標(biāo)。隨著這項(xiàng)技術(shù)的進(jìn)一步完善和推廣，我們有理由期待更加高效、強(qiáng)大的AI視覺生成系統(tǒng)，它們將為創(chuàng)意產(chǎn)業(yè)、教育、娛樂等眾多領(lǐng)域帶來革命性的變化。對(duì)于那些希望深入了解技術(shù)細(xì)節(jié)的讀者，可以通過論文編號(hào)arXiv:2512.16615v1查詢完整的研究?jī)?nèi)容。

Q&A

Q1：對(duì)數(shù)線性稀疏注意力是什么原理？

A：這是一種模仿人類視覺觀察方式的AI技術(shù)。就像我們看地圖時(shí)先看整體輪廓，再看街區(qū)，最后關(guān)注具體建筑一樣，AI先粗略觀察圖像整體，再逐步細(xì)化到具體細(xì)節(jié)。這種分層觀察方式將計(jì)算復(fù)雜度從平方級(jí)降低到對(duì)數(shù)線性級(jí)，大幅提升了處理效率。

Q2：這項(xiàng)技術(shù)能帶來多大的性能提升？

A：在處理256×256像素圖像時(shí)，新方法讓注意力推理速度提升了28.27倍，整體訓(xùn)練速度提升了6.09倍，同時(shí)保持圖像生成質(zhì)量不變。這相當(dāng)于原來需要一小時(shí)的工作現(xiàn)在只需要不到三分鐘就能完成。

Q3：這種技術(shù)何時(shí)能應(yīng)用到實(shí)際產(chǎn)品中？

A：研究團(tuán)隊(duì)已經(jīng)在像素空間圖像生成和PixelFlow等實(shí)際模型中驗(yàn)證了技術(shù)效果，并開發(fā)了高效的GPU實(shí)現(xiàn)方案。目前技術(shù)已經(jīng)比較成熟，預(yù)計(jì)很快就能集成到商業(yè)化的AI圖像生成產(chǎn)品中，讓普通用戶也能享受到更快速的AI圖像生成服務(wù)。