ICICLE: Expanding Retrieval with In-Context Documents

ICICLE: Expanding Retrieval with In-Context Documents Yu-Chen Den*, Yung-Yu Shih*, Zhi Rui Tam, Kuan-Yu Chen, Pu-Jen Cheng, Yun-Nung Chen, Eugene Yang (*Equal Contribution) -- National Taiwan University & Johns Hopkins University Findings of the Association for Computational Linguistics: EMNLP 2026

Setup

Install uv for package management

pip install uv && uv sync

Download data

Use huggingface-hub to donwload prepared dataset

# Stage 1a — Base GR
Abner0803/msmarco_text-with_pseudo_query-100k-gr
Abner0803/nq_text-with_pseudo_query-100k-gr

# Stage 1b — ICL template SFT
Lala8383/msmarco-docquery-pseudo-hn-3shot-v3
Lala8383/nq-item-id-llm-compressed-hardneg-3shot-v4

# Stage 3 — LoRA 100-shot
Lala8383/msmarco-item-id-hardneg-filter-100shot-v4_128k
Lala8383/nq-item-id-llm-compressed-final-few-shot-test-v4

Training

We suggest everyone to use Axolotl to finetune our model

1. Stage-1 Finetune (Qwen3-1.7B-Base -> Generative Retrieval)

axolotl train ./configs/axolotl_qwen3-1.7b_msmarco-100k.yml
axolotl train ./configs/axolotl_qwen3-1.7b_nq-100k.yml

2. In-Context template tuning

axolotl train ./configs/axolotl_icl_msmarco-100k_v3.yml

[!NOTE] We'll need to inference this stage's checkpoint first to build the training data for DPO.

3. DPO (Stage 2 routing calibration)

# Build DPO preference pairs from Stage 1b inference output
python src/construct_dpo_pairs.py \
    --inference_file retrieve_results/<stage1b_eval>.json \
    --test_file data/<icl-3shot-test>/test.jsonl \
    --output_dir data/dpo-v4-copy-full

python src/build_hard_cd_dpo.py \
    --h1success_file   data/dpo-v4-copy-full/dpo_h1success.jsonl \
    --dpo_file         data/dpo-v4-copy-full/dpo.jsonl \
    --output_file      data/dpo-v4-copy-full/dpo_h1success_cd_zs_unseen.jsonl \
    --zero_shot_unseen retrieve_results/zero_shot_unseen.json \
    --strategy         zero_shot_unseen --seed 42

# Train DPO (shared config for MSMARCO + NQ; swap paths for NQ)
axolotl train ./configs/axolotl_stage2_dpo.yml

4. Long Context Extension (Stage 3 LoRA)

# Build 100-shot ICL training data
python -m src.build_train_docquery_v4 --config-name=build_stage3_msmarco
python -m src.build_train_docquery_v4 --config-name=build_stage3_nq

# Train LoRA
axolotl train ./configs/axolotl_stage3_msmarco_lora.yml
axolotl train ./configs/axolotl_stage3_nq_lora.yml

# Merge LoRA adapters into full checkpoints
bash scripts/merge_lora_last_ckpt.sh

Inference

Run inference code with below command when finish training ICICLE

python -m src.inference_icl_with_tag

Citation

If you use this code, please cite the following:

@inproceedings{den2026icicle,
      title={ICICLE: Expanding Retrieval with In-Context Documents},
      author={Den, Yu-Chen and Shih, Yung-Yu and Tam, Zhi Rui and Chen, Kuan-Yu and Cheng, Pu-Jen and Chen, Yun-Nung and Yang, Eugene},
      booktitle={Findings of the Association for Computational Linguistics: EMNLP 2026},
      year={2026}
}