MedSymmFlow: Hybrid Model for Medical Imaging

• MedSymmFlow is a hybrid generative-discriminative model that integrates classification, generation, and uncertainty quantification under a single conditional flow framework for medical imaging.
• It leverages a latent-space formulation and semantic RGB prototype conditioning to enhance multiclass performance and scalability on high-resolution MedMNIST datasets.
• The model employs reverse-time integration for classification and uses prototype distance as an inherent uncertainty measure, facilitating effective selective prediction.

MedSymmFlow is a generative-discriminative hybrid model for medical imaging that is built on Symmetrical Flow Matching and is designed to unify classification, generation, and uncertainty quantification within a single conditional flow framework. It uses a latent-space formulation to scale to high-resolution inputs, conditions the flow on semantic mask representations, and performs classification by reverse-time integration of a learned semantic flow rather than by a separate discriminative head. In the reported evaluation on four MedMNIST datasets, MedSymmFlow matches or exceeds established baselines in classification accuracy and AUC, while its uncertainty estimates are validated by performance improvements under selective prediction (Caetano et al., 25 Jul 2025).

1. Origins and conceptual scope

MedSymmFlow inherits its formal basis from Symmetrical Flow Matching, or SymmFlow, which learns image synthesis and semantic inference as opposing flows within a single model. In SymmFlow, the image variable $x$ and semantic variable $y$ are both perturbed along time-indexed paths, and a joint velocity field is trained to recover their corresponding target velocities. The original SymmFlow formulation was introduced for semantic image synthesis, semantic segmentation, and classification on benchmarks such as CelebAMask-HQ, COCO-Stuff, MNIST, and CIFAR-10, and it emphasized bi-directional consistency and efficient sampling (Caetano et al., 12 Jun 2025).

MedSymmFlow adapts that framework to medical imaging and reorients it toward clinical decision support. The motivation is explicitly tied to the requirement that clinical models provide both accurate predictions and well-calibrated uncertainty, since overconfident wrong predictions can lead to unsafe recommendations, whereas transparent uncertainty supports selective prediction, triage, and referral. A central design claim is that generation, classification, and uncertainty quantification are usually developed in isolation, while MedSymmFlow unifies them through a single conditional generative mechanism (Caetano et al., 25 Jul 2025).

A common misconception is that MedSymmFlow is primarily a conventional classifier with a generative auxiliary module. The reported formulation is instead intrinsically joint: image generation integrates the image component of the flow forward in time, whereas classification integrates the semantic component backward in time. This suggests that its discriminative behavior is not appended post hoc, but emerges from the same learned transport structure that governs synthesis.

2. Flow formulation and learning objective

The method is framed in the language of Flow Matching. Let $p_0(x)$ denote a simple prior such as a Gaussian and $p_1(x)$ denote the data distribution. Flow Matching defines a probability path $\{p_t(x)\}_{t\in[0,1]}$ and a target velocity field $u(x,t)$ satisfying the continuity equation

$$\partial_t p_t(x) + \nabla \cdot (p_t(x) u(x,t)) = 0,$$

with sampling governed by the probability flow ODE

$\frac{d x_t}{d t} = u(x_t,t), \qquad t\in[0,1].$

A standard straight-path instantiation uses

$x_t = (1-t)x_0 + t x_1,$

with $(x_0,x_1)\sim \pi_X$, giving the constant target velocity $y$0 (Caetano et al., 25 Jul 2025).

MedSymmFlow extends this to a joint image-semantic setting. The paired straight-line paths are

$y$1

with target velocities

$y$2

The learned joint velocity field $y$3 has image and semantic components, and training minimizes the symmetrical flow-matching objective

$y$4

The boundary conditions are also explicit: $y$5, $y$6, $y$7 is the class-specific RGB embedding plus uniform noise of amplitude $y$8, and $y$9 is noise (Caetano et al., 25 Jul 2025).

Classification is obtained by reverse-time integration in semantic space:

$p_0(x)$0

The semantic prediction is then decoded by nearest-prototype matching,

$p_0(x)$1

The paper also states what the model does not use: there is no separate cross-entropy classification loss, and no KL regularization on the flow is reported. This is important because it distinguishes MedSymmFlow from standard discriminative classifiers and from hybrid systems that retain a conventional supervised head.

3. Latent-space realization and semantic conditioning

The principal architectural modification is a latent-space formulation, denoted LatMSF, which uses the Variational Autoencoder from the Stable Diffusion stack, specifically sd-vae-ft-mse, to encode both images $p_0(x)$2 and semantic masks $p_0(x)$3 into latent variables $p_0(x)$4 and $p_0(x)$5. The flow network operates in this compressed domain, and decoded outputs are reconstructed back to RGB space. The stated purpose is to make high-resolution medical images tractable at $p_0(x)$6 while preserving semantics (Caetano et al., 25 Jul 2025).

The flow network is a UNet-like architecture with attention. Reported hyperparameters include Channels in $p_0(x)$7, Depth $p_0(x)$8, Channels Multiple $p_0(x)$9, Heads $p_1(x)$0, Head Channels $p_1(x)$1, Attention Resolution $p_1(x)$2, and Dropout $p_1(x)$3. Time $p_1(x)$4 is embedded and injected into the UNet, and the semantic state $p_1(x)$5 conditions both $p_1(x)$6 and $p_1(x)$7 through concatenation or embedding alongside image features and time embeddings. Training uses cosine annealing learning-rate scheduling with warmup; example settings include batch sizes in $p_1(x)$8, epochs up to $p_1(x)$9, learning rates in $\{p_t(x)\}_{t\in[0,1]}$0, and label noise amplitude $\{p_t(x)\}_{t\in[0,1]}$1 depending on dataset and resolution.

The semantic conditioning strategy is a defining contribution. Instead of grayscale mask-conditioning, which the paper describes as collapsing inter-class distances onto one axis, MedSymmFlow maps each class to a unique RGB prototype and perturbs that prototype with uniform noise amplitude $\{p_t(x)\}_{t\in[0,1]}$2 during training. The semantic target is therefore not a one-dimensional label coordinate but a three-dimensional continuous embedding. The reported interpretation is geometric: RGB conditioning expands semantic geometry from one dimension to three, which improves multiclass separation and yields higher AUC and accuracy, particularly on DermaMNIST.

A second possible misconception concerns the phrase “semantic mask conditioning.” In the classification experiments reported for MedSymmFlow, the semantic variable is not an anatomical segmentation mask; it is an RGB class prototype perturbed into a continuous manifold. The paper explicitly notes that, for segmentation or anatomical masks, conditioning would need to integrate spatial semantics, whereas the present evaluation encodes classification labels as RGB prototypes rather than spatial masks.

4. Experimental setting and empirical performance

The evaluation uses four MedMNIST v2 subsets and reports both $\{p_t(x)\}_{t\in[0,1]}$3 and $\{p_t(x)\}_{t\in[0,1]}$4 settings, with the latent formulation used for the higher resolution. The protocol follows MedMNIST accuracy and AUC evaluation, and results for SymmFlow and MedSymmFlow are reported as mean and $\{p_t(x)\}_{t\in[0,1]}$5 over five independent inference runs to capture probabilistic variability (Caetano et al., 25 Jul 2025).

Dataset	Classes	Train/Val/Test
PneumoniaMNIST	2	4,708 / 524 / 624
BloodMNIST	8	11,959 / 1,712 / 3,421
DermaMNIST	7	7,007 / 1,003 / 2,005
RetinaMNIST	5	1,080 / 120 / 400

The baselines are ResNet-18, ResNet-50, AutoKeras, Auto-Sklearn, MedViT-S, and the original SymmFlow with grayscale masks. At $\{p_t(x)\}_{t\in[0,1]}$6, the main empirical pattern is that RGB prototype conditioning improves multiclass performance relative to grayscale SymmFlow. On DermaMNIST, AUC rises from $\{p_t(x)\}_{t\in[0,1]}$7 to $\{p_t(x)\}_{t\in[0,1]}$8, and accuracy rises from $\{p_t(x)\}_{t\in[0,1]}$9 to $u(x,t)$0. On BloodMNIST, accuracy increases from $u(x,t)$1 to $u(x,t)$2. On PneumoniaMNIST, AUC improves from $u(x,t)$3 to $u(x,t)$4, although accuracy slightly declines relative to SymmFlow (Caetano et al., 25 Jul 2025).

At $u(x,t)$5, LatMSF matches or exceeds strong CNN and ViT baselines on several tasks. Reported results are BloodMNIST AUC $u(x,t)$6 and ACC $u(x,t)$7; DermaMNIST AUC $u(x,t)$8 and ACC $u(x,t)$9; RetinaMNIST AUC $$\partial_t p_t(x) + \nabla \cdot (p_t(x) u(x,t)) = 0,$$0 and ACC $$\partial_t p_t(x) + \nabla \cdot (p_t(x) u(x,t)) = 0,$$1; and PneumoniaMNIST AUC $$\partial_t p_t(x) + \nabla \cdot (p_t(x) u(x,t)) = 0,$$2 and ACC $$\partial_t p_t(x) + \nabla \cdot (p_t(x) u(x,t)) = 0,$$3. The comparison singled out in the text is that RetinaMNIST AUC exceeds MedViT-S, with $$\partial_t p_t(x) + \nabla \cdot (p_t(x) u(x,t)) = 0,$$4 versus $$\partial_t p_t(x) + \nabla \cdot (p_t(x) u(x,t)) = 0,$$5, while DermaMNIST is described as close to MedViT-S, which attains AUC $$\partial_t p_t(x) + \nabla \cdot (p_t(x) u(x,t)) = 0,$$6 and ACC $$\partial_t p_t(x) + \nabla \cdot (p_t(x) u(x,t)) = 0,$$7.

The generative side of the model is evaluated qualitatively rather than through a large standalone synthesis benchmark. The reported sample visualizations show realistic high-resolution images with fine details such as skin texture, hair artifacts, and fundus vessels. This suggests that the latent formulation preserves dataset-specific visual semantics sufficiently well for image synthesis to act as a useful inductive component for classification, although the medical evaluation is centered on accuracy, AUC, and uncertainty rather than on FID-style generative metrics.

5. Uncertainty quantification and selective prediction

Uncertainty estimation is integrated into the prediction mechanism rather than added as a separate calibration stage. After reverse-time integration yields a continuous semantic prediction $$\partial_t p_t(x) + \nabla \cdot (p_t(x) u(x,t)) = 0,$$8, the predicted class is the nearest RGB prototype, and the Euclidean distance to that prototype,

$$\partial_t p_t(x) + \nabla \cdot (p_t(x) u(x,t)) = 0,$$9

is used as the uncertainty measure. Larger distance corresponds to lower confidence (Caetano et al., 25 Jul 2025).

The paper characterizes this as “uncertainty for free,” because the uncertainty signal arises directly from the same generative sampling process used for classification. It also notes that, from Monte Carlo sampling, one could compute variance across samples, predictive entropy of prototype posteriors via soft assignment, and mutual information, but the primary proxy used in the study is the prototype distance. The validation of calibration is operational rather than abstract: Accuracy-Rejection Curves rank predictions by uncertainty and filter those with the largest distances.

The reported result is that filtering by prototype distance consistently improves accuracy, confirming that the distance correlates with error likelihood. The strongest effects are observed on PneumoniaMNIST and RetinaMNIST, whereas BloodMNIST shows negligible effect because the baseline accuracy is already high. A more nuanced observation is reported for DermaMNIST at high resolution: overfiltering can remove correct, confident predictions, so the gain diminishes beyond an optimal rejection threshold. This makes the uncertainty measure clinically relevant for selective prediction, but it also indicates that threshold selection remains task-dependent.

A frequent misunderstanding is that MedSymmFlow’s uncertainty behaves like temperature scaling or another post-hoc calibration layer. The reported mechanism is different. The uncertainty score is a geometric property of the recovered semantic embedding, and selective prediction operates directly on that geometry rather than on a recalibrated softmax output.

6. Computational profile, limitations, and related directions

The computational trade-off is explicit. On BloodMNIST, averaged over 100 single-image batches on an RTX 3090 Ti, ResNet-18 at $\frac{d x_t}{d t} = u(x_t,t), \qquad t\in[0,1].$0 trains in 12 minutes with 63.6 MB memory and 0.4 ms classification latency; ResNet-50 at $\frac{d x_t}{d t} = u(x_t,t), \qquad t\in[0,1].$1 trains in 372 minutes with 286.4 MB memory and 14.5 ms latency; MedViT-S at $\frac{d x_t}{d t} = u(x_t,t), \qquad t\in[0,1].$2 trains in 395 minutes with 170.7 MB memory and 29.0 ms latency. By contrast, SymmFlow at $\frac{d x_t}{d t} = u(x_t,t), \qquad t\in[0,1].$3 trains in 178 minutes with 68.3 MB memory, 15.3 ms classification latency, and 360.1 ms sampling latency; MSF at $\frac{d x_t}{d t} = u(x_t,t), \qquad t\in[0,1].$4 trains in 180 minutes with 68.3 MB memory, 15.3 ms classification latency, and 360.5 ms sampling latency; and LatMSF at $\frac{d x_t}{d t} = u(x_t,t), \qquad t\in[0,1].$5 trains in 2,302 minutes with 712.1 MB memory, 68.2 ms classification latency, and 540.1 ms sampling latency (Caetano et al., 25 Jul 2025).

These numbers support a clear interpretation: the model is slower than purely discriminative baselines because ODE integration and VAE encoding-decoding introduce additional cost. The paper presents this as an expected consequence of ODE-based generation and classification rather than as an implementation anomaly. Sample images were produced with an Euler ODE solver using 25 steps, and the text notes that increasing the number of steps can improve fidelity at additional cost.

The main limitations are also explicit. The evaluation is restricted to curated 2D datasets, whereas many clinically important modalities such as CT and MRI are 3D. The method assumes the availability and quality of class prototypes, and the design of semantic encodings matters. Domain shift and generalization to unbalanced datasets, surgical RGB scenes, and cross-domain shifts remain to be tested. For segmentation or anatomical mask conditioning, spatial semantics would need to be incorporated, since the present classification experiments encode labels as RGB prototypes rather than spatial masks. Distillation and reduced-step solvers are suggested as routes to lower inference latency.

Within the broader flow-matching literature, MedSymmFlow occupies a specific niche. SymmFlow provided the general symmetrical formulation on non-medical benchmarks (Caetano et al., 12 Jun 2025); medical Optimal Transport Flow Matching addressed fast medical image synthesis and enhancement without the joint classification objective (Yazdani et al., 1 Mar 2025); and FlowSDF used conditional flow matching with signed distance functions for uncertainty-aware medical segmentation rather than symmetrical joint image-label transport (Bogensperger et al., 2024). MedSymmFlow’s distinctive contribution is therefore the coupling of medical image generation, reverse-time classification, and uncertainty quantification through one conditional velocity field, together with an RGB semantic geometry and a latent-space implementation tailored to $\frac{d x_t}{d t} = u(x_t,t), \qquad t\in[0,1].$6 medical inputs (Caetano et al., 25 Jul 2025).