UniX: Unifying Autoregression and Diffusion for Chest X-Ray Understanding and Generation

TL;DR

UniX tackles the challenge of unifying chest X-ray understanding and generation by decoupling the tasks into an autoregressive understanding branch and a diffusion-based generation branch, then coupling them through cross-modal self-attention. The model uses a three-stage training pipeline and a rigorous data-cleaning process to enable synergistic learning while maintaining parameter efficiency (approximately one-quarter the parameters of large task-specific baselines). Empirically, UniX achieves substantial gains on understanding (Micro-F1) and generation (FD-RadDino) benchmarks, while matching or approaching single-task models with far fewer parameters. This architecture offers a scalable paradigm for combining semantic reasoning and high-fidelity medical image synthesis in a unified framework, with strong practical implications for chest X-ray analysis and data augmentation.

Abstract

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.

Figures (4)

• Figure 1: The quantitative and qualitative results of UniX. Quantitative results show UniX's superiority over existing unified and single-task medical foundation models in understanding and generation. Qualitatively, UniX enables multi-disease X-ray interpretation and high-fidelity medical image generation.
• Figure 2: Model Architecture. UniX comprises two decoupled yet synergistic branches: an autoregressive understanding branch for semantic encoding, and a diffusion-based generation branch for visual synthesis. To enable effective collaboration between them, we introduce a cross-modal self-attention mechanism that allows semantic features to dynamically guide the generation process. Data Processing and Training Pipeline. To fully exploit the potential of this architecture, we design a rigorous data cleaning pipeline and a three-stage training strategy. This strategy progressively freezes the branches during different stages, ensuring efficient knowledge transfer and stable training.
• Figure 3: Demonstration of Data Processing and Report Generation Efficacy. The application of large language models enables the purification of raw data by eliminating extraneous information. This process ensures that the model prioritizes and extracts pertinent information related to disease diagnosis.
• Figure 4: Qualitative Examples from UniX.(A)-(C) illustrate the model's precise control over the attributes of generated findings, including their severity and location. In (D), the model successfully synthesizes a complex radiographic scene containing multiple findings that are consistent with a full clinical report, highlighting its ability to process and integrate extensive contextual information.