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

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 (~46,400 volumes), with guided cropping annotations from ReXGroundingCT.

Hardware

8x NVIDIA H200 GPUs with DDP via HuggingFace Accelerate, bf16 mixed precision.

4-Phase Training Pipeline

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

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.

Per-Class Results (2,984 validation samples)

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 with CT-RATE weights
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

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
• Pre-computed LeJEPA slice embeddings from the 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

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

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.

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:

@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}, 
}