Multilingual Retrieval Transition Heads

Preprint releasing soon |  Precomputed Results: HuggingFace

---

Overview

This repository contains code to identify Retrieval Transition Heads (RTH) — attention heads in multilingual large language models (LLMs) that are responsible from latent-space to target-language transition. We detect RTH via a multilingual Needle-in-a-Haystack (NIAH) paradigm: the needle (a factual statement) is placed in a haystack written in a language close to the internal latent space (like English), and the model is asked to retrieve it in another language (like German). Heads that consistently attend to the correct needle position across language pairs are identified as retrieval transition heads. Our experiments reveal that RTHs are vital for Chain-of-Thought reasoning in multilingual LLMs. Across four multilingual benchmarks (MMLU-ProX, MGSM, MLQA, and XQuAD) and two model families (Qwen2.5 and Llama3.1), we demonstrate that masking a fraction of top RTH induces sharp performance drops. Specifically, for Llama3.1-8B-Instruct, masking the top-25 RTHs results in an average 36.6-point drop in reasoning accuracy (MMLU-ProX, MGSM). A similar trend is observed in extractive QA (MLQA, XQuAD), where RTH masking yields a 9.0 F1-score.

Masking experiments (evaluating the causal effect of RTH on multilingual benchmarks via lm-harness) live on a separate branch. Switch to masking to access that code:

git checkout masking

---

Repository Structure

.
├── main.py                        # Main entry point for RTH detection
├── utils/
│   ├── preprocess.py              # NIAH data preprocessing → Parquet cache
│   ├── inference.py               # Batched inference + head score accumulation
│   ├── models.py                  # Model loader (LLaMA, Qwen, Phi, Mistral, Mixtral)
│   ├── visualize.py               # W&B head-score plots
│   ├── LLM.py                     # LiteLLM wrapper for subword alignment (via hosted LLM)
│   └── reverse_mapping.py         # Maps decoded token attention back to needle positions
├── faiss_attn/source/             # Custom model forward passes that return attention weights
├── tests/                         # Unit tests
├── run/                           # SLURM job scripts (HPC reference)
│   └── translation/               # Scripts for cross-lingual RTH experiments
├── haystack_for_detect/           # Haystack (To be downloaded from HuggingFace)  
├── results/                       # Per-experiment NIAH result JSONs (auto-created)
├── head_scores/                   # Per-head retrieval scores JSON (auto-created)
└── run/CreateVizFromLLMTesting.ipynb  # Visualization notebook

---

Setup

All experiments were performed on NYU Greene HPC Cluster

1. Environment

conda create -n rth python=3.10 -y
conda activate rth

or using venv:

python -m venv rth_env
source rth_env/bin/activate

2. Install Dependencies

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install transformers accelerate
pip install flash-attn --no-build-isolation
pip install wandb numpy pandas tqdm matplotlib
pip install rouge-score rouge-chinese jieba
pip install unbabel-comet
pip install litellm pydantic
pip install vllm          # only needed if hosting the alignment LLM yourself

Note: flash-attn requires a CUDA-capable GPU and matching CUDA/PyTorch versions. See flash-attention for build instructions.

3. W&B Login

Results and head-score heatmaps are logged to Weights & Biases:

wandb login
export WANDB_ENTITY=<your_entity>
export WANDB_PROJECT=rh_translation

---

Data Preparation

The NIAH data is organized by language. Each language directory must contain:

haystack_for_detect/
├── en/
│   ├── needles.jsonl     # one JSON per line: {"needle":..., "question":..., "real_needle":...}
│   ├── part1/            # text files used as haystack background
│   ├── part2/
│   └── part3/
├── de/
│   └── ...
├── zh/
│   └── ...
└── sw/
    └── ...

A data/ directory is used to cache preprocessed Parquet files (created automatically on first run).

---

Running RTH Detection

RTH detection is performed by main.py. The script:

1. Loads the target model with custom attention forward passes (faiss_attn/source/).
2. Preprocesses NIAH test configurations and caches them as Parquet files (data/).
3. Runs batched inference, recording which attention heads attend to the needle position at each decoding step.
4. Computes per-head retrieval scores and saves them to head_scores/.
5. Logs heatmap visualizations to W&B.

Cross-Lingual RTH Detection — Example: Qwen-2.5 7B, English ↔ German

The cross-lingual setup places the needle in language A while the model retrieves it in language B. We run both directions in parallel.

Step 1 (Optional): An alignment LLM is used to map decoded tokens back to needle positions across languages (subword alignment via LiteLLM). You can use any of the following backends — pick whichever is most convenient:

Option A — Gemini or OpenAI API (no GPU required)

# Gemini
export GEMINI_API_KEY=<your_gemini_key>
export model_name="gemini/gemini-2.5-pro"   # default

# — or — OpenAI
export OPENAI_API_KEY=<your_openai_key>
export model_name="gpt-4o"

No server to start; LiteLLM routes the calls automatically.

Option B — Self-hosted open-source model via vLLM

# On a GPU node (e.g. via SLURM srun or screen)
MODEL_PATH=/path/to/Qwen3-30B   # any instruction-tuned model
PORT=8010

vllm serve "$MODEL_PATH" \
  --host 0.0.0.0 \
  --port $PORT \
  --tensor-parallel-size 2 \
  --gpu-memory-utilization 0.85 \
  --trust-remote-code \
  --served-model-name "qwen3-30b"

export HOSTED_VLLM_API_BASE="http://<server_node>:${PORT}/v1"
export model_name="hosted_vllm/qwen3-30b"

For HPC environments, see run/translation/start_qwen.sh for a ready-made SLURM launcher.

Step 2: Run RTH detection on the target model:

export WANDB_ENTITY=<your_entity>
export WANDB_PROJECT=rh_translation

MODEL_PATH=/path/to/Qwen2.5-7B-Instruct
HAYSTACK_DIR=/path/to/haystack_for_detect
DATA_DIR=/path/to/data

# EN haystack & needle, DE target language
CUDA_VISIBLE_DEVICES=0 python main.py \
  --model_path $MODEL_PATH \
  -s 1000 -e 50000 \
  --haystack_dir        $HAYSTACK_DIR/en \
  --needle_lg           en \
  --translated_needle_lg   de \
  --translated_haystack_dir $HAYSTACK_DIR/de \
  --exp_name            qwen25_7b_trans_ende \
  --parquet_dir         $DATA_DIR &

wait

For HPC (SLURM) environments, reference scripts are provided in run/translation/. See run/translation/get_rth_only.sh for a self-contained SLURM job script and run/translation/start_qwen.sh for launching the vLLM alignment server as a separate SLURM step.

Key Arguments

Argument	Description
--model_path	Path to the HuggingFace model
-s / -e	Min / max context length for NIAH sweep (tokens)
--haystack_dir	Directory of background text files + needles.jsonl
--needle_lg	Language code of the source needle (en, de, zh, sw)
--translated_needle_lg	Target language for cross-lingual retrieval
--translated_haystack_dir	Haystack directory for the target language
--exp_name	Experiment name (used for output file naming)
--parquet_dir	Directory to cache / read preprocessed Parquet data
--filter_needle_indices	Comma-separated needle indices to run a subset, e.g. "0,1,2"

---

Outputs

After a run, two directories are populated:

results/
└── <exp_name>/          # Per-(context_length × depth) NIAH result JSONs

head_scores/
└── <exp_name>.json      # Per-head retrieval score lists
                         # Format: {"layer-head_id": [score_1, score_2, ...], ...}

Loading Head Scores Programmatically

import json, numpy as np

with open("head_scores/qwen25_7b_trans_ende.json") as f:
    head_list = json.load(f)

head_score_list = [
    ([int(x) for x in k.split("-")], np.mean(v))
    for k, v in head_list.items()
]
head_score_list = sorted(head_score_list, key=lambda x: x[1], reverse=True)

print("Top RTH:")
for (layer, head), score in head_score_list[:10]:
    print(f"  Layer {layer:2d}, Head {head:2d}  —  score: {score:.3f}")

Visualization

Head-score heatmaps are automatically uploaded to W&B during each run. For offline visualization, open run/CreateVizFromLLMTesting.ipynb and set model_name to the folder name of your results.

---

Precomputed Head Scores

Precomputed head scores for Qwen-2.5 7B, Phi-3.5 3B, and Llama-3.1 8B across languages EN, DE, ZH, SW are available at:

Link — to be added

---

Masking Experiments

To study the causal role of RTH on downstream multilingual benchmarks (evaluated via lm-evaluation-harness), switch to the masking branch:

git checkout masking

That branch contains code to apply attention-head masks derived from the head scores computed here.

---

Citation (Coming soon)

---

Acknowledgements

This codebase builds on Retrieval Head (Wu et al., 2024) and the Needle-In-a-Haystack framework. We thank the Eleuther AI team for providing the lm-evaluation-harness that we have used to perform benchmarking experiments on reasoning benchmarks. This work was supported in part through the NYU IT High Performance Computing resources, services, and staff expertise. The work is partially funded by NSF CAREER award 2443271 and NSF award RI-2521091.