Skip to content

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

🕒 Published (v1): 2026-01-16 18:59 UTC · Source: Arxiv · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

UniX is a unified chest X-ray foundation model that resolves the inherent conflict between visual understanding and generation by decoupling them into an autoregressive branch and a latent diffusion branch, linked via cross-modal self-attention. At 1.5B parameters—one quarter of the competing LLM-CXR—it achieves 46.1% higher Micro-F1 for report generation and 24.2% better FD-RadDino generation fidelity. It matches task-specific single-purpose models on both tasks simultaneously.

Problem

Existing unified medical foundation models (e.g., LLM-CXR, HealthGPT) share parameters across understanding and generation, creating feature interference: understanding demands semantic abstraction while generation demands pixel-level reconstruction. Additionally, discrete autoregressive generation (codebook-based) loses high-frequency structural detail critical in medical images. Neither parameter sharing nor H-LoRA decomposition fully resolves this paradigm mismatch.

Method

UniX adopts a decoupled dual-branch architecture initialized from Janus-Pro:

  • Understanding branch: A VLM with SigLIP-large-patch16-384 visual encoder (1024-dim → 2048-dim via 2-layer MLP) and a 24-layer LM backbone, trained autoregressively with cross-entropy loss over generated report tokens.
  • Generation branch: Latent diffusion operating in a 16× downsampled, 16-channel VAE latent space; minimizes an MSE flow-matching objective over noise-conditioned latent embeddings.
  • Cross-modal self-attention: Performs joint self-attention over a unified sequence \(S = [T_{in}, N]\) (understanding text tokens + generation latent embeddings) using modality-specific Q/K/V projection matrices (\(W^u_{q,k,v}\) vs. \(W^g_{q,k,v}\), selected by binary modality masks). Attention scores are computed globally, allowing understanding features to implicitly condition generation without explicit cross-attention.

Three-stage training on MIMIC-CXR (163K understanding pairs; 237K generation pairs): 1. Stage 1: Freeze generation branch; SFT understanding branch at 384×384 with cleaned reports (DeepSeek used to strip non-diagnostic noise). 2. Stage 2: Freeze understanding branch; pre-train generation branch at 256×256 with REPA loss aligning the 8th-layer hidden states to RadDino image features (weight ratio 0.5). 3. Stage 3: Freeze understanding branch; fine-tune generation branch at 512×512, extended positional encodings, no feature-level supervision.

Key Contributions

  • Dual-branch architecture that structurally decouples autoregressive understanding from diffusion-based generation, eliminating feature interference.
  • Cross-modal self-attention with modality-specific projection matrices enabling dynamic, content-aware semantic conditioning of the generation process.
  • Three-stage progressive training pipeline with REPA alignment for efficient knowledge transfer.
  • DeepSeek-powered data cleaning pipeline to remove non-diagnostic noise from MIMIC-CXR reports, reducing hallucinations.
  • Parameter-efficient unification: 1.5B parameters outperforming the 12B LLM-CXR on both tasks.

Results

  • Understanding (CheXbert Micro-F1-14, uncertain-as-negative): UniX 53.6 vs. HealthGPT 24.2 (+121%) and LLM-CXR (not reported); comparable to LLaVA-Rad (57.3, 7B) at 1.5B.
  • Understanding (Micro-F1-5): UniX 56.6 vs. HealthGPT 25.1.
  • Generation (FD-RadDino, 512×512): UniX 54.022 vs. LLM-CXR 71.243 at 256×256 (−24.2%); on par with Sana (54.68 at 512×512, specialized single-task model, 0.6B params, same fine-tune data).
  • Generation (Alignment Score, 512×512): UniX 0.635 vs. LLM-CXR 0.319.
  • Pathology-specific generation (FD-RadDino per category): UniX wins 13 of 14 pathology categories over LLM-CXR; competitive with Sana across all 14.
  • Ablation: Freezing the understanding branch during generation training is critical—unfreezing without understanding data collapses Micro-F1-14 from 53.2 to 13.9.

Limitations

  • Evaluated exclusively on chest X-rays (MIMIC-CXR); generalizability to other imaging modalities (CT, MRI, pathology) is undemonstrated.
  • Cross-modal self-attention is not directly ablated against simpler conditioning alternatives (e.g., standard cross-attention), so its marginal contribution is unclear.
  • The three-stage frozen-branch strategy prevents end-to-end joint optimization; understanding and generation cannot co-adapt post-Stage 1.
  • Generation is conditioned on free-text reports only; no image-to-image or multi-view synthesis capability.
  • High-resolution (512×512) generation requires extended positional encodings—scalability to larger resolutions is untested.

Relevance to Foundation Models in Medicine

UniX directly addresses the core architectural challenge in medical foundation models: how to unify perception and synthesis without sacrificing either. Its insight—that conflicting objectives require decoupled parameter spaces connected by a learned attention bridge rather than shared weights—offers a transferable design principle for multi-modal medical AI beyond CXR. The result that a 1.5B unified model matches or beats specialized 12B and 0.6B single-task models is a meaningful data point for the scaling and efficiency frontier. The cross-modal self-attention mechanism also demonstrates that generation quality in medical imaging can be improved by conditioning on explicit semantic understanding features, not just text prompts.