Long-Context Language Modeling with Parallel Context Encoding (2024)

Abstract: Extending large language models (LLMs) to process longer inputs is crucial for a wide range of applications. However, the substantial computational cost of transformers and limited generalization of positional encoding restrict the size of their context window. We introduce Context Expansion with Parallel Encoding (CEPE ), a framework that can be applied to any existing decoder-only LLMs to extend their context window. CEPE employs a small encoder to process long inputs chunk by chunk, enabling the frozen decoder to utilize additional contexts via cross-attention. CEPE is efficient, generalizable, and versatile: trained with 8K-token documents, it extends the context window of LLAMA-2 to 128K tokens, offering 10× the throughput with only 1/6 of the memory. CEPE yields strong performance on language modeling and in-context learning. CEPE also excels in retrieval-augmented applications, while existing long-context models degenerate with retrieved contexts. We further introduce a CEPE variant that can extend the context window of instruction-tuned models using only unlabeled data, and showcase its effectiveness on LLAMA-2-CHAT, leading to a strong instruction-following model that can leverage very long contexts on downstream tasks.1

heading: 1 Introduction
text: Enabling long and extensible context is crucial for large language models (LLMs) to effectively perform complex tasks, such as summarizing a book or answering questions with hundreds of retrieved Web pages. However, several challenges limit the ability of LLMs to leverage long context: (1) LLMs and popular positional encodings (Raffel et al., 2020; Su et al., 2021) do not generalize to sequence lengths longer than the lengths seen during training
1Code and models are available at https://github. com/princeton-nlp/CEPE.
(Press et al., 2022), even after additional fine-tuning (Chen et al., 2023, 2024; Peng et al., 2024, inter alia). (2) Transformers (Vaswani et al., 2017)— the predominant architecture of LLMs—incur a quadratic computational cost and a linear memory cost with respect to the input length, making it expensive to use for long sequences. (3) Highquality long-context data, such as long instructionfollowing data, are scarce and difficult to obtain (Wang et al., 2023; Xiong et al., 2023).
Besides directly fine-tuning LLMs on longer inputs, a series of inference-time modification methods have been proposed recently to scale up the effective context window, either by modifying the attention mechanism (Bertsch et al., 2023; Xiao et al., 2024b; Ivgi et al., 2023) or encoding chunks of context in separate forward passes (Shi et al., 2024; Ratner et al., 2023; Lin et al., 2023). While these methods generalize to longer sequences, the model often fails to effectively leverage the extra tokens and can incur larger inference costs.
2588
In this work, we propose an efficient and lightweight solution to extending the context window of LLMs, called Context Expansion with Parallel Encoding (CEPE ). CEPE is applicable to any pre-trained decoder-only LM by adding two components: (1) a small encoder that encodes the long context in chunks, and (2) a cross-attention module that is inserted at each layer of the decoder to attend to the encoder representations (Figure 2). With a careful selection of unlabeled training data, CEPE can leverage not only long-context documents, but also retrieved documents effectively.
CEPE offers several benefits: (1) length generalization: CEPE is not limited by positional encoding constraints as the long context is encoded in chunks, each with its own positional encoding; (2) efficiency: using a small encoder and processing contexts in parallel reduce computational cost. Since cross-attention attends only to the last layer’s representations from the encoder, CEPE requires much less memory compared to decoder-only LMs, which cache the key-value pairs of every token in every layer; (3) reduced training cost: unlike full fine-tuning approaches, we only tune the encoder and the cross-attention while keeping the large decoder LM frozen; augmenting a 7B decoder with a 400M encoder and cross-attention layers (1.4B parameters) can be done with one A100 80GB GPU.
We apply CEPE to LLAMA-2 (Touvron et al., 2023b) and train it on a filtered version of RedPajama (Together, 2023b) for 20B tokens—only 1% of the pre-training budget of LLAMA-2. We first show that CEPE-LLAMA-2, trained with 8K input length, continues to improve perplexity on longer input up to 128K tokens. Then, we apply CEPE to a retrieval-augmented setting, as
our larger context window allows incorporating more retrieved documents. Compared to existing methods, CEPE achieves better performance on both retrieval-augmented language modeling and open-domain question answering. Additionally, we demonstrate that CEPE can effectively leverage more demonstrations for in-context learning (Brown et al., 2020). All the above is achieved with a much lower memory and computational cost than most previous solutions.
Finally, we propose CEPE-Distilled (CEPED), which extends the context window of instructiontuned models, using only unlabeled data. CEPED distills the behavior of the original instructiontuned model to the new architecture through an auxiliary KL divergence loss, which eliminates the need to curate expensive long-context instructionfollowing data (Ivison et al., 2023). We apply CEPED to LLAMA-2-CHAT and show that while preserving their instruction understanding ability, CEPED-LLAMA-2-CHAT can incorporate more context and improve performance on long-text understanding tasks (Shaham et al., 2023).
To conclude, CEPE is a lightweight framework that can extend context windows of any base or instruction-tuned LMs. We hope CEPE can empower future LLM research with cheap and effective long-context abilities.
2 Method: CEPE
We design CEPE to adapt pre-trained LLMs to perform long-context language modeling on sequences with many tokens (e.g., books). For retrieval augmentation, these long contexts may contain a set of retrieved passages instead. We first describe how CEPE modifies the architecture of LLMs to
attend to representations encoded by a small encoder, and then describe how the CEPE modules are trained. Finally, we extend CEPE to CEPED, which expands the context window of instructiontuned models using only unlabeled data.

heading: 2.1 Architecture
text: CEPE augments off-the-shelf decoder-only LMs by (1) adding a small, bidirectional pre-trained encoder model and (2) inserting cross-attention layers between the self-attention and feed-forward layers in every transformer block of the decoder model.
Notation. Given an input context with T tokens x1, ..., xT , we consider the first m tokens x1, ..., xm as the additional context C and the last n = T − m tokens xm+1, ..., xT as the main input X . The additional context is split into chunks C = C1, ..., Ck, which can contain either segments within a long document or a set of retrieved passages. We use Menc to denote the encoder model with hidden dimension denc and Mdec to denote the decoder-only LLM with hidden dimension ddec.
Encoding chunks. We first encode C1, ..., Ck chunk by chunk using the trainable encoder Menc:
ϕi = Menc(Ci) Φ = CONCAT({ϕi}ki=1)
where ϕi ∈ R|Ci|×ddec is the token-wise last layer hidden state from Menc and Φ ∈ Rm×ddec . Note that Menc is bidirectional, which results in more information-rich representations compared to unidirectional ones. While we do not preserve global positions across different chunks, experiments show that CEPE achieves better or comparable performance to full-attention models that do.
Cross-attention modules. In every decoder layer of the transformer, we insert a cross-attention module between the self-attention and feed-forward layers. To construct the cross-attention module, we provide Φ as keys and values, and the hidden states of X as queries. Note that in order for Mdec to attend to Φ, the key and value projection matrices in the cross-attention module also serve as an up-projection that transforms the denc-dimensional Φ into a ddec-dimensional embedding. Figure 2 illustrates the architecture of CEPE.
Efficiency. Menc is much smaller and encodes contexts in parallel to avoid the quadratic complexity of full attention. This enables CEPE to exhibit a substantially higher training and inference speed than if we were to use Mdec to process all
T = m+ n tokens. Moreover, CEPE drastically reduces memory consumption by avoiding caching (m+n)L key-value pairs (L is the number of layers of Mdec) and instead caching only Φ and nL keyvalue pairs. In particular, using a standard decoderonly model requires O((m + n)Lddec) memory whereas CEPE requires O(mdenc + nLddec). In our setting, m ≫ n and ddec ≫ denc, so in practice, we observe a substantial gap: for each additional token, CEPE requires only 1/256 of the memory compared to encoding them in Mdec. Comparison with existing retrieval-augmented LMs. CEPE takes inspiration from retrievalaugmented models such as FiD (Izacard and Grave, 2021), Atlas (Izacard et al., 2022b), and RETRO (Borgeaud et al., 2022) for attending to parallel encodings. Our method differs in the following key aspects: FiD, Atlas, and RETRO are all fullparameter training methods for encoder-decoder models; Atlas and RETRO are pre-training methods that rely on expensive large-scale retrievalaugmented data; FiD is a fine-tuning method for open-domain QA tasks, which requires taskspecific data. In contrast, CEPE is a lightweight framework that can extend the context window of any existing decoder-only model; we only fine-tune an added small encoder and the cross-attention layers; instead of retrieval-augmented data, CEPE only requires a small amount of filtered unlabeled data, as described in the following section; because CEPE is built on top of powerful pre-trained models, it can be easily adapted to a wide range of tasks without explicit fine-tuning.

heading: 2.2 Data
text: We use RedPajama (RP; Together, 2023b) as our training corpus, which is an open-source reproduction of the LLAMA-1 (Touvron et al., 2023a) training data. It contains about 1 trillion tokens from seven domains: ArXiv, Books, C4-RP, CC, Github, StackExchange, and Wiki. We first partition the corpus into three sets of documents: training, test, and retrieval; we leverage the retrieval corpus for retrieval-augmented language modeling.
We design a data mixture with strong long-range dependencies and a diversity of domains. The training data is preprocessed into two subsets, each of which is a collection of 8192-token sequences: (1) In RPtrain-cat, we concatenate documents together to form training sequences; (2) In RPtrain-filter, we keep documents from the Arxiv and Books domains that
have at least 8, 192 tokens and sample sequences within document boundaries.
Our qualitative analysis found that data from the ArXiv and Books domains naturally contain long documents that are especially useful when training long-context models. It is also important to use a mixture of data from all domains to ensure better generalization. Thus, we use a mixture ratio of 2:1 between RPtrain-filter and RPtrain-cat for training. We alternatively consider using retrieval supervision when generating the training data to improve the model’s ability to leverage retrieved documents. However, our ablations found that our unlabeled data mixture transfers well to retrieval-augmented settings at a much cheaper cost (§7.1).

heading: 2.3 Training
text: We use LLAMA-2-7B (Touvron et al., 2023b) as Mdec (originally trained on 4K length), and insert the new cross-attention layers described in Section 2.1 into it. We add a bidirectional encoder Menc with 435M parameters, yielding 1.8B added parameters in the CEPE model.
Encoder. We first pre-train a bidirectional maskedlanguage model on the RedPajama dataset. Menc follows the configuration of RoBERTa-large (Liu et al., 2019b) but shares the vocabulary with Mdec, LLAMA-2. 2 We train Menc for 100K steps with a batch size of 2, 048 and sequence length of 512 tokens. For more details, refer to §A.1.
Cross-attention. We freeze the weights of Mdec and train only the added cross-attention layers as well as fine-tuning Menc using the cross-entropy loss. We first adopt a warmup training stage designed to teach Mdec to copy from Menc through cross-attention—for each position i ≤ T , the objective is to generate xi+1, conditioned on Menc(x1, . . . , xT ) and Mdec(x1, . . . , xi)—Menc and Mdec share the same input. The warmup stage uses 131M tokens from the training set of RP.3
After the warmup stage, we move to the standard training, where each sequence has T = 8, 192 tokens. We use the last 4, 096 tokens as the decoder input X , and chunk the first 4, 096 tokens into k = 16 contexts of |Ci| = 256 tokens each as the encoder input C, and train for 20B tokens. Freezing the decoder allows CEPE to be trained
2Future works applying CEPE to LLAMA-2 models can simply use our pre-trained encoder.
3Preliminary experiments suggest this stage can stabilize training; ablations can be found in §7.1.
on a single A100 GPU, which is a significant reduction in computational cost compared to training Mdec with sequence length T . For more training details and hyperparameters, please refer to §A.2.
2.4 CEPED for Instruction-Tuned Models We extend our method to CEPE-DISTILLED (CEPED) to augment instruction-tuned models with longer context. Instruction-tuned models (Ouyang et al., 2022; Taori et al., 2023; Touvron et al., 2023b) excel in many downstream applications, but their limited context window restricts their performance in tasks that require long documents (Shaham et al., 2023) or a large number of retrieved passages (Gao et al., 2023). It is challenging to extend these models to longer context windows directly through fine-tuning due to the scarcity of high-quality instruction data.
To this end, we propose CEPED, which uses an auxiliary distillation loss to encourage Menc and the cross-attention layers to learn the capabilities of the already instruction-tuned Mdec. This can be especially useful for settings where the fine-tuning data are not open-sourced, which is the case for LLAMA-2-CHAT (Touvron et al., 2023b).
Distillation loss. We design a distillation objective, where the original Mdec acts as the “teacher” and the CEPED model acts as the “student”. We use the input context CONCAT(C, X) with m + n = 4, 096 tokens, which can fit in the context window of LLAMA-2-CHAT, our choice of Mdec. First, we input CONCAT(C, X) to Mdec and save the logits of X as the teacher logits. During training, C and X are used as inputs to Menc and Mdec, respectively, and we minimize the KL divergence between the output logits of X and the corresponding teacher logits as well as the cross-entropy loss. We train the model for 10B tokens from our data mixture. For more details, see §A.3.

heading: 3 Long-context Language Modeling
text: We start by evaluating CEPE on long-context language modeling benchmarks to assess basic LM abilities, and compare with existing approaches in terms of perplexity, memory, and throughput.
Datasets. We evaluate on ArXiv and Books from our RedPajama test split, as well as three longcontext datasets: PG19 (Rae et al., 2020), ProofPile (Azerbayev et al., 2023), and CodeParrot (Wolf et al., 2023). We filter all documents to have at least 32, 768 tokens, and sample 5, 000 sequences
for each dataset. We calculate the perplexity on the last 256 tokens of each sequence. Following Peng et al. (2024), for the experiments with 128K tokens, we filter documents to have at least 131, 072 tokens and evaluate on only 10 sequences.4
Models. We focus on 7B parameters model. Our baseline includes LLAMA-2 and its long-sequence fine-tuned versions: LLAMA-2-32K (Together, 2023a), YARN-64K, and YARN-128K (Peng et al., 2024). LLAMA-2-32K was trained on RP with upsampled ArXiv and Books domains, and a mixture of other data , and YARN was trained on PG19. While YARN-64K and YARN128K were fine-tuned for only 1.7B and 3.4B tokens, respectively, fine-tuning all parameters requires much more memory relative to CEPE. We also evaluate on training-free long-context methods: STREAMINGLLM (Xiao et al., 2024b) and REPLUG (Shi et al., 2023) with LLAMA-2-7B. Although REPLUG was originally evaluated in a retrieval-augmented setting, we found that it also works well in long-context modeling, when view-
4This is due to computational cost and scarcity of long documents. All 10 sequences are from different documents.
ing the long context as retrieved context. See §B for more implementation details.
For CEPE, we put 2K tokens in the decoder when T = 4, 096, and 4K tokens in the decoder in other settings. Additional tokens are split into chunks of 256 tokens and fed into the encoder.
Results. We show the results in Table 1. Compared to the two fully fine-tuned models, LLAMA2-32K and YARN-64K, CEPE achieves either lower or comparable perplexity across all datasets with lower memory usage and higher throughput despite being trained only 8K sequences. Furthermore, CEPE continues to improve on perplexity while maintaining low memory use at 128K tokens, well beyond its training lengths (8K); on the other hand, LLAMA-2-32K and YARN-64K cannot generalize beyond its training length and the memory cost increases significantly. At 128K tokens, we outperform or achieve comparable perplexity with all applicable baselines on all domains. Although CEPE observes higher perplexity than YARN-64K on PG19 and Book at 32K, we note that YARN-64K was fine-tuned on PG19, and thus has a significant domain advantage.
We also outperform REPLUG across all domains while achieving much higher throughput— REPLUG encodes additional context in chunks but each chunk requires a forward pass of the main input, incurring slow speed, due to which we omit REPLUG at longer length. STREAMINGLLM maintains a low memory usage and a reasonable throughput, but the perplexity does not always decrease as the sequence length increases, likely due to the model only using a limited number of cached key-value pairs. Compared to STREAMINGLLM, CEPE can leverage all input tokens and achieve better perplexity with better throughput.

heading: 4 Retrieval-augmented Applications
text: Retrieval-augmented settings naturally benefit from long-context LMs, as models can leverage the ad-
ditional context to include more retrieval results. Thus, we test if CEPE trained on long-context data can transfer to retrieval-augmented settings.

heading: 4.1 Retrieval-augmented Language Modeling
text: Datasets. We use the test and retrieval split of RedPajama described in §2.2 for retrieval-augmented LM evaluation. Each sequence contains 2, 048 tokens, and the first 256 tokens are used as the query to retrieve passages from the retrieval split. The retrieval corpus contains 200M documents of 256 tokens each, and we use Contriever (Izacard et al., 2022a) to retrieve k passages for each sequence.
Models. We evaluate full-context baselines, LLAMA-2, LLAMA-2-32K, and YARN-64K, by simply prepending the retrieved passages to the input sequence. We also evaluate REPLUG, which
runs one forward pass for each retrieved passage and aggregates the results. CEPE uses 2, 048 tokens in the decoder and retrieved passages are fed through the encoder in parallel.
Results. The results are shown in Table 2. CEPE can effectively improve perplexity by using the retrieved contexts, outperforming REPLUG. Notably, CEPE continues to improve perplexity even with k = 50 (trained with k = 16). CEPE transfers well to the retrieval-augmented setting whereas the full-context decoder models degrade.

heading: 4.2 Open-domain Question Answering
text: Given a question and a large corpus of documents, open-domain question answering (QA) requires the model to retrieve relevant passages and generate the answer. A model that can leverage a large number of retrieved passages without being distracted by irrelevant ones is desirable for this task.
Datasets. We adopt three open-domain QA datasets: Natural Questions (NQ; Kwiatkowski et al., 2019; Lee et al., 2019), TriviaQA (Joshi et al., 2017), and PopQA (Mallen et al., 2023). For each question, we use Contriever to retrieve k passages from Wikipedia.5
Models. We compare CEPE with LLAMA-2, LLAMA-2-32K, and REPLUG. For each model, we use two in-context demonstrations. For CEPE, we use 10 passages in the decoder, and all other passages are encoded separately by the encoder. Refer to §C.1 for more details.
Results. The results are shown in Figure 3 CEPE consistently outperforms all models across all datasets and k. Notably, CEPE outperforms LLAMA-2-32K on NQ and TriviaQA by over 3 and 4 EM points, respectively. Furthermore, CEPE does not degrade in performance as the k increases, while other models often perform worse at larger k, as they are sensitive to the large amount of redundant or irrelevant passages (Liu et al., 2024).

heading: 5 In-Context Learning
text: In-context learning (ICL; Brown et al., 2020) is one of the most important emerging qualities of LLMs. In this experiment, we examine whether CEPE can effectively utilize demonstrations from the encoder context to improve performance. Specifically, we use a range of classification tasks that contains a
5Snapshot from 2018-12-20, and each passage is 100 words (Karpukhin et al., 2020).
large number (up to 150) of categories (Ratner et al., 2023), where the model can benefit from additional demonstrations. Following Ratner et al. (2023), we use a test set size of 250 examples for each dataset.
Models. Our baseline is LLAMA-2 with 2 demonstrations. For CEPE, we add additional demonstrations in the encoder. We also compare with an “oracle”, where the LLAMA-2 decoder takes 40 demonstrations. Note that the oracle is significantly more expensive. More details are in §C.2.
Results. The results are shown in Table 3. We first observe that compared to the decoder-only baseline, CEPE can effectively use the additional demonstrations from the encoder context; the performance further increases or remains consistent with more demonstrations in the encoder. However, there is still a large gap to the 40-demonstration oracle. Our hypothesis is that in-context learning requires both query-demonstration interactions and demonstration-demonstration interactions, which CEPE cannot provide. Regardless, CEPE can be always applied on top of the decoder-only model to add additional demonstrations, with little extra computational and memory cost.

heading: 6 Instruction-tuned Models for Long Text Understanding
text: In this section, we show that applying CEPED on LLAMA-2-CHAT produces instruction-following models that can leverage long inputs.
Datasets. ZeroSCROLLS (Shaham et al., 2023) is a collection of zero-shot long-context understanding tasks that require instruction-following abilities. Specifically, we test on NarrativeQA, QASPER, QuALITY, GovReport, SummScreenFD, and QMSum, which all have large validation sets made available. We follow the formats and instructions of Shaham et al. (2023) for each dataset, except the long text is placed before the instructions.
Models. For CEPED, we use 2, 048 tokens in the decoder and put the remaining tokens to the encoder as chunks of 256 tokens. We compare with LLAMA-2-CHAT and LLAMA-2-32K INSTRUCT,6 which was fine-tuned on multi-round conversational data as well as long-context summarization and QA data. We allow the model to generate 1, 024 tokens for the summarization tasks and 50 tokens for the question answering tasks.
6https://huggingface.co/togethercomputer/ Llama-2-7B-32K-Instruct
Results. Table 4 shows that CEPED improves upon LLAMA-2-CHAT with 2K tokens across all tasks. The performance of CEPED improves or remains consistent as we scale up the number of tokens in the context window. Notably, on NQA—the only dataset with a significant number of samples longer than 32K tokens (see Table 11)—CEPED improves upon LLAMA-2-CHAT by 3 points in F1 scores. Furthermore, CEPED outperforms LLAMA-2-32K INSTRUCT on 4 out of the 6 tasks, despite being trained on unlabeled data. We provide qualitative examples and analysis in §C.3.

heading: 7 Ablation Studies
text: We conduct comprehensive ablations to show the effectiveness of our training data mixture, pretraining and fine-tuning the encoder, and the warmup training stage. We also ablate to verify the effectiveness of the KL divergence loss in CEPED. Lastly, we evaluate on Needle in the Haystack (Kamradt, 2024), and analyze the results in §D.

heading: 7.1 Training Settings
text: Training with retrieved documents. Even though CEPE was trained with long documents, it achieves strong performance on retrievalaugmented applications. We also test a different
data strategy: training CEPE with retrieved documents using the training settings from §2.3.
As shown in Table 5, we find that training on retrieved documents (RetDoc) results in slightly stronger performances in the retrievalaugmentation setting, but perform worse on long documents. Augmenting a pre-training corpus with retrieval contexts is extremely computationally and storage expensive. CEPE’s simple long-document data strategy achieves a good balance between efficient training and strong performance on both long-context and retrieval-augmented applications.
Choices of unlabeled data. In Table 5, we show the results of training CEPE with only the fil-
tered documents from the ArXiv and Books domains (RPtrain-filter), and only the concatenated RP documents (RPtrain-cat). We find that training on RPtrain-filter is more beneficial for the longdocument setting and training on RPtrain-cat is better for the retrieval setting, but using a mixture of both leads to a more balanced and generalizable model. Our findings corroborate with the recent work on long-context data engineering (Fu et al., 2024). We also ablate different training strategies for the encoder, and find both the warmup stage and fine-tuning are crucial for strong performance. More details on the training ablations are in §E.1.

heading: 7.2 KL Divergence
text: The key component of CEPED is the KL Divergence loss. To understand the importance of this auxiliary loss, we explore the performance of CEPED when trained without the KL Divergence loss as well as with difference coefficients for each loss. Results are shown in Table 14, and we find that the KL Divergence loss is crucial for summarization tasks and QALT. More details are in §E.2.

heading: 8 Related Work
text: Long-context language models. Many recent works on long-context LMs aim to solve the problem of positional embedding extrapolation in transformers (Peng et al., 2024; Chen et al., 2023). Others fine-tune LMs on longer sequences (Xiong et al., 2023; Chen et al., 2024; Together, 2023a) or compress the context into shorter forms (Yoshida et al., 2020; Choromanski et al., 2021; Chevalier et al., 2023). Notably, several recent papers propose to extend the context window of LMs by modifying the attention mechanism: Xiao et al. (2024b) discover the use of “sink tokens” in sliding windows, and Bertsch et al. (2023); Xiao et al. (2024a) retrieve relevant tokens from a cache instead of attending to all tokens. This results in memoryefficient long-context LMs, but they have diminishing returns with longer contexts, as the same positional embedding may be seen multiple times and they can not fully utilize all tokens. The key advantage of CEPE is that it does not degrade for inputs longer than the training length while achieving great efficiency compared to full fine-tuning approaches. Techniques have been designed for specific applications, such as for in-context learning (Ratner et al., 2023; Hao et al., 2022), but we focus on general long-context language modeling.
Novel architectures and pre-training methods, such as S4 (Gu et al., 2022), RPT (Rubin and Berant, 2023), YOCO (Sun et al., 2024), and Mamba (Gu and Dao, 2023), also extend the context window at greater efficiency. However, pre-training is extremely expensive at scale and these methods cannot leverage existing powerful pre-trained LLMs. It is also unclear if state-space models can achieve comparable performance with transformers (Jelassi et al., 2024).
Retrieval-augmented language models. Augmenting LMs with retrieval has been useful in many applications, such as open-domain question answering. Recently, combining LMs with retrieval systems for more generalized purposes has been explored: Guu et al. (2020); Borgeaud et al. (2022); Izacard et al. (2022b); Min et al. (2023) pre-train LMs with retrieval, and Shi et al. (2024); Lin et al. (2023) use logits interpolation from separate forward passes to incorporate retrieval information.
Our architecture is similar to Atlas (Izacard et al., 2022b), RETRO, and RETRO-Fitting (Borgeaud et al., 2022) , but they use retrieval-augmented data for pre-training, which can be expensive to acquire at a large scale. CEPE only requires fine-tuning on long document data, which are much more efficient to obtain; CEPE is also applicable to any decoderonly LM, allowing us to extend context windows for pre-existing strong models. Finally, these works do not consider long-context language modeling settings in addition to retrieval-augmented settings.

heading: 9 Conclusion
text: We propose CEPE to extend the context window of existing language models. The key idea is to leverage a small encoder and cross-attention modules to process long inputs and achieve low memory and computational complexity. Compared to existing methods, CEPE extrapolates to input lengths well beyond the training length, while remaining efficient and effective. Consequently, CEPE augments pre-trained models to be performant on both long-context and retrieval-augmented applications. We also show that CEPED can be applied to instruction-tuned models with additional contexts using an auxiliary loss with only unlabeled data. We believe that there is still room for improvement through better data to train flexible and robust models. We hope our work can be a useful and accessible tool for the community to study long-context models in diverse applications.
Limitations
One limitation of our work is the focus on LLAMA2-7B. We hope that future work can investigate the applicability of our framework to a wider variety of LLMs of different sizes. Similarly, we only applied CEPED to LLAMA-2-CHAT-7B, but we look forward to other researchers applying it to other instruction-tuned or fine-tuned models.
We also acknowledge that certain hyperparameters are not studied in depth due to training costs – such as the ratio between RPtrain-filter and RPtrain-cat, learning rate, and the size of the small encoder model. We also fixed Contriever (Izacard et al., 2022a) to be the retriever in this work, but studying a greater range of retrievers would be useful.
Ethics Statement
LLMs are known to potentially output harmful and/or offensive language, and the LLAMA-2based models we use in this work are no exceptions. Since these models are trained on internetsize corpora (e.g., RedPajama), it can be difficult and expensive to filter out such offensive language.
Our models are also fine-tuned on RedPajama, which means they may also generate undesirable language. Although addressing this issue in the large-scale pre-training corpus is out of the scope of this work, we hope that future work will carefully resolve possible misuse issues in these models.

heading: Acknowledgements
text: We want to acknowledge Sadhika Malladi, Alexander Wettig, Mengzhou Xia, Dan Friedman, Amanda Bertsch, Adithya Bhaskar, Victoria Graf, and other members of the Princeton NLP group for their useful feedback and discussion. Tianyu Gao is supported by an IBM PhD Fellowship. This work is also gratefully supported by an NSF CAREER Award (IIS-2239290) and Intel.

heading: A Training Details
text: A.1 Pre-training Encoder
The encoder follows the configuration of RoBERTalarge (Liu et al., 2019b) – it has 24 layers, a hidden size of 1024, and 16 attention heads. However, we use the architecture of LLAMA-2, which means that the vocabulary size is different and the attention module contains an additional output projection. We refer to Liu et al. (2019b) and Touvron et al. (2023a) for more details.
We pre-trained the encoder for 100K steps on RP using the masked language modeling objective (Devlin et al., 2019). We used a batch size of 2048 sequences, where each sequence consisted of 512 tokens. The learning rate was set to 10−3 with a warm-up of the first 4% of the steps. We used eight A6000 GPUs with a gradient accumulation of 16. Furthermore, we employed a masking rate of 30% and disabled the next sentence prediction objective. We always replace the token with the [MASK] token if it is masked instead of replacing it with a random token or the original token. Finally, we used the AdamW optimizer (Loshchilov and Hutter, 2019) with β1 = 0.9, β2 = 0.999, and ϵ = 10−8, as implemented by the HuggingFace Transformers library (Wolf et al., 2020).
A.2 Training CEPE
The attention module in LLAMA-2 consists of four projection matrices: key, value, query, and output. In contrast to original transformers (Vaswani et al., 2017), the output projection matrix is used as an additional attention output projection. When we first insert the cross-attention layers into the decoder, we initialize the weights of the key, value, and query projection matrices with the respective weights from the decoder’s self-attention layer in the same transformer block. Furthermore, since the hidden dimension of the encoder is smaller than the hidden dimension of the decoder, d < D, we use only copy the first d rows of the key and value projection matrices from the self-attention module to the cross-attention module. Lastly, the output projection matrix is initialized with all zeros. While we did not investigate this initialization in detail, the intuition is that we want the model should use the tokens from the encoder using a similar mechanism as the decoder uses for its own tokens. However, we want the model to learn the output projection from scratch, as it may be too disruptive to have doubled the number of attention modules.
Then, we employ a warmup initialization method that simply trains the model to copy the input tokens from the encoder to the decoder. Specifically, we use the same inputs X = C for both Menc and Mdec, and X consists of n = 256 tokens. However, for the encoder, we chunk X into k = 4 sequences of 64 tokens to construct C. This step was trained for 4K steps with a batch size of 128 and peak learning rate of 5× 10−4, which totals to 131M tokens. We noticed that the model quickly learned to copy the input tokens from the encoder to the decoder, and the loss was close to zero after just 1K steps. The intuition behind this initialization strategy was to instill a strong inductive bias between the encoder input and decoder outputs. From our early experiments, we found that this initialization strategy helped stabilize the later training.
Finally, we train CEPE for 20K steps with a batch size of 128. We use eight A100 80 GB GPUs with a per-device batch size of 2 and gradient accumulation of 8, which took approximately 750 GPU hours. We also use a peak learning rate of 3×10−4 with a warm-up of 4% of the steps and a cosine learning rate schedule. We use the AdamW optimizer with β1 = 0.9, β2 = 0.999, and ϵ = 10−8.
During the standard training, we also use masking to inject noise into the parallel encodings. Intuitively, additional contexts may not always be in chunks of exactly 256 tokens in practice; for example, in open-domain question answering, the retrieved passages may vary in length. Furthermore, we may not always use exactly k = 16 contexts during inference, so we may want to train CEPE with instances with different k. To this end, for each encoder context Ci ∈ C, we apply masking with a probability of 0.3. When masking, we mask out the entire context Ci with probability 0.1. With probability 0.9, we mask out the suffix tokens x|Ci|−t+1, . . . , x|Ci|, where x1, . . . , x|Ci| are the tokens in Ci and t ∼ U(1, |Ci|). We choose to only mask out suffixes to maintain a natural distribution for encoder inputs. While we did not study the masking rate extensively, we found in preliminary experiments that masking did not hurt perplexity while improving performance on downstream tasks. We leave further explorations on masking in the encoder input for future work.
A.3 Training CEPED We leverage an additional distillation loss for CEPED. For encoder input C and decoder input X ,
we first calculate the logits Mdec(concat(C, X)) by running forward passes with the original model parameterized by Mdec. Due to storage constraints, we only store the top 50 likelihoods and their indices in the vocabulary for each token in X , following (Askell et al., 2021; Bai et al., 2022).
Then, during training, we define the distillation loss as the KL Divergence between the teacher model’s probability distribution and the student model’s probability distribution for the previously stored top 50 tokens. Concretely, our distillation loss is defined as follows:
LKL = DKL(Mdec(S)||MCEPE(C, X))
where S = concat(C, X), Mdec(S) is the probability distribution of the top 50 tokens for X , and MCEPE(C, X) takes C as the encoder input and X as the decoder input and outputs the probability distribution of the same 50 tokens on X .
Although Bai et al. (2022) also uses an additional category that represents the sum of all other tokens’ probabilities, we found that this may cause the KL Divergence to be undefined when the sum of other probabilities is 0. For our main model, we use a coefficient of 2 in front of LKL when adding to the cross-entropy loss to calculate the total loss. We experiment with this coefficient in §7.1.

heading: B Baseline Implementations
text: REPLUG. Although REPLUG (Shi et al., 2024) was introduced as a method to augment language models with retrieval, we found that the technique of interpolating logits from separate forward passes can also transfer well to the long context setting. Among the methods that we compare to, REPLUG uniquely improves performance upon the base model in both the long-context and the retrievalaugmented LM settings. This gives us an additional point of comparison across the two settings.
Following the original authors, we use Contriever (Izacard et al., 2022a) to calculate the scores for each previous context by using the first 256 tokens following the previous context as the query in the long-context setting. We did not include the additional memory and inference time costs of calculating the Contriever scores in our evaluation.
STREAMINGLLM. We follow the implementation of STREAMINGLLM from the original authors7
7https://github.com/mit-han-lab/ streaming-llm
(Xiao et al., 2024b). Specifically, we use their best settings, where we enable the positional shifts and cache 4 sink tokens and 2044 recent tokens.
The original code evaluates the model using a stride of 1 token at a time, where the cache is updated after every token, but this is not feasible for our large-scale evaluation. Therefore, we use a stride of 2048 tokens, and we update the cache after each stride. We show the difference in performance between the two settings in Table 6, and we found that STREAMINGLLM benefits from using a larger stride. We leave future exploration in this direction to future work.

heading: C Evaluation Settings
text: C.1 Open-domain Question Answering
The full results for the open-domain question answering experiments are shown in Table 7. REPLUG only uses up to k = 30 passages due to memory constraints, and LLAMA-2 has a window size of 4096, which limits k to 20. While LLAMA2-32K can use more than k = 60 passages with a context size of 32K, we only use up to 60 passages due to the cost of generation. For each demonstration, we only show the top 1 retrieved passages instead of the top k passages.
C.2 In-context Learning
For our in-context learning experiments, we use the datasets commonly used in previous works (Zhao et al., 2021; Lu et al., 2022; Han et al., 2023; Ratner et al., 2023): SST-2 (Socher et al., 2013), MovieReview (MR Pang and Lee, 2005), AGNews (Zhang et al., 2015), SST-5 (Socher et al., 2013), TREC (Voorhees and Tice, 2000), DBPedia (Zhang et al., 2015), NLU (Liu et al., 2019a), BANKING77 (Casanueva et al., 2020), CLINIC150 (Larson et al., 2019). We follow the prompts used in Ratner et al. (2023) for all datasets. During evaluation, we first calculate the log-likelihood of each option and select the option with the highest likelihood. We sample the in-context learning demonstrations from the training set such that each label has an equal number of demonstrations (except for possible remainders).
Furthermore, we first calculate the accuracy for each dataset using four different metrics: likelihood, likelihood normalized for length, calibrated likelihood, and calibrated likelihood normalized for length. We calibrate using Domain Conditional PMI (Holtzman et al., 2021), but use the empty
k NQ TQA PopQA
LLAMA-2
1 28.37 56.44 27.17 5 31.91 61.08 33.83
10 32.27 62.09 34.67 15 31.19 61.35 33.67 20 30.39 60.31 31.35
LLAMA-2-32K 10 30.64 56.00 32.38 15 31.27 56.98 33.48 20 31.97 57.28 33.75 30 30.66 57.57 33.90 60 30.58 57.03 34.61
REPLUG
5 31.27 61.21 32.40 10 31.52 61.35 32.31 15 30.80 60.89 31.62 20 30.30 60.41 31.11 30 29.78 59.99 30.27
CEPE
10 32.27 62.09 34.67 15 33.30 62.30 34.67 20 33.85 62.26 34.83 30 33.91 62.33 34.85 60 34.07 62.26 34.98
Table 7: Open-domain QA results. We report exact match scores for the Natural Questions(NQ) test set, TriviaQA(TQA) validation set, and PopQA test set. All models use two-shot in-context learning. k is the number of retrieved passages, and CEPE uses the top 10 passages in the decoder and all passages in the encoder.
string as the domain string for all datasets for simplicity. We then choose the metrics that yield the highest score for the LLAMA-2 model in the twoshot setting and apply the same metrics to all other models. The metrics used for each dataset are shown in Table 8. In this work, we did not investigate how to best calibrate CEPE in ICL settings. We leave these explorations for future work.
C.3 ZeroSCROLLS We use a subset of the ZeroSCROLLS (Shaham et al., 2023) with large validation sets: NarrativeQA (Kočiský et al., 2018), Qasper (Dasigi et al., 2021), QuALITY (Pang et al., 2022), GovReport (Huang et al., 2021), SummScreenFD (Chen et al., 2022), and QMSum (Zhong et al., 2021). Specifically,
we use the validation sets of these datasets made available by SCROLLS (Shaham et al., 2022).
However, we follow the same evaluation setup as ZeroSCROLLS, where models are evaluated in the zero-shot setting. We also use the same evaluation metrics as ZeroSCROLLS for each dataset. For the question answering datasets (NarrativeQA, Qasper, and QuALITY), we allow the model to generate up to 50 tokens, and we use greedy decoding.
For the summarization datasets (GovReport, SummScreenFD, and QMSum), we allow the model to generate up to 1, 024 tokens, following the original authors. For SummScreenFD and QMSum, we use greedy decoding, and for GovReport we use nucleus sampling (Holtzman et al., 2020) with a temperature of 1.0 and top-p of 0.95 and a minimum generation length of 10 tokens.
This is because GovReport has a much longer gold summary than the other datasets, and sampling methods are typically used in long-generation settings. Furthermore, greedy decoding may degenerate. The minimum generation length helps prevent trivial outputs, such as empty strings. To account for the randomness in the sampling method, we averaged GovReport performance over 3 seeded runs, and we found that the standard deviation is less than 0.20 ROUGE-L scores in all settings.
Furthermore, we hypothesize that LLAMA-2-
32K INSTRUCT may overfit to specific domains such as GovReport due to being trained on BookSum, a summarization dataset that also has long gold summaries (Kryscinski et al., 2022).
We also show some generation examples in Table 9 and 10. We find that CEPED can especially benefit from the additional contexts in the encoder in the QA datasets, where the answer may be localized to just one small part of the entire input. On the other hand, summarization tasks do not catastrophically fail when the model only has access to only part of the input, as the model can still generate a coherent summary for the provided context, achieving reasonable ROUGE-L scores.

heading: D Needle in the Haystack
text: Needle in the Haystack (Kamradt, 2024) is a synthetic task designed to test in-context retrieval abilities of long-context language models. In this task, a needle (e.g., “The best thing to do in San Francisco is to eat a sandwich and sit in Dolores Park on a sunny day”) is placed in a long text of essays, and the model is tasked with retrieving the needle from the context—that is, the model is expected to answer “What is the best thing to do in San Francisco?” with the needle.
We evaluate CEPE-LLAMA-2, where 2K tokens are used in the decoder, and the rest of the context is used in the encoder. For this experiment, we also scale the cross-attention scores with the ratio between the number of tokens during inference and the number of tokens during training, following Peng et al. (2024). We follow the evaluation setting of Fu et al. (2024), but split the context by sentence boundaries when inputting additional tokens to the encoder. We find that while CEPELLAMA-2 is able to perfectly retrieve the needle in context at the training sequence length (~10K tokens) and at some longer sequences (~14K tokens),
it struggles at other lengths. This is likely due to the training/inference discrepancies on the input length and we will explore training objectives that can mitigate such discrepancies in the future. Other efficient methods, such as StreamingLLM (Xiao et al., 2024b), also struggle with copying tokens that are far from the current context.

heading: E Ablations
text: We show the full results for the ablation studies in Table 12 and 13, and 14. The subsections below describes the ablation settings in more details.
E.1 Training Settings
Training with retrieved documents. In this subsection, we study training CEPE with retrieved documents. Specifically, we pair the training sequences described in §2.2 with retrieved passages from the retrieval split of RP using Contriever. We use the first 256 tokens of the decoder input X as the query and retrieve k = 16 passages to form the additional contexts C. Then, we train CEPE with retrieved passages (RetDoc) using the same training settings as in §2.3.
Encoder training. We also investigate how to best train the encoder. To this end, we train CEPE with (1) freezing the encoder after pre-training and the warmup stage, (2) training with a randomly initialized encoder, and (3) using the pre-trained model without the warmup stage. As shown in Table 5, we find that the copying warmup stage and fine-tuning the encoder during training are both crucial for strong performance.
E.2 KL Divergence The key component of CEPED is the KL Divergence loss. To understand the importance of this auxiliary loss, we explore the performance of CEPED with changes to the loss function in this subsection. Let LCE ,LKL be the cross entropy loss and the KL Divergence loss, respectively. Then, the total loss is L = cCELCE + cKLLKL, where cCE and cKL are coefficients for the cross entropy loss and the KL Divergence loss, respectively. We vary the coefficients cCE and cKL to study the importance of the KL Divergence loss. The results are presented in Table 14. Without the KL Divergence loss, CEPED may still perform well on NarrativeQA and Qasper, but the performance on the summarization tasks and QuALITY may decrease as a result.
Encoder Input C1: We propose a novel pre-training method called BRLM, which can effectively alleviates the distance between different source language spaces. Our proposed approach significantly improves zero-shot translation performance, consistently surpassing pivoting and multilingual approaches. Meanwhile, the performance on supervised translation direction remains the same level or even better when using our method. Related Work In recent years, zero-shot translation in NMT has attracted widespread attention in academic research. Existing methods are mainly divided into four categories: pivot-based method, transfer learning, multilingual NMT, and unsupervised NMT. Pivot-based Method is a common strategy to obtain a source→target model by introducing a pivot language. This approach is further divided into pivoting and pivot-synthetic. While the former firstly translates a source language into the pivot language which is later translated to the target language BIBREF4, BIBREF5, BIBREF12, the latter trains a source→target model with pseudo
Encoder Input C2: , NMT heavily relies on large-scale parallel data, resulting in poor performance on low-resource or zero-resource language pairs BIBREF3. Translation between these low-resource languages (e.g., Arabic→Spanish) is usually accomplished with pivoting through a rich-resource language (such as English), i.e., Arabic (source) sentence is translated to English (pivot) first which is later translated to Spanish (target) BIBREF4, BIBREF5. However, the pivot-based method requires doubled decoding time and suffers from the propagation of translation errors. One common alternative to avoid pivoting in NMT is transfer learning BIBREF6, BIBREF7, BIBREF8, BIBREF9 which leverages a high-resource pivot→target model (parent) to initialize a low-resource source→target model (child) that is further optimized with a small amount of available parallel data. Although this approach has achieved success in some low-resource language pairs, it still performs very poorly in extremely low-resource or zero-resource translation scenario. Specifically, BIBREF8 reports that without any child model training data,
Encoder Inputs [C3, . . . , C17] Omitted...
Decoder Input X: tokens are selected to be masked. Among the selected tokens, 80% of them are replaced with [MASK] token, 10% are replaced with a random BPE token, and 10% unchanged. The prediction accuracy of masked words is used as a stopping criterion in the pre-training stage. Besides, we use fastalign tool BIBREF34 to extract word alignments for BRLM-HA. Experiments ::: Main Results Table TABREF19 and TABREF26 report zero-shot results on Europarl and Multi-UN evaluation sets, respectively. We compare our approaches with related approaches of pivoting, multilingual NMT (MNMT) BIBREF19, and cross-lingual transfer without pretraining BIBREF16. The results show that our approaches consistently outperform other approaches across languages and datasets, especially surpass pivoting, which is a strong baseline in the zero-shot scenario that multilingual NMT systems often fail to beat BIBREF19, BIBREF20, BIBREF23. Pivoting translates source to pivot then to target in two steps, causing inefficient translation process. Our approaches use one encoder-decoder model to translate between any zero-shot directions, which is more efficient than pivoting. Regarding the comparison between transfer approaches, our cross-lingual pretraining based transfer outperforms transfer method that does not use pretraining by a large margin. Experiments ::: Main Results ::: Results on Europarl Dataset. Regarding comparison between the baselines in table TABREF19, we find that pivoting is the strongest baseline that has significant advantage over other two baselines. Cross-lingual transfer for languages without shared vocabularies BIBREF16 manifests the worst performance because of not using source↔pivot parallel data, which is utilized as beneficial supervised signal for the other two baselines. Additional Decoder Input Omitted... You are given a scientific article and a question. Answer the question as concisely as you can, using a single phrase or sentence if possible. If the question cannot be answered based on the information in the article, write "unanswerable". If the question is a yes/no question, answer "yes", "no", or "unanswerable". Question: what are the pivot-based baselines? Answer:
Model Outputs: LLAMA-2-CHAT output: unanswerable. CEPE output with encoder contexts: pivot-based baselines include pivoting and pivot-synthetic. Gold answers: pivoting, pivotingm
Table 9: ZeroSCROLLS generation example on the Qasper dataset. CEPE sees the entire article through the decoder and the encoder, whereas LLAMA-2-CHAT only sees a 2K token window. For brevity, we omit part of the decoder input and only show 2 out of k = 17 encoder inputs for CEPE.
Encoder Input C1: Phoebe: Almost sunrise. Do you think you’re ready to try the window again? Prue: Yeah, yeah, but Abraxas will be ready for us here. We have to take him by surprise, go where we’re most powerful, where we’re most connected. [Cut to the park. Prue, Piper and Phoebe have joined hands around a stone.] Prue, Piper and Phoebe: "Hear now the words of the witches, the secrets we hid in the night, the oldest of Gods are invoked here, the great work of magic is sought." [Cut to Abraxas undoing the spell that gave them their powers.] [Cut back to the girls.] Prue, Piper, Phoebe: "In this night, and in this hour we call upon the ancient power." [Cut back to Abraxas. The pages of the Book of Shadows turn over. He turns back to the spell and continues to read it backwards.] [Cut back to the girls.] Prue, Piper, Phoebe: "Hear
Encoder Input C2: hoebe: Honey, we’ve done a lot of good as witches too. You know that. But that doesn’t mean that bad things still aren’t gonna happen. But just because we can’t help that, doesn’t mean that is our fault. Piper: Or yours. Phoebe: Come here. (They all hug.) Grams’ Voice: The Power of Three. (They look around.) Phoebe: Okay, did you guys hear that? Piper: Mm-hmm. Prue: Grams? Grams’ Voice: The Power of Three. Prue: How is that possible? Phoebe: I don’t know but it’s gotta to be a good thing if we’re all hearing it. Piper: Maybe it means we’re connected. Phoebe: To be as one. Maybe that’s what the Book meant by the Power of One. Piper: The three of us working as one. Prue: Which is what the triquetra represents. (Phoebe looks at her watch)
Encoder Inputs [C3, . . . , C44] Omitted...
Decoder Input X: the secrets we hid in the night, the oldest of Gods are invoked here, the great work of magic is sought." [Cut back to Abraxas. The Book disappears.] [Cut back to the girls. The Book drops out from the sky and lands on the rock.] Phoebe: Prue, Piper... Prue: Keep chanting. Prue, Piper, Phoebe: Hear now the words of the witches, the secrets we hid in the night..." (Abraxas appears above them) Abraxas: What are you doing? How’d you bring me here? Prue, Piper, Phoebe: "The oldest of Gods are invoked here, the great work of magic is sought. In this night, and in this hour we call upon the ancient power. In this night, and in this hour we call upon the ancient power. In this night, and in this hour we call upon the ancient power." (Abraxas is vanquished. Phoebe grabs the Book.) Additional Decoder Input Omitted...
You are given a script of a TV episode. Summarize the episode in a paragraph.
Summary:
Model Outputs: LLAMA-2-CHAT output (ROUGE-L = 19.1): Prue, Piper, and Phoebe find an old Book of Shadows in the attic of their home, which they believe to be the original Book of Shadows. They start to read it, but Abraxas, the demon who wrote the book, appears to them, warning them to stop reading the book because he has returned to Earth and is seeking the power in the book to take over. The three sisters are able to defeat Abraxas and close the book, but the experience causes them to reflect on their lives as witches and their bond as sisters.. CEPE output with encoder contexts (ROUGE-L = 19.3): In this episode of Charmed, Prue is worried about the future because Andy is not coming back and they don’t know what their new life will be like. Prue and Piper have a hard time coming to terms with the fact that their powers are no longer needed and they are losing their sense of purpose and identity. At the same time, they are struggling to cope with their feelings for each other and their relationship. Phoebe, meanwhile, has found a new source of power and is learning to control her abilities in a new way. As the three sisters work together to defeat an evil spirit that is threatening their home, they are able to come to terms with their powers and their relationship and find a new sense of purpose and identity. Gold answers: The sisters face a demon named Abraxas , who steals the Book of Shadows and undoes some of the spells they’ve cast (in reverse order), releasing old vanquished enemies of the Charmed Ones and forcing them to remember the vanquishing spells from memory alone. The sisters have one chance to recapture the Book of Shadows or they will lose their powers forever. The sisters also meet their new neighbors, Jenny and her uncle Dan Gordon . Phoebe and Piper learn through the Wiccan community that because their anniversary of activating their inherent powers falls on an equinox , a wiccan holy day , each of their powers will be more developed and greatly magnified but only temporarily.
Table 10: ZeroSCROLLS generation example on the SummScreenFD dataset. CEPE sees the entire TV script through the decoder and the encoder, whereas LLAMA-2-CHAT only sees a 2K token window. For brevity, we omit part of the decoder input and only show 2 out of k = 44 encoder inputs for CEPE.

heading: Future Work
text: 3): In this episode of Charmed, Prue is worried about the future because Andy is not coming back and they don’t know what their new life will be like. Prue and Piper have a hard time coming to terms with the fact that their powers are no longer needed and they are losing their sense of purpose and identity. At the same time, they are struggling to cope with their feelings for each other and their relationship. Phoebe, meanwhile, has found a new source of power and is learning to control her abilities in a new way. As the three sisters work together to defeat an evil spirit that is threatening their home, they are able to come to terms with their powers and their relationship and find a new sense of purpose and identity. Gold answers: The sisters face a demon named Abraxas , who steals the Book of Shadows and undoes some of the spells they’ve cast (in reverse order), releasing old vanquished enemies of the Charmed Ones and forcing them to remember the vanquishing spells from memory alone. The sisters have one chance to recapture the Book of Shadows or they will lose their powers forever. The sisters also meet their new neighbors, Jenny and her uncle Dan Gordon. Phoebe and Piper learn through the Wiccan community that because their anniversary of activating their inherent powers falls on an equinox , a wiccan holy day , each of their powers will be more developed and greatly magnified but only temporarily. Table 10: ZeroSCROLLS generation example on the SummScreenFD dataset. CEPE sees the entire TV script through the decoder and the encoder, whereas LLAMA-2-CHAT only sees a 2K token window. For brevity, we omit part of the decoder input and only show 2 out of k = 44 encoder inputs for CEPE