🤖 machine learning

CONFLUX: A Latent Diffusion Model for 3D Chest-CT Synthesis with RL Post-Training

CONFLUX is a latent diffusion model for 3D chest-CT synthesis that combines a variational autoencoder with a rectified-flow transformer and reinforcement learning post-training to generate high-fidelity, controllable medical volumes that significantly outperform existing baselines in both image quality and adherence to specified clinical attributes.

Original authors: Max Van Puyvelde, Halil Ibrahim Gulluk, Wim Van Criekinge, Olivier Gevaert

Published 2026-07-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Max Van Puyvelde, Halil Ibrahim Gulluk, Wim Van Criekinge, Olivier Gevaert

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you want to build a library of 3D chest CT scans for doctors to study, but you can't scan enough real patients because of privacy rules or because certain rare diseases are too hard to find. You need a machine that can "dream up" realistic fake scans. But here's the catch: if you ask the machine to make a scan of a 70-year-old man with a specific lung nodule, it shouldn't just make a generic old man's scan; it must actually include that nodule.

The paper introduces CONFLUX, a new AI system designed to do exactly this. Think of it as a highly skilled, controllable "3D artist" that creates chest CT scans based on a specific recipe.

Here is how it works, broken down into three simple stages:

1. The Compression (The "Suitcase")

Creating a full 3D CT scan is like trying to paint a massive mural on a wall the size of a football field. It takes too much memory and time.

What CONFLUX does: First, it uses a "suitcase" (a 3D autoencoder) to compress the massive CT scan into a tiny, high-quality suitcase (a "latent space").
The Analogy: Imagine taking a giant, detailed 3D model of a house and folding it down into a small, flat blueprint. The AI learns to fold it so perfectly that when you unfold it later, the house looks exactly the same. The AI does all its creative work on these tiny blueprints, not the giant houses, which makes it much faster and smarter.

2. The Painting (The "Flowing River")

Once the AI is working on the blueprints, it needs to draw them.

What CONFLUX does: It uses a "Rectified Flow Transformer."
The Analogy: Imagine a river flowing from a chaotic, foggy mountain (random noise) down to a clear, calm lake (a perfect medical scan). The AI learns the exact path of this river. Instead of taking a winding, confusing path, it learns to flow in a straight, efficient line.
The Control: This is where the "recipe" comes in. You tell the AI: "I want a scan of a 60-year-old woman with a heart condition." The AI uses this information to steer the river, ensuring the final lake (the scan) has those specific features. It uses a special "steering wheel" (adaptive layer normalization) to make sure the river flows exactly where you want it to.

3. The Coach (The "Reinforcement Learning")

Even with a good steering wheel, the AI sometimes gets lazy. It might make a scan that looks realistic but forgets to include the specific lung nodule you asked for. It looks like a real scan, but it's missing the point.

What CONFLUX does: The authors added a "post-training" stage using Reinforcement Learning (RL).
The Analogy: Imagine the AI is a student taking a test.
- Before the Coach: The student writes an essay that looks great, but they missed the main question. The teacher (the AI's original training) just says, "Good handwriting," and gives a passing grade.
- With the Coach: A strict grader (a separate AI classifier) checks the essay. If the student asked for a "lung nodule" and the essay doesn't have one, the grader gives a harsh penalty. The student (the AI) then practices again and again, trying to get a better score.
- The Result: The paper calls this "Group-Relative Policy Optimization." It's like a coach telling the student, "You did better than the other students who forgot the nodule, but you can do even better." This training forces the AI to stop guessing and start actually delivering the specific medical findings requested.

What Did They Achieve?

The team tested CONFLUX against other top-tier 3D medical AI models.

Quality: The fake scans were incredibly realistic, scoring much better on "realism tests" than the competition.
Control: This is the big win. When they asked for specific medical conditions, CONFLUX included them much more reliably than before. The "Coach" training reduced the gap between what was requested and what was delivered by nearly 50%.
The Gift: They didn't just keep the AI; they released the model and a massive dataset of 200,000 synthetic chest CT scans. These scans come with the "recipe" (metadata) attached, so researchers can use them to study specific diseases without needing real patient data.

In Summary

CONFLUX is a 3D medical image generator that:

Compresses huge scans into tiny, manageable blueprints.
Paints new scans by flowing from noise to reality, guided by a specific recipe.
Trains itself with a strict coach to ensure it actually follows the recipe, making sure the fake scans are not just pretty, but medically accurate to the request.

It's a tool for creating high-quality, privacy-safe medical data that doctors and researchers can use to study diseases, provided they stick to the specific conditions the AI was trained to recognize.

Technical Summary: CONFLUX

Problem Statement

Controllable generative models for 3D medical images aim to synthesize volumes with specific clinical attributes for tasks such as cohort augmentation, privacy-preserving data sharing, and counterfactual study design. However, achieving high-fidelity, natively 3D synthesis that is simultaneously faithful to requested conditioning remains a significant challenge. Existing models often struggle to guarantee that individual generated samples exhibit the specific clinical findings requested, as standard training objectives (like flow matching) optimize for aggregate distributional realism rather than per-sample attribute realization. This gap is particularly acute in chest CT, which requires high resolution and volumetric coherence while being described by structured radiological metadata.

Methodology

CONFLUX is a three-stage latent diffusion framework designed for chest CT synthesis, integrating a variational autoencoder (VAE), a rectified-flow transformer, and a reinforcement learning (RL) post-training stage.

1. Latent Autoencoder (Stage 1)

To mitigate the computational cost of direct voxel-space generation, a 3D convolutional VAE compresses preprocessed CT volumes ( $216 \times 176 \times 200$ ) into a lower-resolution latent space ( $16 \times 27 \times 22 \times 25$ ).

Architecture: The encoder and decoder utilize group-normalized residual blocks with 3D self-attention at the lowest resolution.
Training: The model minimizes an $\ell_1$ reconstruction loss, a Kullback–Leibler regularizer, and a tri-planar LPIPS perceptual loss.
Normalization: Post-training, the latent space is normalized to near-unit variance to facilitate flow matching.

2. Conditional Rectified-Flow Transformer (Stage 2)

A single-stream transformer generates latent embeddings from Gaussian noise within the learned latent space.

Flow Matching: The model is trained under a rectified-flow objective, learning to predict the constant velocity vector field connecting noise to data latents along straight-line paths.
Conditioning: Structured radiological metadata (18 binary abnormality findings, sex, age decade, and reconstruction kernel) is concatenated into a 42-dimensional vector. This vector modulates the transformer via adaptive layer normalization (adaLN-zero).
Sampling: Generation involves integrating the probability-flow ODE from $t=1$ to $t=0$ using a time-shifted Euler grid, followed by decoding with the frozen VAE decoder.

3. Reinforcement Learning Post-Training (Stage 3)

To address the disconnect between aggregate distribution matching and per-sample attribute fidelity, CONFLUX employs Group-Relative Policy Optimization (GRPO).

Motivation: Standard flow matching does not explicitly reward whether a specific requested finding appears in a single generated volume.
Mechanism: The deterministic ODE sampler is replaced with a stochastic sampler (SDE) to provide tractable log-probabilities for policy gradient updates.
Reward Signal: A frozen, independent classifier trained on real latent representations scores generated volumes based on the agreement between requested and realized findings. The reward is defined as the negative weighted cross-entropy between the classifier's predictions and the conditioning vector.
Optimization: The model is fine-tuned to maximize group-relative advantages (subtracting the mean reward of a group of rollouts per prompt and dividing by the standard deviation). Updates use a clipped PPO surrogate with a KL-divergence penalty against the pre-RL reference policy to prevent drift.

Key Contributions

Controllable 3D Latent Rectified-Flow Model: A native 3D architecture for chest CT that conditions directly on structured radiological metadata via adaLN, achieving synthesis quality competitive with strong 3D diffusion baselines.
RL Post-Training for Faithfulness: The first application of GRPO to a 3D medical flow model. This stage explicitly optimizes for the reliability of requested findings in generated volumes, verified by an independent judge model.
Open Release: The release of the trained model and a $\sim$ 200,000-volume synthetic chest-CT dataset. This dataset is conditioned on real metadata proportions, spanning a wide variety of clinical findings, intended for cohort augmentation and controlled study design.

Results

Synthesis Quality

Evaluated on a held-out validation set ( $N=501$ ) using tri-planar Fréchet Inception Distance (FID), CONFLUX significantly outperforms existing baselines:

Tri-planar FID: 32.3 (CONFLUX) vs. 74.6 (MAISI) and 145.4 (GenerateCT).
Coverage and Density: CONFLUX samples are 3.5 $\times$ denser and cover 6 $\times$ more of the real manifold compared to MAISI.
Diversity: The model maintains high diversity (MS-SSIM 0.507), comparable to real data (0.41), avoiding mode collapse.

Conditioning Faithfulness

Post-training with GRPO improves the agreement between requested and realized attributes as measured by an independent image-space classifier:

Performance Gain: Average Precision (AP) increased from 0.330 to 0.344, and AUROC from 0.684 to 0.699.
Gap Recovery: While the absolute gains appear modest, they recover 47% of the performance gap between the base model and the real-data ceiling (AP 0.360).
Specificity: Improvements were specific to the optimized findings group; unweighted groups (sex, age, kernel) remained unchanged within noise.

Reward Classifier Performance

The classifier used for the RL reward signal achieved a macro AUROC of 0.793 on the real validation split, outperforming CT-CLIP variants and CT-Net on the CT-RATE benchmark, validating its utility as a reward signal despite operating on compressed latents.

Significance and Claims

The paper claims that CONFLUX represents a significant step forward in controllable 3D medical image synthesis by decoupling high-fidelity generation from attribute control. By integrating RL post-training, the authors demonstrate that it is possible to directly optimize for the "faithfulness" of clinical attributes in individual samples, a property not guaranteed by standard flow-matching objectives. The release of a large-scale, conditionally diverse synthetic dataset aims to facilitate research in scenarios where real data is scarce or privacy-constrained. The work positions itself as the first to adapt GRPO for 3D medical flow models, offering a pathway to more reliable and controllable generative tools for clinical study design.