Self-Reflective Reinforcement Learning for Diffusion-based Image Reasoning Generation


Illustration of self-reflective reasoning step. Through self-reflective processes of repeated denoising and re-noising, diffusion models achieve image reasoning generation adhered to physical laws and counterintuitive physical phenomena.

Abstract

Diffusion models have recently demonstrated exceptional performance in image generation task. However, existing image generation methods still significantly suffer from the dilemma of image reasoning, especially in logic-centered image generation tasks. Inspired by the success of Chain of Thought (CoT) and Reinforcement Learning (RL) in LLMs, we propose SRRL, a self-reflective RL algorithm for diffusion models to achieve reasoning generation of logical images by performing reflection and iteration across generation trajectories. The intermediate samples in the denoising process carry noise, making accurate reward evaluation difficult. To address this challenge, SRRL treats the entire denoising trajectory as a CoT step with multi-round reflective denoising process and introduces condition guided forward process, which allows for reflective iteration between CoT steps. Through SRRL-based iterative diffusion training, we introduce image reasoning through CoT into generation tasks adhering to physical laws and unconventional physical phenomena for the first time. Notably, experimental results of case study exhibit that the superior performance of our SRRL algorithm even compared with GPT-4o.

Overview

Overview of the algorithm. It includes two processes: multi-round reflective denoising process and condition guided forward process. These two processes are repeated for K rounds.

Visualization Results

Images related to physical laws

Reasoning generation of images related to physical laws.

Images related to counterintuitive physical phenomena

Reasoning generation of images related to counterintuitive physical phenomena.

Reasoning process

Reasoning generation process of the prompt related to a balance. Initially, the model generates images of a balance tilted left without objects or tilted right with lighter objects on the left and heavier ones on the right, both following physical laws. Eventually, it learns to create images defying logic: a balance tilts left with no objects on the left and a ball on the right.

Comparison with GPT-4o

Example of physical phenomenon related prompts and counterintuitive physical phenomena prompts by GPT-4o and our algorithm.