Long-Short Alignment for Effective Long-Context Modeling in LLMs

We thank Reviewer qzDj for the comments. We will address your questions in the following part.

Q1. Figure 1 (a), the NoPE loss is almost zero. Is anything wrong?

A1. Thank you for your question. The NoPE loss is indeed almost zero (around 1e-5). This observation is also supported by our theoretical results in Appendix C.1.

Q2. The proposed method may degrade the performance for the shorter length, which is proved in Table 2 and Table 3.

A2. Thank you for your feedback. We apologize for the potential misunderstanding in Table 2 and 3. Table 2 and Table 3 do not evaluate performance with respect to sequence length. Instead, they show that our proposed method may experience a slight degradation compared to the baseline when trained for only 50 epochs (8.92 v.s. 8.95 in Table 2, 6.89 v.s. 6.92 in Table 3). Importantly, as shown in the same tables, our proposed method benefits significantly from longer training (100 and 200 epochs), demonstrating the effectiveness with sufficient training.

Q3. The method is sensitive to hyperparameters for the loss misalign, which is proved in Table 6. Is there any way to help choose the hyperparameter?

A3. Indeed, the misalignment loss is a regularization term, and like many regularization techniques, it can influence the training process—a phenomenon observed in prior work [1, 2]. As mentioned in Section 5.4 (Line 412), we suggest using a coefficient $\alpha$ between 0.1 and 0.3 as a default to mitigate the risk of over-regularization while maintaining performance.

[1] Zhao et al, When Will Gradient Regularization Be Harmful? ICML 2024.

[2] Srivastava et al. Dropout: a simple way to prevent neural networks from overfitting. Journal of Machine Learning Research 2014.

Thanks again for your comments, and hope our response could address your concerns. Please let us know if you have additional questions.

Q1. About Delétang et al. 2022. This makes me question if the motivation is valid.

A1. Thank you for your insightful comment. We agree that the difficulty of achieving length generalization in Transformers likely stems from multiple factors. In the related work section, we have acknowledged that positional encoding plays a significant role in this challenge. However, our findings suggest that output misalignment is another important factor, as evidenced by our comparisons.

To clarify, we do not claim that output misalignment is the sole reason for poor length generalization. Rather, we identify it as one contributing factor and show that addressing it leads to improved generalization. While some tasks may not exhibit an expanding output support set with increasing length, our focus is on cases where such a shift does occur, as it presents a clear source of distributional mismatch.

Moreover, our results demonstrate that mitigating output misalignment leads to better length generalization, even if it does not fully resolve the problem. This suggests that output alignment is a meaningful component of the broader challenge, rather than an all-encompassing solution. We will clarify this distinction in the final version.

Q2. On the connection between synthetic task and language modeling task.

A2. We acknowledge that the tasks in the two sections are not the same type. However, our motivation for including Section 3 is not to claim that the two tasks are identical, but rather to provide a controlled setting that isolates one specific challenge in length generalization: output misalignment.

While sequence modeling is more complex, the fundamental issue remains—the output distribution can shift as input length varies, which can harm generalization. Section 3 demonstrates this effect in a simpler setting where the misalignment can be analyzed more clearly. This serves as motivation for Section 4, where we propose methods to mitigate similar issues in the more complex sequence modeling setting. We will refine the discussion to better bridge the gap between the synthetic tasks and sequence modeling.

Q3. The paper would benefit a lot by comparing it with some representative length generalization methods.

A3. Thank you for the suggestion. Since most representative length generalization methods modify RoPE, we compare our method applied on top of EABF against CREAM [1], another method that builds on EABF. The results for 100 steps training are shown in the table below.

We find that our method outperforms CREAM in both LongBench-E score and perplexity, suggesting that output alignment provides additional benefits beyond positional encoding modifications.

[1] Wu et al. An Efficient Recipe for Long Context Extension via Middle-Focused Positional Encoding. NeurIPS 2024.

Q4. Refine and rationalize their claim on which lengths should the consistency be enforced on. The author can provide experiments on something like $[0.95*l_{train},l_{train}]$.

A4. Thank you for pointing this out. We agree that "moderate" would better reflect our intended meaning in L064 and L216. Our goal is to enforce consistency across moderate length differences rather than extremely small ones, as small variations may not provide sufficient regularization. We will refine the wording in the paper to better reflect this insight.

Following your suggestion, we also conduct experiments using the sampling range $[0.95*l_{train},l_{train}]$ (which implies that the sampling range of $l_{extra}$ is $[1, 0.05l_{train}]$) following the setting in Table 7. The results of 200 steps are shown in the table below.

While restricting $l_{extra}$ to $[1, 0.05 \cdot l_{train}]$ still provides some improvement, the broader range $[1, l_{train}/2]$ is more effective. This suggests that keeping consistency over moderate rather than minimal length differences is beneficial.

Q5. I was confused on how the hyper-parameters is selected, could you clarify?

A5. Sure. The batch size is 256. We use SGD as the optimizer and search for the initial learning rate in {5e-5, 1e-4, 5e-4}. The learning rate follows a cosine schedule. The models are trained for 40,000 epochs. The training loss is Mean Squared Error (MSE).

Q6. In Section 3, why are those transformation selected?

A6. To align the output distribution across different input lengths, $f$ is chosen based on the following criteria: (1) $f$ must be a bijection to ensure the existence of $f^{-1}$ (2) $f$ should be a contraction mapping on $[1, +\infty]$ to reduce the discrepancy in outputs across different input lengths.

Also thanks for the minor comments! We will address them in the final version.

We thank Reviewer Ueg1 for appreciating the insight of our paper. We will address your questions in the following part.

Q1. (From Other Strengths And Weaknesses) I believe comparison with other methods in length generalization is needed. Also seeing if this technique can be combined with other techniques to see if it further improves the performance.

A1. Thank you for your suggestion. To clarify, both Table 2 and Table 3 compare our proposed method combined with a length generalization technique against the length generalization technique alone. The results demonstrate that our method consistently improves performance when combined with existing approaches.

To further address your point, we compare our method applied on top of EABF against CREAM [1], another method built on EABF. Additionally, we evaluate whether combining our method with CREAM leads to further improvements. The results for 100-step training are shown below:

These results show that our method outperforms CREAM alone in both LongBench-E score and perplexity. Furthermore, combining our method with CREAM leads to additional improvements. This supports our claim that output alignment plays a crucial role in length generalization and can complement existing approaches.

[1] Wu et al. An Efficient Recipe for Long Context Extension via Middle-Focused Positional Encoding. NeurIPS 2024.

Q2. (From Questions For Authors) In the length generalization experiment, why do you think 1 over square root function works the best? Is it just because the values are in [0,1] range? What about other functions that puts the output distribution in this range? Do they work as well?

A2. We conducted additional experiments using alternative parameterization functions, including $f(x)=1/\log(x+1)$ and $f(x)=1/x$ following your suggestions. The results, presented in the anonymous link https://ibb.co/39M3D8Nq, show that while these functions also improve performance, $f(x)=1/\sqrt{x}$ still performs the best.

The superior performance of $f(x)=1/\sqrt{x}$ is not solely due to its range being within [0,1]. Rather, we hypothesize that it is because, both empirically (Figure 1(b)) and theoretically (Theorem C.1), the test loss is proportional to the square of the input length when no parameterization function $f$ is applied. Therefore, using $f(x)=1/\sqrt{x}$ effectively normalizes the output and mitigates this issue.

Q3. (From Questions For Authors) What positional encoding is used for main results like in table 1? Do you think if changing the positional encoding might affect the results? A comparison of them against each other might be helpful. Especially when the initial positional encoding is not good at length generalization and your approach improve length generalization which means your approach can be stand alone for length generalization.

A3. Thank you for your question. For the main results in Table 1, GPT-J-6B, GPT-NeoX-20B, Llama2-7B, and Qwen-7B-8K use the original RoPE, while Yarn-Llama2-7B-8K and CLEX-Llama-4K use modified RoPE and achieve better length generalization.

Changing the positional encoding can indeed affect the results, as seen in Table 1—models using modified RoPE tend to perform better in terms of length generalization. However, even among models using the same RoPE (e.g., GPT-J-6B, GPT-NeoX-20B, Llama2-7B, and Qwen-7B-8K), we observe varying degrees of length generalization. This suggests that while positional encoding plays a role, it is not the sole factor.

As shown in Table 11, our method does not lead to significant performance gains when using the original RoPE, likely because using original RoPE will lead to slow convergence speed in length generalization task [1, 2]. However, when we use modified RoPE, our method provides notable improvements in length generalization, as shown in Table 1. This suggests that while the choice of positional encoding does impact performance, our method works best when combined with a modified version of RoPE that better supports length generalization.

[1] Chen et al. Extending Context Window of Large Language Models via Positional Interpolation. arxiv 2306.15595

[2] Zhang et al. Extending LLMs' Context Window with 100 Samples. arxiv 2401.07004

Thanks again for your comments, and hope our response could address your concerns. Please let us know if you have additional questions.

Q1. Maybe there is an interesting connection to be discussed with SCAN.

A1. Thanks for highlighting these works! Indeed, SCAN-based methods leverage CFG-like rules to generate training data that reduce context sensitivity, which can help improving out-of-distribution generalization. This is indeed relative to the concept behind our proposed output matching, which ensures semantic consistency across different input lengths. However, there are key differences:

• Different Focus. SCAN-based approaches primarily operates from the input perspective by generating structured data, while our method focuses on output matching and introduces a fresh output perspective.
• Better Efficiency. SCAN-based approaches need manual CFG design, making them time-consuming and less scalable, while our method can offer a more efficient solution with only ~5% additional computation.

We will clarify this distinction and discuss potential connections to these prior works in the revised paper.

Q2. The results in Section 3 did not seem very closely connected to the proposed methods in Section 4.

A2. Thank you for your insightful comments. We clarify that Section 3 is crucial for our paper for the following reasons:

• Fundamental Motivation for Length Generalization Strategies: Section 3 highlights that one of the key challenge in length generalization is that the output distribution shifts when input length changes. This sets the stage for why controlling output alignment is crucial.
• Bridging Task Reformulation and Model-Based Regularization: Both Section 3 and Section 4 aim to reduce output distribution discrepancies across different input lengths. When we have a strong prior about the task, we can reformulate it directly (as in Section 3). However, when such a prior is unavailable, we instead introduce a regularization approach (Section 4) to encourage the model to implicitly learn this property. Without Section 3, the necessity of such regularization may be less clear. By including both, the paper offers a broader perspective on addressing length generalization, rather than focusing solely on a single heuristic.

Thank you again for raising this point, we will make Section 3 more concise to avoid distracting the main contributions of Section 4.

Q3. The proposed terminology of "Output Alignment" is overloaded with various techniques for model post-training.

A3. Thanks for pointing it out! We will use “Output Matching” instead to avoid such confusion.

Q4. It would be good to understand the limitations of the proposed method. The ablations in section 5.3 could be expanded to highlight potential weaknesses of the method. Does this change if the regularization coefficient is increased?

A4. Thanks for your suggestions! We acknowledge the importance of discussing the limitations of our method and have addressed this in Section 5.4 (line 408). Following your suggestions, we conduct additional experiments to further analyze the effect of $\alpha$ based on the ablations in Table 5. Specifically, we extend the range of $\alpha$ to include 0, 0.1, 0.3 and 0.5. The results are shown in the table below.

The results indicate that increasing $\alpha$ leads to a decline in scores for Depth = 0% and 25%, which confirms the potential drawback of excessive regularization. This conclusion is also consistent with the claim in Section 5.4.

Q5. One of the models looks very underfit in Figure 1a, perhaps worth investigating the training configuration.

A5. The model that appears underfit in Figure 1(a) uses Alibi positional encoding. We indeed observe that the Alibi-based model converges more slowly than those using other positional encodings, likely due to the complex inductive bias introduced by Alibi.

Q6. It could be interesting to see if different positional encoding schemes lead to different degrees of "long-short misalignment" in the predictable way.

A6. Thank you for your thoughtful insight. Regarding the connection between positional encoding schemes and long-short misalignment, we observe that most modern LLMs rely on RoPE, often with modifications to its hyperparameters. However, as shown in Table 1, even among models using RoPE, length generalization ability varies significantly. For example, while GPT-J-6B, GPT-NeoX-20B, Llama2-7B, and Qwen-7B-8K all use RoPE, only Qwen-7B-8K demonstrates strong length generalization and low long-short misalignment. Furthermore, Yarn-Llama2-7B-8K and CLEX-Llama-4K, which employ modified versions of RoPE, also exhibit improved generalization and reduced misalignment. These observations suggest that while positional encoding choices influence long-short misalignment, they do not fully determine it.