AdaRAG-CT

AdaRAG-CT

Beyond the Embedding Bottleneck: Adaptive Retrieval-Augmented 3D CT Report Generation<br> _{accepted to ECCV 2026, then desk-rejected for bibliographic inaccuracies}

📄 arXiv · 💻 GitHub · 🤗 Models & Data

This HuggingFace repository holds only the weights and data (data/, results/). The code lives on GitHub — clone it there, then download this bundle into the repo root as shown below.

Contrastive 3D CT embeddings concentrate 90% of their variance in just 2 of 512 dimensions, and scaling the LLM from 8B to 70B gives no gain — the bottleneck is visual, not generative. AdaRAG-CT compensates by retrieving organ-indexed report sentences and adaptively injecting them during generation, lifting Clinical F1 from 0.420 (CT-Agent) to 0.480.

<p align="center"> <img src="./figures/figure1.png" width="620"> </p>

Results (CT-RATE validation)

CT-CHAT reproduced under our unified evaluation protocol. Base (ViSD-Boost + CT-CLIP) is our base model: it takes organ-level ViSD-Boost embeddings plus whole-volume CT-CLIP embeddings as visual input.

Install

conda create -n adaragct python=3.12 && conda activate adaragct
pip install -r requirements.txt

The Llama backbones are not redistributed here — they are gated on HuggingFace and must be requested and downloaded separately: meta-llama/Llama-3.1-8B-Instruct (8B) and meta-llama/Llama-3.3-70B-Instruct (70B). Only training needs them; evaluation and inference run from the released checkpoints alone. Point model.model_path in the training configs at your local copy.

Data & Checkpoints

Everything else — CT embeddings, retrieval indices, and all four checkpoints — is on 🤗 HuggingFace, laid out exactly like this repo, so download straight into the repo root:

# 8B track only: data + Base 8B + AdaRAG-CT 8B  (~45 GB)
huggingface-cli download LiangRenjie/AdaRAG-CT --repo-type model --local-dir . \
    --include "data/*" "results/base_8b/*" "results/adaragct_8b/*"

# everything, including the 70B checkpoints  (~200 GB)
huggingface-cli download LiangRenjie/AdaRAG-CT --repo-type model --local-dir . \
    --include "data/*" "results/*"

This gives you data/ (embeddings, annotations, oracle/retrieval contexts) and results/{base,adaragct}_{8b,70b}/. Four checkpoints: Base 8B/70B (standalone merged models) and AdaRAG-CT 8B/70B (LoRA adapter + projector, loaded on top of the matching base). The released predictions, metrics, and training logs under results/ already ship with this repo.

The pipeline starts from precomputed embeddings — extracting ViSD-Boost organ features and CT-CLIP whole-volume features from raw CT volumes is not part of this release; use the released data/embeddings/ instead.

Evaluate

Score the released predictions directly:

python -m adaragct.eval.cal_metrics results/adaragct_8b/infer_step_2000.jsonl --output metrics.json

Add --compute-llama-score for the LLaMA score, --bootstrap 1000 for 95% confidence intervals.

Or generate predictions, then score:

# AdaRAG-CT (adaptive retrieval)
python -m adaragct.inference.inference_rag --checkpoint results/adaragct_8b/checkpoint_step_2000 --output pred.jsonl

# Base / no-retrieval (same checkpoint, retrieval disabled)
python -m adaragct.inference.inference_rag --checkpoint results/adaragct_8b/checkpoint_step_2000 --no-rag --output base_pred.jsonl

python -m adaragct.eval.cal_metrics pred.jsonl --output metrics.json

Key inference flags: --no-rag (disable retrieval), --text2text (Text2Text retrieval pipeline), --oracle (oracle context), --top-k, --max-retrievals.

On-the-fly retrieval encodes each generated probe sentence with a fine-tuned text encoder (code + indices ship under data/retrieval/; it auto-downloads microsoft/BiomedVLP-CXR-BERT-specialized on first run). For a download-free run, use --oracle (precomputed contexts). To reproduce the exact paper numbers, score the released predictions in results/.

Train

AdaRAG-CT

AdaRAG-CT trains a [RAG] trigger token on top of a frozen base model via LoRA, mixing oracle and retrieved contexts. Every required input — base checkpoint, CT embeddings, and the precomputed oracle/retrieval contexts — comes from the HuggingFace bundle, so training runs directly with no extra preprocessing:

# 8B
python -m adaragct.train.train_rag --config configs/adaragct_8b.yaml
# 70B
python -m adaragct.train.train_rag --config configs/adaragct_70b.yaml

Each config points to data/ (embeddings, oracle_context_top3.jsonl, retrieval_context_top3.jsonl) and a base checkpoint under results/.

Base model

The base model — organ-level ViSD-Boost embeddings plus a whole-volume CT-CLIP embedding projected into a Llama backbone — is trained in two stages by adaragct.train.train_base. Download the released Base checkpoints if you only want to train AdaRAG-CT on top of them; the steps below are for reproducing a Base checkpoint from scratch.

Stage 1 — projector pretraining (train_projector_only: true). The LLM is frozen and only the five projectors (whole-CT + lung/heart/esophagus/aorta) are trained, to align the visual embeddings with the LLM hidden space. Saves checkpoints/step_N/projector.pt.

python -m adaragct.train.train_base --config configs/base_8b_projector.yaml   # 8B
python -m adaragct.train.train_base --config configs/base_70b_projector.yaml  # 70B

Stage 2 — LoRA SFT (freeze_projector: true). Loads a Stage-1 projector through model.pretrain_checkpoint, freezes it, and trains LoRA + embed_tokens/lm_head. Set model.pretrain_checkpoint to the Stage-1 step you want (the released models used step 9000), then:

python -m adaragct.train.train_base --config configs/base_8b.yaml   # 8B
python -m adaragct.train.train_base --config configs/base_70b.yaml  # 70B

Merge. Stage 2 saves a PEFT folder; merging it into the backbone produces the standalone model directory that the AdaRAG-CT configs expect as model.model_path (this is the format the released results/base_8b/checkpoint is in):

python -m adaragct.tools.merge_peft_checkpoint \
    --peft-dir results/train/base_8b/checkpoints/step_5000 \
    --output-dir results/base_8b/checkpoint

Stage-1 projector checkpoints are not released, so reproducing a Base checkpoint means running both stages. adaragct/tools/ also holds convert_base_checkpoint.py / convert_rag_checkpoint.py for converting older single-file .pt checkpoints into the PEFT folder format.

Citation

@misc{liang2026embeddingbottleneckadaptiveretrievalaugmented,
      title={Beyond the Embedding Bottleneck: Adaptive Retrieval-Augmented 3D CT Report Generation},
      author={Renjie Liang and Yiling Ma and Yang Xing and Zhengkang Fan and Jinqian Pan and Chengkun Sun and Li Li and Kuang Gong and Jie Xu},
      year={2026},
      eprint={2603.15822},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2603.15822},
}

Acknowledgements

CT-RATE / CT-CLIP · ViSD-Boost · LLaVA · Self-RAG