CheXmix: Unified Generative Pretraining for Vision Language Models in Medical Imaging

Abstract: Recent medical multimodal foundation models are built as multimodal LLMs (MLLMs) by connecting a CLIP-pretrained vision encoder to an LLM using LLaVA-style finetuning. This two-stage, decoupled approach introduces a projection layer that can distort visual features. This is especially concerning in medical imaging where subtle cues are essential for accurate diagnoses. In contrast, early-fusion generative approaches such as Chameleon eliminate the projection bottleneck by processing image and text tokens within a single unified sequence, enabling joint representation learning that leverages the inductive priors of LLMs. We present CheXmix, a unified early-fusion generative model trained on a large corpus of chest X-rays paired with radiology reports. We expand on Chameleon's autoregressive framework by introducing a two-stage multimodal generative pretraining strategy that combines the representational strengths of masked autoencoders with MLLMs. The resulting models are highly flexible, supporting both discriminative and generative tasks at both coarse and fine-grained scales. Our approach outperforms well-established generative models across all masking ratios by 6.0% and surpasses CheXagent by 8.6% on AUROC at high image masking ratios on the CheXpert classification task. We further inpaint images over 51.0% better than text-only generative models and outperform CheXagent by 45% on the GREEN metric for radiology report generation. These results demonstrate that CheXmix captures fine-grained information across a broad spectrum of chest X-ray tasks. Our code is at: https://github.com/StanfordMIMI/CheXmix.

• The paper introduces an early-fusion unified generative pretraining strategy that jointly models discrete image and text tokens, bypassing projection bottlenecks.
• It achieves state-of-the-art performance in CheXpert classification, image inpainting, and radiology report generation with notable gains in AUROC and PSNR metrics.
• The model’s use of 50% token masking and causal attention demonstrates robustness under high occlusion, offering versatile applications in clinical AI.

CheXmix: Unified Generative Pretraining for Vision-LLMs in Medical Imaging

Introduction and Motivation

Medical imaging analysis using multimodal foundation models (FMs) has historically relied on contrastive learning paradigms adapted to LLMs through LLaVA-style late-fusion architectures. This architectural decoupling imposes a projection bottleneck, which is especially problematic in high-stakes clinical imaging due to the attenuation of fine-grained features. CheXmix proposes an early-fusion unified generative approach, leveraging a joint sequence of discrete image and text tokens rather than adapting off-the-shelf visual encoders via projection layers. The motivation is to directly align image and language modalities within a transformer decoder, thereby maximizing cross-modal information flow and mitigating catastrophic forgetting associated with multimodal instruction tuning.

Model Architecture and Pretraining Strategy

CheXmix extends the Chameleon early-fusion paradigm to the clinical domain, introducing a two-stage multimodal generative pretraining protocol:

1. Stage 1 (S1): Standard autoregressive next-token prediction is performed on sequences mixing chest X-ray VQ-GAN tokens and radiology report tokens, initialized from the RadPhi-2 clinical LLM.
2. Stage 2 (S2): Masked autoregressive training, where 50% of both image and text tokens are masked and the network learns to reconstruct the missing information, enhances the discriminative and generative capacity—integrating strengths of MAEs and MLLMs for robust fine-grained representation learning.

This early-fusion, unified architecture contrasts with state-of-the-art LLaVA-style medical MLLMs (e.g., CheXagent) that require post-hoc instruction finetuning and maintain separate discriminative and generative modules.

The architectural distinction and the flexible downstream utility of CheXmix are highlighted in the model comparison visualization.

Figure 1: CheXmix and CheXagent architectural comparison, contrasting a unified generative backbone with modular, late-fusion LLaVA-style decoupling.

The generative pretraining framework, token masking schedule, and the operational flexibility (classification, inpainting, report generation, retrieval) are presented in the methodological overview.

Figure 2: Overview of CheXmix’s generative pretraining pipeline and the multimodal tasks used for evaluation.

Experimental Protocol

Training data comprises over 550,000 chest X-ray and radiology report pairs from five public datasets, yielding a sequence corpus exceeding 627 million tokens. CheXmix is trained on 1,024-token VQ-GAN image sequences and tokenized text capped at 1,300 tokens—covering both findings and impression sections.

Evaluation spans:

• CheXpert findings classification (discriminative)
• Image inpainting under variable masking
• Radiology report generation (generative, including test-time augmentation)
• Multimodal retrieval (image-to-text, text-to-image)
• External validation on ChestX-ray14 (classification) and ReXGradient (report gen.)

Results

Discriminative Performance: CheXpert Classification

Under no masking, CheXmix (S1+S2) achieves an AUROC of 0.712 for CheXpert findings, outperforming general-domain generative baselines (Chameleon +2.4%, M3AE +3.6%, MAE +7.1%) and its own S1 variant (+7.2%). Robustness to masked/occluded visual input is a key finding: at 40–80% masking, CheXmix (S1+S2) surpasses CheXagent (SigLIP encoder) by up to 8.6% AUROC—establishing superiority in the fine-grained, partial-information regime.

Generative Performance: Image Inpainting

Quantitative metrics (PSNR, MS-SSIM) consistently favor CheXmix (S1+S2), particularly with high masking. At 80% masking, CheXmix delivers a 51.1% PSNR improvement over the text-only generative baseline and a 20% gain over S1. Qualitative reconstructions show perceptual retention of relevant anatomic details at extreme occlusion.

Figure 3: Inpainting improvement with CheXmix pretraining at high masking; anatomical structures are preserved when 80% of tokens are missing.

Figure 4: CheXmix inpainting visualizations for pathologies (COPD and consolidation cases) at high masking rates.

Generative Performance: Radiology Report Generation

On clinical radiology report generation (measured by GREEN and CheXbert), CheXmix (S1+S2) exceeds the performance of CheXagent by 45% on factual consistency (GREEN) while matching CheXbert. The performance drop with increasing masking is modest (only 25% on GREEN from 0% to 80% masking), compared to the catastrophic collapse in variant S1.

Test-time augmentation (TTA), using ensemble generation from disjoint masking, further boosts report fidelity—with up to 16% improvement on CheXbert at 80% masking.

Figure 5: Test-time augmentation mechanism with disjoint masked views yields substantial gains in radiology report quality.

Retrieval and External Validation

CheXmix matches or closely trails specialized contrastive/retrieval encoders (e.g., SigLIP) in both image-to-text and text-to-image retrieval, demonstrating strong cross-modal alignment from generative pretraining.

On external datasets, CheXmix generalizes robustly, outperforming domain-specific and generalist models on ChestX-ray14 and ReXGradient.

Architectural and Training Ablations

Ablation studies confirm:

• Masking during pretraining is essential—50% masking yields optimal discriminative performance, and benefits persist up to 90%.
• Causal attention masking (as opposed to bidirectional masking for image tokens, which is common in general vision transformers) provides consistent improvements in both discriminative and generative settings, contrary to findings in some general-domain unified models.

Implications and Future Directions

CheXmix demonstrates that unified early-fusion generative pretraining—jointly modeling image and text tokens in a clinical context—yields a single, flexible backbone capable of SOTA performance under both discriminative and generative objectives. It eliminates the need for separate modules or instruction finetuning post-projection and showcases robustness to missing or occluded image information, which is common in clinical practice.

These findings indicate that domain-specific early-fusion generative modeling, with masked pretraining, may supersede contrastive and adapter-based architectures for future biomedical foundation models. This paradigm shift is of significant practical value in clinical AI, promising increased accuracy, diagnostic fidelity, and model versatility. From a theoretical perspective, these results provide new evidence regarding the inductive biases required for fine-grained medical representations and the importance of causal generative objectives in specialized domains.

Conclusion

CheXmix validates the efficacy of early-fusion, masked generative pretraining for unified vision-language modeling in the medical domain. It achieves top performance in chest X-ray classification, image inpainting, and clinical report generation, especially under partial information. The unified architecture exhibits strong flexibility and sets a precedent for future medical multimodal FMs, warranting further exploration of early-fusion approaches in broader clinical modalities and high-resolution scenarios.