OpenBMB/SHIFT

OpenBMB/SHIFT

Description: Souce code for paper "SHIFT: Gate-Modulated Activation Steering for Knowledge Conflict Mitigation in Retrieval-Augmented Generation"

Language: Python

License: MIT

Stars: 3

Forks: 0

Open issues: 0

Created: 2026-06-20T04:39:19Z

Pushed: 2026-06-23T05:42:20Z

Default branch: main

Fork: no

Archived: no

README:

📖 Overview

SHIFT is a lightweight framework for resolving knowledge conflicts in retrieval-augmented generation. Instead of directly editing internal neurons, SHIFT adds a small learnable gate module to frozen LLMs, allowing them to adaptively balance retrieved context and parametric knowledge during generation. With fewer than 0.01% trainable parameters, SHIFT improves context reliance while minimizing unintended effects on general model capabilities.

⚙️ Setup

1. Create Conda Environment

conda create --name shift python==3.10.0
conda activate shift
git clone https://github.com/OpenBMB/SHIFT.git
cd SHIFT

2. Install PyTorch

pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124

3. Install Flash Attention

pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

4. Install the rest of the dependencies:

pip install -r requirements.txt

5. Patch vLLM model files This project requires modifications to the vLLM implementations of Qwen3 and LLaMA. After installing the requirements, run the patch script:

bash vllm/patch_vllm.sh

> Note: Please run this script after 'pip install -r requirements.', because installing or reinstalling vLLM may overwrite the patched files.

⚡ Data

Our training data can be downloaded from SHIFT. After downloading, place the files into the dataset folder.

To construct the data from scratch, download the files from MRQA-Shared-Task-2019. Use the downloaded data to synthesize it with FlashRAG.

🔥 Training

GRPO with a single GPU:

python single_gpu.py

GRPO with multiple GPUs:

python multi_gpu.py

📊 Evaluation

For MRQA and ConfiQA:

python eval.py

For MMLU, use lm-evaluation-harness

🎯 Analysis

We also provide the t-SNE visualization pipeline for gates in SHIFT, with corresponding figures available under the figs folder:

python run_batch_tsne.py

For Qwen-3-0.6B:

For Qwen-3-8B:

🎉 Acknowledgement

Our work is built on the following codebases, and we are deeply grateful for their contributions.

• vllm
• nano-aha-moment
• lm-evaluation-harness
• gated_attention
• FlashRAG

🥰 Citation

If you find this work useful, please cite our paper and give us a shining star 🌟

@article{Li2026shift,
title={SHIFT: Gate-Modulated Activation Steering for Knowledge Conflict Mitigation in Retrieval-Augmented Generation},
author={Li, Ruochang and Huang, Pengcheng and Liu, Zhenghao and Yan, Yukun and Xie, Huiyuan and Gu, Yu and Yu, Ge and Sun, Maosong},
year={2026}
url={},
}

📧 Contact

If you have questions, collaboration opportunities, or potential PhD opportunities in the United States, please feel free to email:

ruochangli@gmail.com