[论文解读] RL-I2IT: Image-to-Image Translation with Deep Reinforcement Learning

Jing Hu, Ziwei Luo+8 more|arXiv (Cornell University)|Sep 24, 2023

Brain Tumor Detection and ClassificationNeuroscience被引用 3

一句话总结

本文提出 RL-I2IT，一种用于图像到图像翻译的深度强化学习框架，通过使用低维的‘概念计划’作为元策略，将任务分解为迭代的轻量化步骤，从而稳定训练并提升泛化能力。通过结合辅助学习与跳跃连接的软演员-critic（SAC）方法，该方法在高维连续动作空间中实现了鲁棒、高质量的图像翻译，同时降低了模型复杂度并提升了样本效率。

ABSTRACT

Most existing Image-to-Image Translation (I2IT) methods generate images in a single run of a deep learning (DL) model. However, designing such a single-step model is always challenging, requiring a huge number of parameters and easily falling into bad global minimums and overfitting. In this work, we reformulate I2IT as a step-wise decision-making problem via deep reinforcement learning (DRL) and propose a novel framework that performs RL-based I2IT (RL-I2IT). The key feature in the RL-I2IT framework is to decompose a monolithic learning process into small steps with a lightweight model to progressively transform a source image successively to a target image. Considering that it is challenging to handle high dimensional continuous state and action spaces in the conventional RL framework, we introduce meta policy with a new concept Plan to the standard Actor-Critic model, which is of a lower dimension than the original image and can facilitate the actor to generate a tractable high dimensional action. In the RL-I2IT framework, we also employ a task-specific auxiliary learning strategy to stabilize the training process and improve the performance of the corresponding task. Experiments on several I2IT tasks demonstrate the effectiveness and robustness of the proposed method when facing high-dimensional continuous action space problems. Our implementation of the RL-I2IT framework is available at https://github.com/Algolzw/SPAC-Deformable-Registration.

研究动机与目标

• 为解决单步深度学习模型在图像到图像翻译任务中的局限性，如参数量过高、过拟合及泛化能力差等问题。
• 克服将强化学习应用于图像翻译任务中高维连续状态与动作空间的挑战。
• 通过引入具有低维‘概念计划’的元策略来指导动作生成，提升训练稳定性和模型泛化能力。
• 通过迭代决策实现轻量级网络的高效、渐进式图像翻译，降低计算负担。
• 通过特定任务的辅助学习策略与跳跃连接，实现训练稳定化，保留空间细节。

提出的方法

• 该框架采用三网络架构：规划器、演员与评论家，使用软演员-critic（SAC）实现最大熵强化学习。
• 提出一种新型‘概念计划’，作为状态与动作之间的低维中间表征，使高维动作预测变得可行。
• 规划器利用带有跳跃连接的卷积神经网络（CNN）从当前状态生成概念计划，以保留空间细节。
• 演员结合当前状态与概念计划生成动作（图像变换），其中计划作为条件引导，降低探索方差。
• 应用特定任务的辅助学习损失，通过保留结构与语义细节来稳定训练并提升性能。
• 从规划器的下采样层到演员的上采样层引入跳跃连接，增强梯度流动并保留细粒度图像细节。

实验结果

研究问题

• RQ1 具有元策略与概念计划的强化学习框架是否能有效处理图像到图像翻译中的高维连续动作空间？
• RQ2 低维概念计划在图像翻译任务中如何提升训练稳定性和模型泛化能力？
• RQ3 辅助学习与跳跃连接在多大程度上提升了 RL-I2IT 框架的性能与收敛性？
• RQ4 与单步深度学习模型相比，该方法的迭代式、分步方法在参数效率与翻译质量方面表现如何？
• RQ5 所提出方法是否能在降低模型复杂度的同时，实现更优的鲁棒性与多样化表现，从而在多种 I2IT 任务中取得具有竞争力的结果？

主要发现

• RL-I2IT 框架在多个图像到图像翻译基准测试中达到最先进或具有竞争力的性能，涵盖图像到图像翻译、语义分割与风格迁移等任务。
• 使用概念计划显著缩小了有效动作空间，使轻量级演员网络能够实现稳定训练，并减少过拟合。
• 辅助学习与跳跃连接提升了特征保留能力与梯度流动，带来更快收敛速度与更高品质的生成图像。
• 该方法在多种 I2IT 任务中表现出强鲁棒性，在泛化能力与样本多样性方面优于标准 GAN 与 U-Net 基模型。
• 与单步模型相比，该框架以更少的参数与更低的计算成本实现了高质量翻译，适用于资源受限环境的部署。
• 实证结果表明，结合最大熵的元策略可改善探索能力，减少模式崩溃，实现多样化且逼真的图像生成。