Effective Length Extrapolation via Dimension-Wise Positional Embeddings Manipulation

Thank you for your dedicated review of our DPE. In the following, we will carefully respond to your questions.

---

Response to Question 1: Pseudocode for DPE

We provide pseudocode outlining the three stages of DPE implementation, including:

• Algorithm 1 – Detecting effective relative distance for different dimensions (Detecting).
• Algorithm 2 – Identifying key dimensions (Selection).
• Algorithm 3 – DPE FlashAttention Forward(Scaling)

Algorithm 1: Detecting effective relative distance for different dimensions
Input: dimension groups, group_num, target_length, detecting_lengths
Output: best_scale_factors
s = L / L_pretrain / 2;
Initialize detect_scale_factor = [s, s, ..., s] for each group; 
for group_id in range(group_num) do 
| for detecting_length in detecting_lengths do 
| | detect_scale_factor[group_id] = target_length / detecting_length; 
| | accuracy = niah_detect(detect_scale_factor); 
| | if accuracy > best_accuracy then 
| | | best_scale_factors[group_id] = detect_scale_factor[group_id]; 
| end 
end 

Algorithm 2: Identifying key dimensions
Input: Q, K tensors with shape [L, d], topk parameter
Output: topk_dimensions
Q_norm = norm(Q.reshape(L, d/2, 2), dim=-1); 
K_norm = norm(K.reshape(L, d/2, 2), dim=-1); 
QK_norm = mean(Q_norm * K_norm, dim=0); 
topk_dimensions = topk(QK_norm, k=topk); 

Algorithm 3: DPE FlashAttention Forward
Input: Q, K, V [L, d], window size W, scale_factors, topk_dim
Output: attention output
# <--- Calculating position for dimension groups --->
position = [0, 1, 2, ..., L-1]; 
for group_id in range(group_num) do 
| s = scale_factors[group_id]; 
| Scale_Q_position[group_id] = position // s + W - W // s; 
| Scale_K_position[group_id] = position // s; 
end 
# <--- Calculating frequency for dimension groups --->
Cos, Sin = rotary_emb(V, position);
for group_id in range(group_num) do
| Scale_Q_Cos[group_id], Scale_Q_Sin[group_id] = rotary_emb(V, Scale_Q_position[group_id]);
| Scale_K_Cos[group_id], Scale_K_Sin[group_id] = rotary_emb(V, Scale_K_position[group_id]);
end
# <--- get final frequency --->
Scale_Q_Cos, Scale_Q_Sin, Scale_K_Cos, Scale_K_Sin = Concat_freq_by_dim(); # Scale_Q_Cos: [L, d]
Scale_Q_Cos, Scale_Q_Sin, Scale_K_Cos, Scale_K_Sin = Select_topk_dim(Topk_dim); # Scale_Q_Cos: [L, d]
# <--- Calculating sliding window attention with normal positions --->
Q, K = apply_rotary_pos_emb(Q, K, Cos, Sin);
O_sliding, LSE_sliding = flash_attn(Q, K, V, sliding_window=W);
# <--- Calculating self-attention at the left-bottom triangle with scaled positions --->
N = L - W; 
Q_scale, K_scale = apply_rotary_pos_emb(Q, K, Scale_Q_Cos, Scale_Q_Sin, Scale_K_Cos, Scale_K_Sin);
O_scale, LSE_scale = flash_attn(Q_scale[-N:], K_scale[:N], V[:N]);
# <--- Merge the attention outputs from the sliding window and left-bottom triangle --->
attention_output = merge_attentions(O_sliding, O_scale, LSE_sliding, LSE_scale);

---

Response to Question 2: DPE's performance on generative tasks and effects of selective dimension scaling

We evaluate long-document summarization using the HELMET benchmark to assess DPE’s performance on generative tasks. The results (shown below) demonstrate that DPE maintains strong generative quality.

We also compare selective dimension scaling(DPE) with scaling all dimensions(DPE w/o Topk dimension selection scaling). The results indicate that selective scaling does not negatively affect coherence or factual accuracy. This is because DPE identifies and adjusts only the key dimensions necessary for extrapolation—scaling the remaining dimensions has minimal impact.

---

Response to Question 3: Comprehensive Efficiency Analysis

We evaluate the efficiency of DPE and the standard Flash Attention based on Llama3.1 8B in our Appendix B. To clarify DPE's real‑world applicability, we also introduce other baselines in efficiency tests.

We report statistical results including decoding speed and memory usage. We report the average across 20 runs. Latency (seconds / per token) is averaged over 10 token inferences at each length. Memory consumption corresponds to peak GPU usage during inference.

Since the first two steps of DPE—detecting effective lengths and selecting key dimensions—are performed prior to inference, the only additional cost during inference comes from lightweight computations on RoPE frequencies in the scaling step. As a result, DPE introduces negligible overhead: inference speed per token remains highly comparable (e.g., 2.404s vs. 2.346s at 128k), and peak memory usage is nearly identical (e.g., 53.60 GB vs. 53.64 GB at 128k).

---

Response to Question 4: theoretical intuition and proofs about effective length

The variation in effective-length curves across dimensions is primarily due to RoPE’s use of varying frequencies, while the training context length remains fixed. We discuss this in Lines 125–133.

Beyond that, mid-range frequencies tend to collapse while low and high frequencies extrapolate better. This pattern arises because many key dimensions for extrapolation are concentrated in the mid-frequency range（32-48 dimensions）. We compute the proportion of attention contributions on LLaMA3-8B across frequency ranges (shown in the table below). Most of the contribution concentrates in the intermediate frequency dimensions.

Dear Reviewer X83u,

As this is nearing the end of the author response period, we kindly ask if there is anything additional that you would like us to address. We appreciate your valuable suggestions and feedback.

Sincerely,

Authors

Thank you for your dedicated review of our DPE. In the following, we will carefully respond to your questions.

---

Response to the Weakness 1: proof on the long-context understanding and reasoning capabilities

We agree that many existing long-context benchmarks may not fully capture a model’s long-context understanding and reasoning capabilities. To better demonstrate the effectiveness of DPE in these aspects, we evaluate it on the many-shot in-context learning and summarization tasks on 128k context length from the HELMET[1] benchmark. Results are shown in the table below.

We observe that for models requiring length extrapolation, such as LLaMA3-8B-8K, DPE yields substantial improvements on both many-shot ICL and summarization tasks. For 128k-pretrained models, DPE does not degrade ICL performance and brings slight improvements on summarization. These results support the effectiveness of DPE in enhancing long-context understanding and reasoning.

Detail results on Summ sub-tasks (InfBench_Sum, Multi_LexSum) and ICL sub-tasks (TREC_Coarse, TREC_Fine, Banking77, Clinic150,nlu) are shown below.

[1] Yen H, Gao T, Hou M, et al. Helmet: How to evaluate long-context language models effectively and thoroughly[J]. arXiv preprint arXiv:2410.02694, 2024.

---

Response to Question 1: Qwen-2.5-7B's NIAH Performance

Yes, the RoPE column reflects the direct evaluation of the original model using standard RoPE on the NIAH task. Since single-needle retrieval is no longer challenging for current LLMs, we follow the multi-needle setup introduced by LLaMA 3 [1]. Specifically, we use a 4-needle setting in our NIAH task, where the model must retrieve all 4 target tokens from the context. This makes the task significantly more challenging and better reflects long-context retrieval capability.

[1] Grattafiori et al., The LLaMA 3 Herd of Models, arXiv:2407.21783

---

Response to Question 2: Relation between DPE and Normal RoPE

Whether DPE behaves the same as standard RoPE depends on the detected effective length. If the evaluation context length is within and equal to the effective length, DPE becomes functionally equivalent to RoPE and performs no modification. For example, in Table 2, we evaluate LLaMA3.1-7B (128k), which is tested within its pretraining length. However, our detection reveals that its effective length is actually shorter than 128k. Based on this, DPE modifies the relative distance matrix accordingly and achieves noticeable performance gains.

---

Response to Question 3: The reason for using RULER

The LLaMA-3-8B and other models in Table 2 are evaluated at a 128k context length, while the LLaMA-2-Chat model is evaluated at 4k. We include LLaMA-2-Chat@4k as a classic baseline from the original RULER[1] paper. The purpose of including this baseline is to compare long-context performance with that of a strong baseline model under short-context settings.

RULER is a widely adopted long-context benchmark that allows flexible customization of the test length. By evaluating models at both short and long contexts on RULER, we can measure whether the gap between short and long-context understanding is being closed.

[1] Hsieh C P, Sun S, Kriman S, et al. RULER: What's the Real Context Size of Your Long-Context Language Models?[J]. arXiv preprint arXiv:2404.06654, 2024.

Dear Reviewer Cdu5,

As this is nearing the end of the author response period, we kindly ask if there is anything additional that you would like us to address. We appreciate your valuable suggestions and feedback.

Sincerely,

Authors

Thank you for your dedicated review of our DPE. In the following, we will carefully respond to your questions.

---

Response to Question 1: Previous Training-free Length Extrapolation Strategy

Different from Self-Extend [1] and DCA [2], which uniformly manipulate all rope's hidden dimensions, DPE performs dimension-wise adaptation: it detects the effective length of each RoPE dimension and identifies key dimensions for context extension. In this way, DPE adjusts the pre-trained models with minimal modifications while ensuring that each dimension reaches its optimal state for extrapolation.

InfLLM [3], on the other hand, is an inference-time optimization method based on KV block selection. Since it belongs to a different category from methods that manipulate the relative distance matrix, we do not include it in our discussion.

[1] Jin, Hongye, et al. "Llm maybe longlm: Self-extend llm context window without tuning." arXiv preprint arXiv:2401.01325 (2024).

[2] An, Chenxin, et al. "Training-free long-context scaling of large language models." arXiv preprint arXiv:2402.17463 (2024).

[3] Xiao, Chaojun, et al. "Infllm: Training-free long-context extrapolation for llms with an efficient context memory." arXiv preprint arXiv:2402.04617 (2024).

---

Response to Question 2: DPE's Presudocode

We provide pseudocode outlining the three main steps of DPE implementation, including:

Algorithm 1 – Detecting effective relative distance for different dimensions,

Algorithm 2 – Identifying key dimensions,

Algorithm 3 – DPE FlashAttention Forward.

Algorithm 1: Detecting effective relative distance for different dimensions
Input: dimension groups, group_num, target_length, detecting_lengths
Output: best_scale_factors
s = L / L_pretrain / 2;
Initialize detect_scale_factor = [s, s, ..., s] for each group; 
for group_id in range(group_num) do 
| for detecting_length in detecting_lengths do 
| | detect_scale_factor[group_id] = target_length / detecting_length; 
| | accuracy = niah_detect(detect_scale_factor); 
| | if accuracy > best_accuracy then 
| | | best_scale_factors[group_id] = detect_scale_factor[group_id]; 
| end 
end 

Algorithm 2: Identifying key dimensions
Input: Q, K tensors with shape [L, d], topk parameter
Output: topk_dimensions
Q_norm = norm(Q.reshape(L, d/2, 2), dim=-1); 
K_norm = norm(K.reshape(L, d/2, 2), dim=-1); 
QK_norm = mean(Q_norm * K_norm, dim=0); 
topk_dimensions = topk(QK_norm, k=topk); 

Algorithm 3: DPE FlashAttention Forward
Input: Q, K, V [L, d], window size W, scale_factors, topk_dim
Output: attention output
# <--- Calculating position for dimension groups --->
position = [0, 1, 2, ..., L-1]; 
for group_id in range(group_num) do 
| s = scale_factors[group_id]; 
| Scale_Q_position[group_id] = position // s + W - W // s; 
| Scale_K_position[group_id] = position // s; 
end 
# <--- Calculating frequency for dimension groups --->
Cos, Sin = rotary_emb(V, position);
for group_id in range(group_num) do
| Scale_Q_Cos[group_id], Scale_Q_Sin[group_id] = rotary_emb(V, Scale_Q_position[group_id]);
| Scale_K_Cos[group_id], Scale_K_Sin[group_id] = rotary_emb(V, Scale_K_position[group_id]);
end
# <--- get final frequency --->
Scale_Q_Cos, Scale_Q_Sin, Scale_K_Cos, Scale_K_Sin = Concat_freq_by_dim(); # Scale_Q_Cos: [L, d]
Scale_Q_Cos, Scale_Q_Sin, Scale_K_Cos, Scale_K_Sin = Select_topk_dim(Topk_dim); # Scale_Q_Cos: [L, d]
# <--- Calculating sliding window attention with normal positions --->
Q, K = apply_rotary_pos_emb(Q, K, Cos, Sin);
O_sliding, LSE_sliding = flash_attn(Q, K, V, sliding_window=W);
# <--- Calculating self-attention at the left-bottom triangle with scaled positions --->
N = L - W; 
Q_scale, K_scale = apply_rotary_pos_emb(Q, K, Scale_Q_Cos, Scale_Q_Sin, Scale_K_Cos, Scale_K_Sin);
O_scale, LSE_scale = flash_attn(Q_scale[-N:], K_scale[:N], V[:N]);
# <--- Merge the attention outputs from the sliding window and left-bottom triangle --->
attention_output = merge_attentions(O_sliding, O_scale, LSE_sliding, LSE_scale);

Thank you for your dedicated review of our DPE. In the following, we will carefully respond to your questions.

---

Response to Weakness 1: Technical Choices of Selecting Key Dimensions

We choose k = 48 because not all RoPE dimensions benefit from modification—some already encode positions well and do not require further adjustment. Using k = 64 does not make DPE equivalent to existing baselines. Our method applies dimension-specific scaling to extend each selected dimension to its effective length, unlike prior work that scales all dimensions uniformly. This selective approach leads to a more precise and effective position embedding manipulation.

---

Response to Weakness 2: Baselines in Experiment Results

DPE is a training-free method focused on length extrapolation, so we primarily compare against other training-free baselines such as Self-Extend. We evaluate on models like LLaMA3-8B (8k) and Mistral-v0.2-7B (32k), which require extrapolation to reach 128k. In contrast, models like LLaMA3.1 (128k) and Qwen2.5 (128k) do not require extrapolation, so we do not include extrapolation baselines for them in the main tables. Additional results for these models are provided in Appendix E.

For Table 3, our motivation is to show that DPE performs well not only on retrieval-style tasks like RULER, but also across diverse tasks in the InfiniteBench suite (e.g., Code, Math), demonstrating its broad effectiveness. This experiment setting is also consistent with STRING[1].

[1] An C, Zhang J, Zhong M, et al. Why Does the Effective Context Length of LLMs Fall Short?[J]. arXiv preprint arXiv:2410.18745, 2024.

---

Response to Weakness 3: Inconsistent Task Performance

We also observed that performance on the VT task is not fully consistent with other tasks. The VT task takes the following form:

noises....
VAR X1=12345...... VAR Y1 = 54321 ...... 
VAR X2=X1...... VAR Y2 = Y1 ...... 
VAR X3=X2...... VAR Y3 = Y2 ...... 
Find all variables that are assigned the value 12345. Answer: X1 X2 X3

From our case study, we found that many failure cases produced incomplete answers such as: Answer: X1 <eos>. This suggests that models often terminate generation prematurely after the first answer.

We hypothesize that VT may place higher demands on a model’s instruction-following capabilities, especially over long contexts. Since different models exhibit varying degrees of instruction-following robustness at extended lengths, this may explain the observed inconsistencies across tasks.

---

Response to Question: Mechanical Explanations

For Figure 1, the main reason for differing effective lengths across RoPE dimensions lies in their varying frequencies, while the training context length remains fixed. We discuss this in Lines 125–133. Similar observations have been made in prior work: YaRN [1] and LLaMA3 [2] apply interpolation only to low-frequency dimensions while leaving high-frequency ones unchanged.

For Figure 2, Barbero et al. [3] show that in RoPE-based models, a few dimensions contribute disproportionately to the attention mechanism. Since OOD issues are closely tied to attention behavior, we believe the key dimensions identified for attention also play a significant role in length extrapolation.

[1] Peng et al., YaRN*: Efficient Context Window Extension of* Large Language Models, arXiv:2309.00071

[2] Grattafiori et al., The LLaMA 3 Herd of Models, arXiv:2407.21783

[3] Barbero et al., Round and Round We Go! What Makes Rotary Positional Encodings Useful?, arXiv:2410.06205

---

Response to Typos:

We would like to express our sincere appreciation for your dedicated review of our work. Thank you for pointing out the detailed errors in our paper, we will make modifications based on your suggestions to enhance the quality of our paper. If you have any additional questions or suggestions. We would be delighted to engage in further discussion with you.

Dear Reviewer DyvB,

As this is nearing the end of the author response period, we kindly ask if there is anything additional that you would like us to address. We appreciate your valuable suggestions and feedback.

Sincerely,

Authors