This is the official repository for 📜 Effective Length Extrapolation via Dimension-Wise Positional Embeddings Manipulation.
We introduce Dimension-Wise Positional Embeddings Manipulation (DPE), a training-free and plug-and-play framework for context extension. DPE identifies and selectively manipulates the most influential dimensions of the rotary positional embeddings by analyzing their attention contribution to attention scores. This strategy optimizes the model’s dimensional adaptation with minimal modifications to the pretrained model, enabling each dimension to extrapolate in an optimal manner.
DPE proceeds in three steps: a. Detecting the Effective Relative Distance: We partitioning the model’s hidden dimensions into several groups (in practice, we divide them evenly into eight groups). For each group, we detect the most effective relative distances via a detecting task using NIAH. b. Identifying the Key Dimensions: The key dimensions are identified using their attention contribution. c. Scaling the Position Indices: For each selected dimension, we apply a customized scaling factor derived from its most effective relative distance.
pip install transformers==4.47.0
# We use flash-attn==2.7.0
pip install flash-attn --no-build-isolationmodel_name = "model_name in ./config/dpe_config.yaml" # support llama, qwen, mistral
model_path = "/path/to/your/model"
config = AutoConfig.from_pretrained(model_path)
# DPE config, best settings for all the models can be found in ./config
with open('./config/dpe_config.yaml', 'r') as f:
dpe_config = yaml.safe_load(f)
model_config = dpe_config[model_name]
local_window_size = model_config['local_window_size']
scale_factors = model_config['scale_factors']
dimension_groups_range = model_config['dimension_groups_range']
# If select topk dimensions, set select_topk_dim=True
select_topk_dim = False
topk_dim=48 # number of selected dimensions
selected_dim_path = model_config['selected_dim_path'] # result of identified key dimensions
# apply dpe
replace_with_dpe(scale_factors, local_window_size, dimension_groups_range, select_topk_dim, topk_dim, selected_dim_path)
model = AutoModelForCausalLM.from_pretrained(model_path, attn_implementation="flash_attention_2", trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto")import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
from dpe_for_llama import replace_with_dpe
model_path = "meta-llama/Meta-Llama-3-8B-Instruct" # support llama, qwen, mistral
config = AutoConfig.from_pretrained(model_path)
# DPE
# Settings for Llama-3-8b-Instruct, best settings for all the models can be found in ./config
local_window_size = 1024
scale_factors = [4, 4, 2, 16, 32, 32, 16, 128] # scale size for each dimension group
dimension_groups_range = [0, 8, 16, 24, 32, 40, 48, 56, 64] # dimension range for each group
# If select topk dimensions, set select_topk_dim=True
select_topk_dim = False
topk_dim=48 # number of selected dimensions
selected_dim_path = "dim_select_result/llama3-8b-Instruct/qk_2_norm_selected_dim.pt" # result of identified key dimensions
replace_with_dpe(scale_factors, local_window_size, args.all_dims, select_topk_dim, topk_dim, selected_dim_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, attn_implementation="flash_attention_2", trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto")
prompt = f"There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.\n\nThe pass key is 123456. Remember it. 123456 is the pass key.\n " + \
"The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again. " * 4000 + \
"\nWhat is the pass key?\nThe passkey is "
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
prompt_length = inputs["input_ids"].shape[1]
output = tokenizer.decode(model.generate(**inputs, max_new_tokens=128)[0][prompt_length:], skip_special_tokens=True)
print("[input length]: ",prompt_length)
print("[model output:]", output)We use NIAH (4-needle) to detect the optimal scale factors for our pretrained model.
Generate the test data by running the data generator script. You can modify the MAX_CONTEXT_LENGTH parameter to adjust the data length.
bash ./niah/scripts/run_data_generator.shExecute the NIAH detection script with your specific parameters:
- Set
TEST_MAX_LENGTHto your desired test sequence length - Set
PRETRAIN_LENGTHto your model's pretraining context length
bash ./effective_length_detection/run_niah_detect.shAfter this step, you'll obtain results for different detection groups. By sorting these results, you can identify the optimal dimensions for your model and their corresponding scale factors.
step 1: Sample query and key's hidden states for your model
python ./dimension_selection/sample_qk.pystep 2: Calulate 2-norm attention contribution for sampled hidden states
python ./dimension_selection/qk_2_norm_select_dim.pyAfter this step, we can get ./dimension_selection/result/llama3-8b-Instruct/qk_2_norm_selected_dim.pt
"llama3-8b-instruct":
local_window_size: 1024
scale_factors: [4, 4, 2, 16, 32, 32, 16, 128]
dimension_groups_range: [0, 8, 16, 24, 32, 40, 48, 56, 64]
selected_dim_path : "dim_select_result/llama3-8b-Instruct/qk_2_norm_selected_dim.pt"This section contains the data and code for validating DPE in our paper.
We use NIAH (4-needle) to detect effective relative distance and test the latest models from the open-source community using DPE and RoPE. The haystack consists of Paul Graham's essays, and the needles are 6-digit numbers. To run the evaluation:
cd niah/scripts
bash niah.shThe test and evaluation code is from the official release of RULER which contains diverse sythetic tasks to test the long-context ability of LLMs.
# step 1: generate the test data
cd ruler
python auto_prepare_data.py --model_path /path/to/model --max_length 131072 --temp llama3.1 (or qwen2)
# All tasks
[
"niah_single_1", "niah_single_2", "niah_single_3",
"niah_multikey_1", "niah_multikey_2", "niah_multikey_3",
"niah_multivalue", "niah_multiquery", "vt", "cwe", "fwe", "qa_1", "qa_2"
]This will generate the processed data for the 13 tasks in RULER and save it to the data-jsonl folder.
To test the model on RULER, run the following command:
# step 2: test and evaluate the model
cd scripts
bash ruler_dpe.shThis command will generate a prediction file in the Predictions/task_name/model_name/directory.
If you find our work helpful or relevant to your research, please kindly cite our paper:
@misc{lu2025effectivelengthextrapolationdimensionwise,
title={Effective Length Extrapolation via Dimension-Wise Positional Embeddings Manipulation},
author={Yi Lu and Wanxu Zhao and Xin Zhou and Chenxin An and Chenglong Wang and Shuo Li and Yuming Yang and Jun Zhao and Tao Ji and Tao Gui and Qi Zhang and Xuanjing Huang},
year={2025},
eprint={2504.18857},
archivePrefix={arXiv},
primaryClass={cs.CL},
url={https://arxiv.org/abs/2504.18857},
}