用强化学习构建个性化的二维码

694次阅读  |  发布于1年以前

技术概述

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 由三个部分构成：

第一个 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

Copyright© 2013-2020

All Rights Reserved 京ICP备2023019179号-8