用強化學習構建個性化的二維碼

 

技術概述

AIGC 在圖像生成領域如火如荼,StableDiffusion 加各種 LORA,ControlNet,大家玩得不亦樂乎。但是基於擴散模型的方式,仍然存在很多問題,比如抽卡成功率過低,生成圖像的細節仍需優化。具體到二維碼生成,目前 hugging face 上的幾個 ControlNet 確實可以生成不錯的二維碼和語義融合的圖像,但是往往需要大量嘗試,並且加上後續的一些迭代修改,才能保證生成的圖像能被正確地識別爲想要地二維碼。我們通過強化學習加課程學習的方式,在保證出圖效果的基礎上,將二維碼識別率從 20%提高至 80%。

強化學習簡介

強化學習是機器學習的一種,它是基於讓機器不斷進行嘗試並在嘗試中獲得的反饋信息指導它如何更優地進行決策。基於馬爾可夫決策過程(Markov Decision Process, MDP)是強化學習中一個非常重要的框架。

MDP 體系中包含五個主要部分:狀態(State),動作(Action),策略(Policy),獎勵(Reward)和折扣因子(discount factor)。狀態表示學習系統當前的狀態,動作則表示在某個狀態下可以進行的操作,策略是指系統在每個狀態下選擇不同動作的規則,獎勵是指在特定狀態進行特定動作後獲得的反饋,而折扣因子則是用於控制當前和未來獎勵之間的權重比。具體流程如下圖所示:

圖片

在基於 MDP 的強化學習中,學習的目標是找到一個最佳策略使得從起始狀態到終止狀態累計獎勵最大。強化學習算法,就是通過學習環境反饋的獎勵,不斷地調整策略,最終形成一個最佳的策略。具體過程中,會利用“探索-開發”策略,即在探索新的可能性和利用已知的最佳策略之間找到一種平衡。

這就是基於 MDP 的強化學習的基礎概念,可以在很多問題上得到應用,比如自動駕駛、遊戲 AI、機器人導航等地方都有廣泛的應用。

Stable Diffusion

SD(Stable Diffusion)是最近非常火熱的圖像生成領域的算法,其核心思想仍然逃離不了下面的這個基本的數學定理:

/任何兩個有相同的 support,並且絕對連續的分佈 , ,可以找到對應的 transformation,將  轉化成  的分佈,也可以找到另外的 transformation,將  轉化成  的分佈/

在圖像生成領域,我們可以將  對應成符合一定語義或者條件的圖像集合對應的分佈。 對應成相同 sample space 的高斯分佈。

擴散過程()

對於擴散模型來說,從  到  這一步,使用了經典物理的擴散過程。簡單來說,就是不斷在圖像信號  上一步一步地去疊加高斯噪聲,並且同時降低  信號,經過大量步數的疊加,整個圖像得信噪比降至無限接近 0 的水平,即獲得一組了完全的高斯噪聲分佈 。

可以參考一篇非常不錯的 Blog:https://lilianweng.github.io/posts/2021-07-11-diffusion-models/

去噪過程()

去噪過程,即生成過程是我們使用強化學習主要優化的階段。既然如此我們就使用 MDP 的語言將整個去噪過程建模一個強化學習能夠容易使用的場景。

圖片

簡單來說,我們從一個高斯噪聲  開始,scheduler 採樣獲得一個 latent。

  1. scheduler 將 latent 發送給 Agent,這邊 Agent 由 ControlNet 和 StableDiffusion 的模型構成。
  2. Agent 根據接收到的 latents,和其他的一些 context,例如(positive/negative prompts, conditions)等等預測出當前步數對應的 noise,並將預測的 noise 發送給 scheduler
  3. scheduler 將接受到的 noise,通過去噪的算法進行去噪,得到新的 latents

不斷重複上面三個步驟,直至整個去噪過程完畢,最後生成一個圖像,計算生成圖像對應得 reward,可以是衡量圖像質量得 reward,或者其他個性化的 reward。對於上述的這整個去噪過程,我們就可以用強化學習的各種算法去進行優化。

 

強化學習優化二維碼生成

現狀

當前 hugging face 上公佈了大量的 ControlNet,可以用來生成個性化的二維碼,例如:

  1. https://huggingface.co/monster-labs/control_v1p_sd15_qrcode_monster
  2. https://huggingface.co/Nacholmo/controlnet-qr-pattern-v2

一般的使用流程即上傳自己目標的二維碼,然後進行一些 prompts,最後生成具有給定 prompts 語義的二維碼。但是目前這些 ControlNet 比較嚴重的問題,抽卡成功率極低,生成一些融合效果不錯,並且能被成功識別的概率不會超過 20%。

對此我們使用強化學習課程學習兩大技術手段進行 finetune 已有的 ControlNet,使得最終的模型生成的圖像既可以很好地保證融合的效果和識別成功率。

算法介紹

訓練框架

我們的強化學習優化二維碼生成過程的訓練框架圖如下所示:

圖片

首先整個算法的輸入包括三個部分:

  1. 一張底圖
  2. 底圖對應的 caption
  3. 目標的二維碼

這三個部分構成一個訓練集 然後將三個部分輸入給模型,包括 Diffusion Model 和用於控制的 ControlNet,其中用於 finetune 的 ControlNet 的需要保持在訓練模式。在整個生成過程中將數據保存用於後續強化學習訓練。

Reward 設置

在獲得最終的生成圖像之後,我們將用其計算二維碼生成這個應用場景下對應的 reward。目前算法的 reward 由三個部分構成:

  • 二維碼能否被正確識別
  • 生成圖像和原圖像的近似度(使用 CLIP 模型)
  • 生成圖像和 prompt 的近似度(使用 CLIP 模型)

第一個 reward 是我們這個優化算法最直接的目的,即提高出圖的成功率。但是可以想象只有第一個 reward 肯定是不行的,因爲強化學習很容易學習到,直接輸出輸入的目標二維碼,而忽略底圖和 prompt 所給定的語義信息。這樣即使算法能夠很好地提高二維碼識別成功率,仍然和我們的優化初心背離甚遠。因此我們通過增加第二個 reward 來控制生成圖像和底圖,prompt 之間的語義相似度。這樣可以在提高識別成功率的同時極大程度地去將用戶給定的語義或者底圖信息融合進入二維碼。

課程學習

在架構圖中,我們發現在 ControlNet 那一層多了一個 紅色自循環,這代表我們在訓練過程中加入了課程學習的 scheduler。

衆所周知,在 ControlNet 和 Diffusion Model 結合的時候,增加 ControlNet 的控制權重,那麼生成圖像則會向 ControlNet 的所給定的條件靠攏。對於二維碼生成這個應用來說也不例外,在原有的 ControlNet 下,你只需要瘋狂提高 ControlNet 的權重,那麼也是可以保證生成的圖像的識別率,但是其融合效果就會越來越差。因此,之前我們只能通過大量的抽卡,抽中一個能夠很好 balance 這兩者的圖像。

我們這邊借鑑了課程學習的思想,對於 ControlNet 的控制強度設置了一個 scheduler,保證在整個訓練過程中控制強度從高到低按某種方式遞減。直覺上,這樣很符合強化學習和課程學習結合的想法,先從簡單的任務開始,慢慢增加任務的難度,後續的強化學習可以借鑑過去簡單任務上的經驗。

算法結果

這邊展示幾張在訓練過程中保存的光子二維碼。所列出二維碼沒有進行特別的 Prompt 調優,和後處理,均爲隨機採樣出的 Prompt,直出的圖像。

圖片
圖片
圖片
圖片

下面展示一下整個訓練過程的曲線:

二維碼識別成功率:

圖片

從這張圖上,可以看到隨着課程學習的更新,任務難度增加,二維碼識別會突然降低,但是隨着一段時間的強化學習,二維碼的識別會再次上升。

CLIP 提供的語義 Reward:

圖片
未來工作

通過使用強化學習優化二維碼生成這個工作,展現了強化學習在當下火熱的 AIGC 領域的潛力。我們相信未來強化學習的技術會越來越多作爲補充加入整個藝術製作流程當中,比如:

  1. 通過在人類標註的數據上學習一個 reward 模型來評判生成圖像的質量,並將 RLHF 技術加入擴散模型的調優。
  2. 通過將生成視頻的穩定性等等一些評判指標作爲 reward,以此提高當下視頻生成技術。
  3. 在使用多視角圖重建 3D 模型的時候,一致性往往得不到保障,我們也可以在其中加入強化學習,提高多視角圖的重建效果。

參考文獻

Black, Kevin, et al. "Training diffusion models with reinforcement learning." arXiv preprint arXiv:2305.13301 (2023).

 

 

作者:騰訊遊戲光子工作室羣 Alex

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章