- The paper presents a novel diffusion framework (DADD) that disentangles anatomical and disease features for controlled ulcerative colitis progression synthesis.
- It employs a triple-pathway cross-attention mechanism and a delta steering strategy to modulate disease severity with high fidelity.
- Quantitative results demonstrate improved classification metrics and a 2× inference speedup compared to conventional multi-pass CFG methods.
Disentangled Anatomy-Disease Diffusion (DADD) for Controllable Ulcerative Colitis Progression Synthesis
Introduction
The precise generation of longitudinal medical images that model disease progression while preserving patient-specific anatomy is an unresolved problem in generative medical modeling, hindered by the entanglement of structural and pathological information in learned representations. The task is especially pertinent in ulcerative colitis (UC) endoscopy, where disease severity is scored ordinally via the Mayo Endoscopic Score (MES), but datasets are highly imbalanced and longitudinal sampling is impractical due to the invasiveness of repeat endoscopies. The "Disentangled Anatomy-Disease Diffusion (DADD)" framework addresses these limitations by disentangling anatomy from disease severity in diffusion-based generative models, enabling the synthesis of images at arbitrary target severity levels while maintaining the identity of the underlying anatomy.
Figure 1: Sample labeled images from LIMUC showing progressive stages of ulcerative colitis from healthy mucosa to severe ulceration.
Methodology
Problem Analysis: The Anatomy-Pathology Entanglement
Conventional approaches using diffusion models, such as those employing CLIP-based IP-Adapter conditioning, suffer from entanglement: the extracted features from a reference image encode both anatomical structures and the current pathological state. Attempts to modulate disease severity using embeddings (typically Additive Ordinal Embeddings, AOE) thus fail to achieve monotonic or authentic progression, as disease texture from the reference is inextricably tied to the anatomy features injected into the generator.
DADD Framework
The core innovation in DADD is the explicit identification and removal of disease-correlated features from reference image embeddings prior to conditioning the diffusion denoising process. The architecture comprises:
- Additive Ordinal Embedder (AOE): Produces cumulative embeddings for disease severity, reflecting clinical monotonicity.
- Feature Purifier: Utilizes a cross-attention scheme to query CLIP image tokens with the current disease embedding, identifies pathology-associated channels, and applies a learned gating mechanism to erase these components, producing “purified” anatomy tokens.
- Triple-Pathway Cross-Attention: Rather than naive concatenation, anatomy, disease, and delta (change-of-severity) signal pathways are injected at different layers using trainable, resolution-aware gates. This design ensures that coarse U-Net layers encode global structure and high-resolution layers capture pathological textures.
- Delta Steering: A training-free, direction-aware mechanism that enables single-pass modulation of disease progression at inference via signed, magnitude-proportional delta embeddings, eliminating the need for classifier-free guidance (CFG) and thus reducing inference latency.
Figure 2: Framework overview depicting the flow from input image through purified anatomical encoding, disease embedding, delta calculation, and cross-attention-based generation.
Feature Purification via Cross-Attention
The feature purifier operates by querying image tokens with the source disease embedding. The disease-correlated component is gated dimensionwise by a learned MLP and subtracted, yielding disease-agnostic anatomy representations. This learning occurs under a reconstruction objective where, during training, target and source severity match, preventing information leakage into the delta pathway.
Multi-Pathway Cross-Attention & Frequency-Aware Routing
Each pathway (anatomy, disease, delta) attends separately to the U-Net hidden state queries via distinct projections. Routing gates—either hierarchical (disease at fine scales, anatomy at coarse), inverted, or uniform—control the mixing of anatomical and pathological information at various spatial resolutions.
Figure 3: Triple Pathway Cross-Attention mechanism illustrating parallel attention over purified anatomy, disease, and delta tokens with pathway-specific projections.
Delta Steering for Efficient Modulation
The delta pathway, activated only during inference, is constructed by subtracting the projected source embedding from the target embedding—ensuring zero effect during training but enabling direct, explicit control over the direction and magnitude of severity change at test time. This design yields substantial computational gains relative to two-pass CFG.
Experimental Design
Data and Protocol
Training and evaluation occur on the LIMUC dataset (10,284 endoscopic images, heavily imbalanced over MES grades). All generator variants are fine-tuned from Stable Diffusion v1.4, with VAE and CLIP encoder weights frozen.
Baselines:
- IP-Text: Image generation conditioned on CLIP text encoder outputs (e.g., "MES 2").
- IP-AOE: Conditioning via both IP-Adapter image features and AOE disease tokens, with entanglement unresolved.
Variants: DADD with hierarchical (H), inverted (I), or uniform (U) routing.
Metrics: The assessment includes FID, CMMD, improved precision/recall, downstream ResNet-18 classification, and severity fidelity using an independent regression judge (trained only on real images).
Steering Scale Sensitivity
The delta steering scale (λsteer) controls the strength of severity transformation. Empirical tuning selects λsteer=3 as optimal, balancing pathology modulation and anatomical preservation.
Figure 4: Steering scale selection: Increasing λsteer leads to more pronounced severity transformation but excessive values introduce artifacts.
Results
Quantitative Analysis
DADD-H demonstrates superior downstream discriminative fidelity for synthetic images, achieving 61.68% accuracy and 70.74 QWK in classification tasks using exclusively synthetic data and 81.53 QWK/0.46 RMSE in regression, nearly converging to real-data upper bounds. Notably, DADD-I and DADD-U, while yielding superficially lower (better) FID and CMMD scores, fail to generate clinically discriminative pathology, as evidenced by sub-random-classification and regression results.
This disconnect—where under-modulated synthetic images are nearly indistinguishable from real, anatomically correct but pathologically unchanged references—demonstrates the inadequacy of standard image similarity metrics for conditional progression synthesis. The results validate that spatial-frequency-aware injection of disease information is essential to modulate the relevant pathological subspaces.
DADD also confers a 2× inference speedup via single-pass delta steering compared to CFG-dependent baselines.
Qualitative and Feature Space Examination
Side-by-side visualizations confirm that DADD-H produces monotonic, distinguishable MES transitions while maintaining patient-identifying structural cues; the IP-AOE baseline fails to erase source pathology during backward transitions due to entanglement.
Figure 5: Qualitative comparison: DADD-H shows controllable, monotonic transformations in MES severity while preserving anatomy; IP-AOE demonstrates pathological entanglement.
Feature space UMAP visualizations and low Silhouette Scores indicate strong, global overlap between synthetic and real distributions, validating the realism of DADD-H outputs.
Figure 6: Feature-space UMAP: Real and synthetic distributions largely overlap in the ResNet-18 embedding space, suggesting strong anatomical preservation.
Implications and Future Directions
Practically, DADD enables a new paradigm for clinical data augmentation: generating synthetic, anatomically-consistent images at arbitrary disease states for each patient. This approach rebalances datasets skewed toward mild cases and boosts downstream classification and regression performance, serving as a template for similar applications in other longitudinal, progression-based medical domains.
Theoretically, DADD exemplifies the power of cross-attention-based feature editing for explicit disentanglement, potentially generalizable to other domains where attribute entanglement hinders conditional generative modeling. Its delta steering strategy is a generic, efficient alternative to multi-pass conditioning schemes.
However, DADD's fidelity depends on careful tuning of the steering scale; extreme values reintroduce artifacts and disrupt anatomical identity. Extension to per-class adaptive steering schedules and broader validation on other modalities such as radiology or dermatology are future avenues.
Conclusion
DADD provides an effective, efficient, and interpretable mechanism for controllable synthesis of disease progression in ulcerative colitis endoscopy, addressing a core limitation of previous diffusion-based data augmentation approaches. By explicitly erasing embedded disease features via cross-attention, harnessing frequency-aware routing, and introducing a single-pass, signed delta steering mechanism, DADD achieves high-fidelity, monotonic progression control and nearly matches real-data performance in severity scoring tasks. It sets a technical foundation for anatomically and pathologically disentangled conditional generation—a requirement for next-generation medical AI systems.