[논문 리뷰] RL-I2IT: Image-to-Image Translation with Deep Reinforcement Learning

Jing Hu, Ziwei Luo외 8명|arXiv (Cornell University)|2023. 09. 24.

Brain Tumor Detection and ClassificationNeuroscience인용 수 3

한 줄 요약

이 논문은 이미지 간 번역 작업을 반복적이고 경량화된 단계로 분해하여 훈련 안정성과 일반화 성능을 햖थ기 위해 저차원의 '개념 계획'을 가진 메타정책을 사용하는 딥 강화학습 프레임워크인 RL-I2IT를 제안한다. 소프트 액터-크리틱과 보조 학습, 스케일 연결을 활용함으로써, 고차원 연속 행동 공간에서 모델 복잡도를 감소시키고 샘플 효율성을 향상시키며, 강력하고 고품질의 이미지 번역을 달성한다.

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.

연구 동기 및 목표

• 이미지 간 번역에서 한 단계로 처리하는 딥 러닝 모델의 한계를 해결하기 위해, 높은 파라미터 수, 과적합, 일반화 성능 부족 문제를 해결한다.
• 이미지 번역 작업에서 고차원 연속 상태 및 행동 공간에 강화학습을 적용하는 데 도전하는 문제를 해결한다.
• 행동 생성을 안내하는 저차원의 '개념 계획'을 가진 메타정책을 도입함으로써 훈련 안정성과 모델 일반화 성능을 향상시킨다.
• 경량화된 네트워크를 사용한 반복적 의사결정을 통해 효율적이고 점진적인 이미지 번역을 가능하게 하여 계산 부담을 감소시킨다.
• 작업에 특화된 보조 학습 전략과 스케일 연결을 통해 훈련을 안정화시키고, 공간적 세부 정보를 유지한다.

제안 방법

• 프레임워크는 소프트 액터-크리틱(SAC)을 사용한 최대 엔트로피 강화학습을 위한 세 개의 네트워크 아키텍처(플래너, 액터, 크리틱)를 사용한다.
• 상태와 행동 사이의 저차원 중간 표현으로서의 새로운 '개념 계획'이 도입되어 고차원 행동 예측을 가능하게 한다.
• 플래너는 스케일 연결을 사용한 CNN를 통해 현재 상태에서 개념 계획을 생성하여 공간적 세부 정보를 유지한다.
• 액터는 상태와 개념 계획을 모두 사용하여 행동(이미지 변환)을 생성하며, 계획이 조건부 가이드 역할을 하여 탐색 분산을 줄인다.
• 특수 작업용 보조 학습 손실이 적용되어 구조적 및 의미적 세부 정보를 유지함으로써 훈련 안정성과 성능 향상을 도모한다.
• 플래너의 다운샘플링 레이어에서 액터의 업샘플링 레이어에 이르기까지의 스케일 연결은 기울기 흐름을 향상시키고 세밀한 이미지 세부 정보를 유지한다.

실험 결과

연구 질문

• RQ1 메타정책과 개념 계획을 갖춘 강화학습 프레임워크는 이미지 간 번역에서 고차원 연속 행동 공간을 효과적으로 다룰 수 있는가?
• RQ2 저차원의 개념 계획 사용이 이미지 번역 작업에서 훈련 안정성과 모델 일반화 성능을 향상시키는 데 어떤 영향을 미치는가?
• RQ3 보조 학습과 스케일 연결이 RL-I2IT 프레임워크의 성능 향상과 수렴 속도 향상에 어느 정도 기여하는가?
• RQ4 반복적이고 단계적인 접근 방식은 파라미터 효율성과 번역 품질 측면에서 한 단계로 처리하는 딥 러닝 모델과 비교해 어떻게 다른가?
• RQ5 제안된 방법은 모델 복잡도를 줄이고 다양한 I2IT 작업에서 뛰어난 강건성과 경쟁 가능한 성능을 달성할 수 있는가?

주요 결과

• RL-I2IT 프레임워크는 이미지 간 번역, 세분화, 스타일 전이를 포함한 여러 이미지 간 번역 벤치마크에서 최신 기술 수준 또는 경쟁 가능한 성능을 달성한다.
• 개념 계획의 사용은 효과적인 행동 공간을 크게 줄여주며, 경량 액터 네트워크를 사용한 안정적인 훈련과 과적합 감소를 가능하게 한다.
• 보조 학습과 스케일 연결은 특징 보존과 기울기 흐름 향상을 통해 더 빠른 수렴과 고품질의 생성 이미지를 이끌어낸다.
• 이 방법은 다양한 I2IT 작업에서 뛰어난 강건성을 보이며, 일반화 성능과 샘플 다양성 측면에서 표준 GAN 및 U-Net 기반 모델을 능가한다.
• 한 단계 모델 대비 더 적은 파라미터와 낮은 계산 비용으로 고품질 번역을 달성하여 자원 제약 환경에서의 구현에 적합하다.
• 실증 결과에 따르면, 최대 엔트로피를 갖는 메타정책은 탐색을 향상시켜 모드 붕괴를 줄이고 다양한 실사적인 이미지 생성을 가능하게 한다.