Discovering Latent Graphs with GFlowNets for Diverse Conditional Image Generation

TL;DR

Rainbow is a novel framework that uses GFlowNets to decompose input conditions into diverse latent representations, enabling the generation of multiple plausible and diverse images from a single condition, applicable to any pretrained model.

Contribution

It introduces a GFlowNet-based latent graph approach to explicitly model and generate diverse outputs for uncertain conditions in image generation tasks.

Findings

Improves diversity and fidelity in image synthesis tasks.

Effective in natural and medical image datasets.

Enhances counterfactual image generation.

Abstract

Capturing diversity is crucial in conditional and prompt-based image generation, particularly when conditions contain uncertainty that can lead to multiple plausible outputs. To generate diverse images reflecting this diversity, traditional methods often modify random seeds, making it difficult to discern meaningful differences between samples, or diversify the input prompt, which is limited in verbally interpretable diversity. We propose Rainbow, a novel conditional image generation framework, applicable to any pretrained conditional generative model, that addresses inherent condition/prompt uncertainty and generates diverse plausible images. Rainbow is based on a simple yet effective idea: decomposing the input condition into diverse latent representations, each capturing an aspect of the uncertainty and generating a distinct image. First, we integrate a latent graph, parameterized by Generative Flow Networks (GFlowNets), into the prompt representation computation. Second, leveraging GFlowNets' advanced graph sampling capabilities to capture uncertainty and output diverse trajectories over the graph, we produce multiple trajectories that collectively represent the input condition, leading to diverse condition representations and corresponding output images. Evaluations on natural image and medical image datasets demonstrate Rainbow's improvement in both diversity and fidelity across image synthesis, image generation, and counterfactual generation tasks.