Long-Short Alignment for Effective Long-Context Modeling in LLMs
Authors: Tianqi Du
Haotian Huang
Yifei Wang
Yisen Wang
Venue: ICML 2025 poster (2025)
Keywords: Large language model
length generalization
long-short alignment
Primary area: deep_learning->selfsupervised_learning

Abstract: Large language models (LLMs) have exhibited impressive performance and surprising emergent properties. However, their effectiveness remains limited by the fixed context window of the transformer architecture, posing challenges for long-context modeling. Among these challenges, length generalization — the ability to generalize to sequences longer than those seen during training — is a classical and fundamental problem. In this work, we propose a fresh perspective on length generalization, shifting the focus from the conventional emphasis on input features such as positional encodings or data structures to the output distribution of the model. Specifically, through case studies on synthetic tasks, we highlight the critical role of **long-short alignment** — the consistency of output distributions across sequences of varying lengths. Extending this insight to natural language tasks, we propose a metric called Long-Short Misalignment to quantify this phenomenon, uncovering a strong correlation between the metric and length generalization performance. Building on these findings, we develop a regularization term that promotes long-short alignment during training. Extensive experiments validate the effectiveness of our approach, offering new insights for achieving more effective long-context modeling in LLMs. Code is available at https://github.com/PKU-ML/LongShortAlignment.

Long-Short Alignment for Effective Long-Context Modeling in LLMs
Tianqi Du Haotian Huang" 2
Yifei Wang?
Yisen Wang 14
Abstract
eration, reading comprehension, code synthesis, instruction- following, and commonsense reasoning (Radford et al., 2019; Brown et al., 2020; Chowdhery al., 2023; Tou- vron et al., 2023). Their performance has consistently im- proved by scaling both model and dataset sizes (Kaplan et al., 2020). However, the effectiveness of LLMs remains limited by the fixed context window of the Transformer architecture, posing significant challenges for long-context modeling. With larger context sizes, a model can benefit from more in-context learning examples, greater number of reasoning steps, the ability generate longer coherent texts (Li et al., 2024; Huang & Chang, 2023; Wang et al., 2024b). Nevertheless, training Transformer with long input sequences is often prohibitively slow and memory- intensive, making it crucial to understand and improve how LLMs generalize to longer contexts.
Large language models (LLMs) have exhibited impressive performance and surprising emer- gent properties. However, their effectiveness re- mains limited by the fixed context window of the transformer architecture, posing challenges for long-context modeling. Among these chal- lenges, length generalization--th the ability to gen- eralize to sequences longer than those seen dur- ing training-is a classical and fundamental prob- lem. In this work, we propose a fresh perspec- tive on length generalization, shifting the focus from the conventional emphasis on input features such positional encodings or data structures to the output distribution of the model. Specif- ically, through case studies on synthetic tasks, we highlight the critical role of long-short align- ment-the consistency of output distributions across sequences of varying lengths. Extending this insight natural language tasks, we pro- pose metric called Long-Short Misalignment to quantify this phenomenon, uncovering strong correlation between the metric and length gen- cralization performance. Building on these find- ings, we develop a regularization term that pro- motes long-short alignment during training. Ex- tensive experiments validate the effectiveness of our approach, offering new insights for achieving more effective long-context modeling in LLMs. Code is available at https::/ / gi thub. com/ PKU 2.00MMM//oogghhootttigmemeeee
A classical and foundational subproblem within long- context modeling is length generalization -the ability to generalize from shorter training sequences to longer test sequences (Anil et al., 2022). This remains a major chal- lenge even for large-scale Transformers (Liu et al., 2024). Understanding the mechanisms of length generalization is thus an essential step toward achieving robust and efficient long-context modeling.
There exist two dominant approaches to understanding and improving length generalization. The first is to design bet- ter positional encodings (PE) (Press et al., 2021; Su et al., 2024; Kazemnejad et al., 2023; Peng et al., 2024; Chen et al., 2024; Yang. 2023; Zhang et al., 2024b), which help the model systematically encode tokens across a wide range of positions. By reducing the inductive gap between short training sequences and long test sequences, these encodings can partially improve generalization (Kazemnejad al., 2023). The second is to analyze the underlying mechanisms of Transformer models (Zhou et al., 2024; Lee et al., 2023; Veličković & Blundell, 2021; Nogueira et al., 2021; Dele- tang et al., 2022), such as what algorithms they can simulate or how task design affects generalization. However, we ob- serve that one key aspect has been largely overlooked: the output space of the model. In this work, we argue that output behavior plays central role in determining generalization quality across context lengths.
1. Introduction
Large language models (LLMs) have demonstrated impres- sive abilitics in various tasks such as natural language gen-
'Equal contribution 'State Key Lab of General Artificial Intel- ligence. School of Intelligence Science and Technology, Peking University, China "NUS, Singapore 'MIT CSAIL, USA IIstittte for Artificial Intelligence, Peking University, China. Correspon- dence to: Yisen Wang <yisen.wang @ pku.edu.en>.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
We begin our analysis with synthetic tasks involving length

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
generalization: predicting the mean value and the length of binary sequences. Both empirical and theoretical results reveal a stark contrast-Transfommeer generalize well on the mean prediction task but struggle on the length predic- tion task. The key difference lies in the support set of the output distribution: in the mean prediction task, it remains stable across input lengths, while the length prediction task, it varies with the sequence length. We hypothesize that this misalignment in the output distribution leads to poor generalization in the latter task. To verify this, we propose a reparameterization technique named OutRep that explicitly aligns the output distributions across lengths. Our analyses confirm that this approach significantly improves generalization, supporting our hypothesis.
it comes to algorithmic reasoning (Veličković & Blundell, 2021). Specifically, focus simple tasks and study length generalization on the standard Transformer archi- tecture with causal structure. Related to this, Lee et al. (2023) study how well transformers trained from scratch can learn simple arithmetic tasks, and find that no length generalization is observed. Nogueira et al. (2021) find that partial length generalization on addition is observed only when models reach 3B parameters and when the addition questions are presented in reverse order. Zhou et al. (2024) proposes the RASP Generalization Conjecture that Trans- formers tend to learn a -engtthgenerrlizing solution if there exists a short RASP-L program that works for all input lengths. Liu et al. (2023) discovers that the Transformer will learn shortcuts through the study of various synthetic tasks. Besides these explorations on specific tasks, some works study the impact of different positional encodings on the length generalization of math reasoning tasks (Press et al., 2021; Ontanon et al., 2022; Kazemnejad et al., 2023; Ruoss et al., 2023). More details of these works can be viewed in Appendix A.
Building on this insight, we extend our findings natu- ral language tasks. Although natural language outputs are vector-valued and more complex than scalar outputs in syn- thetic tasks, similar misalignment phenomena arise. For instance, two sequences with the same ending but slightly different lengths should ideally produce similar output dis- tributions. However, models that generalize poorly often produce divergent outputs for such inputs. To quantify this phenomenon, we introduce a metric called long-short mis- alignment, which measures the divergence in output distri- butions via symmetric cross-entropy. Both empirical and theoretical analyses show that this metric correlates strongly with long-conteed performance- more so than traditional training making it a reliable indicator of generaliza- tion. Motivated by this, we incorporate the short mis- alignment metric as regularization term during training. Extensive experiments across both synthetic and natural language tasks validate the effectiveness of this approach.
Long-context Modeling on Natural Language Tasks. A series of works (Sun et al., 2023; Chi al., 2022; 2023; Zhang al., 2024b; Chen et al., 2024; Peng et al., 2024; Yang, 2023; Chen et al., 2023a; Fang et al., 2025b) aim to extend the context size of Transformer-based models during fine-tuning, primarily by modifying positional en- codings. For example, Zhang et al. (2024b) introduces a novel extension to RoPE (Su et al., 2024) which combines adjusting RoPE's base frequency and scaling the attention logits to help LLMs efficiently adapt to larger context win- dow. Chen et al. (2024) generalizes the positional encoding scaling approaches to model the continuous dynamics by ordinary differential equations over the length scaling factor. Wang et al. (2024a) proposes a novel approach designed to narrow the generalization gap and provides length extrapo- lation analysis on the feature gap. Our approach, in contrast, focuses on the model's output space, which identifies the crucial role of long-short alignment in length generalization.
The main contributions of this work are as follows:
We identify the crucial role of output behavior in long- context modeling, with a focus on length generaliza- tion. Both empirical and theoretical evidence shows that misalignment in output distributions across input lengths leads to poor generalization.
We introduce the long-short misalignment metric to quantify this output discrepancy and demonstrate its strong correlation with length generalization ability.
3. A Case Study on Synthetic Tasks: How Long-Short Alignment Affects Length Generalization?
We incorporate this metric into novel regularization term. Extensive experiments validate the effectiveness of this approach in boosting long-context modeling performance.
Long-conter modeling aims to extend the reasoning and generation capabilities of language models to longer input sequences (Fang et al., 2025a; Wu et al., 2025; Kuratov et al., 2024; Zhang et al., 2024b). One fundamental challenge in this setting is length generalization the ability to general- ize from shorter training contexts to longer ones. Prior work has identified several factors that influence length general- ization, including the task type (Zhou al., 2024; Jelassi
2. Related Work
Length Generalization on Synthetic Tasks. Our paper is related to the line of work that seeks to understand the capabilities and limitations of Transformer models when
2

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
Examples for Meon Prediction
Examples Length Prodiction
0
1150
Batmagrto
200%
820
0xt-21999000
23000
600
400
533
39055
220
20
34
0-
10
(a) Length generalization in the mean pre- diction task with different positional encod- ings
(h) Length generalization in the length pre-
(c) Length generalization the length pre- diction task with different reparameteriza- tion function.
diction task with different positional encod- ings
Figure 1: Comparison between the length generalization performance in the mean prediction task and the length prediction task. The training sequence length is uniformly selected from [1, 10 (indicated by the light red area) while the test sequence lengths (on the x-axis) can reach a maximum of 50. In the length prediction task (b), the model struggles with length generalization. Conversely, the model demonstrates significantly better length generalization in the mean prediction task (a) using NoPE (indicated by the orange line). Figure (c) shows the length generalization performance in the length prediction task using different reparameterization functions f=r). All three reparameterized targets improve generalization compared to the origin (blue) target. Among them, f(x) 1/ / (red) performs exceptionally well.
et al., 2023; Nogueira et al., 2021) and the design of posi- tional encodings (Ontanon et al., 2022; Kazemnejad et al., 2023; Ruoss al., 2023). In this work, we propose a fresh perspective by examining the consistency of the model's output distribution across inputs of varying lengths, named loor-short alignment, identifying it as a crucial yet underex- plored factor for effective long-context modeling.
ability in the mean prediction task, as the test loss on longer sequences remains nearly consistent with the loss on shorter sequences. We now provide theoretical analysis of this observation.
Theorem 3.1. In the length prediction task, the length gen- eralization loss Elengthi(;)) (13-) has a quadratic relationship with the predicted length leestr i.e.,
Mean Prediction v.s. Length Prediction. We start from a case study on synthetic tasks: in the mean prediction task, the prediction target is the mean value of the sequence, while in the length prediction task, the target is the length of the sequence. We focus on binary input sequences, where each position in the sequence is filled with 0 or with the same probability, and the decoder-only Transformer (Vaswani et al., 2017), a model widely used in both synthetic tasks (Zhou et al., 2024; Jelassi et al., 2023) and LLMss (Touvron al., 2023; Peng al., 2024), which utilizes a causal mask in the self-attention module to enable nuto-regressive gen- eration. More model details can be found in Appendix B. We train the model on sequences with a maximum length of Irain = 10 and test it on sequences with a maximum length of hest = 50. Figure la and Figure 1b display the test results for both tasks. We observe that regardless of the positional embedding used, the test loss of the length pre- diction task dramatically increases when the test sequence length exceeds 10, the maximum training length. Further- more, the test loss continues to rise as the test sequence length grows, indicating that the model demonstrates very low length generalization ability in the length prediction task. In contrast, the model exhibits strong generalization
Stougth (g) (g)reing (hest)
[ (Xrest) ((1esst) 1
= =E =Ex E 10,13/tms
[1
(1)
= =0 ((Ieett hernin) ))
where glamii is the model trained on sequences with max- imum training length brain Xtest is the testing input with length hest-
However, in the mean prediction task, the length generaliza- tion loss has fixed upper bound:
Emean (g9 Jarain leest) O().
(2)
The full statement and the proof are shown in Appendix C. From both the empirical and theoretical results, is evident that while the mean and length of a sequence all con- vey global information, the model's length generalization ability varies across these tasks. A key distinction lies in the differences in output distribution for each task. In the mean prediction task, where the model generalizes well, the output remains within the fixed range of (0, 1]. regardless
3

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
of sequence length. However, in the length prediction task, where generalization is poor, the support set of the output distribution shifts to single-point set { as the sequence length increases to I. This distinction between the two types of tasks motivates us to consider the importance of long- short alignment for better length generalization ability.
4.1. Long-Short Misalignment: Quantifying the Discrepancy of Output Distributions
In synthetic tasks, measure the discrepancy between the support sets of output distributions to capture the dif- ferences in output across varying sequence lengths. How- ever, in natural language tasks, the model output is a vector go (x) E R! where the dimension is the size of the vocab- ulary V/. This makes it challenging to directly apply the same analysis from synthetic tasks to natural language tasks, Despite this, similar long-short misalignment issues can still be observed in natural language tasks. Specifically, for sequence x and its suffixes x[- and [|-)))) where 1 and l are two lengths and [ means the last 4 tokens of x (i = 1, the model's output is expected to remain con- sistent when I and l2 are similar, because the two suffixes share large overlap in tokens, resulting in similar contextual information. However, we find that models with poor length generalization tend to produce distant output distributions when conditioned on these sequences.
Explicit long-short alignment helps length generaliza- tion. We propose oulput reparameterization (OutRep), a reparameterization technique explicitly improve long- short alignment in synthetic tasks, thereby enhancing the model's 's length generalization ability. In the length predic- tion task, the output distribution for sequences of certain lengths is known. Leveraging this prior knowledge, dur- ing training, we apply a reversible function / TR - to map the support sets of output distributions for sequences of varying lengths into more aligned sets. Instead of using the original target y(x), we train the model on the trans- formed target f (y((). At test time, we apply the reverse function f to the output to recover the original predic- tion. This approach aligns the output distributions across different lengths, which is expected to improve length gen- eralization. We consider the following reparameterization functions: f(x) = Vzz (x) log(#) and [xx) = 1/va. We show the experiment results in Figure Ic. It can be observed that all three reparameterization functions success- fully relieve the poor length generalization ability in the length prediction task. Specifically, the reparameterization function f() = 1/2 has a nearly perfect generalization ability when the length is no more than 35. The rising trend when the test sequence length becomes longer is still slow. These results verify our conjecture on the long-short align- ment that better long-short alignment leads to improved length generalization ability. We add more theoretical results and discussions in Appendix C and Appendix D. In the next section, we will extend these findings to the more practical natural language tasks.
4.1.1. METRIC
To quantitatively explore the relationship between long- short alignment and length generalization ability, we want to first design metric to evaluate long-short alignment. We propose to utilize symmetrical cross-entropy (SCE) loss (Wang et al., 2019) to measure the divergence between output distributions conditioned on two distinct sequences. Consider two input sequences, x and x' with corresponding model predictions y = (x) and y' = [x()'. The SCE loss between these predictions is defined as:
Esculy.y') = - ((x', log(y)) + (y, log(y'))) (3)
where () denotes the inner product and the log function is applied element wise. A lower SCE loss between the two predictions indicates better alignment. To assess over- all long-short alignment, we compute the expectation over sequence lengths / and l2 for a given input x:
4. Long-Short Alignment in Natural Language Tasks
Emisalign (go) Ex,1, (CSCE(90(x(11)) go (x(-)=)) +))
In the previous section, we observed positive correlation between length generalization ability and long-short align- ment in synthetic tasks. Motivated by this finding, in this sec- tion, we aim to extend this investigation to natural language tasks. First, we introduce a metric to quantify long- short alignment in sequence modeling and demonstrate its strong correlation with performance on long-com long-context benchmarks. Building on this insight, we propose incorporating this met- ric as a regularization term during training improve long- short alignment, which can lead to the performance gains detailed in Section 5.
(4)
We refer to this metric as the long-short misalignment, where a lower value signifies less discrepancy of output distributions across different lengths. An illustration of this metrie is shown in Figure 2. In practice, we sample I and la from the interval [/raii/ /2, lerain], where lerain represents the maximum context length used during training.
4.1.2. RESULTS
To evaluate the model's length generalization ability, we use the perplexity on long validation sets (16k length) and the LongBench-E score (Bai et al., 2023b). For perplexity evaluation, we select a subset from the RedPajama- -aokk Book
4

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
Table I: The proposed long-short misalignment metric mmssalign of models, with their log of perplexity (PPL) on 16k-long contexts and LongBench-E score. We also provide Liaain as an additional metric in comparison. We find that missalign correlates better with the long-context benchmark performance.
Model
Ctrain
misalign
log(PPL)
LongBench-E Score
GPT-J-6B (Wang, 2021)
2.1
3.4
9.5
7.8
GPT-NeoX-2008 (Black et al., 2022)
2.3
3.2
9,4
9.7
Llama2-7B (Touvron al., 2023)
1.9
3.4
9,4
8.9
RandomPos (Ruoss al., 2023)
2.0
2.8
8.2
9.2
Yam--lama--2-7B-8 (Peng et al., 2024)
2.0
2.6
3.8
21.2
Qwen-7B-8k (Bai al., 2023a)
1.7
2.8
3.2
24.2
CLEX-LLaMA-4K (Chen et al., 2024)
1.9
2.4
1.8
32.7
Metric
Correlation Coefficient with Metric
Grain
0.62
-0.55
missalign
0.85
-0.85
corpus (Computer, 2023), following the protocol in (Chen et al., 2024). LongBench-Fi E is multitask benchmark that comprehensively evaluates large language models' ability to understand long contexts, with task lengths averaging between 5k and 32k tokens.
where c. (i 0, 1, are constants related to lestt
Specifically, each C (.tuaa) increases as last increases, and the ratio C( |C increases as lest increases. This indicates that the lesting length increases, the alignment loss becomes increasingly significant.
Empirical Results. Table 1 shows the long-short mis- alignment metric Lmisalign for the models along with their corresponding long-context evaluation results. Addition- ally, we include the training loss Ltrain for each model on the RedPajama-Bodd corpus, which reflects the per- plexity on sequences with the maximum training length Arain- Interestingly, while the training loss (i.e., log of perplexity on sequences of length 4rain) shows moder- ate correlation with long-context performance metrics- indicating that lower training loss can contribute im- proved length generalization- the long-short misalignment metric Cmisalign demonstrates a much stronger correlation with long-context performance, as evidenced by its higher absolute correlation coefficient. These findings suggest that Lmisatign is a promising indicator of length generalization ability. However, important note that we do not claim any causal relationship based solely these observations. We will elaborate more on this relationship in Section 5.
Proof Sketch. We decompose the generalization error Egen (go: hestt) of long sequence into the misalignment term between this sequence and set of shorter sequences, along with the prediction error of the shortest sequence. The prediction error of the shortest sequence is related to the model's training loss Lrrain while the misalignment term between long and short sequences increases significantly as the testing length grows. This mainly because the model has a limited capacity to align sequences of differ- ent lengths, so aligning longer test sequences requires more intermediate- length sequences, resulting in higher alignment loss. Therefore, reducing the alignment loss can effectively lower the generalization error when testing with longer se- quences.
The full statement and the proof are shown in Appendix C.2. This theorem highlights the importance of minimizing both Lmisalign and Lirain to achieve lower generalization er- ror. Moreover, the testing length lest increases, reducing Luisalign plays more critical role in improving general- ization performance. The above empirical and theoretical results motivate us to explicitly optimize the long-short mis- alignment metric to enhance the model's ability to handle long-
Additionally, we also provide theoretical support for this ob- servation, extending previous work on autoregressive mod- eling (Zhang et al., 2024a) with a theorem:
Theorem 4.1 (Generalization guarantees for the natural language task). Under some model assumptions, the gen- eralization error Egen (go; host)) with testing length beea is upper bounded by the sum of training loss Lirain (ge) and misalignment metric Lauisaliga i.e.,
long-context context sequences, which will be stated in the follow- ing sections.
4.2. Long-Short Misalignment Metric as Regularization Term
Egen(ghumn (gh hhest) < Oftimee) Emisalign Cof (leen))) Lirsin (5)
Since both empirical and theoretical results indicate a strong correlation between the proposed long-short misalignment
cll (hoo))
5

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
Truncate
X
Probs
Long Text
"He likes"
Tokens
04
6ogiiog
the
nothing
SCE Loss
Probs
Short Text
"He likes"
Tokens
logging
the
nothing
Illustration of long-short misalignment metric Lmisalign- Given two input sequences, where one is truncated version of the other, the long-short misalignment metric is computed by taking the expectation on Symmetrical Cross-Entropy (SCE) loss (Wang et al., 2019) between the model's 's predictions for these two sequences.
Figure 2:
metric and long-context benchmark performance, incor- porate this metric regularization term into the training loss, resulting in the new training loss defined as:
Table 2: Performance of the fine-tuned models using only cross-entropy loss (baseline) and an additional long-short misalignment loss on long-context modeling benchmark, LongBench-E score (Bai et al., 2023b) and perplexity on the 8k-length validation
Ciraia (go) Liraingo a Lmisalign(go))
The fine-tuning sequence length is 4k, exactly the same as the training sequence length. The models finetuned with our proposed loss outperform the baseline across different model adaption strategies.
where Lirain is the original cross-entropy training loss and is the regularization coefficient. Calculating these two losses separately during training can be tome-consuming, as the computation of Emisalign requires forward propagation through two distinct sequences. To address this, we propose an efficient implementation for Lirsin We first sample integer lextra from [1, /rain/ /2] and then sample sequence of length htsin + Textra- The first 4tain tokens form the first input sequence, while the last Itain tokens form the second input sequence. Both sequences can be used to compute Ctrain- The overlap between the two sequences starts at token lextra + and continues token Irain. resulting in an overlap of 4rriin lextra tokens. We can calculate the long-short misalignment loss in the overlapping positions. This implementation requires only two forward propaga- tions for the two sequences, resulting in minimal additional time and resource costs compared to calculating the origi- nal train loss Lerain- detailed Pytorch-like algorithm is provided in Appendix E and an overall illustration can found in Figure 3. We will conduct experiments using the proposed regularization term in the next section show its effectiveness.
Benchmark Training steps
LongBench-BB (t)
Perplexity ())
50
100
200
50
100
200
RedPajama-
Cerii (Baseline)
22.7
23.8
24.7
7.21
6.56
6.12
miinnin (Ours)
25.2
6.89
6.24
-0.5Lmisatiga (Ours)
21.9
23.7
24.7
7.44
7.01
6.54
PG19
Corai (Baseline)
20.2
21.4
22.5
8.92
7.89
7.45
miiiniinn (Ours)
20.7
22.1
25.3
8.95
7.92
7.35
-0.5Cmisslien (Ours)
20.1
22.2
23.6
9.42
8.59
8.21
first examine our proposed training loss in Equation (6) on length generalization tasks, where the model is trained on short sequences and tested longer ones. Additionally, explore its application in another common scenario: long- context learning, where both training and testing involve long sequences. Then we conduct analytical experiments and ablation studies to further understand the impact of our proposed loss.
5. Experiments Natural Language Tasks
5.1. Experiments with Training on Short Sequences
In this section, we conduct extensive experiments to verify the effectiveness of our proposed length alignment loss. We
In this section, we consider the classical length generaliza- tion setting, where the model trained on short sequences
6

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
CE Loss
Sequence
I think
he likes
P1 P2
P3 P4
think
he likes
jogging too
Model
SCE Loss
Long-short Misalignment
he likes jogging too
Q1 Q2
Q3 q4
Sequence 2
CE Loss
Next-token Prediction
Figure 3: Illustration of efficiently calculating the total training loss Lirain This implementation requires only two forward propagations for the two sequences, resulting in minimal additional time and resource costs
Table 3: Performance of the finetuned models using only cross-entropy loss (baseline) and an additional shorth misaligmment loss on long-context modeling benchmark, LongBench-E score (Bai et al., 2023b) and perplexity on the 8k-length validation set. The fine-tuning sequence length is 8k. The models finetuned with our proposed loss outperform the baseline across different model adaption strategies.
ity to understand long contexts, with task lengths averaging between 5k and 32k tokens, which has been adopted by many previous works (Chen al., 2024; Jin et al., 2024) as an effective evaluation metric for long-context model- ing. The results, shown in Table 2, indicate that the model fine-tuned with proposed loss consistently outperforms the baseline model on the LongBench-F. benchmark. The fine-tuned model shows lower perplexity on RedPajama- Book and similar perplexity on PG19. These results support the effectiveness of our proposed loss and the intuition that lower misalignment metric Cmisalign leads to better length generalization ability. For the regularization coefficient a, find that larger values do not always improve perfor- mance, they may interfere with the model's next-word prediction.
Benchmark Training steps
LongBench-Ei (1)
Perplexity (1)
50
100
200
50
100
200
LongQLora
Eurin (Baseline)
21.9
22.1
23.4
6.82
6.41
5.82
+0.1C (Ours)
21.8
23.3
25,8
6.72
6.39
5.77
m0.5mmissiign (Ours)
21.4
23.9
25.1
7.07
6.62
5.92
EABF
Lunin (Baseline)
22.1
22.9
23.6
6.89
6.52
6.01
+0.1Cmisalien (Ours)
23.2
24.0
24.8
6.92
6,43
5.91
mis5eiiss (Ours)
22.5
23.2
23.9
7.14
6.78
6.34
5.2. Experiments with Training on Longer Sequences
(4k-long) and tested longer sequences (at least 5k-long). Due to the high computational cost of pre-training large language models from scratch, most current methods fine- tune open-sourced pre-trained models (Chen et al., 2024; Yang, 2023; Peng et al., 2024). our experiments, we use Llama2-7b (Touvron et al., 2023) as the base model and apply the CLEX (Chen al., 2024) adjustment method. We use two datasets: the RedPajama- Book corpus (Computer, 2023) and PG19 (Rae et al., 2019). The experiments are conducted with context length of 4,096, batch size of 64, and maximum of 200 training steps. For the regularization coefficient a, we test values 0.1 and 0.5.
In this section, we consider the scenario that the model is finetuned on a longer sequence than the training length. We use different model adjustment strategies during the fine-tuning stage, to demonstrate that the proposed length alignment loss can be applied various long-context fine- tuning methods. Our experiments use Llama2-7b as the base model. For model adjustments, consider two ap- proaches: LongQLora (Yang, 2023) and EABF (Zhang et al., 2024b). LongQLora leverages multiple techniques, includ- ing Position Interpolation (Chen et al., 2023a), QLoRA (Dettmers et al., 2024), and Shift Short Attention from Lon- gLoRA (Chen et al., 2023b). Meanwhile, EABF introduces dynamic rescaling mechanism the attention layers and applies higher base frequency for RoPE. The experiments are conducted the RedPajama-Book corpus (Computer, 2023), with a context length of 8,192, a batch size of 64, and a maximum of 200 training steps. For the regularization coefficient a, we test values of 0.1 and 0.5.
We evaluate performance using the LongBench-F score (Bai al., 2023b) and perplexity validation sets made up of sequences of length 192 from the corpus of the respective training dataset. LongBench-Ei is multitask benchmark that comprehensively evaluates large language models' abil-

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
Table 4: The overall evaluation results on BABILong (Kura- tov et al., 2024) with sequence lengths of 4K, 8K, 16K.
Table 6: Ablation study on the regularization coefficient a. The setting is the same as Table 2. We adopt RedPajama- Book (Computer, 2023) as the training dataset. We find it important select moderate value for a.
Evaluation Length
Training Loss
4K
8K
16K
Benchmark
LongBench-E (t)
Perplexity (L)
Cirain (Baseline)
48.2
42.4
37.9
Training steps
50
100
200
50
100
200
Cirain + m.immina
49.1
44.4
40.1
RedPajana- BaiRajannaBook
Table 5: The evaluation results BABILong with different locations of the facts in the QAI task. Input length is 16K.
Liran (Baseline)
22.7
23.8
24.7
7.21
6.56
6.12
+0.1 missalggn
23.1
25.2
26.6
6,89
6.24
5,88
+0.3 +0.3C maisnlign
23.4
25.8
27.1
6.95
6.35
5.98
+ 0.55..... miisnligg
21.9
23.7
24.7
7.44
7.01
6.54
Fact Depth (%)
+ +1.0Cmienlign
18.2
19.4
19.9
16.21
14.12
12.92
Training Loss
0
25
50
75
Lerain (Baseline)
75
64
30
69
Table 7: Ablation study the sampling range. The setting is the same as Table 3. We adopt RedPajama- (Com- puter, 2023) as the training datasets and LongQLora (Yang, 2023) as the model adjustment method. We find it impor- tant to carefully balance the sampling range to optimize the model's generalization to longer contexts.
Gerain + 0. miisalin
73
64
38
74
We evaluate performance using the LongBench-E score and perplexity on validation sets composed of sequences of length 8,192 from the RedPajama-Bov corpus. The results are shown in Table 3. Notably, our methods outperform the baseline across both model adjustment strategies. Specifi- cally, models trained with our proposed regularization term achieve up to 2.4% improvement in the LongBench-E score. Similar to the previous experiments, we observe that excessively large regularization coefficient values may not consistently benefit oon-cootext modeling, which is re- flected in the slightly lower LongBench-FF scores and higher perplexity, indicating that overly strong regularization may disrupt the model's training process.
Benchmark
LongBench-E (1)
Perplexity ())
Training steps
50
100
200
50
100
200
Sumpling range of luxtra
(I) [1, /2] (Current)
21.8
23.3
25,8
6.72
6.39
5.77
(2) I. /4]
21.5
23.2
25.7
6.77
6.29
5.81
(3) Musia/ /4, A.aaaii [2]
21.4
22.7
24.5
6.82
6.47
5.94
(4) [1,au)
18.2
18.9
19.1
15.65
13.59
12.52
early-context facts: When the supporting fact is located at the beginning of the input context (fact depth = 0), our method achieves performance comparable to the bascline. This suggests that despite the form of the regularization po- tentially encouraging the model to neglect earlier contexts, it does not lead to this behavior in practice.
5.3. Experiments on BABILong
To further assess the ability to retrieve and utilize distant in- formation, we conduct extensive experiments on BABILong (Kuratov et al., 2024), a challenging reasoning-in gaasoning-in-a-bie ina--haystaccckk task specifically designed for evaluating long-context capa- bilities. BABILong comprises crestion-answering tasks in which the supporting facts for each question are situated at specific positions within the context. Using the model setup described in Section 5.1, which incorporates CLEX as the adjustment method and RedPajama-Book as the train- ing dataset, all models are fine-tuned for 200 steps. The evaluation is performed on input sequences of lengths 4K, 8K, and 16K, with the overall results summarized in Table 4. The results indicate that our proposed method consis- tently outperforms the bascline across all evaluated lengths. Specifically, our method achieves a performance gain of 2.0% at length 8K and 2.2% at length 16K.
(2) Perfor- mance with middle-context facts: When the supporting fact is positioned in the middle of the context (fact depth = 50 or 75), our method shows considerable improvement over the baseline. This indicates that our approach effec- tively mitigates the "loss-in in-the-middle" phenomenon (Liu et al., 2024), a common challenge in large language models. Together, these results strongly support the effectiveness of our proposed regularization term in enhancing length gen- eralization ability, particularly for tasks requiring attention across diverse positions.
5.4. Ablation Studies
Since we incorporate several hyperparameters such as the regularization coefficient a and the sampling range |4 - l2 in the misalignment metric, we conduct extensive exper- iments to explore how these hyperparameters affect the model performance.
Additionally, we analyze the impact of the supporting fact's position within the input context using the BABILong QA1 task, where each question is associated with single sup- porting fact. The results of this analysis are presented in Table 5, offering two key insights: (1) Performance with
Regularization coefficient a. In addition to the settings already provided in the previous experiments (a 0 as the baseline, a = 0.1, and a = 0.5). we evaluated a = 0.3
8

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
Table 8: Comparison between synthetic tasks and natural language tasks.
Synthetic Tasks
Language Tasks
Output Space
Lj unit ball in RD))
Specific Task
Length/Sum prediction
Next token prediction Exist
Long-Shore Short Misalignment
Exist
Priori on Output Distribution
Explicit and predifined
Implicit and task deskddepndeent
Alignment Technique
Explicit reparameterization
Regularization term across lengths
Does the technique decrease long-short misalignment?
Yes (explicitly)
Yes (implicitly through optimization)
Does the technique improve length generalization?
Yes
Yes
and a = 1.0 under the same experimental conditions as Table 2, using CLEX as the model adjustment. We still adopt RedPajama- as the training dataset. The results shown in Table 6, which reveal the following trend: (1) Performance peaks for values in the range 0.3] in both evaluation metrics. (2) Larger values of a (e.g., 0.5 or a = 1.0) lead to a significant decline in performance, confirming the risks of over-regularization. These findings highlight the importance of selecting a moderate value for a. We suggest using a coefficient a between 0.1 and 0.3 as default mitigate the risk of over-regularization.
6. Discussion
Since our work is initially motivated by phenomena ob- served in synthetic tasks, we provide additional clarification on the relationship between synthetic tasks and natural lan- guage tasks by summarizing their key differences and simi- larities in Table While these two types of tasks differ sig- nificantly in their specific forms and output space, they share common challenge: output distribution misalignment across different input lengths. Our analysis highlights that employing an alignment sechiqueewweethe explicit repa- rameterization in synthetic tasks or regularization in natural language tasks-can effectively mitigate this misalignment. This mitigation directly enhances the model's length gener- alization ability, demonstrating the broader applicability of our approach.
Sampling range. In equation 4, we sample ( and l2 from [ursaa/2 22, hraii] by default to avoid input sequence with sig- nificantly different lengths. This is equivalent to sampling lextra from [1, 4 /2. Here we conduct an ablation study examining how different sampling strategies of lextra affect performance. We consider four sampling configurations: (I) The current strategy, sampling from [1, Arain /2]: (2) Sam- pling from narrower range [1, (rrin /4]; (3) Sampling from narrower range [/raai/4, / /2: (4) Sampling from broader range [1, 4rain] and remove the limit that ( and 22 should be in | 4- We conduct experiments using the same setting as Table 3, using LongQLore for model adjustments and regularization coefficient of 0.1. The results are shown in Table 7: (1) Setting achieved performance comparable to the current strategy, while Set- ting 3 showed slightly inferior performance compared to the current strategy. This suggests that aligning outputs between sequences with moderate length discrepancies effectively supports long-context modeling. (2) Setting 4 yielded significantly worse performance than the current strategy, indicating that encouraging alignment between sequences with large length differences adversely affects the model's long-context capabilities. These results un- derscore the importance of carefully balancing the sampling range in the proposed regularization.
7. Conclusion
In this work, we investigated the challenges of long-context modeling in large language models and introduced fresh perspective by examining the output behavior of models across varying input lengths. We identified that misalign- ment in output distributions across sequences of dif- ferent lengths, which we term long-short misalignment, plays a critical role in limiting length generalization classical and foundational subproblem in long-contex mod- eling. Through both synthetic and natural language tasks, we demonstrated that long-short misalignment is strongly correlated with performance long-context inputs. We further proposed regularization term based on this met- ric to explicitly reduce output divergence during training, Extensive experiments confirm that this approach not only improves length generalization but also leads to more effec- tive long-context modeling. Overall, our work highlights the importance of considering the output space when designing and analyzing models for long-context scenarios, offering new dimension for future research in scalable and effective language modeling.
We also compare our methods with baselines under the same computational time in Appendix F.
9

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
Acknowledgements
large language models. arXiv preprini arX1v:3399/12307 2023b.
Yisen Wang was supported by National Key R&D Program of China (2022D0160030), National Natural Science Foun- dation of China (92370129, 62376010), and Beijing Nova Program (20230484344, 20240484642).
Chi, T.-C., Fan, T.-.., Ramadge, P., and Rudnicky, A. Ker- ple: Kernelized relative positional embedding for length extrapolation. In NeurIPS, 2022.
Chi, T.-.., Fan, T.-H., Rudnicky, A., and Ramadge, P. Dis- secting transformer length extrapolation via the lens of receptive field analysis. In ACL,, 2023.
Impact Statement
We believe there are no direct ethical problems in this work since it primarily focuses on the theoretical analysis and improvement of large language models (LLMs). However, LLMs can generate inaccurate harmful content, and this research does not offer direct solution these issues. Users are encouraged to ensure that the LLMs obey the ethical standards when implementngg the proposed methods.
Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung. H. W., Sutton, C., Gehrmann, S., al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-133 2023.
Computer, T. Redpajama: An open source recipe to reproduce llama training dataset, 2023. URL
References
Ahn, K., Cheng, X., Song, M., Yun, C., Jadbabaie, A., and Sra, S. Linear attention is (maybe) all you need (to understand transformer optimization). In ICLR, 2024.
RedPajama -Dataa
Deletang, G., Ruoss, A., Grau-Moya, J., Genewein, T., Wen- liang, L. K., Catt, E., Cundy, C., Hutter, M., Legg, S., Veness, J., et al. Neural networks and the chomsky hier- archy. In ICLR, 2022.
Anil, C., Wu, Y., Andreassen, A., Lewkowycz, A., Misra, V., Ramasesh, V., Slone, A., Gur-Ari, G., Dyer, E., and Neyshabur, B. Exploring length generalization in large language models. In NeurlPS, 2022.
Dettmers, T., Pagnoni, A., Holtzman, A., and Zettlemoyer, L. Qlora: Efficient finetuning of quantized llms. In NeurlPS, 2024.
Bai, J., Bai, S., Chu, Y., Cui, Z.,, Dang, K., Deng, X., Fan, Y., Ge, W., Han, Y., Huang. F., et al. Qwen technical report. arXiv preprint 2rX9 16609, 2023a.
Fang, I.., Wang, Y., Liu, Z.,, Zhang, C., Jegelka, S., Gao, J., Ding, B., and Wang, Y. What wrong with perplexity for long-contexs language modeling? In ICLR, 2025a.
Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L., et al. Longbench: A bilingual, multitask benchmark for long context under- standing. urXiv preprint arXiv:2308. 14508, 2023b.
Fang, L.,, Wang, Y., Liu, Z., Zhang, C., Jegelka, S., Gao, J., Ding, B., and Wang, Y. What is wrong with perplexity for long-context language modeling? In ICLR, 2025b.
Black, S., Biderman, S., Hallahan, E., Anthony, Q., Gao, L.. Golding. L., He, H., Leahy, C., McDonel,, K., Phang, J., et al. Gpt-neox-20b: An open-source autoregressive language model. arXiv preprint arXiv:2204.06744. 2022.
Huang. J. and Chang, K. C.-C. Towards reasoning in large language models: A survey, In ACL, 2023.
Jelassi, S., d' Ascoli, S., Domingo-Ennicn C., Wu, Y., Li, Y., and Charton, Length generalization in arithmetic transformers. arXiv preprint arXiv:2306. 15400, 2023.
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., al. Language models are few-shot learners. In NewrIPS, 2020.
Jin, H., Han, X.. Yang, J., Jiang, Z., Liu, Z., Chang, C.-Y., Chen, H., and Hu, X. Lim maybe longlm: Selfextend Ilm context window without tuning. In ICML, 2024.
Chen, G., Li, X., Meng, Z., Liang, S., and Bing, L. Clex: Continuous length extrapolation for large language mod- els. ICLR, 2024.
Kaplan, J., McCandlish, S., Henighan, T., Brown, B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. 2020.
Chen, S., Wong, S., Chen, L., and Tian, Y. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.155 15595, 2023a.
arXiv preprint arXiv:
Kazemnejad, A., Padhi, I., Natesan, K., Das, P., and Reddy, S. The impact of positional encoding on length general- ization in transformers. NewrlPS, 2023.
Chen, Y., Qian, S., Tang, H., Lai, X., Liu, Z., Han, S., and Jia, J. Longlora: Efficient fine-tuning of long-context
10

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
Kuratov, Y., Bulatov, A., Anokhin, P., Rodkin, I., Sorokin, D., Sorokin, A., and Burtsev, M. Babilong: Testing the limits of Ilms with long context crasoning-in-a--ayytta NeurIPS, 2024.
Touvron, H., Martin, L.,, Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
Lee, N., Sreenivasan, K., Lee, D., Lee, K., and Papail- iopoulos. D. Teaching arithmetic to small transformers. In ICLR, 2023.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L.,, and Polosukhin, I. Attention is all you need. In NeurlPS, 2017.
Li, T., Zhang, G., Do. Q. D., Yue, X., and Chen, W. Long- context Ilms struggle with long in-context learning. arXiv preprint ar 4rXiv:240.002060 2024.
Veličković, P. and Blundell, C. Neural algorithmic reasoning. arXiv preprint urXiv:2105.0276: 2021.
Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, C. Transformers learn shorteuts automata. In ICLR, 2023.
Wang, B. Mesh-Traane Taahh-TTafffrme-JAAX Model-Parallel Imple- mentation of Transformer Language Model with JAX, 2021.
Liu, N. F., Lin, K., Hewitt, J., Paranjape, A., Bevilacqua, M., Petroni, F., and Liang, P. Lost the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157-173, 2024.
Wang. S., Kobyzev, I., Lu, P., Rezagholizadch, M., and Liu, B. Resonance RoPE: Improving context length general- ization large language models. Findings of ACL, 2024a.
Wang. Y.. Ma, X., Chen, Z., Luo, Y., Yi, J., and Bailey, J. Symmetric cross entropy for robust learning with noisy labels. In ICCV, 2019.
Nogueira, R., Jiang, Z., and Lin, J. Investigating the limita- tions of transformers with simple arithmetic tasks. arXiv 2021.
preprint 1rXiv:2102.1300
Wang. Y., Wu, Y., Wei, Z., Jegelka, S., and Wang, Y. A theoretical understanding of self-correction through in- context alignment. NewrIPS, 2024b.
Ontanon, S., Ainslie, J., Fisher, Z., and Cvicek, V. Making transformers solve compositional tasks. In ACL, 2022.
Wu, Y., Wang, Y., Du, T., Jegelka, S., and Wang, Y. When more is less: Understanding chain-of- length in llms. arXiv preprint arXiv:2502.07266 2025.
Peng, B., Quesnelle, J., Fan, H., and Shippole, E. Yarn: Efficient context window extension of large language models. In ICLR, 2024.
Yang. J. Longglora: Efficient and effective method extend context length of large language models. arXiv preprint arXiv:231 1.04879, 2023.
Press, O., Smith, N., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapo- lation. In ICLR, 2021.
Zhang, Q.. Du, T., Huang, H., Wang, Y., and Wang, Look ahead look around? theoretical comparison between autoregressive and masked pretraining. In ICML, 2024a.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models unsupervised multitask learners. OpenAl blog, 2019.
Zhang, Y., J., and Liu, P. Extending llms' context win- dow with 100 samples. arXiv preprint arXiv:2401.0704 0000044 2024b.
Rae, J. W., Potapenko, A., Jayakumar, S. M., and Lillicrap, T. P. Compressive transformers for long-ram range sequence modelling, arXiv preprint arXiv: 1,5507, 2019.
Zhou, H., Bradley, A., Littwin, E., Razin, N., Saremi, O., Susskind, J., Bengio, S., and Nakkiran, P. What algo- rithms can transformers learn? a study in length general- ization. In ICLR, 2024.
Ruoss, A.. Delétang, G., Genewein. T.. Grau-Moya, J., Csordás, R., Bennani, M., Legg, S., and Veness, J. Ran- domized positional encodings boost length generalization of transformers. In ACI. 2023.
Su. J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 1676127, 2024.
Sun, Y., Dong, L., Patra, B., Ma, S., Huang, S., Benhaim, A., Chaudhary, V., Song. X., and Wei, A length- extrapolatable transformer. In ACI., 2023.

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
A. Full Related Work
Length Generalization on Synthetic Tasks. Our paper is related to the line of work that seeks to understand the capabilities and limitations of Transformer models when it comes to algorithmic reasoning (Veličković & Blundell, 2021). Specifically, we focus on simple tasks and study length generalization on the standard Transformer architecture with causal structure. Related to this, Lee et al. (2023) study how well transformers trained from scratch can learn simple arithmetic tasks, and find that no length generalization is observed. Nogueira et al. (2021) find that partial length generalization on addition is observed only when models reach 3B parameters and when the addition questions are presented in reverse order. Jelassi et al. (2023) study models trained on addition and find strong generalization performance when using a few examples of longer sequences. Zhou et al. (2024) proposes the RASP Generalization Conjecture that Transformers tend to learn a wingth-generaliiing solution if there exists a short RASP-L, program that works for all input lengths. Liu et al. (2023) discovers that the Transformer will learn shorteuts through the study of various synthetic tasks. Besides these explorations on specific tasks, some works study the impact of different positional encodings on the length generalization of math reasoning tasks. Alibi adds a linear bias the attention score to achieve better length generalization performance (Press et al., 2021). Ontanon et al. (2022) studies different settings of positional encodings and identifies Transformer configurations that generalize compositionally significantly better in a diverse set of compositional tasks. Kazemnejad et al. (2023) systematically studies the role of no positional encoding (NoPE) in Transformer with causal structure. Ruoss et al. (2023) proposes a randomized positional encoding scheme that simulates the positions of longer sequences and randomly selects an ordered subset to fit the sequence's 's length.
Long-context Modeling on Natural Language Tasks. A series of works (Sun et al., 2023; Chi et al., 2022: 2023; Zhang et al., 2024b; Chen et al., 2024; Peng et al., 2024; Yang, 2023; Chen et al., 2023a) aim to extend the context size of Transformer- based models during fine-tuning, primarily by modifying positional encodings. For example, Zhang et al. (2024b) introduces a novel extension to RoPE (Su et al., 2024) which combines adjusting RoPE's base frequency and scaling the attention logits to help LLMs efficiently adapt to a larger context window. Chen et al. (2024) generalizes the positional encoding scaling approaches to model the continuous dynamics by ordinary differential equations over the length scaling factor. (Chen et al., 2023a) proposes to extend the context length by slightly modifying RoPE via Position Interpolation (PI) and fine-tuning on small amount of data. Our approach, in contrast, focuses the model's output space, which identifies the crucial role of loog-short alignment in length generalization. Wang et al. (2024a) proposes a novel approach designed to narrow the generalization gap by refining the interpolation of RoPE features for OOD positions and provides length extrapolation analysis on the feature gap.
B. Model Details for Synthetic Tasks
We focus on the decoder-only Transformer (Vaswani et al., 2017), model widely used in both synthetic tasks (Zhou et al., 2024; Jelassi et al., 2023) and LLMs (Touvron et al., 2023; Peng et al., 2024) which utilizes causal mask in the self-attention module to enable auto-regressive generation. We consider several positional encodings: learnable positional encoding (Radford et al., 2019), Alibi (Press et al., 2021), rotary positional encoding (Su et al., 2024) and no positional encoding (NoPE). Since recent works found that by removing the positional encoding, Transformers can trained to be well generalized on length (Deletang et al., 2022; Kazemnejad et al., 2023), we adopt this setting (i.c. NoPE) by default. To provide a clear signal for the model to accomplish the tasks, we add both the begin-of-sentence (BOS) token and the end-of-sentence (EOS) token in the sequence. We apply feed-forward neural network on the hidden state of the last token to generate an output of a real number. We train all of our models on the train distribution from scratch to convergence if possible. For all tasks, the length of training examples is sampled uniformly from length I up to the max training length Itrain We select peree-parameters such as the learning rate for each task based on what is required to fit the training set. At test time, the length of the examples will traverse from 1 to the max testing length lustt
C. Theorems and Proofs
C.1. Theoretical Analysis for Synthetic Tasks
Ahn et al. (2024) suggests that linear Transformers serve realistic abstractions for understanding Transformer optimization and generalization. Therefore, following (Ahn et al., 2024; Zhang et al., 2024a), our analysis is based on the linear attention model. The general form of linear attention is given by:
Attn(x) = QK T V= X XVVV
12

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
where WQ, WK WV are projections, and is the length of input x. In all tasks, will linear attention model with causal mask. Specifically, we normalize the output according to its position, which allows linear attention to perform similarly to dot-product attention. For example, the k-th output will be normalized as follows:
E V.K.V. 1 K
QNK
(8)
k
For the synthetic tasks, add bias terms to Q, K, V to mitigate the impact of in the input. The model based linear attention is defined as follows:
(x(x11))) - Σ QNKTVi.
(9)
and it is trained based on the following target function:
Iariin
Cige; Itrain) = ( Mieann
||0((1(s)) - y(xxx11)
(10)
Itrain
where [|s]] indicates the first i-tokens input x. The optimal model trained on sequence with maximum length of Irain is denoted as glrrain We have the following result:
Theorem C.I. In the length prediction task and the sum prediction, the length generalization loss has a quadratic relationship with the predicted length, i.e..
Elength (g0 guaai lest)) Ex (0,1) (1ob (Xrest) (XXesst)
0 ((leet leraia)))
(11)
Esum (gi lustt) O (Chest lerain) 12)
(12)
However, in the mean prediction task, the length generalization loss has fixed upper bound:
Emean Arrain Lestt) O(1).
(
(13)
Proof. Since the input consists of and 1, the Qx, Kg. Va will each have only two possible values. Thus, we may assume that when the input token is 1, the Qk. KR, Vi are q', K', respectively, and when the input token is O, the Qx. Kg. Vi are d" k", n" respectively. Furthermore, we note that in (9), K; and V; always share the same subscript. Therefore, we can treat them single token and replace them with K, where we define k' k'v' and x" k" We can decompose (10) into hrain separate part:
Lggo; lerain)
E(go: I),
E(go; 1)
Ex [1000((p0) - y(xxx1))
(14)
Terain
Next, we consider each task individually.
(a) First, we study the length prediction task. In this case, v(x(q)) = I. Expanding E(go; 1) we have:
f(go: = Fex
(x190(x14)) - 11000~((())))) v(x(11
2
Σ c (1 + ct (Hidn ((-0"=((
2
It's easy to find that É(go. achieves its minimum 0 if and only if of, q" ne" satisfy the following conditions:
df = q" #0,K' K = K" + 0,
(15)
d''' = d" q'K" q" n
13

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
These also properties that the solution of Oelge; I) holds. Note that the first property holds for any {go: and is independent of I. Since L(go; Irain) is composed of E(go; 1). and each optimal solution of E(go; satisfying d' =" + 0, x' x" + 0, then the global solution for (g(g) leran)) = 0 should also have this property. Therefore, we may assume that g q'n' q"rc' / = q 13" = q²=" = r, and our goal is to find the I' that satisfying (g(g: / = 0, which means that this s minimizes [goo htrain): ). In this case, is easy to find that {ggo = (T - 1)2, which means that Øg( 0)/0r = 2(1 - l). Therefore, we have:
ÖCigoi ltrain)
2 lerain
Lerain
[T (1-
(16)
ar
solving (g( /OT 0 and we have r (Itrain + -1)/2, the yreain satisfying q'n' =q"k' q'1" R² q"=²"
(Irain + 1)/2. Therefore, we have:
Elengeh (gharin,, llest) elghreain (lest) (hest
Itrinn + 1
O ((hest Itrain) )2)
(17)
2
(b) We now study the sum prediction task. In this case, s(x(a1) = Ei=r X, Expanding E(go: we have:
[110(((11)== - 11000~x(xx1111) [(x(1))11]
I(go;t) Ex
2 + +i -i) 1 c (t (1- (-)))(())))))
Σ
(Hr0' (1-0)"''
= 2
It's easy to find that E(go. achieves its minimum 0 if and only if o', ,n' satisfy the following conditions:
q gd==0 = q" + 0, of + 0, n² 0, g'''' = q"k'== =1
(18)
This is similar to the previous conditions (15). Thus, assume that q' n' = q"k' =T with x" 0. In this case, we have:
perform similar analysis above, and we similarly
[goo 1)
1 21
Σ
&C (T - 1)= 1+11 1 1)2
(19)
21
Thus we have:
Ol(gu;))
1+1 1 (T - 1),
(20)
ør
which leads to:
Ög(ge hrain)
Arsii de(go;
(hrain + Hivean )r lerain (4rain + 3)
(21)
Ør
Ltrain
ar
Itrain
2
where Hn is the n-th harmonic number, i.e., Hn Ei= 1/i. Solving ØL(go; (rraii) )/or 0 and we have herai (hraan + /1//(ura + H.ram)) so the o Araai satisfying n'n' = d"r¹ = heraiu (hrain + /)//(urruii + Hisam))
Therefore, we have:
leest
Itrain Atraii + +3) 2(hrain + Hivein)
Esum ( (ghaaai beeat) Elggneain (tsr.)
breat
24
Ltrain + hest
(22)
2
+ const
O (Ureet leruiu) ))
14

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
(c) Finally study the mean prediction task. In this case, v(x(1))= =Ei=1 x./l. Expanding {(go: I) have:
((goil) Ex -
i
+1 -)) + ot- (Hro0''d (11'0""'(100000"000 x))"k=)) 9) 2
Σ o
2
It's easy to find that {goo 1) achieves its minimum 0 if and only if a', , satisfy the following conditions:
g/==0 q" * 0, 0, 0, g'n' = o" 11
(23)
This is similar to the previous conditions (18). In fact, the above properties are irrelevant to I, so for all Egge. I) these properties hold. Therefore, in this case, the optimal solution of Elgs:huaio will satisfy (23), which leads to C(go: Itrain) 0. Under this situation, the length generalization error is:
E (gg (gggnin; (lest) = Elggrrat leesst = O.
(24)
Consider the case where the solution is not optimal, i.e., when g' q"n"" 1 + E, where + 0 small, we can similarly obtain:
Esum (gum (gggeinn -hest) Elggrram : (bees + 1)E² 2hrest
S 82
(25)
In conclusion, we have:
Esum (g) (ofreati :lesst) = O11),
(26)
which completes the proof.
C.2. Theoretical Analysis for Natural Language Tasks
We now focus on natural language tasks. We make some changes to the model following (Zhang al., 2024a). In natural language tasks, the model requires an additional projection W to make the output a probability distribution. That is, the model is modified as follows:
go (x(x((n)
Σ Q.K. V.W, k.
1
(27)
1=1
In this case, each token x,, output go (x) and the objective function yxx) are normalized. Without loss of generality, we may assume that the changes the objective function after truncating the inputs are negligible, i.e.
Pr(v(xja) A s(xp,a)) = 0 (1),
(28)
For simplicity, we use the Lg-norm instead of SCE to measure misalignment, i.e.:
Lmisuliga (go) Ex.1i.2 [I190(x(()a)) go (x
(29)
and we use the La-norm instead of the CE loss the training loss function since these two functions differ only by a constant when the output and target are regularized. Under these conditions, we have the following result:
Theorem C.2 (Generalization guarantees for the natural language task). Suppose that damin is the model trained on sequences with maximum training length lrain- When the testing length beest satisfying heat hrains the generalization loss Egen (9 Jasin Iestt has the following upper bound:
Jrsin hrsst) < C.lleea) Lmisalign (g) (aaaii)) + Cf Ctraingg (rain)) ) + c(hean))
Egen (g
(30)
where Oflemee) (i 0, 1, are constants related to hest. Specifically, the ((frae)) increase the hest increases and the ratio (hhraa) /C (Cflemes) increases as lest increases. This indicates that as the testing length increases, the alignment loss becomes increasingly significant.
15

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
Proof. Suppose that
Itriin 1 <
Iraii
2
lextra
(31)
Lerain
lerain - 4 1 1 2
where E [1, (tain] an arbitrary integer. Let ((k) hraan + -lextra
(A: 0, ... ,N), where Ni
[frest Terain) /lextra we have:
Arrain (bowe)) :)) - y(x -boss:))
||
- (Nr + 2)
(1sfrmin (x(-)rom))) g0
derai (xc (x)_-1NM)1)) ||2
Ni +
Lgrniin ((1-))))))))) go (x) 2
(32)
k=0
temin (x(--remia))) 2(x (bon:))))))) (-)
%
For each l: E (0, No- 1). have:
(x (x(-1(++)))))) Momwn (x1 (04-)0))))) |
96
[(++] ( x -- WVw w
x (W xx & WV
[(0)
T x
Xn X
WV W
1(x+1)
((b)
- Xo w 2 (W
x ((((())) (1
[(++]
[(A)
wv w
(33)
x -1-10+00--000)11)0)) - ((00 - my) (11100000(00011)
+
Xo 2 (WK)T
WV
1(++1)
1(R)
12 1extra (1(4) - 1)2C8
(1(++1)
+ lextra 1(++1)
1 Jumin (x (1000 -n1
+
g6 (x -0) ((k) 9
where Co = [ WV 1 is a constant. The first term in (33) can be upper bounded by:
12 Pextra (T(A) - -1)2c8
12 C3
(34)
((K+1))(k332
412
train
and for the second term in (33), we have:
1 Textra 1(++1) d deniin (x1-1000 1(k)
1 Jemin (00
(2(0)
1)lextra derain (x1-10-1)
((++))((A) 90
(35)
(arain (x
-9 Arrain (x x
((k)
12
2
+ Aerain (x_110+1))
Jerain (x (10) 13))
train
In fact, the second term above represents the alignment error between 10)00) and x - where these two sequences have length difference of lextra-
16

---

Long-Shou Alignment for Effective Long-Context Modeling in LLMs
On the other hand, for the third term in (32), i train/ /2, have:
98 Jarain (xr-mmma::)) (X 1
Areatn (X|-)rrtt:))) To farato (x[- ((thrrmmmm)))) /2):) 1
98
Arraon (X(-(1+laraia/])) /2)1) D0
glrain (x (x-)
(36)
Arsaiu (x(- ( 4eoo: (boon))
+
96
The first term above similarly represents the alignment error between two sequences with length difference of lextra while the last term corresponds to the training error. Additionally, it is easy to derive an upper bound for the second term:
Jusin (x1- -X (1+trrson/2))))
Fo Arsin (x1-11) ||2 s and this upper bound also holds for the first term in (32). Meanwhile, if
90
1 Itrinn we have:
11
08
Jrmin (x(-
- Jerain (X(- 9 train (xl (x1
(37)
|6 (x] y(x (xx-hea)))
This result is similar to (36) except for lacking the second term (36). Overall, we have:
N₂
11²
g (x]-)rea:)) - y x
< (Ni + 2)
Aarain (x(=:1+++)(((x))--))) 11) - Arrain (x)-0((())00)(
+ frran (Xi - 96
Areain ( (1 /2)+)
Atrain
1
Atrain (x(-1i)) - 1xxx 1/3
+ 90
12
(CB +
+ N₁
+
412
(38)
N
99 gerai (x]-1ess:) - y(x (hou: |/
1
s (Ni + 2)
90
Jeran (x)-10+-)) - Arain (((A)
Arsin
+ |00 (x(-1:)) - (oss--hcoi))
12 (CO 1) +
N₁
472 47
(39)
It's easy know that the probability of < 4rain 2 and l 4tain 12 [(train 1)/2]/44 herain and L(Itrain + -2)/2)/lrral respectively. Therefore, by taking the expectation over l and on both sides of the inequality, have:
Egen (griii hrest)
((2)
Clean
+
Itrain
-1
-11)
Emiiaaliaa (gg ( + c(1) Ltrain (rasii))
lerain
2
(40)
Fest (CO- + 1)
6larain (- 1)
C lina
412
24trai
train
where
C E2 (Nr 2], c(2)
E [N + 2)].
(41)
Feent
is easy to verify that C increases as Hest increases. Consequently, the constants Offrret) ((2) + [(uraan 4eomt
1)/2]/4rains Cl. clliaa) = (1) beent and Ceffrent) (2) 12 12 (C) + 1)/((1 + ((()) (Glarain - (-1)/tsinn) // also heam
framt increases as hest increases. Moreover, the ratio Ciltuna) ((llemn) increases as lest increases since the ratio ((2) (c())
increases lest increases, which is easy to verify.
17

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
Maximum Training Length: 10
1000
Maximum Training Length:
175
Maximum Training Length:
Training Length:
Maximum Training Length:
Maximum Training Length: 20
800
600
100
Test loss
75
400
50
25
30
40
10
20
30
40
50
Test Sequence Length
Test Sequence Length
(a) Length generalization in sum prediction task with different maximum training sequence lengths
(h) Length generalization length prediction task with different maximum training sequence lengths
Length generalization performance in the sum prediction and length prediction task with different maximum training sequence lengths. Although increasing the training length helps reduce the generalization error, the overall trend of increasing test loss remains unchanged.
Figure 4:
D. Analysis on the Sum Prediction Task
We also examine the sum prediction task, where the label corresponds the sum of the sequence. Similar to the length prediction task, the output space shifts as the sequence length increases. As result, models struggle with length generalization in this task, shown in Figure 5a. However, by applying the reparameterization technique proposed in Section 3, we observe significant improvement in length generalization, as shown in Figure 5b. These results demonstrate the importance of long-short alignment in length generalization.
E. Pytorch-like Code for Implementation of Cirain
efficient implementation for the total training loss.
2 import torch import random
def SCE (outputl prob, output2_prob)
loss torch. (torch. sum cutputl_prot torch.exp (output2_prob) outputl_prop torch. tozch.exp (outputi_prob) -1)))
extra_len Landem.randing max_len//2)
data get get_eatttaaaeeeeeeeeeeeaeaaaaaat htt.aatceeeegooooaoooooooeeeeee.
11
output mcdel (data [:, :max_len])
2 model (data I -max_ max]]
probl torch.nn.
prob2 wtt.:.m..unntiiona....om ttt.:attteo....ooooooooooaaaaoaa wwwtmat (output.. logits)
softmax sottmax(output/ (wwtpot2. Logitss) logits)
15
Select the overlapped part to calculate the misalign loss
probl probl[;, max_len//2Hextra_len
prob2 pr mxx_len//2maa__eeeeexrraaa
19
20
loss_ce (cutputl. loss + putput2.1oss) /
loss mas__misalggn = SCE (probl, prob2)
23
Loss_total lcss_ce alpha loss_miselign
18

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
Learnable PE
fxl=x
500
NoPE
175
[0X)ssertix)
Rotary
f0x==log(x+1)
Alibi
150
T0x (/o=t///ort0+11122
400
125
300
100
Test
75
200
50
100
30
40
10
30
40
50
Test Sequence Length
Test Sequence Length
(a) Length generalization in the sum prediction task
(h) Reparameterization in the sum prediction task
Figure 5: Length generalization in the sum prediction task. Explicit alignment of output space boosts length generalization performance.
Table 9: Performance of the tne-ttune models using only cross-entropy loss (baseline) and an additional long-short misalignment loss on long-context modeling benchmark, LongBench-E score (Bai et al., 2023b) and perplexity on the 8k-length validation set. The fine-tuning sequence length is 4k, exactly the same the training sequence length. We adopt two datasets: RedPajama-Bode (Computer, 2023) and PG19 (Rae al., 2019). The models finetuned with proposed loss outperform the baseline across different model adaption strategies. The comparison is based on fixed total computation time. (a/b) in the training steps mean steps for the baseline and steps for our method.
Benchmark Training steps
LongBench-E (f)
Perplexity (1)
50/47
100/95
200/190
50/47
100/95
200/190
RedPajama-Bel Book
Liaain (Baseline)
22.7
23.8
24.7
7.21
6.56
6.12
Ltrain + 0.1 Cmisalign (Ours)
23.1
25.1
26.4
6.92
6.27
5.91
Ltrain + 0.5Lmisalign (Ours)
21.7
23.4
24.5
7.62
7.16
6.61
PG19
Ctrain (Baseline)
20.2
21.4
22.5
8.92
7.89
7.45
Ltrain + 0.1 0.1 (Ours)
20.7
22.0
25.1
9.02
8.01
7.39
Ltrain + 0.5.misallgn (Ours)
19.6
22.0
23.3
9.82
8.62
8.29
F. Comparison Under Same Computation Time
To account for the additional computation time required to calculate Lmisalign we will compare our methods with the baseline as in Table and Table based on total computation time. we compare our method with the bascline under equivalent total computation costs. Our method introduces an additional computational overhead of approximately 3% to 5% per step. To ensure fairness, we adjust the number of training steps proportionally. For example, when the baseline is trained for 50, 100, and 200 steps, method is trained for 47, 95, and 190 steps, respectively, achieving comparable total computation times. We observe that the performance trends remain consistent: method continues outperform the baseline under equivalent computation time. This underscores the efficiency of our approach despite the minor additional cost.
G. Additional Experiment on Mean Prediction Task with a Different Dataset Setting
In the synthetic experiments in Section 3. and have an equal probability. The mean value of 50 such samples will nearly obey the normal distribution A (0.5, 0.005). This means predicting 0.5 under length of 50 would yield an approximate test loss of 0.005. However, we would like to clarify that the test loss of NoPE remains around le-5 (indicated by the orange line in Figure (()),, which is two orders of magnitude smaller than 0.005. This indicates that the model predicts the mean values
19

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
Table 10: Performance of the finetuned models using only cross-entropy loss (baseline) and an additional long-short misaligmment loss on long-context modeling benchmark, LongBench-E score (Bai et al., 2023b) and perplexity on the 8k-length validation set. The fine-tuning sequence length 8k. We adopt two kinds of model adjustments: LongQLora (Yang, 2023) and EABF (Zhang et al., 2024b). The models finetuned with our proposed loss outperform the baseline across different model adaption strategies. The comparison is based fixed total computation time. (a/b) in the training steps mean steps for the baseline and b steps for method.
Benchmark
LongBench-E (1)
Perplexity (1) 100/95
Training steps
50/47
100/95
200/190
50/47
200/190
Long@Lora
Liaai (Baseline)
21.9
22.1
23.4
6.82
6.41
5.82
Ltaain + 0.1 mmismlign (Ours)
21.6
23.1
25.7
6.79
6.42
5.74
Liaain + 0.5.misalign (Ours)
21.1
23.6
24.9
7.19
6.65
5.96
EABF
Craai (Baseline)
22.1
22.9
23.6
6.89
6.52
6.01
Craii + 0.1 misalinn (Ours)
23.0
23.6
24.5
7.01
6.48
5.86
Lraai + 0.5.misalign (Ours)
22.2
23.1
23.8
7.32
6.88
6.42
of the sequences with high precision, rather than simply guessing fixed number. Therefore, the conclusion that the model can have good length generalization ability in the mean prediction task is reasonable.
Besides, in order avoid a trivial solution of predicting 0.5, we conduct an additional experiment on the mean prediction task. To build the sample of length I, we first sample the number of 1 uniformly from [0, and then randomly build the sequence. In this way, the mean value of the sequence uniformly spans from 0 to I, avoiding trivial prediction. The experimental results are shown in Figure 6, where the model can still achieve good length generalization, consistent with our previous results.
20

---

Long-Short Alignment for Effective Long-Context Modeling in LLMs
1e-5
1.10
1.05
Test loss
0.95
0.90
0.80
Sequence Length
Figure 6: Length generalization in the mean prediction task with different dataset setting. To build the sample of length (, we first sample the number of 1 uniformly from [0, I] and then randomly build the sequence. this way, the mean value of sequence uniformly spans from to I, avoiding trivial prediction. The model can still achieve good length generalization. which is consistent with our previous results.
21


Reviews:
[{"review_id": "e3voEPCKeq", "reviewer": "Reviewer_KAyJ", "summary": "This paper studies how aspects of the output distributions of LLMs relate to length generalization and performance on long-context tasks. First, the authors show that train-test mismatch in the output space can lead to poor generalization on synthetic tasks, and that reformulating the tasks to reduce this mismatch improves performance. Second, on non-synthetic natural language tasks, the authors propose a metric called \"long-short misalignment\" which measures the divergence between the output distributions of a model when it's original context window is truncated (up to 50%). This measures the model's invariance to this truncation. There is a positive correlation between models that are more invariant (and therefore less sensitive to earlier context) and performance on long-context modeling tasks. The authors then propose a regularizer based on this metric, which when applied during training can improve performance on tasks such as LongBench-E.\n\n## Update after rebuttal\n\nI would strongly suggest including the expanded Table 5 in the revised version, and reducing Section 3 to improve the presentations, as proposed in the response. While still a bit unintuitive to me (I would have thought more tasks require sensitivity to earlier context) the empirical results seem reasonably strong across most benchmarks, and the authors show that hyperparameters can be chosen such that sensitivity to earlier context is not entirely eliminated.\n\nAssuming the pledged changes will be implemented, I will increase my score from 2 to 3.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Q1. Maybe there is an interesting connection to be discussed with SCAN.\n\nA1. 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:\n\n- 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.\n- 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.\n\nWe will clarify this distinction and discuss potential connections to these prior works in the revised paper.\n\n---\n\nQ2. The results in Section 3 did not seem very closely connected to the proposed methods in Section 4.\n\nA2. Thank you for your insightful comments. We clarify that Section 3 is crucial for our paper for the following reasons:\n\n- 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.\n- 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.\n\nThank you again for raising this point, we will make Section 3 more concise to avoid distracting the main contributions of Section 4.\n\n---\n\nQ3. The proposed terminology of \"Output Alignment\" is overloaded with various techniques for model post-training.\n\nA3. Thanks for pointing it out! We will use “Output Matching” instead to avoid such confusion.\n\n---\n\nQ4. 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?\n\nA4. 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.\n\n| $\\alpha$ | Depth=0% | Depth=25% | Depth=50% | Depth=75% |\n| --- | --- | --- | --- | --- |\n| 0 | 75 | 64 | 30 | 69 |\n| 0.1 | 73 | 64 | 38 | 74 |\n| 0.3 | 70 | 62 | 38 | 76 |\n| 0.5 | 61 | 54 | 36 | 72 |\n\nThe 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.\n\n---\n\nQ5. One of the models looks very underfit in Figure 1a, perhaps worth investigating the training configuration.\n\nA5. 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.\n\n---\n\nQ6. It could be interesting to see if different positional encoding schemes lead to different degrees of \"long-short misalignment\" in the predictable way.\n\nA6. 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.", "claims_and_evidence": "The main empirical claims seem to be supported. The proposed regularizer improves performance on benchmarks such as LongBench-E.", "methods_and_evaluation_criteria": "The proposed benchmarks seem reasonable.", "theoretical_claims": "There is a theoretical claim in 4.1, although it is only stated imprecisely in the main paper and I did not review the details in the appendix.", "experimental_designs_or_analyses": "The experimental design seems reasonable.", "relation_to_broader_scientific_literature": "To the best of my knowledge, the proposed objective is novel.", "essential_references_not_discussed": "This is not critical, just a connection that seemed relevant as I was reading the paper:\n\n- There is a large body of work in NLP related to developing new methods and Transformer variants to improve length generalization e.g. on length-based splits of SCAN (https://arxiv.org/abs/1711.00350). A focus in this line of work has been on understanding the relation between the degree of context sensitivity and out-of-distribution generalization. For example, various approaches have injected CFG-like biases that reduce context sensitivity. The methods that the authors propose in the submitted paper seems related in that they attempt to reduce context sensitivity, in this case encouraging models to be invariant to context earlier in the context window. Maybe there is an interesting connection to be discussed.", "other_strengths_and_weaknesses": "Strengths:\n\n- The paper studies how a metric related to how invariant a model's outputs are to truncating the input context predicts performance on various long-context and length generalization settings.\n- Inspired by this, the authors propose a regularizer that encourages model outputs to be invariant to truncating early context. This appears to improve generalization performance on several tasks.\n\nWeaknesses:\n\nMy main concerns are related to the paper presentation, which I think could be considerably improved. I think it would also be useful to better understand the weaknesses of the method, since the proposed regularization (which encourages models to ignore early context) seems like it must surely be harmful in some settings. It would be good to clarify this so practitioners can have better intuition for when to apply the proposed method. \n\n- Section 3 was a bit distracting to me as a reader. It is not surprising that reducing train and test mismatch in the output space is valuable for out-of-distribution generalization, including length generalization. The results did not seem very closely connected to the proposed methods in Section 4, despite the attempt of Table 8 to clarify this. Section 3 relates to a property of the *task*, i.e. re-formulating the task to reduce the train-test distribution shift of the outputs, and is well supported in theory. The methods in Section 4 relate to a property of the *model*, i.e. that it should be largely invariant to earlier context. This is less clearly useful from a theoretical perspective, and therefore a bit surprising that it is effective. I think the paper would be stronger without Section 3, at least in the main paper. \n- The proposed terminology of \"Output Alignment\" is quite confusing. First, \"Alignment\" is overloaded with various techniques for model post-training, e.g. RLHF, and the proposed technique has nothing to do with this. The method is most clearly understood as encouraging invariance of output distributions when the context window is truncated, i.e. encouraging a form of context insensitivity. \n- It would be good to understand the limitations of the proposed method. For example, in cases that require sensitivity to information early in the context window, the proposed regularization should intuitively harm performance. The ablations in section 5.3 could be expanded to highlight potential weaknesses of the method (see questions for authors).", "other_comments_or_suggestions": "* nit: It would be helpful to make the Figures easier to read without color, i.e. use different line shadings or styles in Figure 1. \n* One of the models looks very underfit in Figure 1a, perhaps worth investigating the training configuration.\n* The introduction frames the proposed method as a completely different perspective from prior work investigating alternative positional encodings. However, it seems both lines of work have considered ways to reduce the sensitivity of models to potentially irrelevant context, e.g. the biases of methods such as RoPE that discourage long-range dependencies. It could be interesting to see if different positional encoding schemes lead to different degrees of \"long-short misalignment\" in the predictable way.", "questions_for_authors": "1. Section 5.3 - Intuitively, the proposed regularization should be harmful for cases where necessary context appears early in the context window, so it is surprising that the scores for \"Fact Depth\" 0% and 25% are comparable to the baseline. Does this change if the regularization coefficient is increased?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yp5chSEPUi", "reviewer": "Reviewer_Ueg1", "summary": "Authors introduce a novel perspective on length generalization in large language models (LLMs) by focusing on output alignment rather than conventional input-based approaches like positional encodings. Through synthetic task case studies, the authors demonstrate that models generalize better when output distributions remain consistent across varying sequence lengths. They propose the Long-Short Misalignment metric to quantify output alignment and introduce a regularization loss to enhance length generalization. Extensive experiments on synthetic and natural language tasks validate their approach, offering insights into improving LLMs' performance on longer contexts.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer Ueg1 for appreciating the insight of our paper. We will address your questions in the following part.\n\n---\nQ1. (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.\n\nA1. 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.\n\nTo 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:\n\n|  | LongBench-E | Perplexity |\n| --- | --- | --- |\n| Our method | 24.0 | 6.43 |\n| CREAM | 23.6 | 6.62 |\n| Our method + CREAM | 25.2 | 5.94 |\n\nThese 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.\n\n[1] Wu et al. An Efficient Recipe for Long Context Extension via Middle-Focused Positional Encoding. NeurIPS 2024.\n\n---\n\nQ2. (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?\n\nA2. 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. \n\nThe 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.\n\n---\n\nQ3. (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. \n\nA3. 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. \n\nChanging 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.\n\nAs 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.\n\n[1] Chen et al. Extending Context Window of Large Language Models via Positional Interpolation. arxiv 2306.15595\n\n[2] Zhang et al. Extending LLMs' Context Window with 100 Samples. arxiv 2401.07004\n\n---\n\nThanks again for your comments, and hope our response could address your concerns. Please let us know if you have additional questions.", "claims_and_evidence": "The claims in the paper are generally well-supported by empirical and theoretical evidence.", "methods_and_evaluation_criteria": "Yes, I would rather also see more comparison with other methods for improving length generalization.", "theoretical_claims": "Yes, and particularly I am more interested on empirical results as I believe they are more important in the context of this research.", "experimental_designs_or_analyses": "Yes, the experimental setup is okay, but I would rather see more comparison with other methods.", "relation_to_broader_scientific_literature": "The paper builds on prior work on length generalization in Transformers, which has primarily focused on positional encodings and modeling mechanisms. Unlike these studies, it shifts the focus to output alignment, introducing the Long-Short Misalignment metric as a new predictor of generalization and proposing a regularization loss to enhance performance.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strength:\n\n* focusing on output alignment for length generalization which I believe can be another axes to other methods like positional encodings.\n* Providing synthesis task analysis to bake the intention for following the idea.\n* Generally speaking, I think this is an insightful paper. \n\n\nWeakness:\n* 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.", "other_comments_or_suggestions": "See questions.", "questions_for_authors": "Question1: 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?\n\nQuestion 2: I believe more comparison with other length generalization approaches can help the results. As well if combining those methods with yours helps would be one of my questions. \n\nQuestion 3: 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.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4P0ccl1PYu", "reviewer": "Reviewer_F5vG", "summary": "This paper targets **length generalization** problem for LLMs and proposes to **shift from** conventional perspective of **positional encodings and data structures to the output distribution** of the model. They argue that the consistency of output distributions across sequences with different length correlates well with length generalization performance. They name this consistency **outer alignment** and propose a metric called **Long-Short Misalignment** to quantify it, and further design a **regularization loss** to explicitly enhance outer-alignment. They support their claim both empirically and theoretically.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Q1. About Delétang et al. 2022. This makes me question if the motivation is valid.\n\nA1. 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. \n\nTo 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.\n\nMoreover, 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.\n\n---\n\nQ2. On the connection between synthetic task and language modeling task. \n\nA2. 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**.\n\nWhile 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.\n\n---\n\nQ3. The paper would benefit a lot by comparing it with some representative length generalization methods.\n\nA3. 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.\n\n|  | LongBench-E | Perplexity |\n| --- | --- | --- |\n| Our method | 24.0 | 6.43 |\n| CREAM | 23.6 | 6.62 |\n\nWe find that our method outperforms CREAM in both LongBench-E score and perplexity, suggesting that output alignment provides additional benefits beyond positional encoding modifications.\n\n[1] Wu et al. An Efficient Recipe for Long Context Extension via Middle-Focused Positional Encoding. NeurIPS 2024.\n\n---\n\nQ4. 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}]$.\n\nA4. 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.\n\nFollowing 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.\n\n| Benchmark | LongBench-E | Perplexity |\n| --- | --- | --- |\n| Baseline | 23.4 | 5.82 |\n| $[1, l_{train}/2]$ | 25.8 | 5.77 |\n| $[1, 0.05l_{train}]$ | 23.9 | 5.87 |\n\nWhile 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.\n\n---\n\n\nQ5. I was confused on how the hyper-parameters is selected, could you clarify? \n\nA5. 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).\n\n---\n\nQ6. In Section 3, why are those transformation selected?\n\nA6. 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.\n\n---\n\nAlso thanks for the minor comments! We will address them in the final version.", "claims_and_evidence": "- **Concern 1**: my biggest concern is on the difference between mean prediction and length prediction task. The author points out in L160-164 that (1) for mean prediction task, the output remain in $[0,1]$ and (2) for length prediction task, the output support set grows with longer length, (3) this difference motivates them to consider output alignment. However, **there exists many tasks whose output support set do not grow with longer length, see Delétang et al 2022.** Transformers still struggle on these tasks. This makes me question if the motivation is valid.\n\n- **Concern 2**: On the connection between synthetic task and language modeling task. Using the mean prediction and length prediction task (Section) to motivate Section 4 seems not good, since the former is essentially a **regression task** (predicting continuous target) while the author extend their claim to **sequence modeling** subsequently.\n\nReference:   \nDelétang et al 2022. Neural Networks and the Chomsky Hierarchy.", "methods_and_evaluation_criteria": "- **Evaluation criteria 1**: The author proposes a new perspective of improving length generalization but **does not show clearly how important this perspective is**. While it makes sense to compare the proposed regularization with naive fine-tuning, the paper would benefit a lot by comparing it with some representative length generalization methods. Namely, this can answer the question that should the community considers more on this direction because it improves more than conventional perspectives (e.g. positional encodings)? However, I admit that this might be out of scope and is just out of my personal curiosity.", "theoretical_claims": "I checked the Theorem 4.1 (which should be the only theoretical claim), and did not find a clear problem. However, it is possible that there is a miss.", "experimental_designs_or_analyses": "- **Experimental Design 1**: In Section 4.1, the $l_1$ and $l_2$ are sampled from $[l_{train}/2, l_{train}]$ and $l_{train}$ is the model's training context length. This can be a large range, which defeats the claim (L064, L216) that the $x_{[-l_1:]}$ and $x_{[-l_2:]}$ are only **slightly** different? **The motivation of output distribution being similar should only hold when the length differences are indeed small.** Also, in L431, the author mentioned that aligning output distributions with **moderate** length discrepancies improves. I think the author should be more careful on the consistency of their claims.\n\n- **Minor**:\n    - In L593 (Appendix B), I was confused on how the hyper-parameters is selected, could you clarify?\n    - In Section 3, why are those transformation selected?\n    - Again in Section 3, what loss are used to train the model? My guess is MSE (seen from Equation 8) but this should be more explicitly mentioned.", "relation_to_broader_scientific_literature": "- **Broader literature 1**: The regularization loss to enforce output distribution consistency between lengths broadly connects to **better fine-tuning strategies**.\n- **Broader literature 2**: The idea of keeping output distributions similar for similar inputs (here in terms of length) is distantly related to keeping model output similar on inputs with data augmentations (e.g. a rotated image), which is in general used in **Test-time training literature**.", "essential_references_not_discussed": "I believe the authors have discussed all necessary references.", "other_strengths_and_weaknesses": "Summary of strengths and weaknesses.\n- **Strength**: the idea to study length generalization on output distributions is novel, and the key claims are supported both theoretically and empirically.\n- **Weakness**: the motivation obtained from the synthetic tasks does not make a lot of sense to me, and the connection between synthetic task and language modeling task, as well as their corresponding proposed approaches (i.e. Section 3 and 4) is weak, in the paper's current format.", "other_comments_or_suggestions": "Minor comments:\n\n- **Comment 1**: Presentation can be a bit smoother. For example, the first paragraph of Section 3 can be much shorter since the related work and background was just discussed above.\n- **Comment 2**: The title of Table 7 seems wrong: it should be ablation study on 'sampling range'.", "questions_for_authors": "Most of my questions have been mentioned above. I would be happy to increase my score if the authors can address/clarify **Concern 1&2, Experimental design 1** and **Evaluation criteria 1**. For example, I suggest the following, respectively (the author can provide more suitable clarifications if there are any)\n- **Concern 1**: I don't have a concrete idea on how Concern 1 can be addressed. I let the authors decide.\n- **Concern 2**: discuss more on how explicit reparametrization relates to the regularization loss, and how section 3 connects with section 4. \n- **Experimental design 1**: Refine and rationalize their claim on which lengths should the consistency be enforced on, should the length differences be slight, moderate, or something else? Given $l_{train}$ can be large, maybe the author can provide experiments on something like $[0.95*l_{train}, l_{train}]$.\n- **Evaluation criteria 1**: Add a preliminary experimental comparison between their approach and some representative length generalization approaches (e.g. compare to length-extrapolatable position encodings).", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4OaTXZPGWJ", "reviewer": "Reviewer_qzDj", "summary": "This work suggests matching the long-sequence loss and the short-sequence loss, named output alignment. The author has conducted experiments to prove the effectiveness of the proposed methods.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer qzDj for the comments. We will address your questions in the following part.\n\n---\nQ1. Figure 1 (a), the NoPE loss is almost zero. Is anything wrong?\n\nA1. 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.\n\n---\n\nQ2. The proposed method may degrade the performance for the shorter length, which is proved in Table 2 and Table 3.\n\nA2. 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.\n\n---\n\nQ3. 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?\n\nA3. 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.\n\n[1] Zhao et al, When Will Gradient Regularization Be Harmful? ICML 2024.\n\n[2] Srivastava et al. Dropout: a simple way to prevent neural networks from overfitting. Journal of Machine Learning Research 2014.\n\n---\n\nThanks again for your comments, and hope our response could address your concerns. Please let us know if you have additional questions.", "claims_and_evidence": "The work supports the claims and evidence.", "methods_and_evaluation_criteria": "The proposed methods evaluation criteria make sense. However, there are still some questions:\n* **Figure 1 (a), the NoPE loss is almost zero**. Is anything wrong?\n* **The proposed method may degrade the performance for the shorter length, which is proved in Table 2 and Table 3.**\n* **The method is sensitive to hyperparameters for the loss misalign**, which is proved in Table 6.", "theoretical_claims": "I have checked the proofs.", "experimental_designs_or_analyses": "I have checked the soundness/validity of experimental designs and analyses", "relation_to_broader_scientific_literature": "This works propose", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "M/A", "questions_for_authors": "The proposed method is sensitive to the hyperparameter for loss misalign (proved in Table 6). Is there any way to help choose the hyperparameter?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper studies how the output distributions of LLMs relates to their length generalization and long-context performance. The proposed "long-short misalignment" metric and a related regularization loss provides empirical evidence in both synthetic and real tasks. While some reviewers raised concerns around the motivation and theoretical framing, the authors addressed these well during rebuttal, including new results and clarifications. Overall, this paper offers some new insights into length generalization of LLMs, although the (1) comparisons/connections to other length generalization methods and (2) more clarity of the limitations of the proposed regularization loss could make this paper stronger.
Decision: Accept (poster)