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

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

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

TL;DR¶

BTB3D is a causal convolutional encoder-decoder that produces compact, frequency-aware volumetric tokens for 3D CT scans via 3D Haar wavelet compression and lookup-free quantization. A three-stage curriculum (short subvolume pretraining → overlapping-window tiling → decoder-only fine-tuning) enables scalability to 300+ slice volumes without additional memory overhead. It achieves state-of-the-art on both CT report generation and text-to-CT synthesis, demonstrating that tokenization quality—not LLM scale—is the primary bottleneck in 3D medical VLMs.

Problem¶

Existing 3D medical VLMs are bottlenecked by poor volumetric tokenization: (1) contrastive (CLIP-style) pretraining for vision encoders is inappropriate for radiology because multiple reports can describe the same pathology differently, penalizing semantically valid unpaired samples; (2) slice-wise or single-vector compression loses fine anatomy (small nodules, subtle textures); (3) there is no unified 2D/3D representation, preventing transfer of pretrained 2D features; and (4) training on full high-resolution CT volumes (hundreds of slices, 512×512 spatial) is memory-prohibitive, leading to cascaded generation pipelines with inter-slice spatial discontinuities.

Method¶

BTB3D has three components:

Architecture. Input CT x ∈ R^{D×H×W} is first compressed 2× per axis via a 3D Haar wavelet transform into 8-channel frequency subbands. A 3D CNN encoder with factorized causal convolutions (spatial 1×k×k then temporal k×1×1 with (k-1) zero-padded past slices) downsamples to latent tokens with 8× or 16×16×8 effective compression. Lookup-free quantization (LFQ) binarizes each latent vector dimension-wise, yielding integer token maps with codebook size 2^18 = 262,144, without an explicit codebook embedding. A symmetric causal decoder with transposed convolutions reconstructs the wavelet-domain volume, then the inverse Haar transform recovers the CT. Training loss: L = L_rec (ℓ1) + λ_adv L_adv (3D discriminator) + L_vq (commitment + entropy regularization).

Three-stage training curriculum. 1. Stage 1: Train end-to-end on single 2D slices or 9-slice subvolumes to learn local structure. 2. Stage 2: Overlapping temporal tiling—encode successive 9-slice windows with 1-slice overlap; discard the first token of each window (except the first) to maintain consistency; train encoder + decoder on sequences up to 201 slices. 3. Stage 3: Freeze encoder and codebook; fine-tune only the decoder on full 241-slice volumes to improve long-range anatomical coherence.

Downstream integration. For report generation, BTB3D tokens are linearly projected into LLaMA 3.1-8B (matching CT-CHAT setup). For text-to-CT synthesis, tokens drive a 12-layer transformer with cross-attention to T5-encoded text prompts, trained with flow matching.

Key Contributions¶

Causal 3D convolutional encoder-decoder with 3D Haar wavelet compression enabling unified 2D/3D processing and arbitrary-length inference
Lookup-free quantization (LFQ) with entropy regularization for efficient, collapse-free discrete volumetric tokenization
Three-stage training curriculum that scales from 9-slice local pretraining to full 241-slice volumes without increased memory footprint
Reconstruction-based (rather than contrastive) pretraining that avoids radiology-specific semantic mismatch issues
Two compression variants (8³ for fine-grained tasks, 16²×8 for memory-constrained/global tasks) with demonstrated trade-off characterization

Results¶

Reconstruction (CT-RATE, 3-stage ablation): - 8³ model: Stage 1 PSNR 9.35 dB → Stage 2 23.98 dB (+14.6 dB) → Stage 3 28.17 dB; SSIM 0.206 → 0.697 → 0.760 - 16²×8 model: Stage 3 PSNR 26.75 dB, SSIM 0.749

Radiology report generation (CT-RATE test set, 1,304 patients): - Ours-16 F1 = 0.258 vs. CT-CHAT 0.184, Merlin 0.160, CT2Rep 0.160 (+40% relative over CT-CHAT) - BLEU-1: 0.439 vs. CT-CHAT 0.373 (+18%); BLEU-mean: 0.305 vs. 0.272 (+12%) - Out-of-distribution (RadChestCT, binary labels only): Ours-16 F1 = 0.266 vs. best baseline 0.182 (+46% relative)

Text-to-CT generation (CT-RATE, 512×512×241 volumes): - Mean FID: 2.236 (Ours-8) vs. 9.512 (GenerateCT), 12.592 (MedSyn) — 76.5% reduction - FVD (CT-Net): 3.955 vs. 7.659 (GenerateCT) — 48.3% improvement - CLIP text-image alignment: 24.27 vs. 23.63 (GenerateCT)

Limitations¶

Evaluated exclusively on chest CT (CT-RATE dataset); generalization to other anatomies (abdomen, brain) is undemonstrated
Training is compute-intensive: Stages 1–3 require 64 H100 GPUs; report generation requires 40 H100s
fVLM (a strong CLIP-based 3D VLM baseline) was excluded from comparisons due to unavailable weights/code, leaving the contrastive vs. reconstruction comparison partially incomplete
Tiled inference requires overlapping forward passes and is slower than single-pass encoding
LFQ codebook size (262,144) and entropy regularization sensitivity are not ablated
No evaluation on non-chest pathologies or multi-organ settings

Relevance to Vision-Language Models¶

BTB3D directly addresses the vision encoder bottleneck in medical VLMs—an underexplored axis compared to scaling LLM decoders—and demonstrates that reconstruction-based pretraining outperforms contrastive objectives for domains where one-to-many image-text relationships invalidate CLIP's negative-pair assumption. The causal 3D convolutional design with wavelet compression is a concrete architecture for bridging 2D pretrained features and 3D volumetric data, a recurring challenge in non-natural-image VLMs. The work also validates that discrete tokenization quality determines downstream generation and understanding performance more than backbone scale, which has implications for VLM design beyond medical imaging. The open-source release positions BTB3D as a foundation encoder for future 3D medical VLMs analogous to the role ViT/CLIP encoders play in 2D systems.