--- license: cc-by-nc-sa-4.0 datasets: - ibrahimhamamci/CT-RATE - rajpurkarlab/ReXGroundingCT pipeline_tag: text-generation tags: - university-of-kentucky - medical - radiology - chest-ct - vision-language - report-generation - lejepa - llama - lora language: - en base_model: - meta-llama/Llama-3.2-3B --- # Ker-VLJEPA-3B A multi-modal vision-language model for **automated Chest CT radiology report generation**. Ker-VLJEPA-3B takes pre-computed 1024-d slice embeddings from [Guided-Chest-CT-LeJEPA](https://huggingface.co/IBI-CAAI/Guided-Chest-CT-LeJEPA) as input (not raw CT images) and generates free-text narrative findings reports in the style of a radiologist. > **Dependency:** This model requires [IBI-CAAI/Guided-Chest-CT-LeJEPA](https://huggingface.co/IBI-CAAI/Guided-Chest-CT-LeJEPA) to first extract per-slice embeddings from raw CT volumes. Ker-VLJEPA-3B does not process CT images directly. **New SOTA on CT-RATE:** Macro F1 = **0.429**, surpassing the previous state-of-the-art U-VLM (F1 = 0.414) by +3.6%. Developed by the **Institute for Biomedical Informatics Center for Applied AI (IBI-CAAI)** at the University of Kentucky. ## Model Details | Property | Value | |----------|-------| | **Model Type** | Multi-modal Vision-Language Model (VLM) for Chest CT report generation | | **Language Model** | [Llama 3.2 3B](https://huggingface.co/meta-llama/Llama-3.2-3B) with LoRA adapters (r=64, alpha=128) | | **Input Embeddings** | [Guided-Chest-CT-LeJEPA](https://huggingface.co/IBI-CAAI/Guided-Chest-CT-LeJEPA) (ViT-Large) — **required dependency**, provides 1024-d per-slice embeddings from raw CT | | **Visual Encoder** | Z-Zoned Perceiver — compresses variable-length slice embeddings into 32 fixed visual tokens (3072-d) | | **Visual Grounding** | Flamingo-style gated cross-attention adapters at LLM layers 7, 14, 21 | | **Input** | Pre-computed LeJEPA slice embeddings (num_slices x 1024) — **not raw CT images** | | **Output** | Free-text narrative radiology findings report | | **Model Date** | 03/2026 | | **License** | CC BY-NC-SA 4.0 (inherited from CT-RATE dataset terms) | | **Parameters** | ~3.2B (Llama 3B base) + 1.7 GB LoRA + 320 MB bridge components | ## Architecture Ker-VLJEPA-3B bridges vision and language through a 3-stage architecture: ``` Raw CT Volume │ ▼ ┌─────────────────────────────┐ │ Guided-Chest-CT-LeJEPA │ ViT-Large backbone (separate model) │ (per-slice feature extract) │ Produces 1024-d embedding per slice └─────────────┬───────────────┘ │ (num_slices, 1024) ▼ ┌─────────────────────────────┐ │ Z-Zoned Perceiver │ 32 anatomical zones, each with learned queries │ + Global Self-Attention │ Compresses variable-length slices → 32 tokens │ + JEPA Predictor (→ 3072d) │ Projects into LLM hidden space └─────────────┬───────────────┘ │ (32, 3072) ▼ ┌─────────────────────────────┐ │ Llama 3.2 3B + LoRA │ Embedding grafting: visual tokens replace │ + Cross-Attention Adapters │ <|visual_region|> placeholders in prompt │ @ layers 7, 14, 21 │ Cross-attention: text attends to visual tokens │ + Additive Layer Projectors │ at 3 intermediate LLM layers └─────────────┬───────────────┘ │ ▼ Narrative Report ``` ### Key Architectural Innovations - **Z-Zoned Cross-Attention:** Divides the body axis into 32 spatial zones. Each zone's queries attend only to slices within their anatomical region, enforcing spatial specialization by construction. - **Flamingo-style Cross-Attention Bridge:** Rather than simple additive injection (which the LLM can trivially ignore), text hidden states explicitly attend to visual tokens via cross-attention at layers 7, 14, and 21. - **Warm Bridge Technique:** Bridge components (cross-attention adapters, layer projectors, LoRA) are initialized from a prior converged checkpoint rather than random, yielding +4.4% F1 improvement. - **Frozen Cross-Attention + EWC:** During narrative fine-tuning (Phase 4), cross-attention adapters are frozen to preserve visual grounding, while Elastic Weight Consolidation constrains LoRA drift. ## Intended Uses This model is intended for **research purposes** in medical imaging and radiology. - Automated generation of narrative chest CT findings reports - Research into vision-language architectures for medical imaging - Benchmarking on the CT-RATE evaluation protocol **Not intended for:** Clinical diagnosis, treatment decisions, or any patient-facing application without proper clinical validation. ## Training ### Data Trained exclusively on the train split of [CT-RATE](https://huggingface.co/datasets/ibrahimhamamci/CT-RATE) (~46,400 volumes), with guided cropping annotations from [ReXGroundingCT](https://huggingface.co/datasets/rajpurkarlab/ReXGroundingCT). ### Hardware 8x NVIDIA H200 GPUs with DDP via HuggingFace Accelerate, bf16 mixed precision. ### 4-Phase Training Pipeline | Phase | Objective | Trainable Components | Key Hyperparameters | Result | |-------|-----------|---------------------|---------------------|--------| | **Phase 1** | Visual encoder alignment | Visual encoder + JEPA head | LR=5e-5, BS=32, 20 epochs | F1=0.460, AUC=0.811 | | **Phase 2** | Contrastive bridge (InfoNCE + MMD) | Visual encoder + LoRA + JEPA head | LR=3e-5, BS=64, 30 epochs | F1=0.465, AUC=0.816 | | **Phase 3** | Generative fine-tuning (positive-findings-only text) | JEPA predictor + LoRA + cross-attn + layer projectors | LR=2e-5, BS=8, warm bridge init | Gen F1=0.422 | | **Phase 4** | Raw narrative fine-tuning | **LoRA only** (cross-attn frozen, EWC=100.0) | LR=5e-7, BS=8, 14 epochs | Gen F1=**0.429** | **Phase 3 → Phase 4 innovation:** In Phase 3, the model trains on positive-findings-only text to learn pathology detection without gradient domination from normal-text tokens (~90% of raw reports). Phase 4 then adapts to raw radiologist prose with frozen cross-attention + ultra-conservative LoRA updates. ### Generation Configuration | Parameter | Value | Rationale | |-----------|-------|-----------| | Temperature | 0.6 | Optimal via sweep — balances diversity and accuracy | | Top-p | 0.9 | Nucleus sampling | | Repetition penalty | 1.1 | Prevents degenerate loops | | No-repeat n-gram | 4 | Medical reports need some anatomical repetition | | Max new tokens | 384 | Covers 95%+ of reports | ## Results ### CT-RATE Benchmark (Cross-Method Comparison) All methods evaluated on the CT-RATE validation set (2,984 volumes) using the official RadBERT classifier for 18-class binary label extraction. | Method | Macro F1 | Macro Prec | Macro Rec | Notes | |--------|----------|------------|-----------|-------| | CT-CLIP (Hamamci et al.) | 0.194 | — | — | Zero-shot | | CT-CHAT (Hamamci et al.) | 0.287 | — | — | Fine-tuned | | BTB3D (Song et al.) | 0.354 | — | — | | | U-VLM (Lee et al.) | 0.414 | 0.491 | 0.429 | Previous SOTA | | **Ker-VLJEPA-3B** | **0.429** | 0.389 | **0.524** | | ### Per-Class Results (2,984 validation samples) | Class | Prec | Rec | F1 | Support | |-------|------|-----|-----|---------| | Pleural effusion | 0.574 | 0.789 | **0.664** | 370 | | Arterial wall calcification | 0.642 | 0.660 | **0.651** | 849 | | Coronary artery wall calcification | 0.580 | 0.585 | **0.582** | 752 | | Cardiomegaly | 0.437 | 0.743 | **0.550** | 315 | | Lung nodule | 0.491 | 0.560 | **0.523** | 1,344 | | Lung opacity | 0.545 | 0.453 | **0.494** | 1,173 | | Emphysema | 0.346 | 0.667 | **0.456** | 588 | | Lymphadenopathy | 0.393 | 0.532 | **0.452** | 769 | | Pulmonary fibrotic sequela | 0.384 | 0.505 | **0.436** | 819 | | Atelectasis | 0.369 | 0.496 | **0.423** | 698 | | Consolidation | 0.430 | 0.389 | **0.408** | 576 | | Pericardial effusion | 0.238 | 0.614 | **0.343** | 215 | | Mosaic attenuation pattern | 0.256 | 0.494 | **0.337** | 245 | | Hiatal hernia | 0.257 | 0.397 | **0.312** | 413 | | Peribronchial thickening | 0.271 | 0.354 | **0.307** | 347 | | Medical material | 0.194 | 0.722 | **0.305** | 306 | | Bronchiectasis | 0.238 | 0.328 | **0.276** | 326 | | Interlobular septal thickening | 0.357 | 0.142 | **0.203** | 246 | **Macro (default 0.5 threshold):** F1=0.429, Prec=0.389, Rec=0.524 ### Evaluation Protocol Clinical accuracy follows the **CT-RATE evaluation protocol**, the same methodology used by CT-CLIP, CT-CHAT, BTB3D, and U-VLM: 1. **Generate** free-text narrative reports from all validation CT volumes (temp=0.6, top_p=0.9) 2. **Clean** degenerate text (truncate at onset of repeated chars/unicode) 3. **Strip** negation suffixes before RadBERT (prevents false positive extraction) 4. **Extract** 18 binary abnormality labels using the official CT-RATE [RadBERT classifier](https://huggingface.co/zzxslp/RadBERT-RoBERTa-4m) with [CT-RATE weights](https://huggingface.co/datasets/ibrahimhamamci/CT-RATE) 5. **Compute** macro-averaged F1, precision, recall against ground-truth labels The RadBERT classifier achieves **F1=0.982** on ground-truth reports, confirming it is a reliable extraction tool. ## How to Use ### Prerequisites ```bash pip install torch transformers peft safetensors accelerate pip install flash-attn --no-build-isolation # recommended for fast generation ``` You will also need: - A local copy of [meta-llama/Llama-3.2-3B](https://huggingface.co/meta-llama/Llama-3.2-3B) - Pre-computed LeJEPA slice embeddings from the [Guided-Chest-CT-LeJEPA](https://huggingface.co/IBI-CAAI/Guided-Chest-CT-LeJEPA) backbone ### Repository Contents ``` Ker-VLJEPA-3B/ ├── README.md # This file ├── model.py # Inference-only model class ├── generate_report.py # CLI example └── weights/ ├── visual_encoder.safetensors # Z-Zoned Perceiver + JEPA predictor (43 MB) ├── bridge_components.safetensors # Cross-attn adapters, layer projectors, norms (277 MB) ├── lora_adapters/ # LoRA adapter weights for Llama 3.2 3B (1.7 GB) │ ├── adapter_config.json │ └── adapter_model.safetensors └── tokenizer/ # Tokenizer with <|visual_region|> special token ├── chat_template.jinja ├── tokenizer.json └── tokenizer_config.json ``` ### Python API ```python import numpy as np import torch from model import load_model # 1. Load the model (requires local Llama 3.2 3B) model = load_model( llm_path="/path/to/Llama-3.2-3B", weights_dir="weights", device="cuda", ) # 2. Load pre-computed LeJEPA slice embeddings for a CT volume # Shape: (1, num_slices, 1024) — one 1024-d embedding per CT slice # See IBI-CAAI/Guided-Chest-CT-LeJEPA for how to extract these embeddings = torch.from_numpy(np.load("volume_embeddings.npy")).unsqueeze(0).float() mask = torch.ones(1, embeddings.shape[1]) # 1=valid, 0=padding # 3. Generate a narrative report report = model.generate( slice_embeddings=embeddings, mask=mask, temperature=0.6, max_new_tokens=384, ) print(report) ``` ### CLI ```bash python generate_report.py \ --llm_path /path/to/Llama-3.2-3B \ --embeddings volume_embeddings.npy \ --temperature 0.6 ``` ### End-to-End: Raw CT Volume → Report This example shows the full pipeline, from a raw NIfTI CT volume through LeJEPA embedding extraction to report generation. ```python import nibabel as nib import numpy as np import torch import torch.nn.functional as F import timm from huggingface_hub import hf_hub_download from safetensors.torch import load_file from model import load_model # --- Step 1: Load the LeJEPA visual backbone --- backbone = timm.create_model( "vit_large_patch14_dinov2", pretrained=False, num_classes=0, in_chans=1, img_size=518, dynamic_img_size=True, ) ckpt = load_file(hf_hub_download("IBI-CAAI/Guided-Chest-CT-LeJEPA", "model.safetensors")) backbone.load_state_dict(ckpt, strict=False) backbone.eval().cuda() # --- Step 2: Preprocess and extract slice embeddings --- # HU clipping and normalization must match LeJEPA training exactly CLIP_MIN, CLIP_MAX = -997.0, 888.0 MEAN_HU, STD_HU = -142.39, 360.97 RANGE = CLIP_MAX - CLIP_MIN NORM_MEAN = (MEAN_HU - CLIP_MIN) / RANGE NORM_STD = STD_HU / RANGE vol = nib.load("chest_ct.nii.gz").get_fdata() # (H, W, D) in Hounsfield Units slice_embeddings = [] with torch.no_grad(): for i in range(vol.shape[2]): s = torch.from_numpy(vol[:, :, i]).float().unsqueeze(0) # (1, H, W) s = torch.clamp(s, CLIP_MIN, CLIP_MAX) s = (s - CLIP_MIN) / RANGE s = (s - NORM_MEAN) / NORM_STD # Align to patch size 14 _, H, W = s.shape tH, tW = (H // 14) * 14, (W // 14) * 14 if tH != H or tW != W: s = F.interpolate(s.unsqueeze(0), size=(tH, tW), mode='nearest').squeeze(0) emb = backbone(s.unsqueeze(0).cuda()) # (1, 1024) slice_embeddings.append(emb.cpu()) embeddings = torch.cat(slice_embeddings, dim=0).unsqueeze(0) # (1, num_slices, 1024) mask = torch.ones(1, embeddings.shape[1]) # --- Step 3: Generate the report --- model = load_model("/path/to/Llama-3.2-3B", "weights", "cuda") report = model.generate(slice_embeddings=embeddings, mask=mask) print(report) ``` ## Prompt Format The model uses the Llama 3.2 chat template with 32 `<|visual_region|>` tokens injected into the user message: ``` <|begin_of_text|><|start_header_id|>system<|end_header_id|> You are a radiology reporting assistant. Describe thoracic findings based on the provided CT scan visual features. Report only what you observe.<|eot_id|> <|start_header_id|>user<|end_header_id|> Based on the visual features from this CT scan, describe the thoracic findings. <|visual_region|> x 32<|eot_id|><|start_header_id|>assistant<|end_header_id|> ``` Each `<|visual_region|>` placeholder is replaced at the embedding level with the corresponding visual token from the perceiver encoder. Cross-attention adapters at layers 7, 14, and 21 provide additional grounding throughout generation. ## Limitations - **Domain-specific:** Trained exclusively on Chest CT from CT-RATE. Performance on other body regions, modalities, or datasets is unknown. - **Not clinically validated:** This is a research model. Generated reports should not be used for clinical decision-making. - **Hallucination:** Like all generative models, Ker-VLJEPA-3B can produce findings not present in the scan. The 0.389 precision indicates ~61% of generated positive findings are true positives. - **Degeneration:** At high token counts (>384), generation quality degrades. Reports should be truncated. - **English only:** All training text is in English. ## Citation If you use this model, please cite the CT-RATE dataset: ```bibtex @misc{bumgardner2026curriculumdriven3dctreport, title={Curriculum-Driven 3D CT Report Generation via Language-Free Visual Grafting and Zone-Constrained Compression}, author={V. K. Cody Bumgardner and Mitchell A. Klusty and Mahmut S. Gokmen and Evan W. Damron}, year={2026}, eprint={2603.23308}, archivePrefix={arXiv}, primaryClass={cs.CV}, url={https://arxiv.org/abs/2603.23308}, } ```