DejanPorjazovski AnssiMoisio MikkoKurimo tabular

# Out-of-distribution generalisation in spoken language understanding

###### Abstract

Test data is said to be out-of-distribution (OOD) when it unexpectedly differs
from the training data, a common challenge in real-world use cases of machine
learning. Although OOD generalisation has gained interest in recent years, few
works have focused on OOD generalisation in spoken language understanding
(SLU) tasks. To facilitate research on this topic, we introduce a modified
version of the popular SLU dataset SLURP, featuring data splits for testing
OOD generalisation in the SLU task. We call our modified dataset SLURP For OOD
generalisation, or SLURPFOOD. Utilising our OOD data splits, we find end-to-
end SLU models to have limited capacity for generalisation. Furthermore, by
employing model interpretability techniques, we shed light on the factors
contributing to the generalisation difficulties of the models. To improve the
generalisation, we experiment with two techniques, which improve the results
on some, but not all the splits, emphasising the need for new techniques.

###### keywords:

spoken language understanding, generalisation, out-of-distribution, model
interpretability

##  1 Introduction

Spoken language understanding (SLU) systems are the backbone of human-computer
interaction devices that need to understand the meaning of the utterance
before taking an action. When these models are deployed for real-world use,
their performance should be consistent even when presented with out-of-
distribution (OOD) data that differs from the training distribution in an
unpredictable manner [1].

Multiple types of OOD generalisation capacities are desirable but not
necessarily achieved by neural SLU systems. _Length generalisation_ (see e.g.
[2, 3]) is the capacity to process sequences that are longer (or shorter) than
those seen in the training set. _Out-of-vocabulary_ (OOV) generalisation is
needed when the test set includes words or other units not seen in the
training set [4, 5]. A third type is _compositional generalisation_ (CG) [6,
7], required when test samples combine familiar units in novel ways. For
example, novel combinations of slot types in a slot filling task has been
shown to pose difficulties for SLU systems [8, 9]. More generally, neural
(NLP) systems have been found often to fail in tasks that require CG [10, 11,
12], although they do have some capacity for CG [13]. In addition to the
generalisation types applicable to both text and audio sequences, some types
are specific to audio-based models, such as diverse accents [14], different
age groups [15], and various acoustic environments [16].

The traditional pipeline SLU systems consist of two sub-modules: an automatic
speech recognition (ASR) system that generates transcripts and a separate
text-based natural language understanding (NLU) system that extracts meaning
from the generated transcripts [17, 9]. Previous studies that explore
generalisation in SLU employ a pipeline system, focusing on the transcripts
instead of the original audio [18, 9]. As these studies focus on the text-
based language understanding task, the proposed evaluation sets are mostly
text-only data, and as such, they can not always be applied to the end-to-end
(E2E) SLU models that rely solely on audio [19, 20, 5]. The main reason is
that text is easy to segment due to word boundaries, whereas adding or
removing segments of an audio signal corresponding to certain tokens in the
transcript is challenging. Moreover, the audio models' generalisation criteria
can differ from those of text-based models.

Despite these challenges, a few previous works have evaluated the
generalisation capabilities of E2E SLU models. In [4], a system was trained on
three specific phrases and tested with also a new phrase, assessing the OOV
generalisation. Similarly, in [5], E2E SLU models were tested on unique
wordings not seen in training. The findings revealed a notable increase in
error rates as the number of unseen wordings increased. These studies,
however, predominantly focus on a singular type of generalisation. The authors
of [21] noted that generalisation in SLU tasks may be divided into two types:
generalising to diverse speakers and generalising to diverse phrases (e.g.,
n-grams). These two types may be orthogonal and thus can be assessed
separately. We argue, in the same vein, that there are multiple types of
generalisation that SLU models should ideally achieve. However, our data
splitting methods differ from those in [21] in multiple ways. The n-gram-based
splits in [21] may result in OOV words in the test set, or possibly novel
combinations (n-grams) of familiar words, requiring compositional
generalisation. We distinguish OOV splits from CG splits, allowing for a more
fine-grained evaluation. Moreover, we focus on slightly more abstract units
than n-grams, namely the scenario and action labels. Our methods and resources
therefore aim to complement the previous works, adding to the collection of
resources enabling development of more robust SLU systems.

We introduce OOV, CG, and microphone mismatch splits of the SLURP dataset
[22], which we call SLURP For OOD generalisation (SLURPFOOD). Utilising these
splits, we assess SLU models' capacity for OOD generalisation. As a benchmark,
we report results for pre-trained self-supervised models. We also leverage a
model interpretability method (Integrated Gradients [23]) to determine the
underlying causes contributing to the limited generalisation of the systems.
To improve the generalisation, we explore two techniques: TOPK [24] and
segmented processing. We make the SLURPFOOD splits, the models, and the code
for reproducing all the experiments publicly
available111https://github.com/aalto-speech/slurpfood.

Table 1: Statistics for the proposed split configurations. OOV is out-of-vocabulary, CG is compositional generalisation, and DA is Double-Action. The mic mismatch split is not included due to its similarity to the original splits. The table shows class overlaps (intersections of class sets) between the train, dev, and test subsets for each label type (scenario, action, intent). In the brackets are the similarities $C_{\alpha}$ (see Section 2) of the frequency distributions of labels in the different subsets. |  | Original | Non-OOV | OOV | Non-CG | CG | DA Non-CG | DA CG
---|---|---|---|---|---|---|---|---
Subset size | Train | 50627 | 36287 | 36286 | 29750 | 32651 | 36064 | 36064
Dev | 8690 | 6216 | 6216 | 3302 | 3610 | 6192 | 6192
Test | 13078 | 3444 | 3444 | 1480 | 1480 | 3500 | 3500
Action overlap | Train-Test | 46 | (1.00) | 22 | (0.71) | 2 | (0.32) | 44 | (0.99) | 39 | (0.83) | 39 | (0.97) | 35 | (0.93)
Dev-Test | 45 | (0.99) | 21 | (0.71) | 2 | (0.33) | 43 | (0.98) | 38 | (0.83) | 38 | (0.96) | 37 | (0.92)
Scenario overlap | Train-Test | 18 | (1.00) | 13 | (0.83) | 13 | (0.83) | 18 | (1.00) | 16 | (0.84) | 12 | (1.00) | 12 | (1.00)
Dev-Test | 18 | (1.00) | 13 | (0.81) | 13 | (0.81) | 18 | (1.00) | 16 | (0.85) | 12 | (1.00) | 12 | (1.00)
Intent overlap | Train-Test | 59 | (1.00) | 23 | (0.88) | 0 | (0.00) | 57 | (0.99) | 41 | (0.39) | 27 | (0.74) | 0 | (0.00)
Dev-Test | 58 | (0.99) | 22 | (0.87) | 0 | (0.00) | 55 | (0.97) | 39 | (0.38) | 20 | (0.56) | 0 | (0.00)

##  2 SLURPFOOD

SLURP consists of 72395 recordings of 16520 utterances. Each sentence has 5
types of annotations: transcript, named entities, part-of-speech tags,
scenario, and action. _Scenario_ is a label that indicates the general
situation or context for the sentence, for example `qa' (question answering),
`audio', or `transport'. The _action_ indicates what the speaker wants the
(supposed) system to do (although often not very precisely), for example
`factoid' in the `qa' scenario, `volume_up' in the `audio' scenario, or `taxi'
in the `transport' scenario. There is also a sixth, redundant annotation
called _intent_ , which is the combination of scenario and action. For more
details about SLURP, see [22].

We propose four distinct split configurations to assess generalisation in SLU.
The mismatched acoustic environment splits are applicable to the scenario and
intent classification tasks. The CG splits can be applied to scenario
classification or be adapted for intent classification by predicting the
scenarios and actions separately, for example by using a multi-task model. The
OOV and Double-Action CG (DA CG) splits can be used only for the scenario
classification task because their test sets include some action classes that
are not in the training sets. The statistics of the splits are presented in
Table 1.

OOV splits. To create the OOV splits, we first select a test set, shared
between the OOV and non-OOV portions. In this task, the test subset is sampled
from the original SLURP test set, ensuring there are enough different intents.
To create the OOV training and dev subsets, we removed the intents present in
the test set from the original training and dev subsets. This allows us to
assess the model's OOV generalisation ability when the scenario is seen in
training, but the task (action) is novel. For instance, the training set may
contain `launch super mario' where the scenario is `play' and the action is
`game'. Conversely, the test set could include an instruction `please play the
next episode in the podcast', maintaining the `play' scenario, but in a
combination with a novel `podcast' action.

To create the non-OOV training and development subsets to contrast with the
OOV subsets, we replaced half of the difficult splits with elements containing
intents present in the test set. While adding the new samples, we made sure
that the scenario distribution remained the same as in the OOV split.

CG splits. To generate the CG splits, we utilise the _distribution-based
compositionality assessment_ (DBCA) [11] framework. DBCA is a method to
determine how compositional two datasets are by calculating the divergences of
the distributions of _atoms_ , primitive elements, on one hand, and the
divergences of the distributions of combinations of atoms, called _compounds_
, on the other hand. To create a train-test split that requires compositional
generalisation, the atom divergence should be similar between train and test
splits, while compound distribution should diverge. In our CG splits, atoms
are defined as the scenario and action labels, and the compound in each
utterance is the combination of the two, i.e. the intent.

Similarity between distributions $P$ and $Q$ is calculated in the same way as
in [11], using the Chernoff coefficient
$C_{\alpha}(P\|Q)=\sum_{k}p_{k}^{\alpha}\,q_{k}^{1-\alpha}\in[0,1]$ [25], with
$\alpha=0.5$ for the atom distribution similarity and $\alpha=0.1$ for the
compound distribution similarity. The _divergence_ is defined as
$1-C_{\alpha}$. A greedy algorithm, similar to that by [26], splits the data
aiming to maximise the train-test compound divergence while minimising the
atom divergence. A non-CG split with minimal compound divergence, as well as
minimal atom divergence, is also created to contrast with the CG split. In
Table 1, we can see that for the CG splits, the distributions of the actions
and scenarios are relatively similar between training and test sets, whereas
the distributions of intents diverge.

Double-Action CG splits. We provide a second split configuration for testing
the CG. In this split, the task is to predict the scenario from a combination
of two utterances with different actions but same scenario. Again, for both
the CG and non-CG versions, we used a unified test set and modified the
training and development portions accordingly. The train, development, and
test portions are derived from the original SLURP splits.

For the CG split, we ensured that no combination of the same two actions
present in the test set appears in the training or development splits. To
create the non-CG split, we removed half of the samples from the difficult set
and replaced them with two action combinations present in the test set,
preserving the class distribution as in the CG split.

Microphone mismatch splits. In the original SLURP dataset, the utterances were
captured either through a headset or from a distance, simulating diverse
acoustic environments. We used this information to create splits for assessing
SLU models' ability to generalise to mismatched acoustic environments. More
specifically, we selected only the headset recordings from the official SLURP
train and development splits. To test the models' classification ability in
different acoustic environments, we created two test splits, derived from the
official SLURP test set. The first one contains recordings made with a headset
and the second one contains recordings made without a headset. We also made
sure that the speakers were identical in both versions. The number of samples
in training, development, and test splits is 24209, 4173. and 5552,
respectively. The rest of the statistics are identical to the original splits
(see Table 1).

##  3 Experiments and results

###  3.1 Baseline systems

The primary goal of the proposed splits is to have resources for testing the
generalisation ability of E2E SLU models. To test the generalisation ability
of the models on the proposed splits, we created baseline E2E SLU systems for
each of the splits and trained them on the scenario classification task.

All four split configurations are designed for the sequence classification
task. To do sequence classification, we fine-tuned the English pre-trained
WavLM-base-plus model [27] (90.2M trainable parameters) by adding a
classification layer and training it with the negative log-likelihood loss. We
chose this model due to its superior performance on various speech-related
tasks. To combine the utterances in the DA CG/non-CG splits, we processed them
individually using the WavLM model. The embeddings of the two utterances are
then averaged to get a single vector representation.

All the models were trained with a single V100 GPU for 30 epochs, using a
batch size of 10. In the OOV, CG and Mic mismatch, training for one epoch with
the WavLM model took around 16 minutes, whereas for the DA CG it took 45
minutes. For more details about the hyperparameters, refer to the provided
code repository.

Table 2: Micro F1 scores for the scenario classification task. The results are obtained by averaging the scores from two runs with different seeds. The 95% confidence intervals calculated with the bootstrap method are given in the brackets. Split | WavLM | TOPK | Segment
---|---|---|---
Non-OOV | 85.6 (85, 87) | 84.1 (83, 86) | 85.0 (84, 87)
OOV | 52.7 (51, 55) | 54.1 (53, 56) | 53.4 (52, 55)
Non-CG | 80.9 (79, 83) | 80.9 (79, 83) | 81.8 (80, 84)
CG | 61.1 (59, 64) | 57.6 (56, 60) | 59.4 (57, 62)
DA Non-CG | 88.9 (88, 90) | 88.7 (88, 90) | 89.0 (88, 90)
DA CG | 82.6 (81, 84) | 87.4 (87, 89) | 86.3 (85, 88)
Mic Headset | 85.4 (85, 86) | 85.6 (85, 87) | 85.6 (85, 87)
Mic Other | 69.9 (69, 71) | 63.7 (63, 65) | 69.2 (68, 71)

###  3.2 Results

As an evaluation metric, we used the micro F1 score. Additionally, we
calculated the 95% confidence intervals using the bootstrap method222Ferrer,
L. and Riera, P. Confidence Intervals for evaluation in machine learning
[Computer software]. https://github.com/luferrer/ConfidenceIntervals. The
scenario classification results are presented in Table 2.

On the OOV split, the baseline WavLM model performs significantly worse than
the non-OOV, with a 32.9 percentage point performance degradation. The big
performance gap between the easy and difficult splits suggests that the model
struggles to generalise when presented with context unseen during training.

On the CG/non-CG splits the difference in performance is smaller but still
noticeable. On these splits, the WavLM got a performance decrease of 19.8
points. When using a bigger context, like in the case of the DA CG/non-CG
splits, the gap in performance is more narrow, with a decrease of 6.3.

When evaluating the baseline WavLM model on utterances with the same and
different acoustic environments as those in training we can see again a large
drop in performance. Namely, when presented with utterances that do not use a
headset, the performance drops 15.5 percentage points. Since these splits do
not require any modification to the architecture to perform intent
classification, we trained the baseline on that task. The results showed that
the performance dropped from 79.7% to 42.1% when presented with samples that
do not use the headset.

###  3.3 Why do the models fail to generalise?

Figure 1: Top 3 words that are most frequently the most important word for the
prediction (colour intensity signifies frequency), determined by the IG
method, for the OOV and CG splits. The first two rows in each matrix show the
most important words for the samples that were correctly predicted by the IID
and OOD splits, respectively. The third row shows the most important words for
the incorrectly classified samples. The last row shows the most commonly
confused classes on the OOD split with respect to the true class.

As evident from the results, the WavLM model exhibits performance loss when
evaluated on the OOD splits. This trend shows that even though the model has
seen large amounts of data during pre-training, it still struggles to
generalise. To investigate the reason why the model fails to generalise, we
employed an interpretability technique called Integrated Gradients (IG) to
determine the most important words for each prediction. IG is a technique for
estimating the contribution of the input segments (in our case the word
regions) to the final prediction. The attributions are calculated by creating
interpolation steps from the baseline, in our case zeros, to the input and
integrating along the steps. Then, integrated gradients are obtained by
accumulating the gradients. For more technical details, refer to [23]. To
obtain the most important word for each prediction, we took the word region
with the highest attribution value.

Since we are interested in the contribution of each word to the final
prediction, we need to align the audio signal with the corresponding
transcript. For this, we used an ASR-fine-tuned wav2vec2 model. We collected
the most important word for each classified sample and counted the frequencies
of these words within each scenario class. Figure 1 shows the 3 words most
frequently determined as the most important by the IG method. We see the
models often assign a high attribution value to stopwords, such as `a' or
`the'. Focusing on stopwords may indicate that the model has learned spurious
correlations between input features and output classes. For example, the
models might learn to classify the samples based on syntactic structures
independent of their semantic content, which would lead to an inability to
generalise to novel utterances that are structured differently but have the
same intent, or vice versa.

Comparing the important words for correct predictions in the CG split to those
in the the IID (non-OOD) split (right-hand-side column of matrices, the 1st
and 2nd row in each matrix), there is a tendency that the most important words
for the CG split are semantically more relevant to the class than those for
the IID split. For example, the word `calendar' in the calendar class, or
`song' in music class, `lights' in the IoT class, or `podcast' in the play
class. (Although there are 2 counterexamples: email and transport classes.)
This could mean that in the CG split, the model has to use the semantically
relevant words to predict the class, whereas in the IID split, the model can
more often use stopwords to give correct classification. We take this to mean
the CG split more readily assesses the models' capacity to learn to use
features that help generalising to novel samples.

Comparing the important words for correct predictions to those for incorrect
predictions can also give some clues as to why the model confuses some
classes. For example, the Music class seems to be confused with the Play class
because of the word `song' in the OOV split, and `music' in the CG split.

###  3.4 Towards improving the generalisation

Hoping to improve the generalisation of the WavLM model, we experimented with
two approaches: TOPK [24] and segmented processing. TOPK presents a
straightforward method that preserves the existing architecture without any
modifications. It involves averaging solely the top-k highest losses within
each mini-batch, rather than all individual losses. Through experimentation on
a subset of the data, we explored values of k ranging from 2 to 6 to select
the optimal parameter. The best results were obtained with a k=2, therefore we
use that value in the results reported in Table 2.

For the second approach, we apply segment-based processing by obtaining the
embedding vector from WavLM and splitting it along the temporal dimension in
segments with overlapping windows. Then, we take the mean of each segment to
get a single embedding per segment. Finally, we compute the mean of the
embedding segments and pass it through the output layer. If the utterance is
too short and can not be segmented, we use the whole, without segmenting. To
determine the most optimal segment size, we conducted experiments by training
on a subset of the training data. We experimented with segment sizes of 20,
40, 60, 80, 100, 120, and 140. The stride is always half of the segment size.
We obtained the best results with a segment size of 120 and a stride of 60.

The result presented in Table 2 show that both approaches improve the
generalisation on some of the splits, but not on all of them. For instance,
for the OOV splits, both TOPK and segmented processing produce better results
than the baseline, with TOPK being better. On the CG splits generated with the
DBCA method, both TOPK and segmented processing failed to improve the
generalisation. However, for the DA CG split, both approaches outperformed the
baseline. On the Mic other split, the TOPK approach performs significantly
worse than the segmented processing. One reason for that could be due to the
training set containing only samples recorded with a headset that have very
little variance of the acoustic conditions, so the samples with highest losses
are probably difficult not due to the acoustic environment but due to other
factors instead.

##  4 Limitations

In this study, we mostly focus on the scenario classification task, although
intent classification is another task made possible by the intent labels in
the SLURP data. Similarly, named entity annotations in the SLURP dataset could
also be used to create a different type of CG splits. We aimed to shed light
on the reasons for the poor generalisation results of the models by utilising
the integrated gradients method, and we made some hypotheses based on the most
important words, but these hypotheses are not yet conclusively tested, and
therefore only tentatively suggest weaknesses of the systems.

##  5 Conclusions

We introduced data splits to assess different types of OOD generalisation of
SLU models. These splits systematically test the OOV, CG, and acoustic
environment aspects of OOD generalisation. The empirical findings showed
significant challenges in generalisation for the WavLM model across the OOD
splits, underscoring the necessity for alternative modelling strategies. We
investigated TOPK and segmented processing techniques, revealing improvements
in some types of OOD generalisation, although still failing to improve
generalisation to other types of OOD data. By analysing the most important
words affecting the predictions, we noticed that the models often give high
importance to the stopwords, indicating that they might learn spurious
correlations. Moreover, we found that the important words in general differ
between the OOD split types, providing different challenges to the models,
which is corroborated by the fact that the best results are achieved by
different model types depending on the split. Specifically, the CG split seems
to require the system to focus on semantically relevant words more often than
other splits do. In the future, we plan to use these findings of the
weaknesses of the systems to develop methods for improving the generalisation
of E2E SLU models.

##  6 Acknowledgements

We are grateful for the Business Finland project LAREINA under Grant
7817/31/2022. The computational resources were provided by Aalto ScienceIT.

## References

  * [1] J. Liu, Z. Shen, Y. He, X. Zhang, R. Xu, H. Yu, and P. Cui, ``Towards out-of-distribution generalization: A survey,'' _arXiv preprint arXiv:2108.13624_ , 2021.
  * [2] J. Ray Chowdhury and C. Caragea, ``Monotonic location attention for length generalization,'' in _Proceedings of the 40th International Conference on Machine Learning_. PMLR, 2023, pp. 28 792–28 808.
  * [3] H. Zhou, A. Bradley, E. Littwin, N. Razin, O. Saremi, J. Susskind, S. Bengio, and P. Nakkiran, ``What algorithms can transformers learn? a study in length generalization,'' in _The 3rd Workshop on Mathematical Reasoning and AI at NeurIPS'23_ , 2023.
  * [4] L. Lugosch, M. Ravanelli, P. Ignoto, V. S. Tomar, and Y. Bengio, ``Speech Model Pre-Training for End-to-End Spoken Language Understanding,'' in _Proc. Interspeech 2019_ , 2019, pp. 814–818.
  * [5] E. Palogiannidi, I. Gkinis, G. Mastrapas, P. Mizera, and T. Stafylakis, ``End-to-end architectures for asr-free spoken language understanding,'' in _ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_. IEEE, 2020, pp. 7974–7978.
  * [6] J. A. Fodor and Z. W. Pylyshyn, ``Connectionism and cognitive architecture: A critical analysis,'' _Cognition_ , vol. 28, no. 1-2, pp. 3–71, 1988.
  * [7] D. Hupkes, V. Dankers, M. Mul, and E. Bruni, ``Compositionality decomposed: How do neural networks generalise?'' _Journal of Artificial Intelligence Research_ , vol. 67, pp. 757–795, 2020.
  * [8] S. Broscheit, Q. Do, and J. Gaspers, ``Distributionally robust finetuning bert for covariate drift in spoken language understanding,'' in _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , 2022, pp. 1970–1985.
  * [9] A. Ray, Y. Shen, and H. Jin, ``Compositional Generalization in Spoken Language Understanding,'' in _Proc. INTERSPEECH 2023_ , 2023, pp. 750–754.
  * [10] B. Lake and M. Baroni, ``Generalization without systematicity: On the compositional skills of sequence-to-sequence recurrent networks,'' in _International Conference on Machine Learning_. PMLR, 2018, pp. 2873–2882.
  * [11] D. Keysers, N. Schärli, N. Scales, H. Buisman, D. Furrer, S. Kashubin, N. Momchev, D. Sinopalnikov, L. Stafiniak, T. Tihon, D. Tsarkov, X. Wang, M. van Zee, and O. Bousquet, ``Measuring compositional generalization: A comprehensive method on realistic data,'' in _International Conference on Learning Representations_ , 2020. [Online]. Available: https://openreview.net/pdf?id=SygcCnNKwr
  * [12] Y. Yao and A. Koller, ``Structural generalization is hard for sequence-to-sequence models,'' in _Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing_ , 2022, pp. 5048–5062.
  * [13] M. Lepori, T. Serre, and E. Pavlick, ``Break it down: Evidence for structural compositionality in neural networks,'' _Advances in Neural Information Processing Systems_ , vol. 36, 2023.
  * [14] T. Viglino, P. Motlicek, and M. Cernak, ``End-to-End Accented Speech Recognition,'' in _Proc. Interspeech 2019_ , 2019, pp. 2140–2144.
  * [15] A. Potamianos and S. Narayanan, ``Robust recognition of children's speech,'' _IEEE Transactions on speech and audio processing_ , vol. 11, no. 6, pp. 603–616, 2003.
  * [16] R. Haeb-Umbach, J. Heymann, L. Drude, S. Watanabe, M. Delcroix, and T. Nakatani, ``Far-field automatic speech recognition,'' _Proceedings of the IEEE_ , vol. 109, no. 2, pp. 124–148, 2020.
  * [17] N. Gupta, G. Tur, D. Hakkani-Tur, S. Bangalore, G. Riccardi, and M. Gilbert, ``The at&t spoken language understanding system,'' _IEEE Transactions on Audio, Speech, and Language Processing_ , vol. 14, no. 1, pp. 213–222, 2005.
  * [18] J. Gaspers, A. Kumar, G. Ver Steeg, and A. Galstyan, ``Temporal generalization for spoken language understanding,'' in _Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies: Industry Track_ , 2022, pp. 37–44.
  * [19] D. Serdyuk, Y. Wang, C. Fuegen, A. Kumar, B. Liu, and Y. Bengio, ``Towards end-to-end spoken language understanding,'' in _2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_. IEEE, 2018, pp. 5754–5758.
  * [20] P. Haghani, A. Narayanan, M. Bacchiani, G. Chuang, N. Gaur, P. Moreno, R. Prabhavalkar, Z. Qu, and A. Waters, ``From audio to semantics: Approaches to end-to-end spoken language understanding,'' in _2018 IEEE Spoken Language Technology Workshop (SLT)_. IEEE, 2018, pp. 720–726.
  * [21] S. Arora, A. Ostapenko, V. Viswanathan, S. Dalmia, F. Metze, S. Watanabe, and A. W. Black, ``Rethinking End-to-End Evaluation of Decomposable Tasks: A Case Study on Spoken Language Understanding,'' in _Proc. Interspeech 2021_ , 2021, pp. 1264–1268.
  * [22] E. Bastianelli, A. Vanzo, P. Swietojanski, and V. Rieser, ``SLURP: A spoken language understanding resource package,'' in _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , B. Webber, T. Cohn, Y. He, and Y. Liu, Eds. Online: Association for Computational Linguistics, Nov. 2020, pp. 7252–7262. [Online]. Available: https://aclanthology.org/2020.emnlp-main.588
  * [23] M. Sundararajan, A. Taly, and Q. Yan, ``Axiomatic attribution for deep networks,'' in _International conference on machine learning_. PMLR, 2017, pp. 3319–3328.
  * [24] K. Kawaguchi and H. Lu, ``Ordered SGD: A new stochastic optimization framework for empirical risk minimization,'' in _International Conference on Artificial Intelligence and Statistics_. PMLR, 2020, pp. 669–679.
  * [25] J. Chung, P. Kannappan, C. Ng, and P. Sahoo, ``Measures of distance between probability distributions,'' _Journal of mathematical analysis and applications_ , vol. 138, no. 1, pp. 280–292, 1989. [Online]. Available: https://www.sciencedirect.com/science/article/pii/0022247X89903351
  * [26] A. Moisio, M. Creutz, and M. Kurimo, ``Evaluating morphological generalisation in machine translation by distribution-based compositionality assessment,'' in _Proceedings of the 24th Nordic Conference on Computational Linguistics (NoDaLiDa)_. Tórshavn, Faroe Islands: University of Tartu Library, May 2023, pp. 738–751.
  * [27] S. Chen, C. Wang, Z. Chen, Y. Wu, S. Liu, Z. Chen, J. Li, N. Kanda, T. Yoshioka, X. Xiao _et al._ , ``Wavlm: Large-scale self-supervised pre-training for full stack speech processing,'' _IEEE Journal of Selected Topics in Signal Processing_ , vol. 16, no. 6, pp. 1505–1518, 2022.


# Never Miss A Beat: An Efficient Recipe for Context Window Extension of Large
Language Models with Consistent “Middle” Enhancement

Tong Wu, Yanpeng Zhao, Zilong Zheng
Beijing Institute for General Artificial Intelligence (BIGAI), Beijing, China,
National Key Laboratory of General Artificial Intelligence, Beijing, China
{wutong1, zhaoyanpeng<EMAIL_ADDRESS>Corresponding author.

###### Abstract

Recently, many methods have been developed to extend the context length of
pre-trained large language models (LLMs), but they often require fine-tuning
at the target length ($\gg 4K$) and struggle to effectively utilize
information from the middle part of the context. To address these issues, we
propose Continuity-Relativity indExing with gAussian Middle (CREAM), which
interpolates positional encodings by manipulating position indices. Apart from
being simple, CREAM is training-efficient: it only requires fine-tuning at the
pre-trained context window (_e.g_., Llama 2-4K) and can extend LLMs to a much
longer target context length (_e.g_., 256K). To ensure that the model focuses
more on the information in the middle, we introduce a truncated Gaussian to
encourage sampling from the middle part of the context during fine-tuning,
thus alleviating the “Lost-in-the-Middle” problem faced by long-context LLMs.
Experimental results show that CREAM successfully extends LLMs to the target
length for both Base and Chat versions of Llama2-7B with “Never Miss A Beat”.
Our code will be publicly available soon.

(a) Linear Interpolation

(b) YaRN Interpolation

Figure 1: Results of applying different position interpolation methods to the
“Lost-in-the-Middle” task on CREAM and PoSE [1]. We can see that CREAM
outperforms PoSE [1] at every position, with a particularly improvement in the
middle.

##  1 Introduction

Transformer-based Large Language Models (LLMs) are typically pre-trained with
a fixed context window size, _e.g_., 4K tokens in Touvron et al. [2]. However,
many downstream applications, including in-context learning [3, 4] and LLM
agents [5, 6]. necessitate the processing of significantly longer contexts,
_e.g_., up to 256K tokens. Recent works have proposed promising approaches
that efficiently extend the context window of pre-trained LLMs by
interpolating Positional Encodings (PEs) [7, 8, 9, 10, 11] with a short period
of fine-tuning. Unlike other techniques such as efficient transformer [12, 13]
and memory augmentation [14], PE-based methods do not necessitate alterations
to the model’s architecture or the incorporation of supplementary modules.
Consequently, PE-based methods offer the advantages of straightforward
implementation and rapid adaptation, making them a practical solution for
extending the operational range of LLMs in tasks involving larger context
windows.

Despite the simplicity and effectiveness, existing PE-based methods exhibit
two significant limitations. First, prior approaches, such as positional
interpolation [7], still require fine-tuning on the target context window
size, which imposes a substantial computational overhead [1]. Secondly, though
some PE methods have demonstrated potential in handling extremely long
sequences, as evidenced by low sliding window perplexity scores, their
performance deteriorates notably in “in-the-middle” scenarios [15].
Specifically, when the model is required to accurately retrieve and process
content located in the middle of an extended context, there is a marked drop
in performance on the extended window size (Figure 1 and Figure 3).

These observations and insights underscore a fundamental question: Can we
extend the context window size of pre-trained LLMs efficiently while
simultaneously optimizing their effectiveness in processing "in-the-middle"
content?

To answer the above question, we propose CREAM, namely Continuity-Relativity
indExing with gAussian Middle. CREAM is a novel PE-based fine-tuning recipe
that shows both efficiency in fine-tuning and effectiveness in enhanced middle
content understanding. Our key insights lie in manipulating the positional
indices of long target sequences to produce shorter ones within the pre-
trained context window size (Figure 2).

In Section 2.1, we summarize two crucial ingredients of effective positional
indices: continuity that produces densely connected positional indices and
relativity that reveals the long-range dependencies between fragments. CREAM
is a recipe designed with the best of both worlds by introducing two indexing
strategies for continuity and relativity, respectively (Section 2.2). Besides,
to alleviate the “Lost-in-the-Middle” challenge, we introduce truncated
Gaussian distribution for middle segment sampling, enabling the LLM to
prioritize the information in the middle positions, even when performing
positional interpolation within the pre-trained context window size.

In Section 3, we conduct comprehensive experiments to demonstrate the
efficiency and effectiveness of CREAM. We continually pre-trained on Llama
2-7B with CREAM for a short period and extend the context window size from 4K
to up to 256K. Furthermore, we instruction tuning on Llama 2-7B-Chat with
CREAM for 100 steps and obtain promising results. We highlight our empirical
advantages as:

  1. 1.

CREAM can not only fine-tune within the pre-training context window size, but
also alleviate the issue of the model easily getting lost in the middle.
_e.g_., CREAM-YaRN outperforms PoSE-YaRN [1] by over 20% on average in the
“Lost in the Middle” [15] task.

  2. 2.

CREAM can further be enhanced by integrating novel designs on positional
interpolation frequencies (such as Linear [7], NTK [8], Yarn [9], _etc_.), and
can be extended to context window sizes of up to 256K or beyond.

  3. 3.

CREAM-Chat model requires only 100 steps of instruction-tuning to achieve
nearly perfect performance on the Needle-in-a-Haystack pressure test, and it
outperforms existing strong baselines on LongBench [16].

##  2 Methodology

###  2.1 Preliminaries

#### Problem Formulation.

Given an LLM with a pre-trained context window size $N$, our goal is to unlock
the inference capacity of the LLM on the testing data $\mathcal{D}_{\rm test}$
with an extended context window size $L$ (where $L>N$) by efficiently learning
from a small-scale training data $\mathcal{D}_{\rm train}$ with a maximum
sequence length $N$. We expect the extended model to perform reasonably well
in long-context evaluation.

#### Continuity in Positional Encoding.

Transformer-based language models typically encode positional indices
sequentially as $\\{0,1,\cdots,N-1\\}$. Traditional length extension methods
[7, 8, 9] directly fine-tune on the target length $L$ with an updated
positional index. This approach preserves the continuity of all absolute
positions and learns all position indices within $[0,L-1]$, thereby
successfully extending to the target length. Furthermore, PoSE [1] attributed
their superior performance over RandPos [17] to the ensured continuity of
segments during fine-tuning.

#### Relativity in Positional Encoding.

Relative positional encoding (RPE) [18] has been proposed as an effective
positional encoding method, where only the relative positions between two
tokens are considered. Similar as prior works [17, 1, 19], our work focuses on
rotary positional encoding (RoPE) [20], which is one of the most prominent RPE
methods and has been widely applied to LLMs including the recent Llama family
[21, 2, 22]. In RoPE, only the relative distances between position pairs
$(|j-i|;0\leq i<j\leq L-1)$ are learned during fine-tuning (Appendix A). Due
to this property, we can manipulate the position indices such that all
relative positions between $[0,L-1]$ are learnable within the pre-trained
window size.

Figure 2: Illustration of CREAM position interpolation. The pre-trained
context window is divided into three segments: the head, middle, and tail. To
ensure continuity, we fix the lengths of the head and tail to a small value
$k$. To maintain relativity, we set the lengths of the head and tail to $N/3$.
For the middle part, the start and end position indices are determined via
truncated Gaussian sampling, thereby encouraging the model to pay more
attention to the information in the middle part.

###  2.2 Proposed Recipe: Continuity-Relativity indExing with gAussian Middle
(CREAM)

In the following section, we start by introducing our design of dividing the
context window $N$ to learn relative positional information. Then, we propose
two strategies that target continuity and relativity, respectively. Lastly, we
propose a novel truncated gaussian sampling method to enhance the middle part
of the long context. The overall framework is depicted in Figure 2.

#### Context division.

We first discuss the motivations behind our design of the context length.
First, prior works [23, 24] observed that a significant amount of attention
score is allocated to the beginning tokens of a sequence, which can
potentially encode absolute positional information even without explicit
positional encoding [25]. Secondly, the starting and ending tokens of long
contexts can be treated as two pointers that localize the middle indices with
the help of relative encodings. Therefore, we divide the pre-trained context
window into three segments. The detailed ablation results are shown in Section
3.5.

######  Definition 2.1.

Given the pre-trained context window size $N$ and target extended length $L$,
the position set of $\\{Head,Middle,Tail\\}$ is defined as follows:

Head $\displaystyle=\\{0,1,...,L_{h}-1\\},$ (1) Middle
$\displaystyle=\\{P_{s},P_{s}+1,...,P_{e}-1,P_{e}\\},$ Tail
$\displaystyle=\\{L-L_{t},...,L-2,L-1\\},$ $\displaystyle s.t.{}{}$
$\displaystyle L_{h}+(P_{e}-P_{s})+L_{t}=N,$

where $L_{h}$ and $L_{t}$ denote the length of the head and tail segments,
$P_{s}$ and $P_{e}$ denote the start and end position index of the middle
segment.

The relative positions among the three segments in each sample are calculated
in pairs, _i.e_., $\left\\{|j-i|;\forall
i,j\in\\{Head,Middle,Tail\\}\right\\}$.

The formed relative distance union $D_{r}$ learned by the model is given by:

$\small[0,\max(L_{h}-1,P_{e}-P_{s},L_{t}-1)]\cup[P_{s}-L_{h}+1,P_{e}]\cup[L-L_{t}-P_{e},L-1-P_{s}]\cup[L-L_{t}-L_{h}+1,L-1].$
(2)

Given that not all samples possess the same values for $L_{h}$, $P_{s}$,
$P_{e}$, and $L_{t}$, as fine-tuning progresses, the union $D_{r}$ in Equation
2 can encompass the entire range $[0,L-1]$, facilitating the model to learn
all relative positions within the target length $L$.

#### Two segmentation strategies.

For the sake of continuity, we set the $L_{h}$ and $L_{t}$ to a very small
value $k$, where $0\textless k\ll N$. Specifcally, we use $k=32$ in our
experiments. This choice allows the middle segment to closely approximate the
pre-trained context window. To maintain relativity, we divide $N$ equally into
three parts and fix the $L_{h}$ and $L_{t}$ to $N/3$, enabling the model to
learn as many relative positions as possible. In our fine-tuning process, both
types of examples are sampled with equal probability to maintain balance.

#### Truncated Gaussian Middle Sampling

To better focus the training process on the middle part of the long context,
we introduce a truncated Gaussian function. This approach reduces the interval
overlap in Equation 2 and directs the model’s attention towards the middle
section of the long context. In Appendix B, we provide theoretical
justifications of our truncated Gaussian design, indicating that the
maximization of $|D_{r}|$ holds for middle positions in
$[N,L/2)\cup(L/2,L-N]$.

Formally, given the probability density function (PDF) of a Gaussian
distribution:

$\small
f(x)=\frac{1}{\sigma\sqrt{2\pi}}\exp\left(-\frac{(x-\mu)^{2}}{2\sigma^{2}}\right),$

where $\mu$ is the mean and $\sigma$ is the standard deviation. The
corresponding cumulative distribution function (CDF) is:

$\displaystyle
F(x)=\int_{-\infty}^{x}f(t)\,dt=0.5\left(1+E\left(\frac{x-\mu}{\sigma\sqrt{2}}\right)\right),\quad{}E(z)=\frac{2}{\sqrt{\pi}}\int_{0}^{z}e^{-t^{2}}\,dt,$
(3)

where $E(\cdot)$ is the error function. To calculate the CDF value within the
truncated interval, we use a sufficiently large number (_e.g_. 1000) of
equally spaced $x$ values from the given interval $[1,L/N]$:

$\small x_{i}=1+\frac{(1\times(L/N))\cdot(i-1)}{999},\quad i=1,2,\ldots,1000,$
(4)

By substituting Equation 4 into Equation 3, the cumulative distribution
function (CDF) curve is derived within the truncated interval. For sampling
from this truncated Gaussian distribution, the inverse transform method is
employed, as demonstrated in Equation 5:

$\small\alpha=\text{round}(x_{i-1}+\frac{(x_{i}-x_{i-1})(u-F(x_{i-1}))}{F(x_{i})-F(x_{i-1})}),$
(5)

where $u\sim\text{Uniform}(0,1)$, $\text{round}(\cdot)$ represents rounding to
the nearest integer. Finally, we can get:

$\displaystyle P_{e}$ $\displaystyle\sim\text{Uniform}(L_{h}+\alpha\times
L_{m},(\alpha\times N-1)-L_{t}),$ (6) $\displaystyle P_{s}$
$\displaystyle=P_{e}-L_{m}+1,$

where $L_{m}$ denotes the length of the middle segments. In summary, the
overall sampling flow of our algorithm is presented in Algorithm 1.

Algorithm 1 CREAM sampling algorithm

0: Pre-trained context window size $N$, extended context window size $L$,
training sample size $S$, mean $\mu$, variance $\sigma$ and hyperparameter
$k$.

1: Generate enough $x$ equally spaced according to Equation (4).

2: Substitute $x$ into Equation (3) to derive the truncated Gaussian CDF
$F(x)$.

3: for $i=0$ to $S-1$ do

4: Sample $L_{h}\sim\text{DiscreteUniform}(\\{k,N/3\\})$, and let
$L_{t}=L_{h},L_{m}=N-L_{h}-L_{t}$.

5: Sample $u\sim\text{Uniform}(0,1)$, and substitute it into Equation (5) to
get $\alpha$.

6: Calculate the start and end position ids $P_{s},P_{e}$ of the middle part
according to Equation (6).

7: Get position set
$P_{i}=\\{0,1,\ldots,L_{h},P_{s},\ldots,P_{e},L-L_{t}\ldots,L-1\\}$, where
$|P|=N$.

8: end for

9: return $P=\\{P_{0},P_{1},\ldots,P_{S-2},P_{S-1}\\}$.

##  3 Experiments

###  3.1 Experimental Setup

#### Extended Models

We use Llama-2-7B and Llama-2-7B-Chat [2] as the base models and extend their
pre-trained context window size of 4K to a target context length of 32K. The
extended models are referred to as CREAM-Base and CREAM-Chat, respectively.
Note that, though the target context length is 32K, we do not have to fine-
tune CREAM on 32K token long text (see Section 2.2).

#### Benchmarks

We conduct long-context LLM evaluation of CREAM-Base on LongChat-Lines [26]
and Lost-in-the-Middle [15]. Ideally, fine-tuning should not disrupt what the
base model has learned, so we further evaluate CREAM-Base on the language
modeling task and the evaluation benchmark [27] adopted by Llama2.
Additionally, we assess the CREAM-Chat model with Needle-in-a-
Haystack111https://github.com/gkamradt/LLMTest_NeedleInAHaystack and
LongBench[16]. Unless otherwise specified, we use the linear interpolation to
adapt LLMs to a longer context length.

#### Baselines

As far as we know, RandPos [17] and PoSE [1] are similar to our approach in
that they manipulate position indices to enable fine-tuning on the pre-trained
length for context expansion. Therefore, these two methods serve as the
baselines for our primary comparisons. More details about the experimental
setup can be found in the Appendix C.

###  3.2 Effective Context Window Size Evaluation on CREAM-Base

We evaluate the long-context understanding capabilities of the CREAM-Base
model on two tasks: LongChat-Lines222Passkey retrieval [28] is another similar
task for evaluating long-context LLMs, but it is too simplistic to reflect
model performance at different context window sizes, so we use the dataset
provided by Pal et al. [26], which closely aligns with the task described in
Li et al. [29] [26] (Figure 3) and “Lost in the Middle” [15]
(LABEL:tab:lost_in_the_middle).

Figure 3: Results (%) on LongChat-Lines. Each length consists of 50 samples.
All results are fine-tuned on Llama-2-7B with 4K length data through linear
position interpolation. Refer to Appendix D for ablated results using NTK [8]
and Yarn [9]. Table 1: Results (%) on “Lost in the Middle”. “Position”
indicates the correct answers’ index, and each index comprises 500 samples.
All results are fine-tuned on Llama-2-7B with 4K length data.

Model |  Position (75 keys, $\sim$5K tokens) | AVG |  Position (140 keys, $\sim$10K tokens) | AVG
---|---|---|---|---
0 | 18 | 37 | 54 | 74 | 0 | 34 | 69 | 104 | 139
PoSE-Linear | 99.4 | 24.4 | 37.4 | 47.2 | 46.2 | 50.9 | 95.2 | 8.2 | 7.6 | 13.8 | 18.6 | 28.7
CREAM-Linear | 99.6 | 45.6 | 56.0 | 67.0 | 58.0 | 65.2 | 96.6 | 19.8 | 23.4 | 31.0 | 26.2 | 39.4
PoSE-NTK | 98.6 | 49.6 | 44.6 | 40.2 | 41.4 | 54.9 | 97.6 | 3.4 | 0 | 0 | 27.6 | 25.7
CREAM-NTK | 96.2 | 53.8 | 52.6 | 72.8 | 42.0 | 63.5 | 78.6 | 5.2 | 6.0 | 23.4 | 41.8 | 29.9
PoSE-YaRN | 99.6 | 32.6 | 12.2 | 57.2 | 48.4 | 50.0 | 91.8 | 0.6 | 2.8 | 8.2 | 18.8 | 24.4
CREAM-YaRN | 100.0 | 49.6 | 47.6 | 77.4 | 92.6 | 73.4 | 99.4 | 8.0 | 5.8 | 43.8 | 69.2 | 45.2

####  CREAM-Base performs best in retrieving information from long context of
varying lengths.

We extend the context window size up to 32K and compare CREAM with the Llama
2-7B [2], RandPos [17], and PoSE [1]. As the context window size increases,
the performance of all models drops, but CREAM always performs best except for
the window size of 3.6K (see Figure 3). In terms of the average performance
over all context window sizes, CREAM outperforms PoSE by 16%, demonstrating
its good long-context understanding ability.

####  CREAM-Base alleviates the Lost-in-the-Middle issue.

Lost-in-the-Middle is an observation that LLMs are generally good at
retrieving relevant information appearing at the beginning/end of the input
context [15]. To validate the effectiveness of our middle-focused truncated
Gaussian sampling, we evaluate CREAM and compare it with PoSE on the key-value
retrieval task proposed by Liu et al. [15]. We present results in
LABEL:tab:lost_in_the_middle, where the cyan shading indicates middle
segments. We find that: regardless of the chosen interpolation method, CREAM
always outperforms PoSE by a large margin. _e.g_., CREAM-Linear surpasses
PoSE-Linear by 21.2% when the relevant information is placed at 18.

###  3.3 Long Context Understanding Evaluation on CREAM-Chat

We conduct long-context evaluation of CREAM-Chat on two tasks:

  * •

Needle in A Haystack (Figure 9) This task is a test that places an answer
(_i.e_., Needle) at any position of a long context window (_i.e_., Haystack)
and requires a model to retrieve the correct answer given a question-answer
pair. We follow Wu et al. [19] and use the GPT (GPT-3.5-Turbo-0125) score as
the evaluation metric.

  * •

LongBench (LABEL:tab:longbench2) Bai et al. [16] is a more realistic benchmark
because it covers real-world application scenarios like long-context QA and
summarization. Moreover, it is specifically designed for Chat models.

(a) SkipAlign†

(b) CREAM

Figure 4: Results on Needle-in-a-Haystack. † indicates the results excerpted
from Wu et al. [19]. Both results are instruction-tuned on LLaMa2-7B-Chat with
4K length data. The color gradually changes from deep green to deep red,
indicating the Recall performance decreases from 10 to 1.  Table 2: Results
(%) on LongBench. ∗ indicates results reported by Bai et al. [16].
CREAM-7B-32k is instruction-tuned for 100 steps using 4K length data on
LLaMa2-7B-Chat.

Model |  | Single-
---
Doc QA
| Multi-
---
Doc QA
| Summari-
---
zation
| Few-shot
---
Learning
| Code
---
Completion
| Synthetic
---
Tasks
AVG
Llama2-7B-chat-4k∗ | 24.9 | 22.6 | 24.7 | 60.0 | 48.1 | 5.9 | 31.0
XGen-7B-8k∗ | 24.6 | 20.4 | 24.7 | 56.2 | 38.6 | 5.3 | 28.3
InternLM-7B-8k∗ | 17.4 | 20.2 | 16.1 | 50.3 | 36.4 | 4.5 | 24.2
Vicuna-v1.5-7B-16k∗ | 28.0 | 18.6 | 26.0 | 66.2 | 47.3 | 5.5 | 31.9
LongChat-v1.5-7B-32k∗ | 28.7 | 20.6 | 26.7 | 60.0 | 54.1 | 15.8 | 34.3
CREAM-7B-32k | 34.8 | 31.1 | 27.2 | 65.1 | 50.4 | 7.0 | 35.9

####  CREAM-Chat outperforms SkipAlign in context window expansion.

We visualize the results of CREAM-Chat and the recent SkipAlign in Figure 9.
Clearly, CREAM-Chat beats SkipAlign because the performance of SkipAlign [19]
decreases from the window size of 18K while CREAM-Chat displays a perfect
performance everywhere until from the window size of 29K. Notably, CREAM-Chat
is only fine-tuned for 100 steps.

####  CREAM-Chat makes best use of the extended context window size.

We present results on LongBench in LABEL:tab:longbench2. CREAM-Chat again
surpasses strong baseline models, demonstrating its better use of extended
context size. In terms of the average performance over all tasks, it
outperforms the second best model, i.e., LongChat-v1.5-7B-32k [29] , by 1.6%,
though it is only tuned on a very small amount of data and for only 100 steps.

###  3.4 Language Modeling and Standard Benchmark

Following Chen et al. [7], Zhu et al. [1], Peng et al. [9], we perform the
classic language modeling evaluation, _i.e_., perplexity evaluation, on
GovReport [30] and Proof-pile [31]. Since a lower perplexity does not
necessarily imply better model performance on downstream tasks [11, 32, 33,
34], we further conduct evaluation on the standard natural-language-
understanding (NLU) benchmark [27]. This also lets us know whether fine-tuning
hurts the NLU ability of the pre-trained base model.

Table 3: Perplexity results of GovReport and Proof-pile. Each experiment is
the average perplexity of 50 samples, and all results are based on LLaMa2-7B
fine-tuned on 4K data length.

Model |  GovReport |  Proof-pile
---|---|---
4K | 8K | 16K | 32K | 4K | 8K | 16K | 32K
Original | 3.6 | - | - | - | 4.6 | - | - | -
RandPos-Linear | 8.9 | 7.4 | 6.2 | 5.8 | 12.1 | 11.9 | 11.9 | 12.9
PoSE-Linear | 3.8 | 3.2 | 2.7 | 2.5 | 4.7 | 4.6 | 4.6 | 4.4
CREAM-Linear | 3.8 | 3.2 | 2.7 | 2.5 | 4.7 | 4.6 | 4.5 | 4.4
RandPos-NTK | 4.6 | 4.0 | 3.6 | 4.0 | 5.8 | 5.8 | 6.2 | 7.3
PoSE-NTK | 3.7 | 3.2 | 2.7 | 2.6 | 4.7 | 4.6 | 4.5 | 4.7
CREAM-NTK | 3.8 | 3.2 | 2.7 | 2.7 | 4.7 | 4.6 | 4.5 | 4.7
RandPos-YaRN | 5.0 | 4.4 | 4.0 | 4.6 | 6.4 | 6.5 | 6.8 | 9.1
PoSE-YaRN | 3.7 | 3.2 | 2.7 | 2.5 | 4.6 | 4.6 | 4.5 | 4.4
CREAM-YaRN | 3.7 | 3.2 | 2.7 | 2.5 | 4.6 | 4.6 | 4.5 | 4.4

#### Both CREAM and PoSE demonstrate the lowest perplexity.

We apply different positional interpolation methods to RandPos [17], PoSE [1],
and CREAM and report their perplexities in LABEL:tab:linear_ppl. We find that:
CREAM and PoSE have a similar perplexity in different settings and both
outperform RandPos. This occurs primarily because the position indices used
during RandPos fine-tuning are discontinuous, which creates an inconsistency
with the pre-training stage.

Table 4: Experimental results of standard benchmarks. ∗ indicates results cited from Touvron et al. [2], and all results are based on LLaMa2-7B fine-tuned on 4K data length. Model | Zero-Shot | Few-Shot
---|---|---
WinoGrande | TruthfulQA(mc2) | PIQA | BoolQ | ARC-C | HellaSwag
LLaMa-2-7b-hf∗ | 69.2 | 39.5 | 78.8 | 77.4 | 45.9 | 77.2
RandPos-Linear | 63.3 | 39.3 | 76.5 | 66.6 | 32.0 | 48.5
PoSE-Linear | 68.8 | 38.6 | 77.8 | 76.2 | 47.7 | 77.1
CREAM-Linear | 67.5 | 37.4 | 78.5 | 75.4 | 46.8 | 76.9
RandPos-NTK | 68.7 | 35.9 | 78.6 | 74.8 | 45.5 | 74.4
PoSE-NTK | 68.8 | 38.6 | 77.8 | 76.2 | 47.7 | 77.1
CREAM-NTK | 67.5 | 37.4 | 78.5 | 75.4 | 46.8 | 76.9
RandPos-YaRN | 69.3 | 36.6 | 78.3 | 72.5 | 43.4 | 69.2
PoSE-YaRN | 69.4 | 39.6 | 78.1 | 76.7 | 49.0 | 78.0
CREAM-YaRN | 68.7 | 38.5 | 78.0 | 76.4 | 49.0 | 78.0

####  CREAM has nearly the same NLU abilities as the pre-trained base model.

Ideally, fine-tuning should not adversely affect the original capabilities of
the pre-trained base model. Our evaluation of CREAM confirms this, _i.e_.,
CREAM nearly retains all NLU abilities of the base Llama2-7B (see
LABEL:tab:standard_benchmark). Interestingly, CREAM improves over Llama2-7B on
ARC-C and HellaSwag. This is because these two tasks are few-shot tasks with
longer prompts, necessitating the assistance of long-context understanding.
Again, PoSE [1] performs similar to CREAM and both beats RandPos [17].

#### Extending the context length to 256K.

We push the limit and extend the context length of Llama-2-7B up to 256K.
Following Zhu et al. [1], we evaluate the extended model by calculating the
average perplexity over 20 samples from PG-19 [35] and Book3 [36].333We use
sliding window for calculation, with a window size of 32,768 and a sliding
step size of 4,096. Since the PG-19 test set does have enough samples that are
longer than 256K, we select a subset of samples from the PG-19 training set.

Table 5: Perplexity results of PG-19 and Book3. ∗ indicates results copied
from Zhu et al. [1], and CREAM is based on LLaMa2-7B fine-tuned on 4K data
length.

Model | PG-19 | Book3
---|---|---
64K | 96K | 128K | 192K | 256K | 64K | 96K | 128K | 192K | 256K
PoSE-Linear-128K∗ | 22.47 | 26.77 | 31.18 | - | - | 43.62 | 57.08 | 70.87 | - | -
PoSE-NTK-128K∗ | 14.84 | 29.48 | 34.80 | - | - | 16.04 | 31.42 | 37.00 | - | -
PoSE-YaRN-128K∗ | 10.36 | 10.77 | 11.33 | - | - | 12.30 | 13.07 | 13.81 | - | -
CREAM-Linear-192K | 5.9 | 6.0 | 6.1 | 6.1 | - | 7.6 | 7.7 | 7.8 | 7.8 | -
CREAM-NTK-192K | 5.0 | 5.1 | 5.2 | 5.2 | - | 6.9 | 7.0 | 7.0 | 7.1 | -
CREAM-YaRN-192K | 5.0 | 5.2 | 5.2 | 5.3 | - | 7.0 | 7.1 | 7.1 | 7.1 | -
CREAM-Linear-256K | 7.8 | 8.0 | 8.0 | 8.1 | 8.2 | 10.2 | 10.3 | 10.5 | 10.7 | 10.8
CREAM-NTK-256K | 5.1 | 5.3 | 5.3 | 5.4 | 5.4 | 7.2 | 7.3 | 7.3 | 7.3 | 7.4
CREAM-YaRN-256K | 5.2 | 5.3 | 5.4 | 5.4 | 5.5 | 7.1 | 7.2 | 7.2 | 7.3 | 7.3

We experiment with target context lengths 64K, 96K, 128K, 192K, and 256K and
apply different positional interpolation methods to the extended model (see
LABEL:tab:extremely_long). The results of PoSE [1] in LABEL:tab:extremely_long
are based on fine-tuning LLaMa 1-7B with 2K data length, and are provided for
reference only. Surprisingly, the increase of the target context length brings
little to no perplexity increase, demonstrating the stability of CREAM across
different target context lengths, even when the target context is extremely
long.

###  3.5 Ablation Study

To validate the effectiveness of our modeling choices, we further conduct an
ablation study of three main components of CREAM: truncated Gaussian sampling,
fixed start and end segments, and the trade-off between continuity and
relativity.

(a) Gaussian _vs_. Uniform.

(b) Head _vs_. Tail.

(c) Relativity _vs_. Continuity.

Figure 5: Ablation study of CREAM on LongChat-Lines. The result at each length
is estimated using 50 samples.

#### Truncated Gaussian sampling versus Uniform sampling.

We use truncated Gaussian sampling to encourage CREAM to make better use of
the middle part of the context. As a comparison, we replace it with the
Uniform sampling (see Figure 5(a)). We observe that the Uniform sampling
always leads to worse retrieval performance, suggesting the effectiveness of
the truncated Gaussian sampling.

#### Fixing the head and tail segments is crucial for good retrieval
performance.

We compare our choice of fixing the head and tail segments with three
alternatives: (i) removing both the head and tail segment, (ii) fixing only
the head segment, and (iii) fixing only the tail segment (see Figure 5(b)). We
find that: removing the head and tail segments leads to the worst performance;
it results in a complete failure (_i.e_., zero score) for the context size
32K. Keeping either of the head and tail segments performs slightly better
than removing both but underperforms our default choice of fixing both. We
suppose that this is because fixing both gives rise to better relativity
information, a finding that is consistent with that of Han et al. [23].

#### Maintaining a good balance between continuity and relativity is
necessary.

We encourage continuity by setting the head and tail segment lengths to $k=32$
and elicit relativity by letting $k=N/3$ (see Section 2.2). To balance the two
desired properties, we randomly choose $k=32$ and $k=N/3$ with an equal
probability during fine-tuning. Here we compare three scenarios: (1) enforce
only continuity, (2) enforce only relativity, and (3) balance continuity and
relativity (see Figure 5(c)). We find that balancing continuity and relativity
gives rise to the best performance, thus justifying our modeling choice.

##  4 Related Works

#### Efficient Transformers and Extra Memory

FoT [12] addresses the limitations of local attention in transformers by
integrating memory attention layers, which enable large models to learn from a
wide context while reducing interference. Infini-attention [13] incorporates
compressed memory into the standard attention mechanism and integrates masked
local attention and long-term linear attention mechanisms within a single
Transformer block. LLoCO [14] employs LoRA in conjunction with context
compression, retrieval, and parameter-efficient fine-tuning to learn context
offline. Although these methods can successfully extend the long context
window of LLMs, they either require modifications to the attention mechanism
or the addition of extra modules for assistance. In contrast, CREAM does not
require these operations and can be directly applied to pre-trained model.

#### Positional Interpolation

Chen et al. [7] first proposed extending the context window through positional
interpolation, which linearly reduces the input position indices to match the
original context window size, thereby preventing catastrophic high attention
scores from completely disrupting the self-attention mechanism. Subsequently,
various methods (such as NTK [8], ABF [10], and EABF [11]) emerged that modify
the base frequency of rotary positional encoding to achieve positional
interpolation. YaRN [9] introduced a segmented interpolation method, applying
different positional interpolations to different dimensions. LongRoPE [37]
identifies and utilizes two forms of non-uniformity in positional
interpolation through search, and introduces a progressive expansion strategy
for positiona interpolation. Moreover, CREAM can be combined with any
positional interpolation method.

#### Positional Encoding

RandPos [17] first modified position indices so that the model leverages the
relativity of positions, enabling it to extend to the target length with fine-
tuning over shorter lengths. PoSE [1] then emphasized the importance of
continuous segments, dividing the training length into two parts to further
enhance the interpolation effect. CREAM utilizes both relativity and
continuity, and it also better enables the model to focus on the middle part
of the context.

##  5 Conclusion

We proposed Continuity-Relativity indExing with gAussian Middle (CREAM), a
simple yet effective method to extend the context of large language models.
CREAM achieves a trade-off between continuity and relativity, enabling the
model to exploit positional relativity (_i.e_., fine-tuning within the pre-
trained length), while preserving text continuity (_i.e_., remaining as close
as possible to the pre-trained state). Furthermore, by employing truncated
Gaussian sampling, the model can concentrate more on the middle positions
during fine-tuning. Experimental results demonstrate that CREAM outperforms
other methods on both Base and Chat models and effectively mitigates the issue
of “lost in the middle”.

## References

  * Zhu et al. [2023] Dawei Zhu, Nan Yang, Liang Wang, Yifan Song, Wenhao Wu, Furu Wei, and Sujian Li.  Pose: Efficient context window extension of llms via positional skip-wise training.  In _The Twelfth International Conference on Learning Representations_ , 2023.
  * Touvron et al. [2023a] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al.  Llama 2: Open foundation and fine-tuned chat models.  _arXiv preprint arXiv:2307.09288_ , 2023a.
  * Huang et al. [2023] Xijie Huang, Li Lyna Zhang, Kwang-Ting Cheng, and Mao Yang.  Boosting llm reasoning: Push the limits of few-shot learning with reinforced in-context pruning.  _arXiv preprint arXiv:2312.08901_ , 2023.
  * Li et al. [2023a] Jiaqi Li, Mengmeng Wang, Zilong Zheng, and Muhan Zhang.  Loogle: Can long-context language models understand long contexts?  _arXiv preprint arXiv:2311.04939_ , 2023a.
  * Qian et al. [2023] Chen Qian, Xin Cong, Wei Liu, Cheng Yang, Weize Chen, Yusheng Su, Yufan Dang, Jiahao Li, Juyuan Xu, Dahai Li, Zhiyuan Liu, and Maosong Sun.  Communicative agents for software development, 2023.
  * Zheng et al. [2023] Zilong Zheng, Zixia Jia, Mengmeng Wang, Wentao Ding, Baichen Tong, and Songchun Zhu.  Langsuit$\cdot$e: Controlling, planning, and interacting with large language models in embodied text environments, 2023.  URL https://github.com/bigai-nlco/langsuite.
  * Chen et al. [2023] Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian.  Extending context window of large language models via positional interpolation.  _arXiv preprint arXiv:2306.15595_ , 2023.
  * Peng and Quesnelle [2023] Bowen Peng and Jeffrey Quesnelle.  Ntk-aware scaled rope allows llama models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation., 2023.  URL https://redd.it/14lz7j5.
  * Peng et al. [2023] Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole.  Yarn: Efficient context window extension of large language models.  In _The Twelfth International Conference on Learning Representations_ , 2023.
  * Xiong et al. [2023] Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al.  Effective long-context scaling of foundation models.  _arXiv preprint arXiv:2309.16039_ , 2023.
  * Zhang et al. [2024] Yikai Zhang, Junlong Li, and Pengfei Liu.  Extending llms’ context window with 100 samples.  _arXiv preprint arXiv:2401.07004_ , 2024.
  * Tworkowski et al. [2024] Szymon Tworkowski, Konrad Staniszewski, Mikołaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Miłoś.  Focused transformer: Contrastive training for context scaling.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Munkhdalai et al. [2024] Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal.  Leave no context behind: Efficient infinite context transformers with infini-attention.  _arXiv preprint arXiv:2404.07143_ , 2024.
  * Tan et al. [2024] Sijun Tan, Xiuyu Li, Shishir Patil, Ziyang Wu, Tianjun Zhang, Kurt Keutzer, Joseph E Gonzalez, and Raluca Ada Popa.  Lloco: Learning long contexts offline.  _arXiv preprint arXiv:2404.07979_ , 2024.
  * Liu et al. [2024] Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang.  Lost in the middle: How language models use long contexts.  _Transactions of the Association for Computational Linguistics_ , 12, 2024.
  * Bai et al. [2023] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al.  Longbench: A bilingual, multitask benchmark for long context understanding.  _arXiv preprint arXiv:2308.14508_ , 2023.
  * Ruoss et al. [2023] Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Bennani, Shane Legg, and Joel Veness.  Randomized positional encodings boost length generalization of transformers.  In _Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers)_ , pages 1889–1903, 2023.
  * Shaw et al. [2018] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani.  Self-attention with relative position representations.  In _Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers)_ , pages 464–468, 2018.
  * Wu et al. [2024] Wenhao Wu, Yizhong Wang, Yao Fu, Xiang Yue, Dawei Zhu, and Sujian Li.  Long context alignment with short instructions and synthesized positions.  _arXiv preprint arXiv:2405.03939_ , 2024.
  * Su et al. [2024] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu.  Roformer: Enhanced transformer with rotary position embedding.  _Neurocomputing_ , 568:127063, 2024.
  * Touvron et al. [2023b] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al.  Llama: Open and efficient foundation language models.  _arXiv preprint arXiv:2302.13971_ , 2023b.
  * AI@Meta [2024] AI@Meta.  Llama 3 model card.  2024\.  URL https://github.com/meta-llama/llama3/blob/main/MODEL_CARD.md.
  * Han et al. [2023] Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang.  Lm-infinite: Simple on-the-fly length generalization for large language models.  _arXiv preprint arXiv:2308.16137_ , 2023.
  * Xiao et al. [2023] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis.  Efficient streaming language models with attention sinks.  _arXiv preprint arXiv:2309.17453_ , 2023.
  * Kazemnejad et al. [2024] Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy.  The impact of positional encoding on length generalization in transformers.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Pal et al. [2023] Arka Pal, Deep Karkhanis, Manley Roberts, Samuel Dooley, Arvind Sundararajan, and Siddartha Naidu.  Giraffe: Adventures in expanding context lengths in llms.  _arXiv preprint arXiv:2308.10882_ , 2023.
  * Beeching et al. [2023] Edward Beeching, Clémentine Fourrier, Nathan Habib, Sheon Han, Nathan Lambert, Nazneen Rajani, Omar Sanseviero, Lewis Tunstall, and Thomas Wolf.  Open llm leaderboard.  https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard, 2023.
  * Mohtashami and Jaggi [2023] Amirkeivan Mohtashami and Martin Jaggi.  Landmark attention: Random-access infinite context length for transformers.  In _Workshop on Efficient Systems for Foundation Models@ ICML2023_ , 2023.
  * Li et al. [2023b] Dacheng Li, Rulin Shao, Anze Xie, Ying Sheng, Lianmin Zheng, Joseph Gonzalez, Ion Stoica, Xuezhe Ma, and Hao Zhang.  How long can context length of open-source LLMs truly promise?  In _NeurIPS 2023 Workshop on Instruction Tuning and Instruction Following_ , 2023b.  URL https://openreview.net/forum?id=LywifFNXV5.
  * Huang et al. [2021] Luyang Huang, Shuyang Cao, Nikolaus Parulian, Heng Ji, and Lu Wang.  Efficient attentions for long document summarization.  In _2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2021_ , pages 1419–1436. Association for Computational Linguistics (ACL), 2021.
  * Zhangir Azerbayev [2022] Bartosz Piotrowski Zhangir Azerbayev, Edward Ayers.  Proof-pile.  2022\.  URL https://github.com/zhangir-azerbayev/proof-pile.
  * Hu et al. [2024] Yutong Hu, Quzhe Huang, Mingxu Tao, Chen Zhang, and Yansong Feng.  Can perplexity reflect large language model’s ability in long text understanding?  In _The Second Tiny Papers Track at ICLR 2024_ , 2024.
  * Arora et al. [2024] Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Re.  Zoology: Measuring and improving recall in efficient language models.  In _The Twelfth International Conference on Learning Representations_ , 2024.  URL https://openreview.net/forum?id=LY3ukUANko.
  * Park et al. [2024] Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos.  Can mamba learn how to learn? a comparative study on in-context learning tasks.  In _ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models_ , 2024.  URL https://openreview.net/forum?id=xvr0Hctddy.
  * Rae et al. [2019] Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap.  Compressive transformers for long-range sequence modelling.  In _International Conference on Learning Representations_ , 2019.
  * Presser [2020] Shawn Presser.  2020\.  URL https://twitter.com/theshawwn/status/1320282149329784833.
  * Ding et al. [2024] Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang.  Longrope: Extending llm context window beyond 2 million tokens.  _arXiv preprint arXiv:2402.13753_ , 2024.
  * Loshchilov and Hutter [2018] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_ , 2018.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  _Advances in neural information processing systems_ , 32, 2019.
  * Dao [2023] Tri Dao.  Flashattention-2: Faster attention with better parallelism and work partitioning.  In _The Twelfth International Conference on Learning Representations_ , 2023.
  * Gao et al. [2020] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al.  The pile: An 800gb dataset of diverse text for language modeling.  _arXiv preprint arXiv:2101.00027_ , 2020.
  * Zheng et al. [2024] Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al.  Judging llm-as-a-judge with mt-bench and chatbot arena.  _Advances in Neural Information Processing Systems_ , 36, 2024.

##  Appendix A Relative Positional Encoding in RoPE

We provide a simple background proof on the relative positional encoding
performed by Rotary Position Embedding (RoPE) [20]. Given two embedding
vectors $\bm{x}_{q},\bm{x}_{k}\in\mathbb{R}^{d}$ corresponds to query and key
at positions $(m,n)\in[0,L)$, where $d$ is embedding dimension, their encoding
counterparts can be defined as:

$\displaystyle\bm{q}_{m}$
$\displaystyle=f_{q}(\bm{x}_{q},m)=\mathbf{R}_{\Theta,m}^{d}(\bm{x}_{q},m)$
(7) $\displaystyle\bm{k}_{n}$
$\displaystyle=f_{k}(\bm{x}_{k},n)=\mathbf{R}_{\Theta,n}^{d}(\bm{x}_{k},n)$

where

$\mathbf{R}_{\Theta,m}^{d}=\begin{bmatrix}\cos{m\theta_{1}}&-\sin{m\theta_{1}}&\cdots&0&0\\\
\sin{m\theta_{1}}&\cos{m\theta_{1}}&\cdots&0&0\\\
\vdots&\vdots&\ddots&\vdots&\vdots\\\
0&0&\cdots&\cos{m\theta_{d/2}}&-\sin{m\theta_{d/2}}\\\
0&0&\cdots&\sin{m\theta_{d/2}}&\cos{m\theta_{d/2}}\\\ \end{bmatrix}$ (8)

is the rotary matrix,
$\Theta=\\{\theta_{i}=10000^{-2(i-1)/d},i=[1,2,\dots,d/2]\\}$ is pre-defined
rotation angles. Then the self attention score can be obtained with:

$\displaystyle\bm{q}_{m}^{\rm T}\bm{k}_{n}$ $\displaystyle=\langle
f_{q}(\bm{x}_{q},m),f_{k}(\bm{x}_{k},n)\rangle$ (9) $\displaystyle={\rm
Re}\left[\sum_{i=0}^{d/2-1}\bm{x}_{q[2i:2i+1]}\bm{x}_{k[2i:2i+1]}^{*}e^{i(m-n)\theta_{i}}\right]$
$\displaystyle\coloneq g(\bm{x}_{m},\bm{x}_{n},m-n)$

where $\bm{x}^{*}$ represents the conjugate complex of $\bm{x}$, $g$ is the
derived attention function of RoPE. As seen, RoPE only depends on the relative
distances between and encodes the relative position information.

##  Appendix B Theoretical findings of CREAM design

######  Theorem B.1.

If $N\ll L$, the spanning size $|D_{r}|$ of the relative position union in
Equation 2 reaches its maximum iff. one of the following groups of
inequalities satisfies:

$\max(L_{h}-1,P_{e}-P_{s},L_{t}-1)+L_{h}-1<P_{s}<P_{e}<(L-L_{t})/2,$ (10)

or

$(L+L_{h})/2-1<P_{s}<P_{e}<L-L_{t}-\max(L_{h}-1,P_{e}-P_{s},L_{t}-1),\\\ $
(11)

where $\max|D_{r}|=\max(L_{h}-1,P_{e}-P_{s},L_{t}-1)+2N$.

Proof. Denote four intervals in Equation 2 as $S_{i},i=1,\dots,4$. According
to the inequality of inclusion-exclusion principle for the cardinality of the
union of $n$ sets:

$|D_{r}|=|\cup_{i=1}^{4}S_{i}|\leq\sum_{i=1}^{4}|S_{i}|,$ (12)

where the equality holds iff. all sets are pairwise disjoint. That is

$S_{i}\cap S_{j}=\varnothing,\quad{}\forall i\neq j$ (13)

Given intervals as in Equation 2, we have

$\begin{aligned} \begin{cases}{\rm MAX}<P_{s}-L_{h}+1\\\
P_{e}<L-L_{t}-P_{e}\\\ L-1-P_{s}<L-L_{t}-L_{h}+1\end{cases}\end{aligned}\quad
or\quad\begin{aligned} \begin{cases}{\rm MAX}<L-L_{t}-P_{e}\\\
L-1-P_{s}<P_{s}-L_{h}+1\\\ P_{e}<L-L_{t}-L_{h}+1\\\ \end{cases}\end{aligned},$
(14)

where ${\rm MAX}=\max(L_{h}-1,P_{e}-P_{s},L_{t}-1)$. The above inequalities
can be simplified to Equations 10 and 11.

######  Lemma B.2.

Under mild assumptions that $L-L_{t}\approx L$, $L+L_{h}\approx L$, the
maximization in Theorem B.1 holds for all
$(P_{s},P_{e})\in[N,L/2)\cup(L/2,L-N]$.

Proof. Given that

$\displaystyle\max(L_{h}-1,P_{e}-P_{s},L_{t}-1)+L_{h}-1$
$\displaystyle<\max(2L_{h},N-L_{t},N-L_{m})<N$ (15) $\displaystyle
L-L_{t}-\max(L_{h}-1,P_{e}-P_{s},L_{t}-1)$
$\displaystyle>L-\max(N-L_{m},N-L_{h},2L_{t})>L-N,$

the inequalities in Equations 10 and 11 turns into $[N,L/2)\cup(L/2,L-N]$.

######  Theorem B.3.

If $N\ll L$, when the spanning size $|D_{r}|$ of the relative position union
in Equation 2 reaches its maximum, we denote the coverage area of the middle
segment as:

$S_{m}\coloneq\left\\{x|x\in[P_{s},P_{e}],(P_{s},P_{e})\in\left\\{\operatorname*{arg\,max}_{(P_{s},P_{e})}|D_{r}|\right\\}\right\\}$
(16)

thus, we have:

$L\geq S_{m}+L_{h}+L_{t}>L-N/2$ (17)

Furthermore, as $\frac{N}{L}\rightarrow 0$, we have:

$L_{h}+S_{m}+L_{t}\rightarrow L$ (18)

##  Appendix C Experimental Details

#### Model Hyperparameters

We fine-tune all models by optimizing the causal language modeling objective.
A learning rate of $2\times 10^{-5}$ with a linear scheduler is adopted,
incorporating 10 warm-up steps. We use the AdamW [38] optimizer with the
hyperparameter configurations specified by PyTorch [39]. To speed up fine-
tuning, we resort to DeepSpeed 444https://github.com/microsoft/DeepSpeed ZeRO
stage 1 and Flash Attention-2 [40]. We perform fine-tuning on two A100-80G
GPUs with a total batch size of 32 and run inference on a single A100-80G GPU.
For CREAM-Base, we fine-tune it for 1,000 steps on a dataset derived from Pile
[41]; for CREAM-Chat, we fine-tune it for 100 steps on ShareGPT [42]. To
ensure fair comparison, we follow the fine-tuning and inference configurations
established by Zhu et al. [1].

#### Datasets and Training Cost

For training the Base model, we directly utilize The Pile data provided by Zhu
et al. [1], and select samples with token lengths exceeding 4K. For training
the Chat model, we filter the ShareGPT data from public
datasets555https://huggingface.co/datasets/Aeala/ShareGPT_Vicuna_unfiltered.
Specifically, we used the Vicuna prompt template to sequentially concatenate
the ShareGPT data until each data point comprises at least 4K tokens. Then, we
select 3.2K data points to train for 100 steps. Particularly, during the
instruction tuning process, we mask the USER part and allow the model to
calculate the loss only on the ASSISTANT part. We utilize two A100-80G
machines with a global batch size of 32, fully utilizing the available memory.
Running 1,000 steps for the Base model takes approximately 6 hours, while
running 100 steps for the Chat model takes approximately 2 hours.

##  Appendix D LongChat Lines Results

The interpolation methods using NTK and Yarn are presented in Figures 6 and 7.
As can be seen, CREAM performs the same as the Linear method for
interpolation, still outperforming other methods. The result of NTK at 26K-32K
is zero, which is due to the inherent properties of NTK, a finding that is
aligns with Zhu et al. [1].

Figure 6: Results (%) on LongChat-Lines. Each length consisting of 50 samples.
The above are the results of using NTK interpolation on the Llama 2-7B model.
Figure 7: Results (%) on LongChat-Lines. Each length consisting of 50 samples.
The above are the results of using Yarn interpolation on the Llama 2-7B model.

##  Appendix E LongBench Subtasks Results

The results of each subtask in Table LABEL:tab:longbench2 are shown in Tables
LABEL:tab:subtask1 and LABEL:tab:subtask2.

Table 6: Experimental results (%) of the LongBench subtasks selected in Zhang
et al. [11]. † indicates results quoted from Zhang et al. [11]. Len represents
the context length during fine-tuning. All results are based on Llama 2-7B.

Model | Num / Len | Singl-Doc QA | Multi-Doc QA | Summarization | Few-shot Learning | AVG |
---|---|---|---|---|---|---|---
NQA | QAPR | MFQA_en | HPQA | WMQA | MSQ | GR | QMSM | MNWS | TREC | TRVQA | SMSM |
PI† | 3.5K / 16K | 20.1 | 30.4 | 45.3 | 26.1 | 30.1 | 9.9 | 28.1 | 23.7 | 26.6 | 68.0 | 84.9 | 42.5 | 36.3 |
NTK-By-Parts† | 15.9 | 31.1 | 40.1 | 25.4 | 26.6 | 7.2 | 26.7 | 22.4 | 26.9 | 68.5 | 82.8 | 42.9 | 34.7 |
Yarn† | 20.3 | 28.9 | 42.8 | 27.8 | 30.7 | 7.2 | 27.4 | 22.5 | 26.8 | 66.0 | 85.6 | 42.6 | 35.7 |
ABF† | 24.6 | 32.8 | 45.6 | 35.1 | 30.3 | 15.2 | 30.8 | 23.0 | 27.4 | 71.0 | 84.7 | 42.7 | 38.6 |
EABF† | 21.9 | 31.0 | 47.1 | 40.1 | 32.7 | 15.1 | 32.3 | 23.0 | 27.1 | 70.5 | 86.7 | 42.0 | 39.1 |
CREAM | 3.2K / 4K | 23.0 | 34.6 | 46.8 | 42.2 | 33.7 | 17.4 | 30.4 | 24.3 | 26.8 | 69.5 | 84.0 | 41.9 | 39.6 |

Table 7: Experimental results (%) of the LongBench subtasks.

Model | Singl-Doc QA | Multi-Doc QA | Summarization
---|---|---|---
NQA | QAPR | MFQA_en | HPQA | WMQA | MSQ | GR | QMSM | MNWS
Llama2-7B-chat-4k∗ | 18.7 | 19.2 | 36.8 | 25.4 | 32.8 | 9.4 | 27.3 | 20.8 | 25.8
XGen-7B-8k∗ | 18.0 | 18.1 | 37.7 | 29.7 | 21.1 | 10.3 | 27.3 | 20.5 | 26.2
InternLM-7B-8k∗ | 12.1 | 16.7 | 23.4 | 28.7 | 22.8 | 9.0 | 9.7 | 15.9 | 22.8
Vicuna-v1.5-7B-16k∗ | 19.4 | 26.1 | 38.5 | 25.3 | 20.8 | 9.8 | 27.9 | 22.8 | 27.2
LongChat-v1.5-7B-32k∗ | 16.9 | 27.7 | 41.4 | 31.5 | 20.6 | 9.7 | 30.8 | 22.7 | 26.4
CREAM | 23.0 | 34.6 | 46.8 | 42.2 | 33.7 | 17.4 | 30.4 | 24.3 | 26.8

Table 8: Experimental results (%) of the LongBench subtasks. Model | Few-shot Learning | Code Completion | Synthetic Tasks
---|---|---|---
TREC | TRVQA | SMSM | PC | PR_en | LCC | RBP
Llama2-7B-chat-4k∗ | 61.5 | 77.8 | 40.7 | 2.1 | 9.8 | 52.4 | 43.8
XGen-7B-8k∗ | 65.5 | 77.8 | 25.3 | 2.1 | 8.5 | 38.6 | 38.6
InternLM-7B-8k∗ | 52.0 | 77.8 | 21.2 | 3.0 | 6.0 | 44.1 | 28.8
Vicuna-v1.5-7B-16k∗ | 71.5 | 86.2 | 40.8 | 6.5 | 4.5 | 51.0 | 43.5
LongChat-v1.5-7B-32k∗ | 63.5 | 82.3 | 34.2 | 1.0 | 30.5 | 53.0 | 55.3
CREAM | 69.5 | 84.0 | 41.9 | 3.0 | 11.0 | 52.0 | 48.7

It is noteworthy that, to provide further evidence of the efficacy of our
model, we have specifically chosen 12 tasks from the four categories outlined
in Zhang et al. [11] for comparison purposes. As delineated in Table
LABEL:tab:longbench1, we are able to attain superior performance on LongBench
in comparison to EABF [11], even with shorter training lengths and less data.

##  Appendix F Limitations

When extending the context beyond the pre-trained length, there is an
inevitable loss of information due to position interpolation, particularly
when fine-tuning is restricted to the pre-trained length. However, in
comparison to previous methods such as RandPos [17] and PoSE [1], CREAM has
effectively mitigated the issue of “Lost-in-the-Middle” by introducing
truncated Gaussian sampling. Additionally, as discussed in reference Liu et
al. [15], decoder-only models are prone to inherently exhibiting a U-shaped
performance curve on this task. Therefore, completely solving this problem
remains challenging.

##  Appendix G Loss Curve

(a) CREAM-192K Training Loss

(b) CREAM-192K Validation Loss

(c) CREAM-256K Training Loss

(d) CREAM-256K Validation Loss

Figure 8: Fine-tuning loss curve based on Llama 2-7B.The black line represents
Linear interpolation, the pink line represents NTK interpolation, and the cyan
line represents YaRN interpolation.

(a) RandPos Training Loss

(b) RandPos Validation Loss

(c) PoSE Training Loss

(d) PoSE Validation Loss

(e) CREAM Training Loss

(f) CREAM Validation Loss

Figure 9: Fine-tuning loss curve based on Llama 2-7B. The black line
represents Linear interpolation, the pink line represents NTK interpolation,
and the cyan line represents YaRN interpolation.