Discovering Latent Graphs with GFlowNets for Diverse Conditional Image Generation
Authors: Bailey Trang
Parham Saremi
Alan Q. Wang
Fangrui Huang
Zahra TehraniNasab
Amar Kumar
Tal Arbel
Li Fei-Fei
Ehsan Adeli
Venue: NeurIPS 2025 poster ()
Keywords: generative models
diffusion model
diversity
gflownet.
Primary area: deep_learning
TL;DR: Generating diverse plausible images with GFlowNets and diverse condition latent representation in conditional 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 \modelnamenospace, a novel conditional image generation framework, applicable to any pretrained conditional generative model, that addresses inherent condition/prompt uncertainty and generates diverse plausible images. \modelname 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 \modelnamenospace’s improvement in both diversity and fidelity across image synthesis, image generation, and counterfactual generation tasks.

Discovering Latent Graphs with GFlowNets for Diverse Conditional Image Generation
Bailey Trang', Parham Saremi 4,5 Alan Q. Wang?, Fangrui Huang'. Zahra TehraniNasab Amar K,mar
Tal Arbel 4.5 Li Fei-Fei Ehsan Adeli
1 Dept. of Computer Science, Stanford University, Stanford, CA, USA Dept. of Psychiatry and Behavioral Sciences, Stanford University, Stanford, CA, USA 'Dept. of Biomedical Data Science, Stanford University, Stanford, CA, USA *Center for Intelligent Machines, McGill University, Montreal, QC, Canada 'MILA- Quebec AI institute, Montreal, QC, Canada (trangn, alanqu, fangruih, feifeili, eadeli)@stanforr edu oparham.sarent zahra. tehraninasab, amar. kumma]]@aiil megill tal. ar arbel arbee@mcgill.ca
Abstract
Capturing diversity is crucial in conditional and prompt-based image generation, particularly when conditions contain uncertainty that can lead to multiple plau- sible outputs. To generate diverse images reflecting this diversity, traditional methods often modify random seeds, making difficult to discern meaningful differences between samples, or diversify the input prompt, which is limited verbally interpretable diversity. We propose Rainbow, novel conditional image generation framework, applicable to any pretrained conditional generative model, that addresses inherent condition/prompt uncertainty and generates diverse plau- sible 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 latent graph, parameterized by Generative Flow Networks (Gtsow into the prompt repre- sentation computation. Second, leveraging GFlowNets' advanced graph sampling capabilities 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 natural image and medical image datasets demonstrate Rainbow's improve- ment in both diversity and fidelity across image synthesis, image generation, and counterfactual generation tasks.
Introduction
Conditional image generation produces novel images that adhere to given input prompts or conditions like text [32 75). In real-world scenarios, an input prompt has inherent ambiguity, which may correspond to multiple plausible output images 7 40
76 85], especially when prompts are abstract, high-level information. For example, an input text prompt describing a "sunset scene" could map to many valid images, differing in factors such season, light control, and overall ambiance. Similarly, in medical imaging, brain magnetic resonance images (MRIs) two patients of the same age and gender may nevertheless display variability in the structure of brain regions and patterns of intensity
"Corresponding author.
We refer "prompts" and "conditions" interchangeably.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

despite having identical conditions due to subject-specific and medical scanner- scanner-specific details. In both cases, failing to address inherent uncertainty and capturing diversity in generative models can lead to suboptimal decision-making misinterprettaioos, and generation collapse, where limited and uniform outputs fail to represent the necessary variability
13 18
73
Previous attempts at generating diverse images in conditional image generation models, such as using GANs [9], diffusion models [10] 224,, and latent diffusion models [2] (LDMs), can be categorized
into two main approaches: (1) Traditional methods typically rely on randomness; for example, repeating the generation process with different random seeds or varying the random noise on the same seeds in diffusion models (25] 34 49, [S]] to create multiple outputs. While these methods can produce non-identical images, they often fail to capture true diversity of choices and may exhibit inherent biases; (2) Another line of work involves diversifying and adding details to the input prompt verbally using a pretrained Large Language Model (LLLM) such as ChatGPT 22 60
Although this approach can enhance the richness of the generated content, it is confined to text-based conditions and relies on external L.LM models and their own biases. Consequently. these strategies may not adequately address the inherent uncertainty of conditional image generation tasks. In addition, more versatile approach is needed to handle multiple condition types. For instance, generating medical images conditioned on age, sex, diseases, other medical details can enrich datasets in fields where data collection is costly and time-comeen consuminn,,, such as in 3D brain MRI or chest X-ray datasets.
Addressing these limitations, we introduce Rainbow, novel conditional image generation framework designed to produce diverse and plausible images. Rainbow can be integrated into any pretrained conditional image generative model. The primary idea is to create multiple images simultaneously that capture uncertainty by collectively reflecting the input condition. To achieve this, we aim to generate diverse condition representations that encapsulate various aspects of the uncertainty inherent in the input condition within the latent space. Each representation produces a distinct output image while the pretrained generative models remain frozen or minimally modified. As result, Rainbow delivers a range of high-quality images that comprehensively interpret the input prompt.
To achieve diverse condition latent representations that collectively reflect the input condition, we first construct graph structure, called the latent graph, within the latent representation computation. Next, we utilize Generative Flow Networks (GFlowNets) [ to sample diverse trajectories over the graph collectively representing the input condition. Specifically, GFlowNets are designed to capture uncertainty in tasks with multiple possible outputs (multiple modes) by sampling diverse high-quality intermediate representations (e.g., trajectories over a graph) that lead to varied outputs, each representing one possible optimal outcome (one mode of the solution space). GFlowNets have been applied to many contexts, including molecule generation [4], gene regulatory networks 3 48), and dropout masks [42). In Rainbow, trajectories generated by GFlowNets collectively capture diverse interpretations of the input condition, leading to diverse condition latent representations, which are subsequently used to produce diverse output images.
Our contributions include:
First, we introduce Rainbow, novel conditional image generation framework that captures uncertainty and produces diverse images.
Second, by discovering the diversity in condition latent representations, Rainbow is appli- cable to any pretrained conditional generative model, regardless of the condition type, and addresses the limitation of relying on randomness during generation.
Third, our experiments on text- and non-text-based conditions across natural images and medical images (brain MRIs and chest X-rays) demonstrate Rainbow's 's improved capacity to capture uncertainty, generate diverse and plausible images, and benefit downstream tasks.
2 Preliminary
2.1 Generative Flow Networks (GFlowNets)
GFlowNets is probabilistic model that samples diverse high-quality objects f.e. diverse trajectories of node/edge through a graph, where the likelihood of generating a trajectory x is proportional to an unnormalized probability or reward (())
e B(x), with denoting the expectation of some
quantity of interest associated with 4 5 44
GFlowNets samples sequence of actions that modify

---

a compositional trajectory (i.e., adding one edge to a trajectory), starting from a universal initial state and continues through successive modifications dictated by a trainable policy until it reaches a terminal state or achieves specific graph sparsity. This policy is trained so that the probability of terminating the trajectory x at particular final state proportional to the reward R(x))
Specifically, GFlowNets operate on a graph G (S, A), where S is the set of states and A is the set of actions (transitions). The objective is to model a nonnegative flow F A 120. which defines the unnormalized likelihood of taking action to transform from state to s' sampling. the flow F needs to satisfy certain constraints, such as the flow matching constraints [3].
To ensure correct
Flow Matching Constraints is the core principle of Gets which enforces that for any intermediate states, the incoming flow equals the outgoing flow. For any state the state flow F(s) is defined as the total flow through state 8, and the edge flow F(s' + 8) is the flow along transitions s' 8; subsequently, the flow matching constraints is formulated as F(s) E(s F(s' 5) = -(E(sss) F(s + ). The goal is to train the GelowNets model so that the state flow at any terminal state 8T that obtains object & is proportional to the reward F F(8r) x R(x).
2.1.1 Detailed Balance Objective
Detailed Balance Objective [5] (DBO) is one of the training objectives for GFlowNets, along with other approaches such as flow matching (5] and trajectory balance objectives [44]. DBO enforces consistency between forward and backward transitions while aligning terminal states with a reward function R(x). Let 8; denote the state at step i, where Su is the initial state and s, is the terminal state that yields object after steps; DBO defines the forward policy Pr(s;|s;_1:' 0) parameterized by 0 is the probability of transitioning to state 8; from 8;- 1, while the backward policy Pa Pn(#i-1 - 18:;; 0) models the reverse transition; The state flow (s).). scalar function, estimates the unnormalized likelihood of passing through state 8. Subsequently, the DB loss combines two critical terms:
(2,, R(x)) Σ log (0(8i-))P(( - 2 + (log Fo(8n P) |8n
(1)
Fo (B(s)))P(S
R(x)
Specifically, the first term ensures conservation of flow between consecutive states s;-1 and st. By minimizing the squared log-ratio of forward and backward transitions, the loss enforces that the forward probability to transform 8;-1 to 8; equals the backward probability to revert 8; back to 8:-1- In addition, by weighting with their respective flows, this term guarantees that the net flow through any state transition balanced. The second term aligns the final state 8,," 's flow with reward R(x).
2.1.2 Capturing Diversity with GelowNets
GFlowNets promote diversity through three interconnected mechanisms grounded in their flow- matching constraint foundation. First, the terminal state flow aligmment FSTT) x R(x) enforces proportional sampling where candidates x (associated with terminal states sy.) are generated with R(x), preserving all reward modes unlike reinforcement learning's 's arg max R(x) objective 4 Second, the flow conservation constraint ensures global balance: at every non-terminal state E S. incoming flows from predecessors equal outgoing flows to successors, preventing preferential routing to dominant modes while maintaining non-zero probability for all viable paths (5) Finally, the stochastic forward policy Pr(s"|ss; 0) F(s->") derived from F(2)
probability p(x) f(rr)
normalized edge flows F(s + s"), enables amortized trajectory generation. Unlike Markov Chain Monte Carlo's 's (MCMC) local random walks, this allows direct jumps between distant modes (e.g.. structurally distinct molecular graphs with comparable R(x)) through single-pase sampling via Pr (4, bypassing MCMC's iterative transitions [44]. Together, these mechanisms ensure diverse high-reward candidates are sampled proportionally to their rewards while preserving exploration capacity across disconnected regions of the solution space.
2.2 Latent Diffusion Models
Training an LDM consists of two stages. In the first stage, an autoencoder is trained which learns to map each image X to lower-dimensional latent embedding z. Let E, and D, denote the image encoder and decoder making up the autoencoder, respectively. In the second stage, a (conditional) diffusion model is trained on the optimized latent embeddings (X). The generative process of
3

---

Step
Initial representation
Discovering divene condition representations
Generating diverse ouiput images
"Sunses scase
Hotey
Condition Encoder
? Graph Decoder Q
Decoder Dr
Cc
é Gruphs
Latent Diffusion Model
Generator QGEN
represemation
graphs correspondin conditin represcttations
Dormspndingg
Div Divense ouzput amages
Graph Decodar Qu
Compleoe goop
cubbedding
Poeling layer
Edaz 1-b
Edge 0.3
Edgr 5-
New condition represettatinn
Figure 1: Rainbow operates by transforming an input condition into diverse images. Initially, employs the pretrained condition encoder to derive initial representation of the input condition (which contains uncertainty about locations or objects with the given prompt in this example). Then, graphs generator produces multiple trajectories over graph that reflect the input condition. These graphs are encoded into new latent condition representations. New condition representations and latent noisy image processed through the Latent Diffusion Model to acquire denoised image latents, which are subsequently decoded into diverse output images.
an LDM takes in a noisy latent sampled from some prior distribution p(z) and iteratively denoises to produce generated sample 2n by 20 LDM z, c), where c is the condition. Finally, the denoised latent passed through the image decoder D, by X D (2o) to obtain the synthesized image.
During training, noise added the latent representation to create noisy latent image z,, where denotes the diffusion timestep. The model predicts the noise added to the latent image, minimizing the difference between the predicted noise and the actual noise at every timestep f, as described in Equation 2 where ew is the neural backbone that performs tine-conditioned denoising of the latent embedding. Typically, eu is implemented as a coneeonniiional UNet [6] 65).
LDM N0.0.))l Ew (,,
d, c)
3 Rainbow
Rainbow captures the inherent uncertainty in conditional image generation and produces diverse yet realistic images that reflect the input condition. The objective is to decompose the input condition into diverse latent representations, distinct yet jointly interpret the same condition. Each new latent representation is then processed to generate output image, enabling various images corresponding to the input condition. To achieve diversity in latent representations, Rainbow conducts latent graph in latent representation computation and utilizes the GelowNets 4 to sample diverse high-quality trajectories over the graph, which are then decoded into condition latent representations to produce diverse output images.
produce
3.1
Rainbow's Inference Pipeline
Let M denote the number images to be generated, and we assume that models are fully trained. Figure I visualizes the inference process of Rainbow in three main steps.
First, the input condition (e.g input prompt) denoted as R (e.g. 77 fo24), via learned condition encoder &c. Next, the graph generator QGEN takes as input the initial condition embedding e and outputs a set of M distinct trajectories over the graph. Each generated graph is then transformed into new condition representation by the graph decoder model QD. as described in Equation Finally, a latent diffusion model generates images from the set of condition embeddings é @l:M and a set of noisy latents z I:M sampled independently
encoded into condition initial latent representation,

---

from the prior to generate M images X I:M, described in Equation 4
CI: Qu(Quux(c))
(3)
DI(DDD e')), i M]
(4)
3.2 Rainbow's Training Details
This section describes the training strategy to obtain diverse condition latent representations el:M e R.MSS. from the initial c E R detailed algorithm provided in Appendix C We assume the existence of pretrained LDM. During Rainbow training, we freeze condition encoders &c, image encoder ε,, image decoder D,, and Unet model; the graph generator QGEN and graph decoder QD are trained from scratch. We present Rainbow training progress into three stages as detailed below.
Stage 1: Discovering Diverse Graph Representations. We construct undirected graph, G* with N nodes, which yields N(N - 1)/2 non-self-loop edges. Inspired by previous works [47 80) that learn the underlying connections of variables the latent space for greater interpretable context exploration, edge embeddings in Rainbow are randomly initialized. During training, Rainbow assigns interpretable meaning edges without being constrained by pre-defined edge semantics.
We design QUEN as a GFlowNets model that iteratively predicts edges to be added to each of the set of M trajectories over 9", while maintaining a fixed per-graph sparsity, p. At each step, the GFlowNets predicts M edges, adding each edge to one of the M trajectories. This process is repeated for S steps. Specifically, the GFlowNets terminate the edge sampling process once p is reached, and no explicit terminal states are defined (as in ) The total number of edges S is calculated as
S = (1 p)- NN With this design, the number of states in the GFlowNets equals the number of edges plus S 1, which includes one initial state and S states for adding S edges.
In the initialization, we create set of M trajectories T {Tg=D . =o where s represents the state index within the GFlowNets framework. We initialize each trajectory with special starting element, the edge index 0. At state = S. each trajectory in T is expected to be filled with edge indices in the range to N(N-1)
2
At states i with 1 < i S S. input the previous trajectory T and the condition e to predict probability of edges to add to the current M trajectories, one edge for each, as described in Equation 55 In Rainbow, the edge prediction strategy is derived from DBO, which is introduced in Section 2.1.1 with further details in Appendix C After reaching the final state = S, we obtain M trajectories L. each filled with edges.
TI:M TE = QGEN (T.L 11 c).
(5)
Decoding Graphs into Condition Representations. utilized to decode each trajectory into condition representation of shape Sc, denoted as è 2I:M E S... As visualized in Figure [ Qn is designed with three key steps.
Stage 2:
graph decoder model Qp is
First, for each trajectory T QD encodes the sequence of edge indices (of shape R into sequence of edge embeddings of shape R S Subsequently, these edge embeddings are passed through RNN ensure the order correlation of edges. Finally, the output of the RNN is processed by a projection pooling layer to map the sequence from R Sxda to the desired shape S..
The final condition representations are computed as a convex combination between the diverse representations and the original condition representation by a blending factor r::
1:M = + - 7)c.
(6)
Stage 3: Getting reward and computing losses. After obtaining cl:M we perform the diffusion process introduced in Section 2.2 with the added noise cu and get M predicted noise e Subsequently, to evaluate how good the sampled M trajectories, corresponding to M predicted noise, the reward function R(TEMM). defined as the exponential of the negative MSE, as in Equation 7
Our training objective combines two loss terms. First, the GFlowNets loss, Lors, follows the DBO introduced in Section 2. with reward function R(T.L.M, Second, the diffusion denoising loss,
5

---

"Sunset scene with mountain"
"Sunset scene with mountain in specific season"
SD3.5 12
SD2-1
CADS 67
PAG 84
Rainbow (Ours)
Figure 2: Comparison of multiple images generated by basclines and Rainbow. The baseline methods tend to produce images with repetitive layouts and primarily drawing art styles, failing to capture the uncertainty of "season". In contrast, Rainbow generates a variety of sunset scenes, showcasing diverse light levels, grass colors, and effectively capturing different seasons.
LLDM. (mentioned in Section 2.2 computes the mean squared error (MSE) between the added noise and the predicted noise 2l:00
More specifically, CarN is to train the graphs generator QGPN. ensuring the diversity of sampled trajectories as well as the alignment to reward; meanwhile, LLOM to optimize the graph decoder model Qn. The total loss Liotal is weighted combination:
( -
(7)
CUIN = CED(E R() Lixa = aLGFN + BLLDM.
(8)
(9)
(a) Base prompt Spring edges (b) Base prompt + Winter edges
Figure 3: Images generated by Rainbow seasonal edges with the base prompt "Sunset scene with mountain". Most objects and layouts are consistent between the images, with noticeable season- specific details in the second image, such spring flowers and winter snow.
4 Experiment
We conduct experiments to investigate the following hypotheses. Hi: Utilizing diverse graphs facilitates generating diverse images; Hz: Latent graphs be extracted into meaningful and interpretable patterns; H3: Improved ability to capture diversity enhances the performance of downstream tasks. Reproducibility details are in Appendix D
4.1 Experiment Setup
Natural Images We use the Flickr30k dataset 83 which includes about 30k images paired with captions describing daily-life scenes, which contain uncertainty object choices styles. We build our Rainbow on top of the pretrained Stable Diffusion v2-1-base (SD2-1) with frozen pretrained VAE 35 image encoder/decoder. CLIP (56) text encoder, and Unet model. We evaluate the results against SD2-1, Stable Diffusion v3-medium (SD3.5), CADS [67 a recent sampling strategy enhances diversity in the image-generation task, and pretrained checkpoint of PAG [84]- recent work that improves diversity in the text-to-image to-image task by prompt diversifying with GFlowNets. In comparisons, both Rainbow and CADS utilize SD2-1's pretrained encoder-decoder and diffusion models. Rainbow's 's graph generator module includes M = 40,
20, and S = 32.
3D Brain MRIs We curate dataset of about 27k datapoints for training with diagnosed disease from the following datasets: ADNI 53 ABCD Study [3 78) HCP [77), PPMI (55), and AIBL. [4]. This task contains uncertainty in anatomical details such as ventricle sizes. Our setting employs

---

10.5
1.2
10.0
9.5
1.7
Quality (FID
Quality
9.0
8.5
SD3.5
PAG
2.5
RadEdit
8.0
Rainbow
S01 5
7.5
CADS
Rainbow
26.0 27.0 27.5 Diversity (VS
28.5 29.0
6.2
6.4
6.8
7.0
7.2
Diversity (VS
(a) Quantitative comparison on Natural Images
(b) Quantitative comparison on Chest X-rays
Figure 4: Quantitative analysis diversity and image quality of SD-based Rainbow. Rainbow consistently outperforms SD baselines in diversity with higher Vendi Score (VS) across domains and image quality with higher Inception Score (IS) in natural images and FID in chest X-rays.
demographic input conditions on age and binary sex (0: male, female) and fine-tunes the Medical Open Network for Artificial Intelligence (MONAI) [8]'s optimized 3D LDM along with Rainbow training. We benchmark Rainbow against LDM and a GAN-based 337 baselines. Rainbow's graph generator module includes M = 8, N 8, and S
Chest X-rays We use the CheXpert dataset 30 which contains 170k training images. This dataset contains diversity in medical devices (such as chest tubes and wires), diseases (such as pneumonia and pleural -ffuaaleffusion) and anatomical details. We implement Rainbow on top of frozen parameters of finetuned Stable Diffusion v1.5 (SD1.5) by previous work (36) for chest X-ray data. We generate 2D chest X-ray images based on text prompt conditions, e.g., "Chest X-ray showing Support Devices".
In addition to the finetuned SD1.5, 15, we include RadEdit (52), model trained from scratch multiple
chest radiology data such as CheXpert [29]. MIMIC-CXR [31], and NIH-CXR (79 data for image
editing tasks (more details at Appendix D.3 in the result comparison. Rainbow's 's graph generator
module includes M 10, = 20, and
33.
4.2 Experiment Results
4.2.1 Diverse Images Generation Results
Investigating Hi, analyzed the generated images by baseline models and the proposed Rainbow in scenarios where the input condition contains uncertainty.
Natural Images Qualitatively, Figure compares the generated images using two prompts, with over 40 generations displayed. SD3.5 consistently generates repetitive layouts with backlit mountains, an orange sky, and ambiguous seasons. Additionally, SD2-1 offers broader variety of objects and layouts but adheres to a drawing art style for the first prompt and predominantly generates spring scenes with green grass, lacking scasonal diversity for the second prompt. CADS marks diverse objects, yet produces unclear or winter-dominater images for the second prompt. PAG with one base image and diversified prompts does not introduce significant edits in this case and produces repeated layouts and many unclear season. Conversely, Rainbow produces images with diverse objects and light tones and effectively captures seasonal elements such as lush spring greenery and golden autumn foliage, even in the first prompt. For the second prompt. Rainbow demonstrates balanced seasons and clear seasonal features. We quantify the diversity and quality of generations in 60 prompts from the
FID score
Sex Classification Accuracy
Age MAE 1
Real data
le-5
96%
3.55
Synthesis Conv. Conv. sex Synthesis Cony. age
Conv. sex
Synthesis Conv.
Conv. sex
Random
48%
32.17
GAN 19
2.3329
63%
29.23
LDM 8
0.3288 0.3570 0.3590
68%
68
67% 73%
21.52
20.35
19.43
Rainbow (Ours) 0.3149 0.3375 0.3285 79 %
80 %
13.73
18.59
16.40
Table 1: Quantitative evaluation on 3D Brain MRIs. "Conv." stands for "Converted". The lower
FID score, the higher accuracy, and the lower mean absolute error (MAE in years) indicate better performance. Rainbow outperforms baselines across tasks and metrics.
7

---

(a) Actual samples from dataset
(b) Generations by Rainbow
(c) Generations by baseline L.DM
Figure: 5:
Comparison MRI image generations for 63-year-old male individual. Compared to
actual samples from males aged 63-65, Rainbow captures greater diversity in details like ventricle sizes, while the baseline LDM generates images with less variation.
Generations by RadEdit
(b) Generations SD1.5
Generations by Rainbow
Figure 6: Comparison of chest X-ray image generations for the prompt Chest X-ray with support devices. Rainbow is able to provide high-quality images while generating a more diverse set of medical devices compared the baselines (a) and (b).
COCO Validation set [4], with 40 images per prompt. We use the Inception Vendi score (VS) 16 to evaluate diversity and the Inception score (IS) [70] to assess image quality. As shown in Figure Rainbow outperforms the baseline in both diversity (higher VS) and image quality (higher IS). In
addition, observe that all methods produce images relevant to the prompts with similar CLIP
score [57), which is approximately 30.3. Numeric results are in Table 3 classifier-free seree-ggudnccce scale
affects discussed in Figure 18. full 40 generations in Figures 12 16
all in Appendix E
Brain MRIs Figure showeases the generated brain MRI images conditioned on 65- -5-year-old male individual. Unlike younger age groups (e.g. 10-20 years old) with characteristic small ventricles, or the older age groups (e.g. 70+ years old) with large ventricles 17 71, the 60s age range includes wide variety of ventricle patterns [61], as visualized in Figure with actual samples. While all models generated high-quality images, Rainbow effectively captures diverse ventricle patterns with varying ventricle sizes, whereas the baseline LDM tends to produce similar ventricle regions across different samples. We provide full axial, coronal, and sagittal views for this experiment in Figure 28 and visibility of structures and details in generated 3D MRIs in Figure 27 in Appendix
In addition, we quantify brain MRI generations obtained from multiple tasks: image synthesis (with 200 conditions balanced on age and sex) and counterfactuals on age (three age shifts at 10, 40, and 80 years) and sex (flipping binary sex). For each task, we report FID score [72] and performance on age and sex prediction by pretrained classifier models (details at Appendix D.2p. As shown in Table Rainbow outperforms LDM and GAN baselines across metrics and tasks with lower FID, higher sex accuracy, and lower age MAE. To justify age and sex prediction models, we report "Real" results that were tested on real data and "Random" results that were evaluated on random outputs. Figure 29 in Appendix provides counterfactual generations.
Chest X-rays Figure 4p quantifies generations by Rainbow and baselines using FID and VS. Rainbow achieves higher VS, indicating greater diversity than the finetuned SD model, while also improving image quality with lower FID score. Both Rainbow and SD outperform the RadEdit. qualitative comparison, images are generated using the prompt "Chest X-ray showing support devices". where Rainbow generates a more diverse set of medical devices, such as pacemakers, in all generations, while baselines do not show any devices in some images. All models achieve similar CLIP scores of 33.5. Additional results including generations, Figure 22 and numeric results, Table 4 outlined in Appendix

---

SD2-1
Rainbow
Images by Rainbow with the base prompt "Chest X-ray with no significant find- ings" (top) and being appended "supporting devices" edges (bottom) with devices added.
Figure 7:
Figure 8: Diversity image editing. Using the DAC 744 editing pipeline, given the cat (left) and "A cat wearing a wool cap", Rainbow captures diverse cap colors, while SD2-1 generates white caps.
4.2.2 Latent Graph Interpretation
Investigating Hg, we explore sets of edges that present for 4 seasons in Figure Section 4.2.1 From images generated by Rainbow with the two prompts, we first cluster images based observable seasonal features, e.g. snow for winter. We then extract the 10 most frequently added edges for cach season when having "in specific season" in the prompt. Subsequently, we append these 10 extra edges to the original trajectories of the first prompt and generate new images. Figure presents the effect of manipulated graphs on the newly generated images. We can observe the addition patterns of colorful flowers in Figure 3h and snow in Figure 3p. Although edges in the latent graph are not predefined, Rainbow can implicitly learn capture specific context and group edges into meaningful features. For images with four seasons' edges, see Figure 17 in Appendix E
We apply the same approach to chest X-rays explore the set of "support devices" edges by extracting the 10 most frequently added edges when changing the prompt from "Chest X-ray with no significant finding" to "Chest X-ray showing support devices" (visualized in Figure in Appendix Ep.. Figure shows the transformation with added "support devices" edges into trajectories with the appearance of medical devices in the generated images.
4.2.3 Performance on Downstream Tasks
Exploring H3, we conduct two downstream tasks: image editing and counterfactual generation.
We conduct an image-editing task with the natural-image domain that modifies the input image based on a prompt. We compare Rainbow to SD22- using the DAC image-editing pipeline (74]. Figure [ shows the results of editing an image of a cat to add wool cap. DAC with SD2-1 consistently produces images with a white cap, while Rainbow generates caps in multiple colors. This demonstrates Rainbow's 's ability to capture uncertainty (the cap color) and generate diverse samples. For the full 40 generations, see Figure 19 in Appendix
For 3D Brain MRIs illustrated in Figure 9 we perform the age prediction task using training data that includes synthesized data from Rainbow and the LDM baseline. We include 1600 generations arranged by age from 0 to 100 for both sexes, with each condition generating 8 samples. Real data is incorporated in specific proportions. Figure F visualizes that models trained with synthesized data generated by Rainbow achieved better performance. Both models performed best and better than models trained with fully real data at 50%. Specifically, at 50%, both models outperformed those trained with only 1600 real data points (100% dashed line).
4.3 Ablation studies
We assess the effect of varying the number of trajectories M. As shown in Figure TO (left) for 3D brain MRI, a decrease in M leads to drop in performance, with significant performance decrease when M 1. However, even with M
Rainbow still performs better than the baseline LDM. A similar pattern is observed in Figure I0 (right) for chest X-ray data; the lowest performance corresponds to the lowest number of graphs, given the same sparsity, with M = 10 yielding the best performance on diversity assessment. Models for synthesizing chest X-rays are partially trained for 16,500 steps (out of 24,000 steps for fully trained models) with fixed N Additional ablation results varying the number of nodes provided in Appendix B.2

---

30 Brain Experiment
Chedt X-Ray Experiment
0.70
9.0
100%
8.5
LDM
0.65
7-5
Rainbow
0.60
7.0
Quality
P-10
0.50
6.0
5.0
5.5
Cartr
CaTy:
0.50
p-30t
Percentage of Real Samples
25% 50%
humber of trajcctries
Diversity
Figure 9: Age Prediction Task trained on synthesized data. Rainbow achieves the lowest MAE (Mean Absolute Error in years).
Figure 10: Rainbow over number of trajectories (M) show- ing that high values of M yield better performance, which MAE (Mean Absolute Error in years) on the 3D brain MRI and high diversity (VS score) in the Chest X-ray experiments. p is the sparsity in Rainbow. "Conv." stands for "Converted".
5 Related Work
Conditional Image Generation with Diffusion Models have been driven by diffusion models, which iteratively denoise random inputs into coherent samples while outperforming GANs in stability and fidelity [26]. Key innovations include classifier-guidee sampling [12), which steers generation using gradient signals from pretrained classifiers, and latent diffusion models (LDMMs) 63]. which operate
in compressed latent spaces to enable high-resolution synthesis. Text-to-image models like DALL-E 2 (59) and Stable Diffusion (63] leverage large-scale multimodal pretraining to align textual prompts with visual concepts. while extensions like Palette [68] enable fine-graines control through spatial conditioning. These frameworks highlight the versatility of diffusion processes, with applications spanning artistic creation [69]. medical imaging, and beyond.
Diversity in Conditional Generation Balancing diversity and fidelity remains core challenge in conditional generation. GAN-based approaches address mode collapse via mode-seeking regular- ization [46] self-conditioned clustering (43) while diffusion models inherently trade off diversity and quality through their noise schedules Methods like ControlNet [86] enhance controllability by injecting spatial constraints (c.g.. edges, depth maps) into diffusion processes, whereas mutual information regularization (87) improves statistical dependency between latent codes and outputs. Adversarial training with semantic-guided negative sampling [9] further refines diversity GANs, while category- conssssettn objectives (27) optimize photorealism and variation simultaneously. These advances collectively enable richer, more varied outputs without sacrificing semantic alignment.
Generative Flow Networks (GFlowNets) 4 offer a paradigm for sampling compositional objects (e-g., molecules, graphs) with probabilities proportional to a reward function, prioritizing diversity over the single-mode convergence of RL. Recent work extends this to sequential domains via recurrent architectures (50), demonstrating their capacity to model temporal dependencies in tasks like program synthesis. Our work adapts GFlowNets to knowledge graph generation, using the Trajectory Balance objective (45) to align forward edge-addition policies with backward inference while preserving order- dependent semantics through RNNs. To our knowledge, this is the first integration of GFlowNets with latent diffusion models.
Conclusion
We introduced Rainbow, a novel conditional image generation framework that captures uncertainty and produces diverse, plausible images. Rainbow constructs latent graph in latent representation computation and leverages Generative Flow Networks to sample diverse trajectories over the graph, thereby enhancing the diversity of the condition latent representations and outputting diverse images that collectively interpret the input condition. Our experiments across natural and medical images not only demonstrate that Rainbow outperforms existing baselines in generating diverse, plausible images, but also highlight Rainbow's 's flexibility in adapting to any condition type. We discuss limitations and directions for future research in Appendix B
10

---

Acknowledgement
This work was supported in part by National Institutes of Health grant AG089169. This study was also supported by the Stanford School of Medicine Department of Psychiatry and Behavioral Sciences Jaswa Innovator Award, and the Stanford HAI Hoffman- Yee Award. Part of the data used in the preparation of this article was obtained from the Australian Imaging, Biomarker & Lifestyle Flagship Study of Ageing (AIBL) database (aibl.csrro.aug) with preprocessing and harmonization performed by Mind Data Hub at Stanford University. We thank Mohammad Abbasi from Stanford Translational (STAI) Medicine and Mental Health Lab, Stanford University, for his contribution would like thank to preprocessing all the 3D brain MRI datasets used in this work. Additionally,
Natural Sciences and Engincering Research Council of Canada, the Canadian Institute for Advanced Research (CIFAR) Artificial Intelligence Chairs program for grants; Mila Quebec AI Institute, Google Research, Calcul Quebec, and the Digital Research Alliance of Canada for providing grants and computing resources.

---

References
[1] Elwaleed Mustafa Alasar, Mohamet Awad Omer, and Ghada Abdel EI Erahman Sakin. Morpho- metric assessment of aging impact in cranial/ventrriless volumes and ct/mri imaging systems parameters. American Journal of Public Health Research, 7(4):157-160, 2019.
[2] Pietro Astolfi, Marlene Careil, Melissa Hall, Oscar Mañas, Matthew Muckley, Jakob Verbeek, Adriana Romero-Soriano, and Michal Drozdzal. Consistency-diversity--eee conditional image generative models. ArXiv, abs/2406. 10429, 2024.
pareto fronts of
[3] Lazar Atanackovic, Alexander Tong, Bo Wang, Leo J. Lee, Yoshua Bengio, and Jason Hartford. Dyngfn: towards bayesian inference of gene regulatory networks with gflownets. In Proceedings of the 37th International Conference on Neural Information Processing Systems, NIPS '23, Red Hook, NY, USA, 2023, Curran Associates Inc.
[4] Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J. Hu, Mo Tiwari, and Emmanuel Abu- Hussein. Flow network based generative models for non-iterative diverse candidate generation. Advances in Neural Information Processing Systems (NeurIPS). 34:27381-27394, 2021.
[5] Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J. Hu, Mo Tiwari, and Emmanuel Bengio. Gflownet foundations. J. Mach. Learn. Res., 24(1), March 2024.
[6] Aurelia Bustos, Antonio Pertusa, Jose-Maria Salinas, and Maria De La Iglesia- Vaya. Padchest: large chest x-ray image dataset with multi-label annotated reports. Medical image analysis, 66:101797, 2020.
[7] Holger Caesar, Jasper R. R. Uijlings, and Vittorio Ferrari. Coco-stuff: Thing and stuff classes in context. 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1 1209-1218 2016.
[8] Manuel Jorge Cardoso, Wenqi Li, Richard Brown, Nic Ma, Eric Kerfoot, Yiheng Wang, Benjamin Murrey, Andriy Myronenko, Can Zhao, Dong Yang. Vishwesh Nath, Yufan He, Ziyue Xu, Ali Hatamizadeh, Wentao Zhu, Yun Liu, Mingxin Zheng, Yucheng Tang, Isaac Yang, and Andrew Feng. Monai: An open-source framework for deep learning in healthcare, 11 2022.
[9] Miriam Cha, Youngjune Gwon, and H. T. Kung. Adversarial learning of semantic relevance in text-to-image synthesis. ArXiv, abs/181 11200888., 2018
[10] Minshuo Chen, Song Mei, Jianqing Fan, and Mengdi Wang. An overview of diffusion models: Applications, guided generation, statistical rates and optimization. ArXiv, abs/2404.07771, 2024.
[11] Joseph Paul Cohen, Joseph D. Viviano, Paul Bertin, Paul Morrison, Parsa Torabian, Matteo Guarrera, Matthew P Lungren, Akshay Chaudhari, Rupert Brooks, Mohammad Hashir, and Hadrien Bertrand. CiecXXRayViiion A library of chest X-ray datasets and models. In Medical Imaging with Deep Learning, 2022.
[12] Prafulla Dhariwal and Alex Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-877 8794, 2021.
[13] Mischa Dombrowski, Weitong Zhang, Sarah Cechnicka, Hadrien Reynaud, and Bernhard Kainz. Image generation diversity issues and how to tame them, 2024.
[14] Kathryn A Ellis, Ashley I Bush, David Darby, Daniela De Fazio, Jonathan Foster, Peter Hudson, Nicola T. Lautenschlager, Nat Lenzo, Ralph N. Martins, Paul Maruff, Colin Masters, Andrew Milner, Kerryn Pike, Christopher Rowe, Greg Savage, Cassandra Szoeke, Kevin Taddei, Victor Villemagne, Michael Woodward, and David Ames. The australian imaging, biomarkers and lifestyle (aibl) study of aging: methodology and baseline characteristics of 11 12 individuals recruited for a longitudinal study of alzheimer's disease. International Psychogeriatries, (((4))672-687, 2009.
12

---

[15] Anders M. Fjell, Lars T. Westlye, Inge Amlien, Thomas Espeseth, Ivar Reinvang, Naftali Raz, Ingrid Agartz, David H. Salat, Douglas N. Greve, Bruce Fischl, Anders M. Dale, and Kristine B. Walhovd. Minute effects of sex the aging brain: A multisample magnetic resonance imaging study of healthy aging and alzheimer's disease. Journal of Neuroscience, 29(27):877--8783, 2009.
[16] Dan Friedman and Adji B. Dieng. Vendi score: A diversity evaluation metric for machine learning. arXiv preprint arXiv:2306. 13177, 2023.
[17] Shohei Fujita, Susumu Mori, Kengo Onda, Shouhei Hanaoka, Yukihiro Nomura, Takahiro Nakao, Takeharu Yoshikawa, Hidemasa Takao, Naoto Hayashi, and Osamu Abe. Characteriza- tion of brain volume changes in aging individuals with normal cognition using serial magnetic resonance imaging. JAMA Network Open,
1110:2218815--211115555 06 2023.
[18] M. Gerstgrasser, R. Schaeffer, A. Dey, R. Rafailov, H. Sleight, J. Hughes, T. Korbak, R. Agrawal, D. Pai, A. Gromov, D. A. Roberts, D. L.. Yang, and D. L. Donoho. Is model collapse inevitable? breaking the curse of recursion by accumulating real and synthetic data. arXiv preprint arXiv:2404.01413, 2024.
[19] Ian Goodfellow, Jean Pouget-Abadig, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Y. Bengio. Generative adversarial networks. Advances in Neural Information Processing Systems, 3, 06 2014.
[20] Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Improved training of wasserstein gans, 2017.
[21] Ruben C Gur, Raquel E Gur, Warren Obrist, Beryl Skolnick, Martin Reivich, and Thomas Schlaepfer. Sex differences in brain aging: A quantitative magnetic resonance imaging study. Psychiatry Research, (1(:::12 135, 1995.
[22] Yaru Hao, Zewen Chi, Li Dong, and Furu Wei. Optimizing prompts for rext-to-image generation. In Neural Information Processing Systems, 2023.
[23] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. In Advances in Neural Information Processing Systems (NeurIPS), pages 6626 00666663,455 2017.
[24] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc.
[25] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages
Curran Associates, Inc., 2020.
[26] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-666 6851, 2020.
[27] Tao Hu, Chengjiang Long, and Chunxia Xiao. Crd-cgan: Category-consistent and relativistic constraints for diverse text-to-image generation. ArXiv, abs/2107..1351111 2021.
[28] Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4700-4708, 2017.
[29] Jeremy Irvin, Pranav Rajpurkar, Michael Ko, Yifan Yu, Silviana Ciurea-Ileus, Chris Chute, Henrik Marklund, Behzad Haghgoo, Robyn Ball, Katie Shpanskaya, et al. Chexpert: A large chest radiograph dataset with uncertainty labels and expert comparison. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 590-597, 2019.
[30] Jeremy Irvin, Pranav Rajpurkar, Michael Ko, Yifan Yu, Silviana Ciurea- Chris Chute, Henrik Marklund, Behzad Haghgoo, Robyn Ball, Katie Shpanskaya, Jayne Seekins, David A. Mong, Safwan S. Halabi, Jesse K. Sandberg, Ricky Jones, David B. Larson, Curtis P. Langlotz,
13

---

Bhavik N. Patel, Matthew P. Lungren, and Andrew Y. Ng. Chexpert: a large chest radiograph dataset with uncertainty labels and expert comparison. In Proceedings of the Thirty- Third AAAI Conference on Artifiial Intelligence and Thirty-First Innovative Applications of Artificial Intelligence Conference and Ninth AAAI Symposium on Educational Advances Artificial 1AAAA//19////11/9 AAAI Press, 2019.
Intelligence,
[31] Alistair EW Johnson, Tom J Pollard, Nathaniel R Greenbaum, Matthew Lungren, Chih- ying Deng, Yifan Peng, Zhiyong Lu, Roger G Mark, Seth J Berkowitz, and Steven Horng. Mimic-cxr-jpg, large publicly available database of labeled chest radiographs. arXiv preprint 2019.
[32] Minguk Kang and Jaesik Park. Contragan: contrastive learning for conditional image generation. In Proceedings of the 34th International Conference Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc.
[33] Nicole R. Karcher and Deanna M. Barch. The abed study: understanding the development of risk for mental and physical health outcomes. Neuropsychophammaclogg, 46(1):131-142, 2021.
[34] Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. ArXiv, (bs//2206000364, 2022.
[35] Diederik Kingma and Max Welling. Auto-encoding variational bayes, 2022.
[36] Amar Kumar, Anita Kriz, Mohammad Havaei, and Tal Arbel. Prism: High-resolution precise counterfactual medical image generation using language-guide stable diffusion. https//aazxivorr/ass/255....0 00196, 2025.
[37] Gihyun Kwon, Chihye Han, and Dae-shik Kim. Generation 3d brain mri using auto-encoding generative adversarial networks. In Dinggang Shen, Tianming Liu, Terry M. Peters, Lawrence H. Staib, Caroline Essert, Sean Zhou, Pew-Thian Yap, and Ali Khan, editors, Medical Image Computing and Computer Assisted Intervention- MICCAI 2019, pages 118-126, Cham, 2019. Springer International Publishing.
[38] Generative Al Lab. Understanding model collapse: A hidden threat in generative ai, 2024.
[39] Adrien LeCoz, Stéphane Herbin, and Faouzi Adjed. Explaining an image classifier with generative model conditioned by uncertainty, 2024.
[40] Minhyeok Lee and Junhee Seok. Estimation with uncertainty via conditional generative adversarial networks. Sensors, 21(18), 2021.
[41] Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft coco: Common objects in context. In ECCV, 2014.
[42] Dianbo Liu, Moksh Jain, Bonaventure Dossou, Qianli Shen, Salem Lahlou, Anirudh Goyal, Nikolay Malkin, Chris Emezue, Dinghuai Zhang, Nadhir Hassen, Xu Ji, Kenji Kawaguchi, and Yoshua Bengio. Gflowout: Dropout with generative flow networks, 2023.
[43] Steven Liu, Tongzhou Wang, David Bau, Jun- Yan Zhu, and Antonio Torralba. Diverse image generation via self-conditioned gans. In 2020 IEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), pages 14274-14283, 2020.
[44] Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment gflownets. In Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 5955-5967. Curran Associates, Inc., 2022.
[45] Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment gflownets. Advances in Neural Information Processing Systems (NeurIPS), 35:5955-5966 2022.
14

---

[46] Qi Mao, Hsin- Lee, Hung- Yu Tseng, S. Ma, and Ming-Hsuan Yang. Mode seeking generative adversarial networks for diverse image synthesis. In 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 1429-1437, 2019.
[47] Trang Nguyen, Amin Mansouri, Kanika Madan, Nguyen Duy Khuong, Kartik Ahuja, Dianbo Liu, and Yoshua Bengio. Reusable slotwise mechanisms. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[48] Trang Nguyen, Alexander Tong, Kanika Madan, Yoshua Bengio, and Dianbo Liu. Causal inference in gene regulatory networks with gflownet: Towards scalability in large systems, 2023.
[49] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In Marina Mcila and Tong Zhang, editors. Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 1162-8171. PMLR, 18-24 Jul 2021.
[50] Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Gflownet-em for sequential structured prediction. International Conference on Muchine Learning (ICML), 2023.
[51] Wei Peng, Tian Xia, Fabio De Sousa Ribeiro, Tomas Bosschieter, Ehsan Adeli, Qingyu Zhao, Ben Glocker, and Kilian M. Pohl. Latent 3d brain mri counterfactual, 2024.
[52] Fernando Pérez-Gareía, Sam Taylor, Pedro P Sanchez, Boris van Breugel, Daniel C Castro, Harshita Sharma, Valentina Salvatelli, Maria TA Wetscherek, Hannah Richardson, Matthew P Lungren, et al. Radedit: stress-testing biomedical vision models via diffusion image editing. In European Conference on Computer Vision, pages 358-376. Springer, 2024.
[53] Ronald C. Petersen, Paul S. Aisen, Laurel A. Beckett, Michael C. Donohue, Anthony C. Gamst, Danielle J. Harvey, Jr. Jack, Clifford R., William J. Jagust, Leslie M. Shaw, Arthur W. Toga, John Q. Trojanowski, and Michael W. Weiner. Alzheimer's discase neuroimaging initiative (adni): clinical characterization. Neurology, 74(3):201-209, 2010.
[54] Walter H. L. Pinaya, Petru-Daniel Tudosiu, Jessica Dafflon, Pedro da Costa, Virginia Fernan- dez, Parashkev Nachev, Sebastien Ourselin, and M. Jorge Cardoso. Brain imaging generation with latent diffusion models, 2022.
[55] E. Pulliam and A. B. Singleton. The parkinson's 's progression markers initiative (ppmi): Study design and protocol. Movement Disorders, 26(9); 1453-140,,, 2011.
[56] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision, 2021.
[57] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Phil Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. ICML, 2021.
[58] Marianne Rakic, Hallee E. Wong, Jose Javier Gonzalez Ortiz, Beth Cimini, John Guttag, and Adrian V. Dalca. Tyche: Stochastic in-context learning for medical image segmentation, 2024.
[59] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2004.06125, 2022.
[60] Royi Rassin, Aviv Slobodkin, Shauli Ravfogel, Yanai Elazar, and Yoav Goldberg. Grade: Quantifying sample diversity in text-to-image models, 2025.
[61] Naftali Raz, Ulman Lindenberger, Karen M Rodrigue, Kristen M Kennedy, Denise Head, Adrienne Williamson, Cheryl Dahle, Denis Gerstorf, and James D Acker. Aging, sexual dimorphism, and hemispheric asymmetry of the cerebral cortex: Replicability of regional differences in volume. Neurobiology of Aging, 25(3):377-399) 2004.
[62] Robin Rombach, A. Blattmann, Dominik Lorenz, Patrick Esser, and Björ Ommer. High- resolution image synthesis with latent diffusion models. 2022 IEFF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 10674-10685, 2021.
15

---

[63] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[64] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusson models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 10684-10695, June 2022.
[65] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmeniation, 2015.
[66] Mihaela Rosca, Balaji Lakshminarayanan, David Warde-Farley, and Shakir Mohamed. Varia- tional approaches for auto-eneoding generative adversarial networks, 2017.
[67]
Seyedmorteza Sadat, Jakob Buhmann, Derek Bradley, Otmar Hilliges, and Romann M. Weber. CADS: Unleashing the diversity of diffusion models through condition-annealed sampling. In The Twelfth International Conference on Learning Representations, 2024.
[68] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Raphael Gontijo Lopes, et al. Palette: Image-to Amage-to-image -mage diffusion models. arXiv preprint arXiv:21 1.11111..866. 2022.
[69] Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(4):4713 472))4266777 2022.
[70] Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Rudford, and Xi Chen. Improved techniques for training gans. NeurlPS, 2016.
[71] Rachael I. Scahill, Chris Frost, Rhian Jenkins, Jennifer L. Whitwell, Martin N. Rossor, and Nick C. Fox. A longitudinal study of brain volume changes in normal aging using serial registered magnetic resonance imaging. Archives of (eur)logy.60(7)99899-9 60(7)9999-994, 07 2003.
[72] Maximilian Seitzer. pytorch-fid FID Score for PyTorch. htttps://githut www.s/wiw.oommmmmiiiiiaa... pytorch-fid August 2020. Version 0.3.0.
[73] I. Shumailov, Z. Shumaylov, Y. Zhao, N. Papernot, R. Anderson, and Y. Gal. Ai models collapse when trained on recursively generated data. Nature, 625:1-5, 2024.
[74] Xue Song, Jiequan Cui, Hanwang Zhang, Jingjing Chen, Richang Hong, and Yugang Jiang. Doubly abductive counterfactual inference for text-based image editing. 2024 IEEL/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 9162-91179 2024.
[75] David Stap, Maurits J. R. Bleeker, Sarah Ibrahimi, and Maartje ter Hoeve. Conditional image generation and manipulation for user-specified content. ArXiv, abs/2005.0499 2020.
[76] G. Turinici. Diversity in deep generative models and generative ai. arXiv preprint 27Xiv2220.009573, 2023.
[77] David C. Van Essen, Stephen M. Smith, Deanna M. Barch, Timothy E.J. Behrens, Essa Yacoub, and Kamil Ugurbil. The wu-minn human connectome project: An overview. Neurolmage, 80:62-79 79, 2013. Mapping the Connectome.
[78] N. D. Volkow, G. F. Koob, R. T. Croyle, R. L. Balster, A. R. Childress, and C. R. Schuster. The abed study: Enhancing the understanding of adolescent brain and cognitive development. Neurolmage, 163:1-5, 2019.
[79] Xiaosong Wang, Yifan Peng, Le Lu, Zhiyong Lu, Mohammadhadi Bagheri, and Ronald Sum- mers. Chestx-ray8: Hospital-seale chest x-ray database and benchmarks on weakly-supervised classification and localization of common thorax diseases. In 2017 IEEE Conference Computer Vision and Pattern Recognition (CVPR), pages 3462-3471, 2017.
[80] Elizabeth Wilson, Mika Satomi, Alex McL.ean, Deva Schubert, and Juan Felipe Amaya Gonzalez. Embodied exploration of latent spaces and explainable ai, 2024.
16

---

[81] Jiang Xu, Shotai Kobayashi, Shuhei Yamaguchi, Ken-ichi lijima, Kazunori Okada, and Kazuya Yamashita. Gender effects on age-related changes in brain structure. American Journal of Neuroradiology, 21(1):112-118, 2000.
[82] Mengping Yang and Zhe Wang. Image synthesis under limited data: A survey and taxonomy, 2023.
[83] Peter Young, Alice Lai, Micah Hodosh, and Julia Hockenmaier. From image descriptions to visual denotations: New similarity metrics for semantic inference over event descriptions. TACL, 2:67-78, 2014.
[84] Taeyoung Yun, Dinghuai Zhang, Jinkyoo Park, and Ling Pan. Learning to sample effective and diverse prompts for text-to-image generation, 2025.
[85] Guiyu Zhang, Huan ang Gao, Zijian Jiang, Hao Zhao, and Zhedong Zheng. Ctrl-u: Robust conditional image generation via uncertainty-aware reward modeling, 2024.
[86] Lvmin Zhang and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. ArXiv, abs/2302.05543, 2023.
[87] Zhiwen Zuo, Ailin Li, Zhizhong Wang, Lei Zhao, Jianfeng Dong, Xun Wang, and Meng Wang. Statistics enhancement generative adversarial networks for diverse conditional image synthesis. IEEE Transactions on Circuits and Systems for Video Technology, 34:6167-6180, 2024.
17

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope, we have included baselines and experiments that address all claims made.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer:
[Yes]
Justification:
Yes, the paper discusses the limitations of the work performed by the authors,
have addressed it in
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
18

---

Answer: [Yes]
Justification: Yes, the paper provides mathematical formulation for GFlowNet and demon- strating how the algorithm can help Rainbow to enhance diversity.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Yes, the paper fully discloses all the information needed to reproduce the main experimental results, including parameter table and figures for hyperparameter tuning in the Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
19

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: Yes, the paper provides open access the data and specifies that the code will provided upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurlPS code and data submission guidelines (https://nips.cm
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
http:/.cccpubllc/guuuceee Codeeubmii ssi criionPPoiccy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes, the paper specifies all the training and test details, including data splits, hyperparameters, selection criteria, and type of optimizer, clearly in both the main text and the Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the
the paper to level of detail
that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Yes, the paper reports appropriate quantitative and qualitative results in the experiments to substantiate all claims.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
20

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes, the paper provides sufficient information the computer resources
needed to reproduce the experiments, including details about the GPU used and the training duration.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Yes, the research conducted the paper conforms, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Yes, discussed in Appendix A
Guidelines:
21

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negatíve societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have properly credited the creators and original owners of the code, pretrained models, and checkpoints used in our paper, and have explicitly mentioned and respected the license and terms of associated with these assets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
22

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcoee com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
We will contribute the source code and pretrained model from this work research community, and will provide detailed documentation of the training process and usage instructions alongside the assets.
Justification:
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/code/moodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not include any crowdsoureing experiments or research with human subjects. Instead, we have used the Flickr30k dataset, which contains images from human daily life, and we have provided relevant details in the dataset introduction.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not include any crowdsoureing experiments or research with human subjects.
23

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question:
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: Our method development does not involve LLM.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our L.LM policy for what should or should not be described.
24

---

A Broader Impacts
In this work, we presented Rainbow, a method that generates diverse images. These types of methods can play a crucial role in advancing AI models by enhancing robustness, reducing biases, and improving performance in various downstream tasks. By ensuring greater diversity in generated data, these models help mitigate biases that arise from underrepresented groups, building more equitable and generalizable Al systems. Additionally, generative models can be leveraged for counterfactual image generation, enabling the exploration of alternative scenarios for scientific and medical applications. However, while diversity is valuable, the realism of generated natural images raises concerns about misinformation and the misuse synthetic content. In medical applications, counterfactual images must be rigorously validated by domain experts before being used to train autonomous systems, as incorrect or misleading data could lead to severe clinical consequences. Expert validation ensures that these images maintain diagnostic fidelity, preserving patient safety and the integrity of Al-driven medical decision-makinn making.
B Limitations and Future Work
Although we introduced significant advancements in Rainbow and demonstrated experimental im- provements, some limitations suggest potential areas for future work. One limitation is the higher computational resources required for training Rainbow, as we update M trajectories in parallel. Future research could focus optimizing the training process alleviate these computational demands. Although working on latent graph reveals a high level of flexibility for Rainbow to be applicable any kind of condition, another area for improvement is the interpretation of latent graphs in Rainbow. This would aid in enhancing the latent graphs' interpretability more automatically.
In terms of directions for future exploration, that require diversity and the ability to manage uncertainty, such as text generation, recommendation systems, and decision-makinn latent graphs, which could capture uncertainties across multiple tasks. This expansion could lead to the creation of a foundational world model capable of addressing uncertainties in a variety of applications. One crucial analysis to add is anatomical plausibility tests and checks before making use of these images for any clinical application.
direction is to extend Rainbow other domains
tasks. Another promising direction is to scale Rainbow even larger
C Further Computation Details
C.1 Conditional Image Generation with Latent Diffusion
During Training The goal is to generate output image X given input image I R$ with shape Sj and a condition C. Initially, the input image is encoded into a latent representation zí (I) using an encoder E,. The latent code z E R represents the underlying structure of the image in a lower-dimensional latent space, where L, denotes the dimensions of this latent space.
The latent code z!f is then subjected to forward diffusion process, which iteratively adds noise over steps:
q(2) z = -N(z)) Varzi-1s (1 - ar)).,
where ae is a variance scheduling parameter that controls the amount of noise added at each step.
Concurrently, the condition C is encoded into its own latent representation c = (C) using an encoder Ec. This latent representation e E R&c encapsulates the conditioning information needed for the generation process, where Le denotes the dimensions of the conditional latent space.
During the reverse diffusion process, the objective is to reconstruct the original latent representation z from the noisy latent code z/1 guided by the condition latent representation c. This reverse diffusion is generally modeled using neural network cg, which predicts the noise added at each timestep of the diffusion process in a commonly used formulation:
Cit
col2l, (, o.n,
01
VI
Cr.
25

---

where n ~ (O, I) and are scaling factors for the noise at step t. Note that there are variations in diffusion modeling where different parameter schedules or noise prediction methods might be employed.
Finally, the reconstructed latent code 2 is decoded back into the image space produce the final output image X D(z)). where D is the decoder function that maps the latent representation back to the high-ding high-dimensiooal image space, yielding the generated image X.
Image Generation from Input Condition Along with the encoded latent representation from input condition C, a latent image žy is sampled from a Gaussian prior distribution, typically zr (0, I). The reverse diffusion process, conditioned on iteratively refines Zr until obtaining Zg-
This final latent code Zo is then decoded using D to produce the output image X.
Counterfactual Generation from Input Image and Condition The input image I encoded into its latent representation 2/ = (T). The counterfactual condition C' is encoded into c' (C") The latent code z is perturbed with noise and reverse diffusion, guided by c', is applied generate new latent code zo. Finally, this modified latent code Zo is decoded using D to obtain the counterfactual image X.
C.2 GelowNets Training
The Algorithm describes the training process Rainbow to iteratively construct diverse trajectories over the latent graph using GFlowNets and the Detailed Balance (DB) objective. The process is divided into three phases: initialization, iterative edge sampling, and loss computation.
C.2.1 Initialization Phase
The algorithm initializes with the input as the initial condition representation E SR Configuration includes the number of parallel graphs M, the number of nodes N, and the sparsity p. The total number of edges S is calculated S = (1 - is initialized with a special starting edge index O. Two masks forskrd maskk E R and backward mask E MR- created to enforce valid transitions during sampling. These masks control which edges available to be reached in the forward and backward paths. The log- likelihhoood difference tensor II_dif (S++))xM is initialized to track state transitions for the DB loss.
A set of M trajectories T E R
C.2.2 Iterative Edge Sampling Phase
For each step i (1...., "i....,S,, S), the algorithm performs the following operations:
Graph Encoding and Forward/Backwart Probabilities: The current trajectories T.L:M are encoded into latent representations repy E IR xh,, using the graph decoder Q These concate- nated with the repeated condition form combined representation rep R MRMA The forward predictor MLPFW computes log-forwe log-forward probabilities lor_forward E RMxx and log-flow values flo x1 masked to exclude already added edges. The backward predictor MLP Bw computes log- log-backward probabilities bog_backward E TR. M xn, masked to restrict invalid backward transitions.
Edge Sampling and Mask Updates: Edges are sampled from log_forward using a multinomial distribution, and the trajectories T,EM are updated with the new edges. The log- lekelloldood differences the sampled edges. For > 1,
[ accumulate the log-flow and log-forward probabilities 11 fiff 1] is updated with backward probabilities to ensure transition consistency. The masks forward mask and backward_maad are dynamically adjusted to reflect added edges.
C.2.3 Finalization and Loss Computation Phase
After S steps, the completed trajectories TI:M are decoded into condition representations 2I:M using the graph decoder QD, blended with the original condition e via the factor as in Equation The diffusion denoising process generates rewards log rewarrdardd E Mx1 which are incorporated into II difff The DB loss LDB is computed as the mean squared error of f._diff.
LDB
26

---

Algorithm Rainbow Training Pipeline
1: Input: Intial condition representation R
2: M: Number of graphs that are computed parallelly
N: Number of nodes in graph
4: P: Sparsity
5: n N(N -) 0.5 (1 - p): Total number of edges in undirected graphs with nodes
6: S: Number of edges in the final graph
7:
8: GFlowNets Architecture
9: hig he: dimension the encoded graph yi and the encoded condition c, respectively
10: 1== bi + ha: GFlowNets hidden size
11: Q : hh
D Condition encoder model
12: Qá :. - DR*
> Instant graph decoder model (Used during edges sampling) 13: MLPFw R + R D Forward probability and flow predictor, dimensions for forward
probability, and the last dimension for state flow
14: MLP R R
Backward probability predictor
15:
16: GelowNets computation flow for transforming state i to state
17: Inputs: 71:M E
Mxx d' hR
18: repy = Q L:M
E Mxh,,
> Encoding graphs from state i-1
19: rep = concatenaterrepp: d' c'rrepaa(MM))
P Concatenate encoded graphs and the repeated
encoded condition
20: pred =
b Preparing for forward probability and flow prediction 21: log for_fordrdd = log softmax(pred[::>
wwed=MLPPpg MMMFFFFw (rep)
-1 forward_mnee mask inf) > Getting forward
probability that with added edges excluded
22: log_flow = pred(: -1
Getting log flow b Getting
23: log_backward = log_ trg_softmax((co [MLP BBw(pp)) backward_mask * inf))
backward probability among added edges
24: Outputs: log_forward E Radx bog_backwardd E RMxx flow E RMx1
25:
26: Step O. Initialization
27: T1:bd <- 0,11
Empty M undirected graphs with starting special edge index backward mask +- 1Mxx First state for the initial state and S states for adding S edges,
28: forward forward_mask + OMxnn
b Initialize masks
11 fiff <- ((S+))xM
11 [_diff(() not touched
30:
Graphs generator training pipeline using GFlowNets and Detail-balance loss
32: c' = "-Q (c)
for in S
Loops of execution
34:
log lor_forward, log backward, log_flow = QGEN ITIM (1
35:
edges +- multiommialloy_forrraaa E MR Action as edges are added into trajectories
& Sampling edges as actions
36:
T + edges
37:
38:
Updating flows
39:
U1 diffli flow + log loog._.forurrd.gghhhetettttnn
40:
if then
41:
U1 [idiffli 1] log flo_ log_
42:
end
43:
forward mask actions backward, maskk - actions
Updating the forward mask > Updating the backward mask Completing the last turn
44:
if
S then
46:
M-M
I:M
47:
= Dn (TLM) * + c* (1-7) & Decode done graphs into latent condition shape Performing Diffusion denoising conditioned on al:M and get log neward e R as in
al.M
48:
Equation 8
49:
[1di diff(s - = Nog_rewards
50:
end if
51: end for
52: LDB [l_dif [ (e.maa)()
Optimization step
27

---

D
Experiment Setup
Table 2
indicates hyperparameters used in this work for all experiments.
D.1 Hyper-parameter
Parameter Name
Natural Images
3D Brain MRIs
Chest X-rays
Learning Rate
le-5
25e-7
le-5
Pretrained-LDM epochs
80
Rainbow epochs
20
5
Batch Size
1
8
at
1 (Freeze Unet)
0.2
(Freeze Unet)
0.8
Training image shape
3 x 256 x 256
1 160 X 192 x 176
512 x 512
Training condition type
Text prompt
Age and binary sexes
Text prompt
Training condition shape
Dynamic length
2 dimensions
Dynamic length
Latent image shape
x 64
32 40 48 x 44
64 64
Latent condition shape Se
77 x 1024
256
77 x 1024
Inference image shape
512 x 512
160 x 192 x 176
512 x 512
Encoder - Decoder
VAE
VAE nodes
VAE
Graph Size N
20 nodes
20 nodes
Number of Graphs M
40
8
10
Sparsity P
0.83
0.70
0.82
Num. Edges S
32
8
33
Use RNN
Yes
Yes
Yes
Edge Embedding dim
512
128
512
Latent dimhhg = be
1024
1024
1024
Blending factor 7
0.5
0.5
0.5
Table 2: Model Architecture and Parameter Indications
D.2 Evaluation Metrics
Natural Images
To evaluate Rainbow on natural images, we use Inception Score (IS) and Inception Vendi Score (VS). For both metrics, we use the feature extraction model from pre-trained Pytorch Inception-v3.
3D Brain MRIs We use Fréchet Inception Distance 23 (FID) evaluate the feature quality of synthetic images. FID is widely adopted metric for assessing the similarity between the feature distributions of real and generated images. images should exhibit feature distributions similar to those of real images when passed through a pre-trained neural network. FID is computed by calculating the Fréchet distance between two multivariate Gaussian distributions fitted to the feature vectors of real and generated images.
based on the premise that high-quality synthetic
For feature extraction, complex 3D structures and patterns inherent in volumetric data. The model trained on diverse set of 23 medical imaging datasets, enabling it generalize effectively across various medical image types, such as MRI and CT scans. The feature vectors are extracted from the final convolutional layer (conv seg). with a dimensionality images, making it ideal for evaluating perceptual similarity between the real and synthetic images.
3D ResNet50, which is particularly well-suited for capturing the
2048. This layer captures high-level semantic features of the
Lower FID values indicate that the distributions of real and generated images more closely aligned, suggesting that the synthetic images are of higher quality.
To evaluate the faithfulness of capturing these conditions in our generated samples, we train CNN- based age regressor and sex classifier on the real data (i.e., the same data that is used to train our proposed model and all baselines). The architectures for these CNNs can essentially be seen as the
encoder half typical UNet (65), consisting of 4 downsampling levels and convolutional blocks per level, with each block consisting of convolution layer, batchnorm layer, and ReLU layer. The
28

---

age regressor is trained to minimize MSE loss, and the sex classifier is trained to minimize binary cross-entropy.
Chest X-rays To evaluate Rainbow on chest X-rays, we use FID and Vendi Score similar to the previous modalities. For feature extraction,
pre-trained DenseNet-121 28 model from the TorchXray Vision library [[] which is trained on multiple chest X-ray datasets such as CheXpert [29], NIH-CXR [79,, PadChest (6), and MIMIC-CXR [31]. The feature vectors used for calculating the metrics are extracted from the last layer (before the classifier head) with dimensionality of 1024.
D.3 Baselines
Natural Images We use 4 baselines: Stable Diffusion v2-1-base (SD2-1) Stable Diffusion v3-medium CADS [67], which is sampling strategy that anneals the conditioning signal by adding scheduled, monotonically decreasing Gaussian noise to the conditioning vector during inference to balance diversity and condition alignment, and PAG [84), novel approach that frames prompt adaptation as probabilistic inference problem utilizing GelowNet for the generation of diverse, high-quality prompts.
Brain MRIs The GAN based model is from 37 This 3D GAN model addresses both image blurriness and mode collapse problems by leveraging a a-GAN [66] that combines the advantages of Variational Auto-Encoder (VAE) and GAN with an additional code discriminator network. The model also uses the Wasserstein GAN with Gradient Penalty (WGAN-GP) loss (20 to lower the training instability.
The standard LDM is based on 54 does not leverage any latent graphs.
63
It can be viewed as identical to Rainbow, except that it
Chest X-rays We consider two chest X-ray baseline models. The first model is RadEdit 552, a latent diffusion model developed by Microsoft Health. This model is trained on 487,680 frontal view chest X-rays of multiple datasets such as MIMIC-CXR [31, NIH-CXR [79]. and CheXpert [29]. The second baseline fine-tuned Stable Diffusion v1.5 model 33 the CheXpert [29 dataset.
D.4 Training Time and Computation Sources
For general- ceneral-domain experiments, training was conducted on single NVIDIA H100-80GB GPU, completing in 12 hours.
The brain MRI experiment utilized 4 NVIDIA H100-80GB GPUs paired with 32 CPU cores over days pretraining, followed by continued Rainbow training the same hardware configuration for an additional 24 hours.
The chest X-ray experiments utilized 4 NVIDIA A100-80GB GPUS paired with 24 CPU cores. This model was finetuned on 512 x 512 chest X-rays and prompt pairs for 15 hours with an overall batch-size of 8.
E
Addition Experiment Results
E.1
Numeric Evaluation Results
Table
presents numeric results for Figure 4a.
Table presents numeric results for Figure 4p.
E.2 Chest X-Ray Ablation Studies
Figure shows the results for the full ablation study across three parameters N, M, and p for the chest X-ray model. The ablation setup is same as discussed in the main text, where each model is trained on single GPU for 16500 steps. can be observed that the models with more diversity (higher VS) have higher sparsity values, and the models with lower FID have lower sparsity. However,
base
medium
29

---

Model name
18
CLIP
Pixcel VS f
Inception VS
SD3.5
7.37
30.29
3.75
25.09
SD2-1
8.49
30.27
3,74
25.63
CADS
9.46
30.27
3.88
27.85
PAG
9.93
30.21
3,90
26.92
DDIM
8.44
30.28
3.91
26.62
Rainbow
10.45
30.32
3.94
28.90
Table Quantitative comparison natural-image experiment. An upward arrow "f" indicates that higher value corresponds to better performance. Rainbow consistently outperforms SD baselines diversity (higher Vendi scores), image quality (higher score), and prompt context delivery (higher or comparable CL.IP score)
Model name
FID
CLIP
VST
RadE
10.28
33.58
6.12
SD1.5
1.32
33.51
6.44
Rainbow
1.27
33,45
7.16
Table 4: Quantitative comparison of Chest X-Ray experiment. An upward arrow "1" indicates that higher value corresponds better performance. Rainbow is increasing the VS and achieves higher diversity while improving the generation quality by lowering FID. All models have close CLIP similarity scores.
this does not imply that sparsity parameters are the single factor affecting quality and diversity. In general,
squares and triangles in the upper right section of the plot, showing that M 10 and M 15 are overall better than the M = models. Furthermore, blue and orange are also more prevalent in the upper-right section of the plot indicating that models with N 25 do not perform as good as models with fewer nodes.
E.3 Full of 40 Generations for General-domain Experiments
Figure 19 provides the 40 generations that support Figure X in the main text.
Supporting Figure in the main text, Figure 12 present full 40 generations with two prompts by SD3.5, Figure 13 present whole 40 generation with two prompts by SD2-1, Figure [4 present full 40
Color: N
0.70
N=14
N=20
Quality (FID (t)
0.65
N=25
0.60
Marker: M
Selected Model
M-5
M=10
0.55
M=15
0.50
6.6
6.7
Diversity (VSt)
Figure 11: Ablation studies on the Chest X-Ray model by varying the number of nodes N, the number of graphs M, and the sparsity parameters p. Different colors show the values (blue=15, orange=20. green=25). Different shapes show the M values (circle=5, square=10, triangle=15). and the size of the shapes shows the sparsity (biggest=0.8g middle=0.88 sm.4lest=0.64), The final selected model is shown using the arrow labeled as 'Selected Model'.
30

---

present full 40 generation with two prompts by PAG, and Figure present whole 40 generation with two prompts by Rainbow.
generation with two prompts by CADS, Figure 15
Supporting Figure in the main text, Figure presents full 40 generations with sets of seasonal edges.
Figure 20 provides qualitative and quantitative comparison of generated images between models with additional prompts.
Supporting Figure 6 edges.
the main text, Figure 22 presents full 40 generations with sets of seasonal
"Sunset scene with mountain" by SD3.5
(b) "Sunset with mountain in specific season" by SD3.5
Figure 12: SD3.5 generates repeated layouts and objects in both prompts, producing unclear seasons or dominated winter aesthetics in the second prompt.
E.4 Addition prompts for Text-to-image task
Figure 20 provides a comparison of diverse natural-image generation between models 4 additional prompts.
E.5
3D Brain MRI Fidelity Analysis
As visualized in Figure 27 Rainbow maintains sharper anatomical details across all age groups while avoiding artifacts. LDM intoduucs subtle distortions, particularly age-sensitive regions:
Young (14 vs.16): Rainbow preserves fine textures like developing white matter tracts. LDM's output shows mildly blurred cortical layers.
31

---

"Sunset scene with mountain" by SD2-1
(b) "Sunset with mountain in specific season" by SD2-1
Figure 13: SD2-1 generates diverse layouts and objects but heavily prioritizes specific style in the first prompt, while producing ambiguous or spring porinn-dominated seasons in the second prompt.
Middle-aged (45 vs.44): Rainbow retains small vessels and tissue gradients. LDM exhibits "smeared" edges around ventricles.
Elderly (75 vs.72): Rainbow realistically renders age-related atrophy (e.g., widened sulci). LDM generates unnaturally smooth brain surfaces, masking thinning cortex.
E.6 Qualitative Results for 3D MRI Counterfactual Generation Task
For brain MRI, we perform a counterfactual generation task. Given the factual image and condition. the task is to generate a counterfactual MRI based on counterfactual condition age or sex. T the numerical evaluation of the classification of sex and age based on the generated counterfactuals. Rainbow achieves higher performance, which underscores its effectiveness. Figure 29 compares the counterfactual generations. Both Rainbow and LDM generate correct patterns, such as smaller ventricles at younger age 10)) evident from the red regions in the difference plot, and larger ventricles and cortex at an older age (do(ago=80)) shown by the green regions. Additionally, sex conversion from male to female shows smaller ventricles and partially larger cortex regions. However, the baseline LDM exhibits some artifacts. These findings are consistent with previous studies the effects of age and
MRI characteristics
15
21 81
32

---

"Sunset scene with mountain" by CADS
(b) "Sunset with mountain in specific season" by CADS
Figure 14: CADS generates diverse layouts and objects, but prioritizes specific art style in the first prompt. While CADS produces more defined seasonal characteristics in the second prompt, including generations with clear spring environments, some outputs retain ambiguity or disproportionately favor winter season.
33

---

"Sunset with mountain" by PAG
(b) "Sunset scene with mountain in a specific season" by PAG
Figure 15: Generations by PAG
34

---

(a) "Sunset scene with mountain" by Rainhow
(b) "Sunset scene with mountain in specific season" by Rainbow
Figure 16: Rainbow generates diverse layouts, objects, and seasonal environments with high compo- sitional flexibility, achieving balanced variation across spring, summer, autumn, and winter visual details.
35

---

"Sunset scene with mountain" + Spring edges
(b) "Sunset scene with mountain" + Summer edges
(c) "Sunset scene with mountain" " Autumn edges
(d) "Sunset scene with mountain" + Winter edges
Figure 17: Comparison of images generated by Rainbow with manipulated trajectories with 10 extra edges specific to each season. We observe that objects and layouts are consistent between the images, with clear season-specific details such as colorful spring flowers and green grass, hot-toned sky for summer, autumn yellow leaves, and snow for winter.
36

---

CFG
IS t VS
"Sunset scene mountain"
7.5
1.44
3.80
SD3.5
5.0
1.42
3.69
3.0
1.43
3.78
7.5
1.96
4.29
SD2-1
5.0
1.92
4.53
3.0
2.07
5.24
"Sunset scene mountain specific season"
7.5
1.20
3.06
SD3.5
5.0
1.40
3.54
3.0
1.46
3.85
7.5
1.88
4.50
SD2-1
5.0
1.85
4.69
3.0
1.69
5.32
Figure 18: Comparison of the effect of classifier-free guidance (CFG) scale SD baselines across different prompts. For each prompt, 40 images are generated to compute metrics, and 10 are displayed. Metrics include the Inception Score (IS) for image quality and the Inception Vendi Score (Vendi) for diversity assessment. Across models, CFG scales, and prompts, there is no consistent pattern suggesting that reducing CFG yields more diverse images. Specifically, we can observe that decreasing the CFG only affects the detail level of the objects (more realistic, sharper) without influencing the context or object choices. Therefore, reducing CFG does not improve diversity in context, which is the major strength of the proposed Rainbow.
37

---

(a) DAC SD2-1
(b) DAC Rainbow
Figure 19: Diversity in mmage-editing task. Given the original image of a cat (left) and the editing prompt "A cat wearing wool cap", the baseline mostly produces white caps, while Rainbow captures diverse color choices of caps.
38

---

Vendi
CLIP
forest with animals"
SD3.5
3.60
33.33
SD2-1
4.77
33.59
CADS
5.05
33.42
PAG
4.84
31.76
Rainbow
6.09
33.66
"Photo of San Francisco"
SD3.5
3.37
30.08
SD2-1
4.77
30.37
CADS
4.90
29.57
PAG
4.64
29.56
Rainbow
5.46
29.49
Photo of brownies and lemonade'
SD3.5
3.96
34.88
SD2-1
4.95
35.27
CADS
5.50
34.90
PAG
4.63
30.82
Rainbow
5.79
33.39
grayscale tiger"
SD3.5
32.57
SD2-1
1.94
32.31
CADS
2.23
32.63
PAG
2.65
34.87
Rainbow
32.76
Figure 20: Diversity and quality across Rainbow and SD baselines. Each prompt generates 40 images per model. Notably, Rainbow captures seasonal variations in the "forest with animals" prompt and multiple perspectives of San Francisco beyond seas and buildings in SD baselines. For "brownies and lemonade". Rainbow generates diverse relevant objects on the table, and for "tiger". provides multiple views, showcasing enhanced versatility and realism. The SD baselines tend to generate similar layouts and contexts across prompts, demonstrating less diversity.
39

---

Prompt: "A glass mug dropped the ground"
(b) Prompt: "A cow in a desert during a vibrant sunset"
Figure 21: Examples of cases that Rainbow perform not good.
40

---

(a) Generations by RadEdit
(b) Generations by SD1.5
(c) Generations by Rainbow
Figure 22: Generations with the prompt "Chest X-ray showing support devices." Rainbow generates more diverse set of medical devices compared to the SD model and RadEdit, while maintaining image quality comparable to the SD model.
41

---

(a) Generations by RadEdit
(b) Generations by SD1.5
(c) Generations by Rainbow
Figure 23: Generations with the prompt "Chest X-ray showing Cardiomegaly". Rainbow shows more diversity in the anatomy of the generated chest X-ray, while SD mostly generates left and right lungs that are similar to each other. Furthermore, Rainbow provides diversity in the location of the generated support devices.
42

---

(a) Generations by RadEdit
(b) Generations by SD1.5
(c) Generations by Rainbow
Figure 24: Generations with the prompt "Chest X-ray with no significant findings". Rainbow shows more diversity in the anatomy of the generated chest X-ray, while SD has less variation in the anatomical structure of the lungs.
43

---

findings Xnyy showing support
X-my showing atelectads
showing
Node Indox
Nede Indox
02
18 16
14 10
10 16 Node Index
Node Index
Node Index
Node Index
066
004
Node Indox
Indox
Node Indox
20 12 32 10
14
Node
Node Index
Node
Figure 25: Heatmaps of edges across different prompts. In the first row, we present heatmaps of edges across 40 trajectories per prompt, with corresponding prompts labeled. the bottom row, we illustrate the difference in edge distribution compared to the baseline prompt "Chest X-ray with no significant findings". We analyzed and extracted the most frequently added edges and the 10 most frequently removed edges for the difference heatmap. Our observations reveal that (1) some edges consistently appear in most prompts with a significant proportion, and (2) it is evident that certain edges are representative of specific contexts. Particularly, in the difference heatmap, we can see certain edges are added with high proportion, approximately 60%.
(a) Base prompt
Base prompt suport devices edges
Figure 26: Comparison of images generated by Rainbow with the base prompt "Chest X-ray with no specific findings" and with "support device" edges. By adding the edges corresponding the entity "support devices" to the latent graph, we' re able to modify the images with support devices. This demonstrates that Rainbow latent graphs encode structured and interpretable knowledge, and that manipulating these graphs enables fine-grained, concept-specific image editing.
44

---

Samples generated by Rainbow
Age 14, Sex 0
Age 44, Sex 0
Age 77, Sex
Samples generated by LDM
Age 14, Sex 0
Age 44, Sex
Age 77, Sex
Figure 27: Brain MRI fidelity comparison
45

---

(a) Actual samples of male patients in age range from to
(b) Generated MRI by Rainbow
(c) Generated MRT by baseline L.M
Figure 28: Multiple generations per input condition comparison. The figure displays MRI images showcasing 8 samples generated from single input condition. Given the input condition of 65-year-old male, both Rainbow and the baselime LDM can generate plausible MRI images. However, compared actual samples from males aged 60 65, is evident that Rainbow captures a greater diversity in details, such as varying ventricle sizes. In contrast, the baseline LDM tends to generate images with consistently similar ventricle sizes, demonstrating less diversity.
46

---

Original Age (Male)
Original
Ag 10
DI of
Age 20
Dif
Age T0
OIF
Age 00
DM of
Older Ages
Onigrall
Age
Age
Age
60
Samole
10
Figure 29: Counterfactual 3D Brain MRIs by Rainbow and the difference between the original image and generated counterfactuals
Original information: Age 44 (Mile)
Younger Ages Spe = 28
Ape 64
Sample
02e 10 vs 44
Samrole
Samole
20 v3 44
Srmpl
oye - 60 45
Oider Ages
Ape
Trmpds
60 vo 44
Sample 3
EPR = 28
Sanok
t 80 44
Iraple
- 50 46
-1
Figure 30: Counterfactual 3D Brain MRIs by baseline LDM and the difference between the original image and generated counterfactuals
47


Reviews:
[{"review_id": "twlMudhpz8", "reviewer": "Reviewer_VBdx", "summary": "This paper introduces a framework named Rainbow, designed to enhance the diversity of images produced by conditional generative models. The core problem it addresses is that a single input condition or prompt often has inherent uncertainty and can correspond to multiple plausible visual outputs, which standard generative models fail to capture, often producing repetitive results.\n\nRainbow is proposed as a model-agnostic layer that can be integrated with any pretrained conditional generative model, such as a Latent Diffusion Model (LDM), without altering its core components. The method operates in a multi-step process:\n\n* Initial Representation: An input condition (e.g., a text prompt or a set of attributes) is first encoded into an initial latent representation using a standard condition encoder.\n\n* Latent Graph Discovery with GFlowNets: The framework's key innovation is the use of Generative Flow Networks (GFlowNets). A GFlowNet is trained to sample a set of diverse trajectories over a latent graph structure. Each trajectory represents a different plausible interpretation of the uncertainty contained within the initial condition.\n\n* Diverse Condition Generation: These diverse graph trajectories are then decoded into a set of distinct condition representations. Each new representation captures a unique aspect of the input condition's ambiguity.\n\n* Diverse Image Generation: Finally, each of these diverse condition representations is fed into the frozen, pretrained generative model (e.g., an LDM) to produce a corresponding output image. The result is a set of varied images from a single initial prompt.\n\nThe authors evaluate Rainbow on both natural image (Flickr30k) and medical imaging (3D Brain MRI, Chest X-ray) datasets. The results are compared against baselines like Stable Diffusion and other diversity-enhancing methods using metrics for diversity (Vendi Score) and quality (Inception Score, FID). The experiments aim to show that Rainbow can generate a more diverse set of high-quality images that better reflect the inherent ambiguity of the input conditions.", "questions": "1. **Reward design** – The current GFlowNet reward is the exponential of the final denoising MSE, implicitly tied to modes the frozen LDM already renders well. Have you experimented with composite or alternative rewards (e.g., CLIP‐based semantic distances, pairwise contrastive terms) to reduce this potential bias, and if so how did they affect diversity and quality?\n\n2. **Scaling the latent graph** – Reported models use up to *N = 20* nodes and *M = 40* trajectories. What practical or theoretical bottlenecks did you encounter when trying larger *N* or *M*, and do you foresee algorithmic modifications (e.g., sparse batching, trajectory pruning) that would let Rainbow scale to higher‐resolution video or multimodal inputs without prohibitive memory growth?\n\n3. **Automated graph interpretability** – Qualitative edits (e.g., “season” edges, “support‐device” edges) suggest that some latent edges align with human concepts, yet Appendix C lists graph interpretation as future work. Could you outline planned or preliminary methods (attention tracing, concept activation maps, etc.) for systematically mapping edges to semantic attributes?\n\n4. **Evaluation beyond automated metrics** – Diversity and fidelity are assessed with Vendi, IS, FID, and CLIP similarity. Are there plans to incorporate human studies or downstream‐task evaluations (e.g., radiologist scoring, prompt relevance surveys) to validate that Rainbow captures semantically meaningful uncertainty rather than pixel‐level variation?\n\n5. **Clinical deployment pathway** – For the 3-D MRI and chest-X-ray settings, Appendix C notes the need for anatomical-plausibility checks before use in clinical pipelines. Could you elaborate on specific validation protocols (e.g., expert grading, quantitative anatomical metrics) and on whether constraints or priors could be integrated into Rainbow itself to enforce medical realism during sampling?", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n*Methodological Novelty: The paper introduces a novel framework by coupling Generative Flow Networks (GFlowNets) with latent diffusion models. This approach, which samples multiple graph trajectories to capture multimodal uncertainty, is a notable contribution to the field of conditional image generation.\n\n* Consistent Performance Across Domains: The method demonstrates strong and consistent empirical gains across three distinct domains: natural images, 3D brain MRIs, and 2D chest X-rays. The quantitative improvements in both diversity (Vendi Score) and fidelity (IS, FID) over strong baselines validate the effectiveness of the approach.\n\n* Latent Graph Controllability: A significant strength is the demonstration that the learned latent graphs are not just a passive component but can be actively manipulated. The experiments showing that injecting \"season-specific\" or \"support-device\" edges leads to targeted semantic edits suggest that the graph encodes meaningful, editable concepts.\n\n* Transparent Experimental Design: The paper provides a high degree of transparency, detailing hyperparameters and outlining the training algorithm with pseudocode. The inclusion of ablation studies on key parameters like N, M, and ρ further enhances the reproducibility and clarity of the work.\n\n## Weaknesses\n\n* Opaque Edge Semantics: While the latent graph is shown to be controllable, its components lack inherent interpretability. The mapping from graph edges to human-readable concepts is not automatic and must be inferred through post-hoc analysis. This \"black box\" nature of the graph's semantics is a core limitation, as it makes it difficult to understand or directly control the specific type of diversity being generated.\n\n\n* Indirect Reward Signal for GFlowNet: The GFlowNet is trained to generate diverse graphs based on a reward signal derived from the final denoising error of the diffusion model. This is a highly indirect objective. The GFlowNet is not explicitly guided to find graphs that correspond to semantically meaningful variations, but rather graphs that result in low-error reconstructions. This could potentially bias the model toward modes of variation that are easy for the LDM to render, rather than ones that fully capture the true underlying uncertainty in the prompt.", "was_revised": true, "final_justification": "My initial score reflected concerns about the indirect reward signal and the lack of human-centric evaluation. The rebuttal addressed these issues by providing new, targeted experiments, including a test of an alternative reward function and a human study with expert evaluators. These additions have resolved my main reservations and strengthened the paper's claims. My score has been raised to recommend acceptance.", "rebuttal": "1. **Reward design**: The current GFlowNet reward is the exponential of the final denoising MSE, implicitly tied to modes the frozen LDM already renders well. Have you experimented with composite or alternative rewards (e.g., CLIP‐based semantic distances, pairwise contrastive terms) to reduce this potential bias, and if so how did they affect diversity and quality?\n\n\n\n\nThank you for your comment regarding reward design.\n\nWe acknowledge that incorporating diversity and quality evaluation into the reward is a promissing approach. However, utilizing metrics such as VS, IS, or CLIP poses challenges, as these metrics require evaluating generated images using full-step diffusion, while our training typically employs only one diffusion step. Implementing training with full diffusion steps would significantly increase the time and memory requirements, making it challenging to effectively incorporate these rewards.\n\nWe have conducted your suggestion on the pair-wise contrast (PWC) reward and provide a comparison between the original Rainbow model, which uses a reward based on the exponential of the Mean Squared Error (MSE) of the noise, and an updated version that incorporates pair-wise contrast into the reward function. We calculate the reward based on the exponential of two components: the negative MSE of the noises and the negative pair-wise MSE of the decoded trajectories. \n\nQuantitatively, compared to the original Rainbow model, the version with the additional PWC reward achieves comparable image quality as indicated by the IS score, higher diversity as reflected in the VS score, but shows a significant decrease in CLIP score, as represented in the table below. It's important to note that both VS and IS scores are computed without considering the text prompt.\n\n\nIntuitively, the introduction of the PWC reward aims to encourage greater diversity among the decoded trajectories. However, this emphasis on rewarding larger distances may lead to confusion in the model, as it has to balance the diversity of the generated outputs while maintaining contextual similarity.\n\n| | IS | VS | CLIP | Avg. inference time/image |\n|----------           |----------|----------|----------|----|\n| SD2-1                  | 8.49    | 25.63   | 30.27   |  0.39s  |\n| **Rainbow - original** | **10.45**| **28.90** | **30.32**  | 0.40s   | \n| Rainbow - PWC       | 9.43   | 29.81   | 27.47   |  0.40s  |\n\n\n\nQualitatively, we observe two key trends: first, the pairwise MSE of decoded trajectories under the PWC reward increases significantly compared to the original version. Second, images generated using the PWC approach sometimes include irrelevant objects or omit concepts mentioned in the input prompt. For instance, when prompted with \"Sunset scene with mountains,\" the generated images may lack the mountains and sky, featuring only greenery or even introducing unrelated elements like an office.\n\nWhile we agree that integrating PWC is a promising approach, we recognize that further investigation and fine-tuning are necessary to effectively utilize and control this method.\n\n\n\n\n\n\nThank you for your insightful suggestion!\n\n\n-------\n\n2. **Scaling the latent graph** – Reported models use up to  _N = 20_  nodes and  _M = 40_  trajectories. What practical or theoretical bottlenecks did you encounter when trying larger  _N_  or  _M_, and do you foresee algorithmic modifications (e.g., sparse batching, trajectory pruning) that would let Rainbow scale to higher‐resolution video or multimodal inputs without prohibitive memory growth?\n\n\n\n**Regarding hyperparameter selection for GFlowNet**, specifically for the parameters M, N, and ρ, our choices were based on several factors, including computational resources, the need for diversity, task plausibility, and results from fine-tuning.\nWe adjusted the value of N by using a higher value for natural images to allow greater freedom in diversity and concept choices, while for medical datasets, we opted for a lower N value to balance diversity and ensure that the images still follow a domain-specific template. We maintained a relatively high graph sparsity, above 0.6; higher sparsity results in fewer sampled edges. This strategy enabled us to maximize the number of graphs given our available computational resources.\nWe fine-tuned the combinations of M, N, and ρ for optimal performance in terms of both diversity and image quality. Notably, we observed that all combinations produced similar CLIP scores for prompt delivery evaluation. The fine-tuning results for chest X-rays are provided in Figure 11 in the appendix (referred to in line 313).\n\n**Regarding scaling to higher-resolution video or multimodal inputs**, **the inference time is nearly the same between Rainbow and the base model such as SD2-1**, as the burden is on the generation part, such as the UNet. The primary difference in inference time lies in the additional time required for Rainbow to decompose the input into multiple representations, rather than in the image generation phase itself. And even when we scale the number of images to be generated, the total inference time remains the same, as the forward time of the Rainbow module to generate representations in parallel is the same with M = 40, 20, or 10—around 0.104s/prompt.\n\nFor instance, given the same computational resources, if we can generate 10 high-resolution images in parallel in a batch with a baseline model like Stable Diffusion, we can achieve the same with Rainbow using the same generative model. In scenarios where the available computational resources limit us to generating fewer images per batch at a time, e.g., a batch of 2, the SD baseline needs to generate images 5 times with a batch of 2 to get 10 images; Rainbow can still process multiple decomposed condition representations in parallel in a batch of 10, and then concurrently generate images in a batch of 2.\n\nHope these answers clarify your questions.\n\n-----------\n3. **Automated graph interpretability** Appendix C lists graph interpretation as future work. Could you outline planned or preliminary method? \n\nThank you for your question regarding automated graph interpretability.\n\nFirst, we believe that computation on the latent space is crucial in aiding diversity and understanding the relationships between data points. By effectively interpreting the latent space, we can enhance the model's capacity to generate diverse outputs that still align with desired characteristics.\n\nSecond, the following are some hybrid approaches to improve interpretability in this context:\n\n\n    1.  Pre-defined Clusters: We could establish pre-defined clusters and train the model on specific groups of data. This could involve focusing on specific categories, such as images containing people, images with animals in the prompts, or MRIs corresponding to specific age ranges. This targeted training could provide clearer insights into how different clusters influence the generated outputs.\n\n    2.  Integration with Domain-Specific Knowledge Graphs: We could train the latent graph alongside a domain-specific knowledge graph derived from the dataset. For instance, if a specific node in the predefined knowledge graph is activated, we can open relevant edges in the latent graph that connect to corresponding data points. This would create a more structured approach, allowing for the incorporation of prior knowledge into the model's decision-making process.\n    \n\nBy employing these strategies, we aim to enhance the interpretability of our latent graphs, providing more meaningful insights into how different features and relationships within the data contribute to the overall performance of the model.\n\nHope this help with your consideration.\n\n\n---------------\n4. **Evaluation beyond automated metrics** Are there plans to incorporate human studies or downstream‐task evaluations (e.g., radiologist scoring, prompt relevance surveys) to validate that Rainbow captures semantically meaningful uncertainty rather than pixel‐level variation?\n\n**To address your concerns regarding human evaluation**, we present here the results of our human evaluation on the generated chest X-rays. In our study, we conducted human evaluations using 8 prompts, with 12 images generated per prompt, and invited three expert evaluators, all of whom are medical doctors, assessed the realism of the images by selecting which model—either the baseline SD or Rainbow—exhibited higher realism based on the given prompt descriptions. The results showed that **75% of the evaluators voted in favor of the Rainbow model**.\n\nWe are expanding the human evaluation across all three datasets in our future work, including increasing the number of evaluators involved.\n    \n--------------\n5. **Clinical deployment pathway** Could you elaborate on specific validation protocols (e.g., expert grading, quantitative anatomical metrics) and on whether constraints or priors could be integrated into Rainbow itself to enforce medical realism during sampling?\n\nRegarding validation after training, we can conduct human evaluation to select the best model, as mentioned in the answer to question 4.\n\nRegarding priors to integrate into Rainbow training and benefit the sampling results, one approach could be to utilize predefined domain-specific causal graphs for the effects of age and sex on brain MRI, or diseases on chest X-Ray. Another approach could integrate ideas from our response to question 3.\n\nThank you for your insightful comment."}, {"review_id": "VesoTzBOgz", "reviewer": "Reviewer_75uy", "summary": "Rainbow is a new method for generating diverse images from text prompts or other conditions. The main idea is that when given a prompt like \"sunset scene with mountain\" there are many valid ways to interpret it - different seasons, lighting, perspectives, etc. Current methods struggle to capture this diversity, which was shown in different works. Rainbow addresses this by using Generative Flow Networks (GFlowNets) to create multiple graph trajectories that represent different interpretations of the input condition. Each trajectory gets decoded into a different latent representation, which then generates different image from the trajectory.\nThe authors test this on three datasets: natural images, brain MRIs, and chest X-rays. Results show Rainbow generates more diverse images while maintaining quality compared to baselines like Stable Diffusion.\nThe key innovation is connecting GFlowNets (which are good at sampling diverse solutions) with diffusion models (which generate high-quality images). The learned graphs also seem interpretable - certain edges correspond to concepts like \"winter\" or \"devices.\"", "questions": "1. How did you decide on using graph structures to represent uncertainty? Were other structured representations considered or related works can be improved for that then?\n2. What is the computational overhead of Rainbow compared to standard diffusion models - can you add this to the paper to make it clear? How does inference time scale with the number of graphs M?\n3. When given a very specific prompt with little ambiguity (like a red circle on white background), does Rainbow still try to force diversity?  How does it handle low-uncertainty conditions?\n4. Are the interpretable edges (e.g. seasons) consistent across different training runs or did you select the best examples?\n5. Is there a principled way to choose M, N, and $\\rho$ for a new domain, or does it require extensive hyperparameter search -- not sure I feel the intuition here?", "limitations": "The method requires training additional components top of pre-trained models, which adds complexity.\nNo principled way to set hyperparameters - each domain uses very different values\nThe approach may not scale well to high-resolution images due to memory requirements of maintaining M parallel trajectories\nThe interpretability of edges is only demonstrated through cherry-picked examples; no systematic evaluation of interpretation stability\nAssumes all input conditions have meaningful uncertainty to exploit, which may not always be true\nThe connection between graph sparsity and semantic diversity is assumed rather than proven", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 2, "significance": 3, "originality": 4, "strengths_and_weaknesses": "### Strengths:\n- Novel combination of GFlowNets with diffusion models that hasn't been explored before\n- Strong experimental results across three very different domains (natural images, brain MRI)\n- The learned graph structures show some interpretability - can identify edges corresponding to seasons or medical devices\n- Comprehensive evaluation with multiple metrics (IS, VS, FID, CLIP scores)\n- Works with any pre-trained conditional diffusion model without major modifications\n\n### Weaknesses:\n\n- Computational cost seems high (M=40 parallel trajectories) but no timing comparisons provided\n- Hyper parameter selection (nodes, graphs, sparsity) appears arbitrary and domain-specific\n- The method might be forcing diversity even when the prompt is very specific\n- Unclear how stable the interpretable edges are across different training runs\n– minor thing  hard to distinguish (trajectories vs graphs) as notation", "was_revised": true, "final_justification": "Authors produced nice results with innovative approach and addressed all issues during rebuttal, I would like to save my positive score given all discussions.", "rebuttal": "1.  **How did you decide on using graph structures to represent uncertainty? Were other structured representations considered or related works can be improved for that then?**\n\nGraph structure was the only representation we considered. Driven by the objective of generating multiple output images while capturing a shared context, GFlowNet was chosen early on and GFlowNet primarily operates on graph structures.\n\nThe key design choice was **how to sample the graph**. If we add one **node at a time**, the resulting graph resembles a **causal graph**, where edges tend to connect all nodes. In contrast, adding one **edge at a time** yields a structure more like a **knowledge graph**, potentially with multiple disconnected subgraphs.\n\nWe chose the **edge-by-edge** sampling strategy and found that it performed better in practice.\n\nHope this answers your question.\n\n------------\n2.  **What is the computational overhead of Rainbow compared to standard diffusion models - can you add this to the paper to make it clear? How does inference time scale with the number of graphs M?**\n\nThank you forthe suggestion.. We can certainly include this information in the paper for clarity.\n\nIn terms of inference time, we collect the inference time of Rainbow and SD2-1 when generating 40 images per prompt, taking average of 60 prompts. We see that the average time per image for Rainbow is approximately 0.40 seconds, which is comparable to the standard diffusion model, SD2-1, which has an average inference time of 0.39 seconds. The forward time for the Rainbow module, that decompose original prompt representation into 40 prompt representations, itself is 0.104 seconds per prompt when generating 40 trajectories parallelly.\n\nImportantly, varying the number of trajectories (M) from 10 to 40 does not significantly affect the total generation time, as all the trajectories are predicted as output from the GFlownet at once. This parallel processing allows us to maintain efficient inference times, even when exploring a larger number of trajectories, without introducing substantial delays.\n\nWe appreciate your suggestion and will ensure this information is clearly presented in the revised version of the paper. Thank you!\n\n-----------\n3.  **When given a very specific prompt with little ambiguity (like a red circle on white background), does Rainbow still try to force diversity? How does it handle low-uncertainty conditions?**\n\nThank you for your insightful comment.\n\n\nWe have tried this prompt with both the SD2-1 baseline and Rainbow. Quantitatively, both models achieve comparable scores on Inception Score (IS), Visual Similarity (VS), and CLIP evaluation. Qualitatively, we observe that both models interpret the prompt creatively with red objects incorporating circular shapes, such as a red ball on a textured surface, a red circle appearing on a decorative plate, and red dots arranged in patterns against various backgrounds contains white color. Notably, none of the models produce images with a simple white background and a plain red circle.\n\nThis is an interesting and helpful test case to understand models performance. We appreciate your insightful comment and will investigate similar prompts in the future.\n\n----------\n4.  **Are the interpretable edges (e.g. seasons) consistent across different training runs or did you select the best examples?**\n\n\n\nYes, in different runs with varying model configurations and seeds, we can extract a consistent set of edges for features such as \"seasons\" and \"devices\" using the same edge frequency extraction method. The edge indices for each feature may differ across training seeds, but the presence of these interpretable features remains stable. This is because we do not predefine meanings for edges; instead, we allow the model to learn from the data during training. Consequently, the model assigns different sets of edges for features, but the underlying features themselves are reliably learned.\n\nRegarding to selecting samples, we generate and visualize 40 images for the natural images per model. We have provided all 40 generations in the appendix. For the visualization in Figure 2, we select images generated by Rainbow follow Spring-Summer-Autumn-Winter seasons.\n\n-----------------\n5.  **Is there a principled way to choose M, N, and   ρ  for a new domain, or does it require extensive hyperparameter search -- not sure I feel the intuition here?**\n\n**Regarding hyperparameter selection for GFlowNet**, specifically for the parameters M, N, and ρ, our choices were based on several factors, including computational resources, the need for diversity, task plausibility, and results from fine-tuning.\nWe adjusted the value of N by using a higher value for natural images to allow greater freedom in diversity and concept choices, while for medical datasets, we opted for a lower N value to balance diversity and ensure that the images still follow a domain-specific template. We maintained a relatively high graph sparsity, above 0.6; higher sparsity results in fewer sampled edges. This strategy enabled us to maximize the number of graphs given our available computational resources.\nWe fine-tuned the combinations of M, N, and ρ for optimal performance in terms of both diversity and image quality. Notably, we observed that all combinations produced similar CLIP scores for prompt delivery evaluation. The fine-tuning results for chest X-rays are provided in Figure 11 in the appendix (referred to in line 313).\n\nWe hope this answer clarifies your question."}, {"review_id": "eCjIZ1A9vV", "reviewer": "Reviewer_Jz2s", "summary": "This paper introduces \"Rainbow,\" a novel framework designed to enhance the diversity of conditional image generation models. The core problem addressed is the inherent ambiguity in input conditions (e.g., text prompts), which can correspond to multiple plausible image outputs. The authors' key idea is to decompose a single input condition into multiple, diverse latent representations. This is achieved by introducing a latent graph structure and employing Generative Flow Networks (GFlowNets) to sample a set of diverse trajectories (subgraphs) over this graph. The method's efficacy is demonstrated across multiple domains, including natural image generation (Flickr30k), 3D medical imaging (Brain MRIs), and 2D medical imaging (Chest X-rays). The results show consistent improvements in both diversity (measured by Vendi Score) and image quality (measured by IS and FID) compared to strong baselines like various versions of Stable Diffusion, CADS, and PAG. The paper also provides compelling evidence for the interpretability of the learned latent graphs and the positive impact of the generated data on downstream tasks.", "questions": "1. Regarding the latent graph interpretation (Section 4.2.2): How were the sets of edges corresponding to specific concepts (e.g., \"spring edges\") identified? Was this an automated clustering process based on edge frequency, or was there manual selection involved after clustering the output images? Please clarify the procedure.\n\n2. Could you elaborate on the choice of the RNN in the graph decoder ($Q_D$)? Since the final graph is an unordered set of edges, is the sequential processing of the RNN critical? Have you experimented with non-sequential architectures like a simple pooling or a Transformer encoder?\n\n3. In Table 1, the performance of the age/sex classifiers on real data is provided as a reference. What is the performance of these classifiers on the training set of the real data? This would help contextualize the generalization gap and better interpret the results on the synthetic data.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths\n\nNovel and Sound Methodology: The primary strength of this work lies in its innovative integration of GFlowNets with latent diffusion models. Using GFlowNets to explicitly model a distribution over diverse graph structures that represent conditional uncertainty is a principled and powerful approach. It directly tackles the mode-seeking behavior of many generative models and provides a more structured way to explore the solution space than simple noise manipulation. The technical execution, including the use of the Detailed Balance objective and the design of the graph generator and decoder, is well-reasoned and clearly explained.\n\nComprehensive Empirical Evaluation: The authors have conducted an extensive set of experiments across three distinct and challenging domains. The visual comparisons (e.g., Figure 2, \"Sunset scene with mountain\") are convincing. Rainbow clearly generates a wider variety of seasons, lighting, and compositions compared to the baselines, which often produce repetitive or stylistically limited outputs. The quantitative analysis (Figure 4, Table 1, Table 3) supports the claims. Rainbow consistently achieves a better trade-off between diversity (higher Vendi Score) and fidelity (higher IS, lower FID) than the baselines.\n\nModel-Agnostic and Flexible Framework: An advantage of Rainbow is its applicability to any pretrained conditional generative model. By operating on the condition representation and keeping the core generative model frozen, it serves as a versatile plug-in. This is impressively demonstrated by its successful application to different Stable Diffusion versions for natural images and chest X-rays, and a MONAI-optimized LDM for 3D brain MRIs.\n\n\nWeaknesses\n\nComputational Overhead: The paper acknowledges that the method requires higher computational resources due to the parallel generation of M trajectories (graphs). While the results justify the cost, a more detailed analysis of the trade-off between the number of graphs (M) and performance versus computational cost would be beneficial. The ablation study (Figure 10) shows that performance degrades as M decreases, but the cost-benefit analysis isn't fully fleshed out.\n\nClarity on Graph Hyperparameters: The choice of the latent graph size (N nodes) and sparsity (ρ) seems somewhat arbitrary. The ablation study explores these parameters for the Chest X-ray experiment, but a more in-depth discussion on the intuition behind selecting these values for different domains would strengthen the paper. For instance, why was N=20 chosen for natural images but N=8 for brain MRIs? Is there a principled way to set these, or is it purely empirical?\n\nLimited Discussion on Failure Cases: While the presented results are strong, the paper would be more complete with a discussion of failure modes. Are there types of conditions or prompts where Rainbow struggles to generate diversity or maintain fidelity? For example, how does it handle highly complex prompts with many entangled concepts?\n\nDifference with Subspace-based Methods: It seems that the core goals of rainbow and subspace-based representation learning methods are very similar. Why not use subspace methods to learn different semantic representations?", "was_revised": false, "final_justification": "", "rebuttal": "1.  **Regarding the latent graph interpretation (Section 4.2.2): How were the sets of edges corresponding to specific concepts (e.g., \"spring edges\") identified? Was this an automated clustering process based on edge frequency, or was there manual selection involved after clustering the output images? Please clarify the procedure.**\n\n\nThank you for your question. The process is automatic. We rank the top 10 added edges when the prompt changes from _\"sunset scene with mountain\"_ to _\"sunset scene with mountain in a specific season\"_, and then append these edges to the edge set of the first prompt.\n\nA similar procedure was used in the chest X-ray experiment shown in Figure 7, with the corresponding edge heat map provided in Figure 24 (referred to in line 287).\n\nHope this answers your question.\n   \n   \n\n\n-----------------\n2.  **Could you elaborate on the choice of the RNN in the graph decoder ( QD)? Since the final graph is an unordered set of edges, is the sequential processing of the RNN critical? Have you experimented with non-sequential architectures like a simple pooling or a Transformer encoder?**\n\nFirst, we would like to clarify that our edge sampling process is **order-dependent**—we sample edges **one at a time**, and the choice of each edge can influence the selection of subsequent ones. Therefore, maintaining the order of edges is important, and a sequential model like an RNN helps capture this dependency during decoding.\n\nSecond, we would like to mention that our current design **does include a pooling layer after the RNN**, allowing the model to aggregate information across all sequentially sampled edges before making the final predictions.\n\nThird, we conducted an ablation study comparing Rainbow with RNN (RNN -> Pooling) and without the RNN (directly to Pooling) in the graph decoder in the brain MRI experiment. The results are shown below:\n\n| Method | Synthesized Sample (Age MAE)  ↓| Counterfactual on Age (Age MAE) ↓ | Counterfactual on Sex (Age MAE) ↓ |\n| --- | --- | --- | --- |\n| Rainbow w RNN | **19.87** | **15.40** | **14.36** |\n| Rainbow w/o RNN | 23.91 | 20.42 | 19.71 |\n\nWe can see that removing the RNN leads to a noticeable drop in performance across all evaluation settings. This suggests that modeling sequential dependencies among sampled edges is beneficial for generating more accurate and consistent outputs.\n    \n\n\n------------------\n3.  **In Table 1, the performance of the age/sex classifiers on real data is provided as a reference. What is the performance of these classifiers on the training set of the real data? This would help contextualize the generalization gap and better interpret the results on the synthetic data.**\n\nThank you for your question.\n\nWe’d like to provide more context on the \"Real data\" results in Table 1. We sampled MRIs and corresponding conditions from the real dataset to match the age and sex distribution of the generated samples shown in the bottom rows of the table. This ensures a fair comparison when evaluating classifier performance.\n\n\n\n\n\n\n------------------\n4. **Computational Overhead: The paper acknowledges that the method requires higher computational resources due to the parallel generation of M trajectories (graphs). While the results justify the cost, a more detailed analysis of the trade-off between the number of graphs (M) and performance versus computational cost would be beneficial. The ablation study (Figure 10) shows that performance degrades as M decreases, but the cost-benefit analysis isn't fully fleshed out.**\n\nThank you for your question. \n\nWe would like to clarify that across all datasets in our experiments, varying the number of trajectories does not significantly affect the total generation time, as all the trajectories are predicted as output from the GFlownet at once. This parallel processing enables efficient inference, even when generating a larger number of trajectories, without introducing substantial delays.\n\n\nFor more detail, we did an analysis on natural image generation, we measured inference time for Rainbow and SD2-1 when generating 40 images per prompt (averaged over 60 prompts). The average time per image for Rainbow is approximately 0.40 seconds, which is comparable to SD2-1 at 0.39 seconds. The forward pass of the Rainbow module—which decomposes the original prompt into 40 prompt representations—takes only 0.104 seconds per prompt. This forward duration is similar when we generate 40, 20, and 10 graphs.\n\nHope this clarify your concerns.\n\n---------------\n5. **It seems that the core goals of rainbow and subspace-based representation learning methods are very similar. Why not use subspace methods to learn different semantic representations?**\n\nThank you for the insightful suggestion.\n\nRainbow is our initial attempt at decomposing a single condition into multiple structured representations. We see it as a first step that opens up promising directions for future research and improvement.\n\nWe agree that subspace-based approaches are meaningful and can offer benefits, particularly in terms of interpretability. However, training with predefined subspaces can introduce several challenges—for example, determining how to cluster data into subspaces, deciding the number of subspaces, and defining their semantic scope. In contrast, Rainbow offers a more flexible framework by allowing these representations to emerge naturally during training and enabling interpretability through post-hoc analysis of the learned graphs.\n\nWe see both approaches as complementary and believe exploring their integration could be a fruitful avenue in future work.\n\n---\n\n\n\n\n6. **why was N=20 chosen for natural images but N=8 for brain MRIs? Is there a principled way to set these, or is it purely empirical?**\n\nThank you for your question. \n\n**Regarding hyperparameter selection for GFlowNet**, specifically for the parameters M, N, and ρ, our choices were based on several factors, including computational resources, the need for diversity, task plausibility, and results from fine-tuning.\nWe adjusted the value of N by using a higher value for natural images to allow greater freedom in diversity and concept choices, while for medical datasets, we opted for a lower N value to balance diversity and ensure that the images still follow a domain-specific template. We maintained a relatively high graph sparsity, above 0.6; higher sparsity results in fewer sampled edges. This strategy enabled us to maximize the number of graphs given our available computational resources.\nWe fine-tuned the combinations of M, N, and ρ for optimal performance in terms of both diversity and image quality. Notably, we observed that all combinations produced similar CLIP scores for prompt delivery evaluation. The fine-tuning results for chest X-rays are provided in Figure 11 in the appendix (referred to in line 313) for reference.\n\nWe hope this answer clarifies your question.\n\n\n--------------\n7. **Are there types of conditions or prompts where Rainbow struggles to generate diversity or maintain fidelity? For example, how does it handle highly complex prompts with many entangled concepts?**\n\nThank you for your question.\n\nFirst, we would like to note that even very detailed prompts, such as _\"A serene forest scene at sunrise, with golden rays filtering through the trees, vibrant wildflowers in the foreground, and a gentle stream flowing, reflecting the colors of the morning light.\"_, often still contain uncertainty in object types and layout choices. This suggests that capturing diversity remains valuable, as it can be challenging to express every visual detail purely through text.\n\nSecond, we have observed some cases where both the baselines and Rainbow struggle to accurately reflect certain prompt details. For example:\n\n-   _\"A glass falls to the ground\"_ → models may miss the causal implication (e.g., the glass breaking).\n-   _\"#num of cats standing evenly horizontally on the street\"_ → models often struggle to generate the exact number of cats and align them evenly.\n    \nIt is worth noting that images that are unable to convey specific implications can also stem from limitations inherent to the diffusion model itself. We appreciate this important observation and will include some of these failure cases in the revised version of the paper."}, {"review_id": "95S0BiIxol", "reviewer": "Reviewer_PWCV", "summary": "Rainbow is a plug-and-play framework that adds explicit diversity to any pretrained conditional image generator. Instead of varying random seeds or paraphrasing prompts, Rainbow decomposes the input condition into multiple latent representations. Concretely, it inserts a latent graph—parameterised and sampled by a Generative Flow Network (GFlowNet)—into the prompt-encoding stage. Each GFlowNet trajectory captures a different plausible interpretation of the condition; feeding these trajectories through the frozen generator yields a batch of diverse, yet still faithful, images. Experiments on natural and medical datasets show simultaneous improvements in diversity (e.g., higher LPIPS / dispersion scores) and fidelity (lower FID) for image synthesis, conditional generation and counterfactual tasks, demonstrating Rainbow’s generality and effectiveness.", "questions": "How do graph depth, branching factor, and node-feature type affect both diversity and fidelity?\nCould you provide an ablation that varies these hyper-parameters?", "limitations": "see weaknesses", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Stenghs:\n1. Rainbow can be bolted onto any pretrained conditional generator without retraining the base model, making it easy to adopt across diffusion, GAN and autoregressive backbones.\n2.  Instead of random seeds or ad-hoc prompt edits, the method explicitly decomposes the condition into multiple latent routes, yielding diversity that is semantically tied to ambiguities in the original prompt.\n3. Leveraging GFlowNets’ trajectory-sampling provides a mathematically grounded way to cover high-reward—but diverse—modes of the latent graph, avoiding mode collapse common in MCMC or pure random seeding.\n\n\nWeaknesses:\n1. Comparisons are mostly against “different seeds” or prompt-paraphrasing. State-of-the-art diversity techniques—e.g., DDIM stochasticity schedules, classifier-free guidance sweeps, Multimodal MCTS, or Latent Diffusion Ensembles—are missing, making it hard to gauge Rainbow’s real advantage.\n2. Sampling diverse trajectories may lower prompt adherence or introduce artifacts. The paper reports FID and a single diversity metric but provides no human evaluation or CLIP-based alignment score to show that each variant still obeys the original condition.\n3. GFlowNets are known to be sensitive to reward scaling and exploration temperature. The paper does not include convergence curves, variance across runs, or guidelines for hyper-parameter selection, which may hinder reproducibility.\n4. The abstract states that each trajectory “captures an aspect of the uncertainty,” but no qualitative or quantitative analysis is given to link individual graph paths to human-interpretable factors (e.g., colour palette, viewpoint).", "was_revised": true, "final_justification": "I would like to thank the author for the rebuttal points, i will increase my rating.", "rebuttal": "1.  **Comparisons are mostly against “different seeds” or prompt-paraphrasing. State-of-the-art diversity techniques—e.g., DDIM stochasticity schedules, classifier-free guidance sweeps, Multimodal MCTS, or Latent Diffusion Ensembles—are missing, making it hard to gauge Rainbow’s real advantage.**\n\nThank you for your thoughtful review. We would like to take this opportunity to clarify our choices regarding baseline comparisons and to provide additional details about our methodologies.\n\nFirst, we would like to clarify that we do not solely rely on againsting “different seeds” or prompt-paraphrasing baselines; we also compared our method against baselines that intentionally trying to address the diversity of output images. Specifically, in our natural images experiment, we incorporated the following key baselines:\n\n    1.  **Traditional LDM:**  With this baseline, the \"diversity\" is not considered\n    2.  **CADS:**  With this baseline, the \"diversity\" is taken into account by employing a sampling strategy that anneals the conditioning signal. This is achieved by adding scheduled, monotonically decreasing Gaussian noise to the conditioning vector during inference.\n    3.  **PAG:**  This baseline considers \"diversity\" by not only incorporating prompt-paraphrasing but also by sampling from an unnormalized density function, enabling the generation of both high-quality and diverse prompts.\n\nIn conclusion, our baselines cover a wide range of recent techniques aimed at addressing diversity effectively.\n\nSecondly, in Figure 18 in the Appendix, we also provide comparisons with variations in classifier-free guidance to highlight the benefit of Rainbow in generating diverse images in contextual level.\n\nThirdly, it is worth noting that one of the novel ideas of Rainbow is to capture diversity by decomposing the condition into diverse representations. Consequently, our baselines include methods that generate a single output for uncertain input conditions, effectively highlighting the advantages of Rainbow in achieving greater diversity.\n\nAdditionally, we include a comparison with the DDIM scheduler (implementation provided by stable-diffusion-webui) in the table below, demonstrating that Rainbow consistently outperforms baselines in both image quality and diversity, as well as ensuring obeying original prompt. About qualitative analysis, generations by DDIM have the same remaining limiation as other baselines, that is, images are non-identical but similar in context, not capture the diverse representation of the input prompt's uncertainty. For instance, generation with \"Sunset scene with mountain in a specific season\" produce almost all spring season over 40 images. We will add this additional analysis to the revised version of the manuscript.. \n\n\n| **Model name**    | **IS ↑** | **CLIP ↑**  | **Pixcel VS ↑**  | **Inception VS ↑**  |\n|-------------------|----------|-------------|------------------|---------------------|\n| **SD3.5**         | 7.37     | 30.29       | 3.75             | 25.09               |\n| **SD2-1**         | 8.49     | 30.27       | 3.74             | 25.63               |\n| **CADS**          | 9.46     | 30.27       | 3.88             | 27.85               |\n| **PAG**           | 9.93     | 30.21       | 3.90             | 26.92               |\n| **DDIM**          | 8.44     | 30.28       | 3.91             | 26.62               |\n| **Rainbow**       | **10.45**| **30.32**   | **3.94**         | **28.90**           | \n\n\nFinally, we apologize for not being able to find any studies on MCTS or Latent Diffusion Ensembles specifically applied to image generation tasks. If you have any references or further details regarding these approaches, it would be greatly appreciated if you could share them, as we are eager to explore relevant techniques to enhance our research.\n\n\n\n\n---------------\n\n2.  **Sampling diverse trajectories may lower prompt adherence or introduce artifacts. The paper reports FID and a single diversity metric but provides no human evaluation or CLIP-based alignment score to show that each variant still obeys the original condition.**\n\nThank you for your insightful feedback. \n\n**To address your concerns regarding obeying original condition**, we would like to clarify that we have provided evaluations based on the CLIP score for experiments with prompt-based conditions. Specifically:\n\n-   **For natural images**, we present the CLIP score in line 249, indicating that Rainbow achieves significantly higher diversity while preserving or slightly increasing the CLIP score at around 30.3. This result demonstrates the method's ability to maintain prompt adherence. Further details can be found in Table 3 in the appendix (we also copied this table and present it in the answer for your first question above).\n    \n-   **For chest X-ray images**, we include a CLIP evaluation in line 273, where we report that all models achieve similar CLIP scores at 33.5. This indicates that Rainbow improves diversity over the baseline while still adhering to the original prompts. The full CLIP scores are provided in Table 4 in the appendix.\n\n\n**To address your concerns regarding human evaluation on adherence to original conditions**, we present the results of our human evaluation on the generated chest X-rays. In our study, we conducted human evaluations using 8 prompts, with 12 images generated per prompt. Three evaluators, all of whom are medical doctors, assessed the realism of the images by selecting which model—either the baseline SD or Rainbow—exhibited higher realism based on the given prompt descriptions. The results showed that **75% of the evaluators voted in favor of the Rainbow model**.\n\nWe are expanding the human evaluation across all three datasets in our future work, along with increasing the number of evaluators / experts involved.\n\n\n\nWe hope this clarifies how our approach balances diversity and prompt adherence.\n\n\n\n\n\n-----------\n\n3.  **GFlowNets are known to be sensitive to reward scaling and exploration temperature. The paper does not include convergence curves, variance across runs, or guidelines for hyper-parameter selection, which may hinder reproducibility.**\n\nThank you for your insightful comment.\n\nRainbow is built on a pretrained image encoder-decoder and U-Net models, contributing to its stability and reducing sensitivity during GFlowNet training. We will ensure to include reward and loss curves in the final version of the paper to enhance the clarity of our findings.\n\n**Regarding hyperparameter selection for GFlowNet**, specifically for the parameters M, N, and ρ, our choices were based on several factors, including computational resources, the need for diversity, task plausibility, and results from fine-tuning.\nWe adjusted the value of N by using a higher value for natural images to allow greater freedom in diversity and concept choices, while for medical datasets, we opted for a lower N value to balance diversity and ensure that the images still follow a domain-specific template. We maintained a relatively high graph sparsity, above 0.6; higher sparsity results in fewer sampled edges. This strategy enabled us to maximize the number of graphs given our available computational resources.\nWe fine-tuned the combinations of M, N, and ρ for optimal performance in terms of both diversity and image quality. Notably, we observed that all combinations produced similar CLIP scores for prompt delivery evaluation. The fine-tuning results for chest X-rays are provided in Figure 11 in the appendix (referred to from line 313).\n\nWe hope this addresses your concerns.\n\n---------------\n4.  **The abstract states that each trajectory “captures an aspect of the uncertainty,” but no qualitative or quantitative analysis is given to link individual graph paths to human-interpretable factors (e.g., colour palette, viewpoint).**\n\nWe appreciate your observation and acknowledge the importance of linking graph trajectories to interpretable features. We would like to clarify that Rainbow demonstrating clear benefits in improving diversity, both quantitatively and qualitatively, our analysis does show that the graph can be interpreted to human-understandable features such as \"seasons\" and \"support devices\" in Section 4.2.2. This supports our claim that each trajectory \"captures an aspect of the uncertainty\".\n\nIn addition, we agree that we need one step further to establish a more direct connection between individual graph paths and specific interpretable factors. We plan to address this direct interpretability in the next paper for a big upgraded version of Rainbow as part of our future research.\n\nWe hope this answer clarifies your concerns!\n\n\n-----------\n5. **How do graph depth, branching factor, and node-feature type affect both diversity and fidelity? Could you provide an ablation that varies these hyper-parameters?**\n\nReffering to the fine-tuning results presented in Figure 11 in the appendix, we observe the effects of the number of nodes (N), sparsity ρ that affect the number of edges on sampled trajectories, and the number of trajectories (M) on diversity (as measured by the VS score) and fidelity (as measured by the FID score) as follows:\n\n-   Configurations with M=5 and a sparsity level of ρ=0.64, regardless of the value of N, tend to produce lower scores in diversity and exhibit variability in image quality.\n-   Configurations with N=20 and ρ=0.82 (a mid-range value), regardless of the value of M, tend to stabilize image quality while also generating highly diverse outputs.\n-   Sparsity levels of ρ=0.88, independent of the other configurations, result in instability, producing both low and high scores in terms of diversity and quality.\n   \nBased on these fine-tuning results, we selected the combination that achieves a balance between image quality and diversity by choosing the configuration with the highest ratio of VS^3/FID, which is N\\=20, M\\=10, and ρ\\=0.82.\n\nWe hope this answer clarifies your question."}]

Author rebuttal:

Meta-review: The paper proposes a method for generating diverse images from a fixed trained conditional generative model, by decomposing the conditioning signal into diverse latents, and then integrating graphs over these latents into prompt representations using GFlowNets. 

The method achieves good results on both natural and medical images, is new, and creative. 

There were some concerns around baselines, performance, inference time, and presentation, but these were addressed satisfactorily in the author/reviewer discussion. 

Overall, all the reviewers are positive about the paper, I agree and recommend acceptance.
Decision: Accept (poster)