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

AI summary

Overview

Research area: Computer vision and vision-language modeling for 3D medical imaging, specifically chest CT report generation and text-conditioned 3D CT synthesis.

Technical level: Advanced. The paper assumes familiarity with encoder-decoder architectures, discrete latent representations, wavelet transforms, contrastive pretraining, diffusion/flow-matching generation, and clinical NLP metrics.

Scope: The paper introduces BTB3D, a causal convolutional encoder-decoder that produces compact volumetric tokens for 3D CT, trained with a three-stage curriculum, and evaluates it on reconstruction, radiology report generation, and text-to-CT synthesis.

What This Paper Is About

Vision-language models for 3D medical images lag behind their 2D counterparts because 3D scans contain hundreds of slices, paired 3D image-text data is scarce, and existing vision encoders are pretrained with contrastive objectives that are poorly suited to radiology. Slice-wise tokenization also blurs fine anatomy, so diagnostic detail is lost. The paper's goal is to fix the tokenization and reconstruction side of the pipeline — producing better volumetric tokens — rather than relying on larger language backbones alone.

Key Contributions

1. A causal convolutional encoder-decoder (BTB3D) that unifies 2D and 3D training and inference. Causal 3D convolutions let the model process scans slice by slice, in a manner analogous to a temporal sequence, so it can scale to arbitrarily long scans and remain compatible with pretrained 2D features.

2. Frequency-aware volumetric tokenization. A 3D Haar wavelet transform compresses each CT, and lookup-free quantization converts the encoder output into discrete integer tokens without any codebook or embedding lookup, improving speed and memory efficiency for large volumes.

3. A three-stage training curriculum consisting of (i) local reconstruction, (ii) overlapping-window tiling, and (iii) long-context decoder refinement. The model learns from short slice excerpts yet generalizes to scans exceeding 300 slices without additional memory overhead.

4. State-of-the-art results on two downstream tasks: clinical F1 improvement of 40% over CT2Rep, CT-CHAT, and Merlin for report generation, and a 75% FID reduction plus halved FVD compared to GenerateCT and MedSyn for text-to-CT synthesis, producing anatomically consistent 512 × 512 × 241 volumes.

Main Findings

• Three-stage training drives reconstruction quality. On full-volume reconstruction, the 8×8×8 variant rises from PSNR 9.350 / SSIM 0.206 / MSE 0.117 at Stage 1, to 23.980 / 0.697 / 0.005 at Stage 2, to 28.166 / 0.760 / 0.001 at Stage 3. The 16²×8 variant goes from 11.067 / 0.353 / 0.079 to 23.808 / 0.700 / 0.005 to 26.750 / 0.749 / 0.002. Stage 2 yields the largest single improvement: PSNR increases by over 14 dB, SSIM triples, and MSE drops an order of magnitude.

• Report generation on CT-RATE improves. The 16×16×8 variant reaches F1 0.258, precision 0.260, recall 0.260, CRG 0.370, BLEU-1 0.439, and BLEU-mean 0.305, versus CT-CHAT at F1 0.184, BLEU-1 0.373, BLEU-mean 0.272. This is described as a 40% relative F1 improvement over CT-CHAT, with BLEU-1 up 18% and BLEU-mean up 12%. Merlin shows high precision but low recall (under-reporting), while CT2Rep and CT-CHAT show high recall but lower precision (hallucinated abnormalities).

• Out-of-distribution generalization holds on RadChestCT. The 16×16×8 model achieves F1 0.266, precision 0.272, recall 0.329, versus Merlin and CT-CHAT both at F1 0.182 and CT2Rep at 0.133 — a 46% relative improvement over the best baselines. Text-based metrics are omitted because only binary labels are available.

• Text-conditional CT synthesis improves substantially. The 8×8×8 variant reaches mean FID 2.236 (axial 2.479, sagittal 2.166, coronal 2.062), FVD with CT-Net 3.955 and with I3D 325.51, CLIP Text-Img 24.270 and Img-Img 88.352. GenerateCT scores mean FID 9.512 and FVD 7.659 (CT-Net), MedSyn mean FID 12.592 and FVD 13.927. The paper reports mean FID dropping from 9.51 to 2.24 (a 76.5% reduction) and FVD-CT-Net improving by 48.3%.

• Compression is a trade-off. The 8×8×8 variant excels at fine-grained tasks such as text-to-CT synthesis, while the 16×16×8 variant suits memory-constrained settings and high-level semantic tasks such as report generation.

• Contrastive pretraining is argued to be the wrong fit. Because paired reports typically describe an entire scan, training on cropped or partial slices introduces semantic misalignment, and in radiology multiple reports can describe the same conditions differently, so penalizing unmatched pairs may degrade medical semantics.

Methodology in Plain English

The pipeline converts a CT volume into a compact sequence of discrete tokens and back again.

Compression step. Each scan is passed through a 3D Haar wavelet transform. Every non-overlapping 2×2×2 voxel block yields 8 subband coefficients — one low-frequency and seven high-frequency components. The channel dimension grows from 1 to 8, but the volume size shrinks by a factor of 8, cutting memory and compute. For 2D single slices, the same 3D transform applies, producing a low- and high-frequency pair, so 2D and 3D training share one architecture.

Encoder. Residual blocks use factorized 3D convolutions that separate spatial and temporal processing: a 1×k×k spatial convolution followed by a k×1×1 temporal convolution, with causal padding of (k−1) zero slices in the past and none in the future. This makes encoding strictly causal — the token at index t only sees slices ≤ t. Downsampling uses strided convolutions. The 8×8×8 configuration applies two stride-2 convolutions per axis (4× reduction), which combined with the 2× wavelet downsampling gives effective 8× compression; the 16×16×8 variant adds a spatial stride-2 layer for 16× spatial and 8× temporal compression.

Quantization. Lookup-free quantization binarizes each feature vector by taking the sign of each element, packs the results into integers, and forms a token map. The number of codes K = 2^d, with token dimension 18 and codebook size 262,144. No codebook or embedding lookup is needed, and an entropy regularization term encourages uniform code usage to prevent code collapse.

Decoder. A mirror-image causal decoder with transposed convolutions reconstructs the wavelet-domain volume, and the inverse Haar transform produces the final CT. Training losses are reconstruction (L1, chosen over L2 for sharper detail), adversarial (a 3D discriminator applied directly in the CT domain rather than the wavelet domain, for stability), and quantization. Perceptual losses such as VGG-based features are deliberately omitted because natural RGB features mismatch grayscale medical images.

Three-stage training. Stage 1 trains everything end-to-end on single slices or 9-slice subvolumes to learn local structure. Stage 2 uses overlapping 9-slice windows stepping by 8 slices; in each window the first token is discarded and the second kept (both are kept in the first window), producing the latent sequence [z₁¹, z₂¹, z₂², z₂³, …, z₂ᵀ] with T = ⌊(D−1)/8⌋. This overlapping scheme maintains temporal consistency at lower memory cost than one-shot encoding, and the decoder processes all tokens in a single forward pass. Stage 3 freezes the encoder and codebook and fine-tunes only the decoder to model long-range anatomical dependencies. For inference, the paper offers one-shot and tiled strategies, adopting tiled inference for all downstream experiments due to its robustness and consistency.

Setup. Training and in-domain evaluation use CT-RATE: 25,692 chest CT scans from 21,304 patients, with 20,000 patients in training and 1,304 in test. Volumes are converted to Hounsfield Units and clipped to [−1000, 1000]. Stage 3 resamples to 0.75 × 0.75 × 1.5 mm and crops/pads to 512 × 512 × 241. Pretraining ran on 64 NVIDIA H100 GPUs with DDP and mixed precision (Stage 1 batch size 8 for 9-slice subvolumes and 40 for single slices; Stages 2 and 3 batch size 1 with 201 and 241 slices). Report generation uses LLaMA 3.1-8B with the CT-CHAT setup, trained for 40,000 iterations with DeepSpeed ZeRO-3 on 40 H100 GPUs. Generation uses a 12-layer transformer (1024 hidden size, 16 heads) with flow matching, [7,7,7] windowed self-attention, [2,2,2] patching, and T5v1.1-base prompt encoding on 16 H100 GPUs for 1500 epochs.

Why This Matters

Research impact. The paper argues that precise three-dimensional tokenization, rather than larger language backbones alone, is essential for scalable vision-language modeling in 3D medical imaging. It also offers a concrete critique of contrastive pretraining in the medical domain, and releases an open-source codebase at https://github.com/ibrahimethemhamamci/BTB3D.

Real-world applications:

• Automated radiology report generation from chest CT, reducing radiologist documentation load.
• Text-guided synthesis of anatomically consistent CT volumes for data augmentation where paired 3D data is scarce.
• Educational and simulation use, generating example volumes from clinical prompts for training.
• Surgical and treatment planning, where high-fidelity synthetic volumes can supplement limited real scans.

Industry relevance. Co-authorship spans the University of Zurich, NVIDIA, Imperial College London, and FAU Erlangen-Nürnberg, indicating direct commercial interest in efficient 3D medical foundation models. The lookup-free quantization and causal design target memory efficiency, which matters for deploying long-volume models on constrained hardware. The paper notes that its method requires substantial computational resources, which may limit accessibility.

Future Directions

• Extending beyond chest CT. The authors acknowledge BTB3D is evaluated only on chest CT because large-scale paired datasets for other anatomies and modalities are lacking; applying it to MRI or other body regions is an open problem.
• Clinical validation. The paper states that expert reader studies and prospective clinical trials are necessary to assess real-world safety and efficacy, beyond the clinical and generative metrics used here.
• Reducing compute requirements. Training on 64 H100 GPUs for pretraining and 40 H100 GPUs for report generation limits accessibility; more efficient training or inference is a natural next step.
• Comparing against newer encoders. fVLM was excluded because neither its report generation codebase nor model weights were available, and MAISI was excluded because it is not designed for text-conditional generation — leaving head-to-head evaluation against these methods unresolved.

Target Audience

Researchers and engineers working on 3D medical imaging, vision-language models, and generative models for clinical data; medical AI practitioners interested in report generation or synthetic CT for augmentation; and machine learning scientists focused on tokenization, discrete latent representations, or long-sequence volumetric modeling. Readers without a background in deep generative architectures will find the methodology section dense, though the high-level framing of why 3D VLMs lag 2D ones is accessible.

Authors’ 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