Source-linked AI summary
UniX: Unifying Autoregression and Diffusion for Chest X-Ray Understanding and Generation
Ruiheng Zhang, Jingfeng Yao, Huangxuan Zhao, Hao Yan, Xiao He, Lei Chen, Zhou Wei, Yong Luo, Zengmao Wang, Lefei Zhang, Dacheng Tao, Bo Du
TL;DR
Medical foundation models struggle to unify semantic understanding with pixel-level generation because shared architectures and discrete generation can compromise these objectives. UniX decouples them into autoregressive and diffusion branches linked by cross-modal self-attention, supported by staged training. It reports stronger dual-task performance with fewer parameters and competitiveness with dedicated single-task models.
Problem
Unified medical models struggle with conflicting semantic-abstraction and pixel-reconstruction objectives, while discrete generation loses fine-grained image detail.
Method
UniX uses decoupled autoregressive understanding and diffusion generation branches connected by cross-modal self-attention that dynamically guides generation with semantic features.
Results
46.1% improvement in understanding performance (Micro-F1) and 24.2% improvement in generation performance (FD-Raddino) are achieved using only a quarter of LLM-CXR’s parameters.
Takeaways & Limitations
UniX delivers stronger understanding and generation than prior unified models while reaching competitiveness with dedicated single-task medical foundation models.
Takeaways & Limitations
Continued fine-tuning can reduce loss while key performance metrics decline, indicating overfitting to specific patterns rather than general principles.
Abstract
from arXiv · showhide
Despite recent progress, medical foundation models still struggle to unify visual understanding and generation, as these tasks have inherently conflicting goals: semantic abstraction versus pixel-level reconstruction. Existing approaches, typically based on parameter-shared autoregressive architectures, frequently lead to compromised performance in one or both tasks. To address this, we present UniX, a next-generation unified medical foundation model for chest X-ray understanding and generation. UniX decouples the two tasks into an autoregressive branch for understanding and a diffusion branch for high-fidelity generation. Crucially, a cross-modal self-attention mechanism is introduced to dynamically guide the generation process with understanding features. Coupled with a rigorous data cleaning pipeline and a multi-stage training strategy, this architecture enables synergistic collaboration between tasks while leveraging the strengths of diffusion models for superior generation. On two representative benchmarks, UniX achieves a 46.1% improvement in understanding performance (Micro-F1) and a 24.2% gain in generation quality (FD-RadDino), using only a quarter of the parameters of LLM-CXR. By achieving performance on par with task-specific models, our work establishes a scalable paradigm for synergistic medical image understanding and generation. Codes and models are available at https://github.com/ZrH42/UniX.
1. Introduction
UniX addresses the conflict between semantic understanding and pixel-level medical image generation through decoupled branches connected by cross-modal guidance. Its data and training strategies support efficient dual-task modeling, with reported gains over prior unified approaches.
- Unified medical models face conflicting objectives because understanding requires semantic abstraction while generation requires pixel-level reconstruction.
- UniX separates understanding and generation into autoregressive and diffusion branches, respectively, to reduce task interference and preserve high-fidelity synthesis.
- Cross-modal self-attention dynamically injects understanding features into the diffusion process, linking semantic reasoning with image generation.
- A rigorous data-cleaning pipeline and stagewise optimization train the understanding branch first, then pre-train and fine-tune the generation branch while progressively freezing branches.
- 46.1% improvement in understanding performance (Micro-F1) and 24.2% improvement in generation performance (FD-Raddino) are achieved using only a quarter of LLM-CXR’s parameters.
2. Related Work
Medical foundation models have developed along separate understanding and generation tracks, while unified models seek to combine both within one architecture. UniX uses autoregressive understanding, diffusion generation, and cross-modal interaction to address the limitations of shared and discrete approaches.
- Single-task Medical Foundation Models: Single-task understanding models support diagnosis, question answering, and report generation, commonly combining visual encoders with autoregressive language models.
- Single-task Medical Foundation Models: Single-task generation models use diffusion approaches for synthesis, super-resolution, and inpainting, with synthetic data potentially supporting visual understanding.
- Unified Medical Foundation Models: Unified models typically share a Transformer backbone across understanding and generation, but parameter sharing creates conflicts between abstraction and detail preservation.
- Unified Medical Foundation Models: Discrete generation compresses continuous pixels into fixed codebooks, discarding high-frequency details and limiting fidelity for fine-grained pathological patterns.
- Unified Medical Foundation Models: UniX combines autoregressive and diffusion paradigms in decoupled branches, with cross-modal self-attention dynamically allowing understanding features to guide generation.
3. Method
UniX separates chest X-ray understanding and generation into autoregressive and diffusion branches, then links them with cross-modal self-attention for coordinated processing. A staged optimization strategy progressively trains the branches while preserving their specialized roles.
- Dual-branch architecture: UniX uses an autoregressive branch for semantic abstraction and a diffusion branch for high-fidelity chest X-ray synthesis.The generation branch operates in a VAE-encoded latent space and reconstructs images from semantics extracted by the understanding branch.
- Understanding via autoregression: The understanding branch models multimodal comprehension autoregressively for visual semantics, report reasoning, and causal visual-textual interpretation.It trains on a multimodal sequence and computes cross-entropy over output textual tokens.
- Generation via latent diffusion: The generation branch estimates a target velocity field in latent diffusion using mean-square error on noisy latent variables.Semantic embeddings from the understanding branch condition generation, enabling disease-specific synthesis and improved lesion localization.
- Cross-modal self-attention: Cross-modal self-attention performs bidirectional information flow over a unified sequence, allowing semantic features to modulate generation and generative states to feed back.Unlike conventional cross-attention, interactions are learned through joint attention scores rather than explicit conditioning.
- Cross-modal self-attention: The shared attention operation retains distinct parameter spaces for understanding and generation tokens while synchronizing both branches.The mechanism is intended to improve the fidelity and clinical consistency of generated images.
- Three-stage training pipeline: UniX progressively aligns the branches through three training stages, freezing one branch while training or pre-training the other.Stage 1 trains understanding with paired images and reports; later stages freeze understanding and train generation, including high-resolution fine-tuning.
4. Experiments
UniX is evaluated on chest X-ray understanding and generation using cleaned data, staged training, and comparisons with unified and task-specific models. It shows strong understanding, generation, and pathology-specific performance, while training dynamics reveal plateauing and overfitting risks.
- Experimental Setup: Experiments use MIMIC-CXR data with cleaned image–report pairs for understanding and separate generation-branch processing following CheXGenBench.The understanding split contains 163,344 training and 2,365 test pairs; the generation split contains 237,387 training and 4,352 test pairs.
- Evaluation: UniX evaluates report reliability with CheXbert F1 and generation using FD-RadDino, KD-RadDino, alignment, and PRDC metrics.The PRDC metrics assess generation accuracy and diversity, while Alignment Score measures image–text consistency.
- Training Dynamics: Understanding performance improves rapidly before plateauing, while continued fine-tuning lowers loss but causes key metrics to decline through overfitting.Generation-branch mean squared error also decreases rapidly early and then slows, with sharper metric gains between 50K and 75K pre-training steps.
- Understanding Results: UniX improves understanding over unified models, outperforms similarly scaled single-task models, and approaches larger single-task models.The comparison emphasizes performance relative to parameter count rather than a single reported numerical score.
- Generation Results: UniX improves image quality and image–text alignment over LLM-CXR, consistently outperforms single-task models, and performs on par with Sana.The generation benchmark also reports strong accuracy and diversity, with UniX matching or slightly exceeding Sana on some metrics.
- Pathology-Specific Results: UniX achieves top performance in 13 of 14 pathology categories and remains competitive with Sana on fine-grained lesion generation.The reported results emphasize fidelity to subtle pathological cues and clinically relevant details across diverse diagnostic conditions.
5. Ablations
Ablations show that keeping the understanding branch fixed during generation training protects comprehension and accelerates generation. Updating both branches creates trade-offs, especially without semantic supervision.
- Data Processing: Data cleaning removes non-diagnostic report content so training targets emphasize clinically relevant visual–textual associations.The described noise includes underscores, technical metadata, and conversational fillers.
- Optimization Strategy: Freezing the understanding branch during generation training yields rapid generative gains without harming understanding accuracy.The ablation compares generation-only updates with joint updates under different data mixtures.
- Optimization Trade-offs: Unfreezing the understanding branch without understanding data severely degrades comprehension and provides no generation benefit.Adding understanding data partially mitigates comprehension loss but slows generative learning because both objectives share updates.
6. Conclusion
UniX combines architectural decoupling with cross-modal coordination for chest X-ray understanding and generation. Its dual-branch design and self-attention mechanism improve both tasks while using fewer parameters than prior unified models.
- UniX combines autoregressive understanding with diffusion-based generation in a dual-branch architecture that decouples the two tasks and prevents mutual interference.The autoregressive branch handles understanding, while diffusion-based generation addresses visual synthesis.
- Cross-modal self-attention aligns the branches and dynamically uses understanding features to guide generation.
- UniX delivers stronger understanding and generation performance than prior unified medical foundation models while using fewer parameters, reaching competitiveness with dedicated single-task models.