Reliable Decision Making via Calibration Oriented Retrieval Augmented Generation

This repository provides the implementation of Reliable Decision Making via Calibration Oriented Retrieval Augmented Generation (Accepted by NeurIPS 2025).

Abstract

Recently, Large Language Models (LLMs) have been increasingly used to support various decision-making tasks, assisting humans in making informed decisions. However, when LLMs confidently provide incorrect information, it can lead humans to make suboptimal decisions. To prevent LLMs from generating incorrect information on topics they are unsure of and to improve the accuracy of generated content, prior works have proposed Retrieval Augmented Generation (RAG), where external documents are referenced to generate responses. However, previous RAG methods focus only on retrieving documents most relevant to the input query, without specifically aiming to ensure that the human user's decisions are well-calibrated. To address this limitation, we propose a novel retrieval method called Calibrated Retrieval-Augmented Generation (CalibRAG), which ensures that decisions informed by RAG are well-calibrated. Then we empirically validate that CalibRAG improves calibration performance as well as accuracy, compared to other baselines across various datasets.

---

✅ Prerequisites

All experiments were conducted on a single NVIDIA RTX A6000 GPU.

# Clone the repository
git clone https://github.com/chaeyoon-jang/calibrag.git
cd calibrag

# Install Git LFS
git lfs install

# Pull large files managed by Git LFS
git lfs pull

# Set up the environment
cd code
conda create -n calibrag python=3.9 -y
conda activate calibrag
pip install -r requirements.txt

---

✅ Synthetic Data Generation

1. Create Open-ended Questions

python -m experiments.api --data_dir ./data/dev/raw --type oe

2. Generate RAG Data

1. Download Preprocessed Passage Data

   wget https://dl.fbaipublicfiles.com/dpr/wikipedia_split/psgs_w100.tsv.gz
   

2. Run Retrieval

   sh scripts/retrieve_dev.sh
   

3. Generate LLM Outputs

For baselines,

sh scripts/base_calibrag_lm_outputs.sh

For CalibRAG,

sh scripts/calibrag_lm_outputs.sh

4. Simulate Human Decision-Making

sh scripts/calibrag_decision.sh

5. Evaluate Results

For baselines,

sh scripts/base_train_api.sh

For CalibRAG,

sh scripts/calibrag_api.sh

📂 Finalized Dataset: The finalized dataset is available in the data/dev directory within the repository after cloning with git-lfs.

---

✅ Training Methods

1. CT-LoRA

python -m experiments.train.train_calibration_tune \
  --model_name="Meta-Llama-3.1-8B-Instruct" \
  --batch_size 2 \
  --gradient_accumulation_steps 2 \
  --uc_type "ct"

2. CT-Probe

python -m experiments.train.train_classifier_tune \
  --model_name "Meta-Llama-3.1-8B-Instruct" \
  --batch_size 4

3. CT-Ling (Sampling)

python -m experiments.train.train_calibration_tune \
  --model_name="Meta-Llama-3.1-8B-Instruct" \
  --batch_size 2 \
  --gradient_accumulation_steps 2 \
  --uc_type "ling"

4. CT-Number (Sampling)

python -m experiments.train.train_calibration_tune \
  --model_name="Meta-Llama-3.1-8B-Instruct" \
  --batch_size 2 \
  --gradient_accumulation_steps 2 \
  --uc_type "number"

5. CalibRAG Training

python -m experiments.train.train_calibrag \
  --model_name="Meta-Llama-3.1-8B-Instruct" \
  --batch_size 2 \
  --gradient_accumulation_steps 2 \
  --with_lora True

---

✅ Test Data Generation

1. Create Open-ended Questions

    python -m experiments.api --data_dir ./data/test/raw --type oe

2. Generate RAG Data

   sh scripts/retrieve_test.sh

3. Produce LLM Outputs with Uncertainty

For baselines,

sh scripts/base_eval_lm_outputs.sh

For CalibRAG,

sh scripts/calibrag_eval_lm_outputs.sh

4. Simulate Human Decision-Making for Testing

For baselines,

sh scripts/base_eval_decision.sh

For CalibRAG,

sh scripts/calibrag_decision.sh

5. Evaluate Results

For baselines,

sh scripts/base_eval_api.sh

For CalibRAG,

sh scripts/calibrag_api.sh

---

✅ References

https://github.com/facebookresearch/contriever
https://github.com/tatsu-lab/linguistic_calibration
https://github.com/activatedgeek/calibration-tuning
https://github.com/esteng/pragmatic_calibration

✅ Citation

@inproceedings{jang2025reliable,
  title={Reliable Decision-Making via Calibration-Oriented Retrieval-Augmented Generation},
  author={Jang, Chaeyun and Cho, Deukhwan and Lee, Seanie and Lee, Hyungi and Lee, Juho},
  booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
  year={2025}
}