Prompt-Guided Image Editing with Masked Logit Nudging in Visual Autoregressive Models

Amir El-Ghoussani, Marc Hölle, Gustavo Carneiro, Vasileios Belagiannis

View Original ↗
AI 導讀 technology AI 重要性 4/5

MLN 演算法在視覺自迴歸模型中實現免反演編輯,1024px 圖像處理僅需 1.6 秒。

  • 採用 Logit Nudging 機制在機率空間平衡原圖特徵與新提示詞。
  • 透過分析源與目標提示詞的交叉注意力差異,精準生成修改遮罩。
  • 結合量化修正技術降低重建誤差,1024px 編輯耗時僅需 1.6 秒。

傳統擴散模型在編輯 1024px 高解析度圖像時,往往需要等待數十秒並受困於反演誤差。最新發表於 arXiv 的研究提出 Masked Logit Nudging (MLN) 方法,能在不依賴反演的視覺自迴歸 (VAR) 模型架構下,將 1024px 圖像編輯時間大幅壓縮至 1.6 秒,同時在多項標準測試中達到頂尖的感知品質。

擺脫擴散模型反演限制,探討視覺自迴歸的挑戰

擴散模型 (Diffusion Models) 仰賴反演技術 (Inversion,推導出生成原圖所需的初始雜訊) 來進行圖像編輯。然而,反演誤差會在採樣過程中不斷累積,導致生成的圖像經常出現局部結構扭曲或整體排版走樣的問題。即便後續引入注意力圖等中間特徵來控制編輯範圍,依然耗費大量運算資源且容易失效。

視覺自迴歸 (Visual Autoregressive, VAR) 模型,如 SWITTIInfinity,透過從粗到細 (coarse-to-fine) 的多尺度預測策略,解決了傳統自迴歸模型逐個預測 token 造成的低效問題。這使得 VAR 模型具備了直接在潛在空間中高速生成完整特徵圖的能力。

即便擁有速度優勢,要在 VAR 架構中實現基於提示詞的圖像編輯 (Prompt-Guided Image Editing) 仍面臨技術瓶頸。現有基於 VAR 的編輯方法,要不是依賴容易產生累積誤差的偽反演 (pseudo-inversion),就是受限於特定分詞方案。這促使研究團隊著手開發一種無需反演、且適用於多種架構的新型編輯機制。

Logit Nudging 機制:機率空間引導生成軌跡

為了解決編輯精準度的挑戰,研究團隊提出 Masked Logit Nudging (MLN),其核心運作於模型的 Logit (神經網路未經歸一化的原始預測輸出) 空間。該機制不直接修改源圖像的底層編碼,而是介入模型解碼時的機率預測步驟。

在給定原圖與目標提示詞後,模型會固定最初幾層低解析度的 token 映射表,並自迴歸地生成剩餘的高解析度映射。為了防止目標提示詞過度改變圖像原有結構,MLN 採用了借鑒自無分類器引導 (Classifier-Free Guidance, CFG) 的概念,但做出關鍵性修改。

不同於傳統 CFG 在資料空間中進行插值,MLN 是在機率空間進行微調。它將目標提示詞引導的軟性預測與來源圖像的 One-hot 原始 token 進行動態平衡。透過隨著解析度尺度遞減的推力 (Nudging Strength),系統能在低解析度階段強制維持整體版面佈局,並在高解析度階段允許局部細節依據新提示詞進行語意轉換。

提取交叉注意力差異,精準定位圖像編輯遮罩

單純依賴 Logit 微調仍可能對不相關的背景區域造成意外更動。為了嚴格限制目標提示詞的影響範圍,研究團隊設計了一套專屬的空間遮罩策略,透過分析源提示詞與目標提示詞之間的交叉注意力 (Cross-Attention) 差異來決定修改區域。

系統會將編碼後的來源 token 分別送入配備源提示詞與目標提示詞的 Transformer 解碼器,並記錄前向傳遞過程中的注意力啟動狀態。比較兩次傳遞中不同區域對文字的響應後,模型計算出絕對差異值,藉此抓出語意發生改變的精準位置。

研究發現 Transformer 的中間層(第 3 到 27 層)能提供最清晰的注意力結構。系統隨後將這些差異映射圖進行正規化,並取前 80%(於 512px 實驗中)的高差異像素生成二值化的編輯遮罩。在這個遮罩內部套用 Logit 微調,而在遮罩外部則強力保留源圖像特徵,確保修改只發生在真正與提示詞變更相關的區域。

引入量化修正技術,消除特徵離散化的重建偽影

在進行未套用編輯的純圖像重建測試時,研究人員觀察到編碼與解碼過程會累積不可忽視的量化誤差。當連續的圖像特徵被強制映射到有限字典中的離散代碼向量 (Codebook Vectors) 時,會產生微小的重建落差,並在多尺度解碼中逐層放大。

簡單地將這些殘差 (Residuals) 加回最終特徵圖會產生明顯的視覺偽影。因為原始殘差包含了解碼器無法正確解析的非流形 (Off-manifold) 方向,強行加入只會破壞畫面。

為此,團隊開發了量化修正 (Quantization Refinement) 步驟。他們採用「投影-更新」的迭代循環,在保留編輯區域不受影響的前提下,將殘差重新投影回模型學習到的代碼本嵌入空間。這種做法逐步過濾掉無法解碼的雜訊成分,大幅提升了未編輯區域的像素級保真度與整體畫面真實感。

PIE 基準測試:1024px 圖像編輯僅需 1.6 秒

PIE-Benchmark(包含 700 張涵蓋物件替換、風格轉換等場景的標準測試集)上,MLN 展現了壓倒性的效能優勢與速度。

1024×1024 的高解析度測試中,MLN 達到了最低的 LPIPS(感知圖像塊相似度),以及最高的 CLIP 分數(評估圖像與文本語義對齊程度),同時單圖編輯耗時僅 1.6 秒。相較之下,採用擴散模型或流匹配 (Rectified Flows) 的現有方案處理同等解析度往往需要漫長的等待時間。

512×512 的實驗中,MLN 同樣以 0.82 秒的速度創下紀錄,並在背景保留指標(PSNR、MSE)上顯著勝過現有的 VAR 編輯方法 AREdit 與 VARIN。此外,這套機制展現了優異的架構通用性,能無縫套用在 SWITTI 與 Infinity 等不同的網路中,無需任何額外的微調訓練。

透過機率空間微調與交叉注意力遮罩,MLN 讓自迴歸模型實現了秒級高解析度精準圖像編輯。

Abstract

We address the problem of prompt-guided image editing in visual autoregressive models. Given a source image and a target text prompt, we aim to modify the source image according to the target prompt, while preserving all regions which are unrelated to the requested edit. To this end, we present Masked Logit Nudging, which uses the source image token maps to introduce a guidance step that aligns the model's predictions under the target prompt with these source token maps. Specifically, we convert the fixed source encodings into logits using the VAR encoding, nudging the model's predicted logits towards the targets along a semantic trajectory defined by the source-target prompts. Edits are applied only within spatial masks obtained through a dedicated masking scheme that leverages cross-attention differences between the source and edited prompts. Then, we introduce a refinement to correct quantization errors and improve reconstruction quality. Our approach achieves the best image editing performance on the PIE benchmark at 512px and 1024px resolutions. Beyond editing, our method delivers faithful reconstructions and outperforms previous methods on COCO at 512px and OpenImages at 1024px. Overall, our method outperforms VAR-related approaches and achieves comparable or even better performance than diffusion models, while being much faster. Code is available at 'https://github.com/AmirMaEl/MLN'.