By Samuel Yeh and Sharon Li.

This codebase is the official implementation for the paper, "How Retrieved Context Shapes Internal Representations in RAG." It covers the main experiment pipeline:

1. Generate question-only answers and split them into correct / incorrect.
2. Export questions for retrieval in a separate project.
3. Categorize retrieved documents as relevant, distracting, or random.
4. Generate and judge answers for retrieval settings: relevant, distracting, random, relevant_3_distracting, relevant_3_random, and all_20.
5. Extract last prompt-token representations for each setting.
6. Plot single-doc last-layer PCA.
7. Plot multi-doc last-layer PCA.
8. Plot single-doc PCA across layers.
9. Plot cosine similarity versus response correctness.

Supported datasets are trivia, nq, pop, and strategy.

Supported models are gemma3-27B, llama4-17B, qwen3-80B, plus their -base variants.

Run commands from the repository root.

python generate_question_only.py --dataset trivia --model gemma3-27B --root .
python prepare_retrieval_data.py --dataset trivia --root .

We use the MassiveDS retrieval-scaling pipeline to retrieve documents. After exporting queries with prepare_retrieval_data, run retrieval with RulinShao/retrieval-scaling; see that repository for setup and execution details. The downstream scripts expect the retrieval output as a JSONL file with ctxs.

After retrieval writes the JSONL file, categorize documents:

export AZURE_OPENAI_ENDPOINT="..."
export AZURE_OPENAI_API_KEY="..."
python categorize_retrieved_documents.py \
  --dataset trivia \
  --root . \
  --retrieved-jsonl retrieved_results/post_processed/dedup_merged_trivia_top1000.jsonl

Generate retrieval answers, judge them, and extract representations:

python generate_with_retrieval.py \
  --dataset trivia \
  --model gemma3-27B \
  --root . \
  --retrieved-jsonl retrieved_results/post_processed/dedup_merged_trivia_top1000.jsonl

python extract_prompt_representations.py \
  --dataset trivia \
  --model gemma3-27B \
  --difficulty correct \
  --root .

Plot scripts are under plot_scripts, or use the shared CLI:

python plotting.py single-last --root .
python plotting.py multiple-last --root .
python plotting.py single-diff-layer --root .
python plotting.py sim-vs-response --root .