近日，華盛頓大學(xué)的李坤?。↘unjun Li）、楊程言（Cheng-Yen Yang）、黃正能（Jenq-Neng Hwang）以及新加坡國(guó)立大學(xué)的陳子亙（Zigeng Chen）共同發(fā)表了一項(xiàng)突破性研究，標(biāo)題為《使用尺度感知KV緩存壓縮實(shí)現(xiàn)內(nèi)存高效的視覺自回歸建?！罚∕emory-Efficient Visual Autoregressive Modeling with Scale-Aware KV Cache Compression）。這篇論文于2025年5月26日在arXiv預(yù)印本平臺(tái)發(fā)布，編號(hào)為2505.19602v1，研究代碼已在GitHub（https://github.com/StargazerX0/ScaleKV）開源。

想象一下，你正在使用一款能根據(jù)文字描述生成精美圖片的AI工具。當(dāng)你嘗試生成高分辨率圖像時(shí)，系統(tǒng)卻提示"內(nèi)存不足"或運(yùn)行極為緩慢——這正是當(dāng)前視覺AI面臨的一大挑戰(zhàn)。視覺自回歸模型（Visual Autoregressive Model，簡(jiǎn)稱VAR）雖然在圖像生成質(zhì)量上取得了顯著進(jìn)步，但在處理高分辨率圖像時(shí)會(huì)消耗大量?jī)?nèi)存資源。例如，使用目前最先進(jìn)的Infinity-8B模型生成1024×1024分辨率的圖像時(shí)，僅KV緩存（一種存儲(chǔ)中間計(jì)算結(jié)果的內(nèi)存空間）就需要消耗高達(dá)85GB的顯存！這使得普通用戶幾乎無法在個(gè)人設(shè)備上運(yùn)行這類模型。

而這篇研究正是針對(duì)這一內(nèi)存瓶頸提出了創(chuàng)新解決方案。研究團(tuán)隊(duì)開發(fā)了名為"ScaleKV"的技術(shù)，能將模型所需內(nèi)存減少90%，同時(shí)幾乎不影響生成圖像的質(zhì)量。這就像是將一個(gè)原本需要專業(yè)服務(wù)器才能運(yùn)行的軟件，通過巧妙的壓縮技術(shù)，讓它能在普通家用電腦上流暢運(yùn)行。

那么，ScaleKV是如何實(shí)現(xiàn)這一"魔法"的呢？這項(xiàng)突破背后的原理與我們?nèi)粘Ｉ钪械脑S多優(yōu)化方式類似，讓我們一起來了解這個(gè)既簡(jiǎn)單又巧妙的解決方案。

一、視覺自回歸模型：高質(zhì)量與高內(nèi)存需求的矛盾

我們先來了解什么是視覺自回歸模型（VAR）。如果說傳統(tǒng)的圖像生成AI就像一位畫家一次性在畫布上完成整幅作品，那么VAR則更像是一位從粗略草圖逐步精細(xì)化的藝術(shù)家，先勾勒出大致輪廓，再逐步添加細(xì)節(jié)。

具體來說，VAR模型采用了"下一尺度預(yù)測(cè)"的創(chuàng)新方法，相比傳統(tǒng)的"下一個(gè)像素點(diǎn)"或"下一個(gè)圖像塊"預(yù)測(cè)，這種方法允許模型從低分辨率逐步生成高分辨率圖像。想象一下，它首先創(chuàng)建一個(gè)模糊的16×16小圖像，然后是更清晰的32×32版本，接著是64×64，以此類推，直到最終生成完整的高分辨率圖像。這種"粗到細(xì)"的生成方式不僅提高了效率，還顯著改善了圖像質(zhì)量和模型的零樣本泛化能力（即在未見過的場(chǎng)景中的表現(xiàn)）。

然而，這種方法帶來了一個(gè)嚴(yán)重的挑戰(zhàn)：內(nèi)存消耗問題。在生成過程中，模型需要記住之前所有尺度的信息，這就像一位畫家需要不斷參考自己之前繪制的所有草圖版本。隨著圖像尺度的增加，需要存儲(chǔ)的信息量呈指數(shù)級(jí)增長(zhǎng)。對(duì)于1024×1024分辨率的圖像，模型需要處理超過10,000個(gè)圖像令牌（token），而存儲(chǔ)這些中間計(jì)算結(jié)果的KV緩存（Key-Value Cache）可能消耗高達(dá)85GB的顯存！

這導(dǎo)致的結(jié)果是：即使是最先進(jìn)的VAR模型，如Infinity-8B，也很難在消費(fèi)級(jí)硬件上生成高分辨率圖像，嚴(yán)重限制了這些模型在實(shí)際應(yīng)用中的推廣。

二、關(guān)鍵發(fā)現(xiàn)：不是所有層都需要完整記憶

研究團(tuán)隊(duì)通過深入分析VAR模型的工作方式，發(fā)現(xiàn)了兩個(gè)關(guān)鍵的性質(zhì)，這為內(nèi)存優(yōu)化提供了突破口。

首先，他們觀察到不同的Transformer層對(duì)歷史信息的依賴程度存在顯著差異。這有點(diǎn)像一個(gè)團(tuán)隊(duì)中的不同成員：有些人需要掌握整個(gè)項(xiàng)目的所有歷史細(xì)節(jié)才能工作（我們稱之為"起草者"或Drafters），而另一些人只需要關(guān)注當(dāng)前手頭的具體任務(wù)（我們稱之為"精細(xì)器"或Refiners）。

研究團(tuán)隊(duì)通過可視化注意力模式（模型關(guān)注信息的方式）發(fā)現(xiàn)，某些層展現(xiàn)出分散的注意力模式，廣泛關(guān)注來自多個(gè)尺度的上下文信息。就像一位藝術(shù)總監(jiān)需要掌握整個(gè)創(chuàng)作過程的大局觀。而另一些層則主要關(guān)注當(dāng)前處理的圖像部分，就像一位專注于精細(xì)繪制某個(gè)局部細(xì)節(jié)的畫師。

其次，他們發(fā)現(xiàn)這些注意力模式會(huì)隨著生成尺度的增加而演化。在早期尺度，所有層都需要廣泛收集信息；而在后期尺度，"精細(xì)器"層會(huì)變得更加專注于局部細(xì)節(jié)，幾乎不需要參考早期尺度的信息。

這些發(fā)現(xiàn)挑戰(zhàn)了現(xiàn)有的緩存管理策略。傳統(tǒng)方法要么對(duì)所有層應(yīng)用相同的緩存分配（就像給團(tuán)隊(duì)中的每個(gè)人分配相同大小的工作空間），要么基于位置進(jìn)行緩存削減（如僅保留最近的信息）。然而，VAR模型需要一種能夠同時(shí)考慮層特定需求和尺度依賴特性的自適應(yīng)分配策略。

三、ScaleKV：智能壓縮KV緩存的創(chuàng)新方案

基于上述發(fā)現(xiàn)，研究團(tuán)隊(duì)開發(fā)了ScaleKV（Scale-Aware KV Cache）框架，這是一種針對(duì)視覺自回歸模型的KV緩存壓縮技術(shù)。它的核心思想很簡(jiǎn)單：不同的神經(jīng)網(wǎng)絡(luò)層有不同的記憶需求，因此應(yīng)該分配不同大小的"工作內(nèi)存"。

ScaleKV的工作原理可以分為三個(gè)主要步驟：

第一步是識(shí)別"起草者"和"精細(xì)器"層。研究團(tuán)隊(duì)引入了一個(gè)名為"注意力選擇性指數(shù)"（Attention Selectivity Index，ASI）的指標(biāo)，用于量化每一層的注意力模式。這個(gè)指標(biāo)考慮了兩個(gè)關(guān)鍵因素：一是該層對(duì)當(dāng)前圖像地圖的關(guān)注程度，二是該層對(duì)歷史序列的注意力集中度。簡(jiǎn)單來說，高ASI值表明該層要么強(qiáng)烈關(guān)注當(dāng)前圖像地圖，要么對(duì)特定歷史信息展現(xiàn)高度選擇性，這表明它是一個(gè)"精細(xì)器"。相反，低ASI值表明該層在前綴上下文中分布注意力更廣泛，這是"起草者"的特征。

第二步是建立緩存預(yù)算分配策略。在確定了"起草者"和"精細(xì)器"后，ScaleKV實(shí)現(xiàn)了一種高效的預(yù)算分配策略，確?？們?nèi)存消耗與均勻預(yù)算分配相同，同時(shí)為"精細(xì)器"實(shí)現(xiàn)尺度依賴的削減。具體來說，"精細(xì)器"的緩存預(yù)算會(huì)隨著尺度k的增加而線性減少：Br(k) = Br(0) - δ·k，其中δ控制精細(xì)器預(yù)算衰減率。節(jié)省下來的內(nèi)存隨后重新分配給"起草者"，確保Bd(k) >> Br(k)，以便與每一層的計(jì)算需求保持一致。

第三步是KV緩存選擇。在確定了每一層的緩存預(yù)算后，ScaleKV實(shí)現(xiàn)了一種高效的令牌選擇策略，決定應(yīng)該保留哪些特定的KV狀態(tài)。對(duì)于每個(gè)令牌地圖rk，它首先將地圖分割成N個(gè)區(qū)塊，從每個(gè)區(qū)塊中選擇中心令牌形成一個(gè)觀察窗口W。這種采樣方法確保了跨令牌地圖的空間覆蓋，同時(shí)維持最小的內(nèi)存占用。然后，基于這些令牌與觀察窗口的注意力交互評(píng)估其余令牌的相對(duì)重要性。

這個(gè)過程非常像一個(gè)聰明的編輯團(tuán)隊(duì)：先確定哪些成員需要全面了解項(xiàng)目歷史（起草者），哪些只需專注于當(dāng)前任務(wù)（精細(xì)器）；然后根據(jù)每個(gè)人的工作性質(zhì)分配不同大小的資源；最后，對(duì)于每個(gè)人的工作資料，保留那些最重要、最常被參考的部分，丟棄那些幾乎不會(huì)用到的內(nèi)容。

四、實(shí)驗(yàn)驗(yàn)證：效果超乎預(yù)期

研究團(tuán)隊(duì)在兩個(gè)不同規(guī)模的VAR模型上評(píng)估了ScaleKV的性能：Infinity-2B和Infinity-8B，以驗(yàn)證其在不同規(guī)模模型上的通用性。他們?cè)O(shè)置了三種內(nèi)存預(yù)算約束：原始KV緩存大小的4%、10%和20%，以模擬不同資源限制環(huán)境下的部署場(chǎng)景。

在MS-COCO 2017驗(yàn)證集上的實(shí)驗(yàn)結(jié)果令人驚喜。ScaleKV在不同內(nèi)存預(yù)算下均顯著優(yōu)于所有基線方法，F(xiàn)ID（Fréchet Inception Distance，越低越好）、LPIPS（學(xué)習(xí)的感知圖像補(bǔ)丁相似度，越低越好）和PSNR（峰值信噪比，越高越好）指標(biāo)都取得了顯著改善。

在最受限的預(yù)算（4%）下，ScaleKV比下一個(gè)最佳基線在Infinity-2B和Infinity-8B上分別降低了31.2%和48.5%的FID。隨著預(yù)算增加，這種優(yōu)勢(shì)進(jìn)一步擴(kuò)大，在20%預(yù)算下，ScaleKV在兩個(gè)模型上分別實(shí)現(xiàn)了1.82和1.45的FID分?jǐn)?shù)，相比所有競(jìng)爭(zhēng)對(duì)手都有實(shí)質(zhì)性的改進(jìn)。LPIPS結(jié)果進(jìn)一步驗(yàn)證了這些發(fā)現(xiàn)，ScaleKV在20%預(yù)算下分別實(shí)現(xiàn)了0.08和0.06的分?jǐn)?shù)，而PyramidKV（最接近的競(jìng)爭(zhēng)對(duì)手）則為0.11和0.10，表明與原始輸出的感知相似度更高。

更令人印象深刻的是，ScaleKV不僅保持了像素級(jí)一致性，還保留了語義理解能力。在GenEval和DPG這兩個(gè)評(píng)估感知質(zhì)量和語義對(duì)齊的基準(zhǔn)測(cè)試中，使用僅10%原始KV緩存的ScaleKV壓縮模型表現(xiàn)出色。對(duì)于Infinity-2B，ScaleKV壓縮后的模型表現(xiàn)甚至略微超過了原始模型（GenEval分?jǐn)?shù)從0.725提高到0.730）；對(duì)于Infinity-8B，性能幾乎保持不變（GenEval分?jǐn)?shù)為0.790 vs 0.792，DPG分?jǐn)?shù)為86.49 vs 86.61）。

這種性能保持尤為顯著，因?yàn)镮nfinity模型在這些基準(zhǔn)上已經(jīng)優(yōu)于大多數(shù)現(xiàn)有方法，包括更大的模型如DALL-E 3和Emu3-8.5B。而ScaleKV壓縮后的Infinity-8B僅需8.5GB的KV緩存內(nèi)存，相比原始的85GB是一個(gè)巨大的減少。

五、內(nèi)存效率與時(shí)間成本分析

研究團(tuán)隊(duì)對(duì)Infinity-8B模型的推理過程進(jìn)行了全面的內(nèi)存消耗分析。他們發(fā)現(xiàn)，由于序列長(zhǎng)度顯著延長(zhǎng)，Infinity模型的KV緩存是最大的內(nèi)存消費(fèi)者，大約需要模型解碼操作所需內(nèi)存的10倍。ScaleKV大幅降低了KV緩存內(nèi)存需求，將其壓縮到原始模型的10%。

更重要的是，隨著批量大小增加，ScaleKV的內(nèi)存節(jié)省變得更加顯著。研究團(tuán)隊(duì)能夠使用不到80GB的總內(nèi)存生成批量大小為16的圖像，而原始模型僅KV緩存就需要170GB內(nèi)存。

雖然ScaleKV主要是為了提高內(nèi)存效率而開發(fā)的，但它也通過減少?gòu)埩吭L問和傳輸操作帶來了顯著的推理加速。隨著圖像分辨率增加，由于令牌序列中的指數(shù)級(jí)增長(zhǎng)，推理延遲大幅增加。ScaleKV在單個(gè)NVIDIA H20 GPU上實(shí)現(xiàn)了高達(dá)1.25倍的加速，隨著分辨率增加，性能提升變得更加明顯。這些結(jié)果展示了ScaleKV在資源受限環(huán)境中部署的潛力，以及將VAR模型擴(kuò)展到超高分辨率（如4K）的可能性，這在以前受到內(nèi)存瓶頸和推理延遲的限制。

六、研究意義與未來展望

這項(xiàng)研究的意義遠(yuǎn)超過技術(shù)層面。通過解決VAR模型的內(nèi)存瓶頸問題，ScaleKV為高質(zhì)量圖像生成的民主化鋪平了道路。

首先，它大大降低了硬件要求，使普通消費(fèi)者和小型創(chuàng)作者也能在個(gè)人設(shè)備上運(yùn)行先進(jìn)的圖像生成模型。想象一下，即使是普通筆記本電腦也能生成專業(yè)級(jí)質(zhì)量的1024×1024圖像，這將使更多人能夠訪問和利用這些創(chuàng)新技術(shù)。

其次，減少的內(nèi)存占用意味著更低的能源消耗，有助于更可持續(xù)的AI部署。在當(dāng)前對(duì)AI能源消耗日益關(guān)注的背景下，這一點(diǎn)尤為重要。

第三，通過使超高分辨率生成（如4K）成為可能，ScaleKV為專業(yè)內(nèi)容創(chuàng)作、醫(yī)學(xué)成像和科學(xué)可視化應(yīng)用創(chuàng)造了新機(jī)會(huì)。這些應(yīng)用場(chǎng)景通常需要極高的圖像細(xì)節(jié)和質(zhì)量，而這在以前受到計(jì)算資源限制。

這項(xiàng)研究也為未來工作指明了方向。研究團(tuán)隊(duì)指出，雖然ScaleKV在不同容量的模型上展示了強(qiáng)大的壓縮性能，但在更大規(guī)模的VAR模型（如20B參數(shù)以上）上的評(píng)估將提供關(guān)于方法可擴(kuò)展性的額外見解。此外，ScaleKV作為一種后訓(xùn)練KV緩存壓縮解決方案，依賴于預(yù)訓(xùn)練的VAR模型并反映原始模型的輸出。因此，如果原始VAR模型的基線質(zhì)量不令人滿意，那么使用這種方法實(shí)現(xiàn)高質(zhì)量結(jié)果可能會(huì)面臨挑戰(zhàn)。

總結(jié)來說，ScaleKV代表了視覺自回歸模型領(lǐng)域的一個(gè)重要進(jìn)步，通過實(shí)現(xiàn)顯著的內(nèi)存優(yōu)化而不犧牲生成質(zhì)量，它為高分辨率圖像生成的廣泛應(yīng)用開辟了新的可能性。這項(xiàng)研究提醒我們，有時(shí)最有效的創(chuàng)新不是開發(fā)更大的模型，而是讓現(xiàn)有模型更加高效和可訪問。