1. 一眼概覽

MaskUNet 提出了一種基于可學(xué)習(xí)掩碼的參數(shù)篩選機(jī)制，在不更新預(yù)訓(xùn)練U-Net參數(shù)的前提下，有效提升了擴(kuò)散模型的圖像生成質(zhì)量和下游泛化能力。

2. 核心問題

當(dāng)前擴(kuò)散模型在不同時(shí)間步使用相同U-Net參數(shù)生成結(jié)構(gòu)和紋理信息，限制了模型的表達(dá)靈活性。該研究聚焦于：如何在不更改預(yù)訓(xùn)練U-Net的參數(shù)下，提升其對不同時(shí)間步和樣本的適應(yīng)性，以生成更高質(zhì)量的圖像？

3. 技術(shù)亮點(diǎn)

• 參數(shù)掩碼機(jī)制：提出可學(xué)習(xí)的二值掩碼，對預(yù)訓(xùn)練U-Net的參數(shù)進(jìn)行篩選，使其在不同時(shí)間步與樣本中發(fā)揮最大效能；
• 雙重優(yōu)化策略：設(shè)計(jì)基于訓(xùn)練（使用擴(kuò)散損失）與免訓(xùn)練（使用獎(jiǎng)勵(lì)模型）的兩種掩碼優(yōu)化方法，適應(yīng)不同場景需求；
• 廣泛驗(yàn)證：在COCO及多個(gè)下游任務(wù)（圖像定制、關(guān)系反轉(zhuǎn)、文本轉(zhuǎn)視頻）中驗(yàn)證，展示優(yōu)越性能和強(qiáng)泛化能力。

4. 方法框架

圖片

MaskUNet方法如下圖流程所示：

• 引入掩碼機(jī)制：對預(yù)訓(xùn)練U-Net參數(shù)施加時(shí)間步和樣本相關(guān)的二值掩碼，實(shí)現(xiàn)參數(shù)選擇性激活；
• 訓(xùn)練方式一：帶監(jiān)督學(xué)習(xí)：通過MLP生成掩碼，聯(lián)合時(shí)間嵌入與樣本特征進(jìn)行訓(xùn)練，目標(biāo)函數(shù)為擴(kuò)散損失；
• 訓(xùn)練方式二：免訓(xùn)練優(yōu)化：借助獎(jiǎng)勵(lì)模型（如ImageReward與HPSv2）指導(dǎo)掩碼更新，無需額外訓(xùn)練掩碼生成器。

該機(jī)制不修改原U-Net參數(shù)結(jié)構(gòu)，而是通過靈活的掩碼動(dòng)態(tài)激活權(quán)重，從而提升模型表達(dá)能力。

5. 實(shí)驗(yàn)結(jié)果速覽

圖片

在COCO 2014和COCO 2017兩個(gè)文本到圖像的零樣本生成任務(wù)中，MaskUNet相較于原始的Stable Diffusion 1.5與LoRA方法，在圖像質(zhì)量指標(biāo)（FID）上均有顯著提升。例如，在COCO 2014數(shù)據(jù)集上，MaskUNet將FID分?jǐn)?shù)從12.85降低至11.72，COCO 2017上則從23.39降至21.88，表現(xiàn)出更強(qiáng)的生成能力。同時(shí)，在圖文一致性方面（CLIP分?jǐn)?shù)）與其他方法持平，說明MaskUNet在不影響語義對齊的前提下，顯著增強(qiáng)了圖像質(zhì)量。

在多個(gè)下游任務(wù)如DreamBooth圖像定制、Textual Inversion新概念學(xué)習(xí)、ReVersion關(guān)系圖像生成以及Text2Video-Zero文本轉(zhuǎn)視頻中，MaskUNet均展現(xiàn)出更強(qiáng)的個(gè)性化表達(dá)能力與細(xì)節(jié)還原能力，進(jìn)一步驗(yàn)證了其作為通用增強(qiáng)組件的實(shí)用價(jià)值。

6. 實(shí)用價(jià)值與應(yīng)用

MaskUNet方法適用于文本生成圖像、視頻生成、圖像定制、關(guān)系表達(dá)等任務(wù)，尤其在無需大規(guī)模參數(shù)更新的資源受限場景下表現(xiàn)出色，適合作為輕量級增強(qiáng)模塊嵌入現(xiàn)有擴(kuò)散框架中。

7. 開放問題

? 掩碼機(jī)制在跨模態(tài)生成（如音頻到圖像）任務(wù)中是否同樣有效？

? MaskUNet是否可以與LoRA等參數(shù)高效微調(diào)方法協(xié)同工作以實(shí)現(xiàn)更強(qiáng)性能？

? 如何進(jìn)一步壓縮掩碼生成模塊的計(jì)算量，使其適用于移動(dòng)端或邊緣設(shè)備？