Better Tokens for Better 3D: Advancing Vision-Language Modeling in 3D Medical Imaging

TL;DR

BTB3D tackles the core bottleneck of 3D vision–language modeling in radiology by introducing a causal 3D convolutional encoder–decoder that produces compact, frequency-aware volumetric tokens through 3D Haar wavelet compression. A three-stage training curriculum scales learning from short subvolumes to full-length CT volumes, enabling robust reconstruction and long-context decoding without prohibitive memory costs. Empirically, BTB3D sets new state-of-the-art in radiology report generation and text-conditioned 3D CT synthesis, achieving significant gains in clinical metrics (e.g., up to a 40% improvement in F1 over strong baselines) and generative fidelity (substantial FID/FVD reductions). The approach highlights that precise 3D tokenization and scalable training, rather than solely larger language backbones, are essential for effective multimodal 3D medical imaging, with open-source code to foster reproducibility.

Abstract

Recent progress in vision-language modeling for 3D medical imaging has been fueled by large-scale computed tomography (CT) corpora with paired free-text reports, stronger architectures, and powerful pretrained models. This has enabled applications such as automated report generation and text-conditioned 3D image synthesis. Yet, current approaches struggle with high-resolution, long-sequence volumes: contrastive pretraining often yields vision encoders that are misaligned with clinical language, and slice-wise tokenization blurs fine anatomy, reducing diagnostic performance on downstream tasks. We introduce BTB3D (Better Tokens for Better 3D), a causal convolutional encoder-decoder that unifies 2D and 3D training and inference while producing compact, frequency-aware volumetric tokens. A three-stage training curriculum enables (i) local reconstruction, (ii) overlapping-window tiling, and (iii) long-context decoder refinement, during which the model learns from short slice excerpts yet generalizes to scans exceeding 300 slices without additional memory overhead. BTB3D sets a new state-of-the-art on two key tasks: it improves BLEU scores and increases clinical F1 by 40% over CT2Rep, CT-CHAT, and Merlin for report generation; and it reduces FID by 75% and halves FVD compared to GenerateCT and MedSyn for text-to-CT synthesis, producing anatomically consistent 512*512*241 volumes. These results confirm that precise three-dimensional tokenization, rather than larger language backbones alone, is essential for scalable vision-language modeling in 3D medical imaging. The codebase is available at: https://github.com/ibrahimethemhamamci/BTB3D

Figures (6)

• Figure 1: (a) Stage 1: A 9-slice subvolume is compressed via a wavelet transform, then encoded causally using two stride-2 temporal convolutions. The encoder processes the input strictly causally by prepending zero-padded slices. Tokens are decoded back to the wavelet domain by a symmetric causal decoder. (b) Stage 2: To scale to long CT volumes, we introduce overlapping temporal tiling (retaining only the second token from each window) to ensure consistent representation. Stage 3 follows the same scheme but trains only the decoder to refine long-range anatomical reconstruction.
• Figure 2: Reconstruction results of our 8$\times$8$\times$8 model across axial, coronal, and sagittal planes after each training stage. Progressive improvements demonstrate the effectiveness of our three-stage strategy, with Stage 2 providing the largest gains in anatomical fidelity and inter-slice consistency.
• Figure 3: Example generations for a clinical prompt: “Chest CT scan of a 63-year-old male: No findings compatible with pneumonia were detected. Mild sequelae changes are observed in both lungs. Atherosclerotic changes noted, with slight increase in the calibration of vascular structures in the mediastinum. Hepatosteatosis. Hiatal hernia.”. We show one representative slice per anatomical plane. Ground-truth volumes are omitted, following standard practice in text-to-image generation.
• Figure 4: Qualitative reconstruction results across axial, coronal, and sagittal planes for two BTB3D variants: 16$\times$16$\times$8 (top) and 8$\times$8$\times$8 (bottom). The figure shows outputs after Stage 1 (short-volume training), Stage 2 (overlapping-window training), and Stage 3 (decoder refinement), compared to the ground truth. The progressive improvements highlight the effectiveness of our three-stage training strategy. Stage 2 yields the largest gain in anatomical fidelity and inter-slice consistency, while Stage 3 further sharpens structural details such as lung fissures and vascular boundaries.
• Figure 5: Example of radiology report generation for the same 3D chest CT scan using our BTB3D method (both 16$\times$16$\times$8 and 8$\times$8$\times$8 variants) compared to baseline models (CT-CHAT, Merlin, CT2Rep) and the ground truth report. Key phrases from the ground truth are highlighted and matched across model outputs using consistent colors to indicate alignment. Our BTB3D framework, especially the higher compression rate variant, produces more detailed, clinically relevant, and accurate radiology reports, showing superior coverage of anatomical structures and abnormalities.