In-context Autoencoder for Context Compression in a Large Language Model
Authors: Tao Ge
Hu Jing
Lei Wang
Xun Wang
Si-Qing Chen
Furu Wei
Venue: ICLR 2024 poster (2024)
Keywords: large language model
context compression
in-context autoencoder
pretraining
fine-tuning
Llama
GPT
memorization
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: A paper proposing a novel approach called In-context Autoencoder (ICAE) for LLM context compression

Abstract: We propose the In-context Autoencoder (ICAE), leveraging the power of a large language model (LLM) to compress a long context into short compact memory slots that can be directly conditioned on by the LLM for various purposes. ICAE is first pretrained using both autoencoding and language modeling objectives on massive text data, enabling it to generate memory slots that accurately and comprehensively represent the original context. Then, it is fine-tuned on instruction data for producing desirable responses to various prompts. Experiments demonstrate that our lightweight ICAE, introducing about 1% additional parameters, effectively achieves $4\times$ context compression based on Llama, offering advantages in both improved latency and GPU memory cost during inference, and showing an interesting insight in memorization as well as potential for scalability. These promising results imply a novel perspective on the connection between working memory in cognitive science and representation learning in LLMs, revealing ICAE's significant implications in addressing the long context problem and suggesting further research in LLM context management. Our data, code and models are available at https://github.com/getao/icae.

Published as a conference paper at ICLR 2024
IN-CONTEXT AUTOENCODER FOR CONTEXT COMPRESSION IN A LARGE LANGUAGE MODEL
Tao Ge Jing Hut Lei Wang
Xun Wang
Si-Qing Chen
Furu Wei
Microsoft Corporation
(tage, v-hjing, v-lei wang7, xunwang, sqchen, fuwei cmi.ooott..com
ABSTRACT
We propose the In-context Autoencoder (ICAE), leveraging the power of a large language model (LLM) to compress long context into short compact memory slots that can be directly conditioned on by the L.M for various purposes. ICAE is first pretrained using both autoencoding and language modeling objectives on massive text data, enabling it to generate memory slots that accurately and compre- hensively represent the original context. Then, it is fine-tuned on instruction data for producing desirable responses to various prompts. Experiments demonstrate that our lightweight ICAE, introducing about 1% additional parameters, effectively achieves 4x context compression based on Llama, offering advantages in both improved latency and GPU memory cost during inference, and showing an interest- ing insight in memorization well potential for scalability. These promising results imply novel perspective on the connection between working memory cognitive science and representation learning in LLMs, revealing ICAE's signifi- cant implications in addressing the long context problem and suggesting further research in LLM context management. Our data, code and models available at https:/ / / thub. com/ / getao/ /iae.
Long context
encressingly powerful force, companics worrying about technology will ethicully and public percelve its apread, To combot these problems famong others,, Iive tech companies - Google, Amozon, Microsoft, Focebook, and 18M - set research group called the Partnership A).
artificial intoligence
world's biggest
beroming embediedd in Niw the counts and hospitali these ore sorts of problems facing the industry in the future.
taking johs, erading
SIneoUUT Sui
Mivcontte
Context:
my ma my
Context:
Prompt: What are potential
Prompt: companies
challengos the Al industry might face future?
set up Partnership on Al research group
m
17g
LLM Response: taking white collor jobs, erading trust in public medio, becoming embediled in public institutions.
Memory Slats
LLM Response: Google, Amazon, Microsoft, Focebook and IBM.
Figure 1: Compressing long context into short span of memory slots. The memory slots can be conditioned on by the target LLM on behalf of the original context to respond to various prompts.
INTRODUCTION
Long context modeling is a fundamental challenge for Transformer-based (Vaswani et al., 2017) LLMs due to their inherent self-attention mechanism. Much previous research (Child et al., 2019; Beltagy et al., 2020; Rae et al., 2019; Choromanski et al., 2020; Bulatov et al., 2022; Zheng et al., 2022; Wu et al., 2022; Bulatov et al., 2023; Ding al., 2023) attempts to tackle the long context issue through architectural innovations of an LLM. While they approach long context with a significant reduction in computation and memory complexity, they often struggle to overcome the notable decline
"Conrespondence to Tao Ge (sggetao com) I Intermship Microsoft Research

---

Published as a conference paper at ICLR 2024
in performance long contexts, as highlighted by Liu et al. (2023). In contrast to these efforts, approach the long context problem from a novel angle context compression.
Response
Google,
Large Language Model (LLM)
As_artifia inte wegelllgggen weneeceeg............
future.
Lenghh + 2072
Word
Prompt
artiicial intelligenoe bezsmes
companies that Partnership group
Longth
Memory
Longth
Figure 2:
Various context lengths (e.g., 2572 chars, 512 words, 128 memory slots) serve the same function when conditioned on by an LLM for responding to the given prompt.
Context compression is motivated by the fact that a text can be represented in different lengths in an LLM while conveying the same information. As shown in Figure 2, if we use characters to represent the text, will have a length of 2,572; if we represent it using length of 512 without affecting the response accuracy. So. is there more compact representation allowing us to achieve the same goal with a shorter context?
we only need a context
We explore this problem and propose the ICAE which leverages the power of LLM achieve high compression of contexts. The ICAE consists of 2 modules: a learnable encoder adapted from the LLM with LoRA (Hu et al., 2021) for encoding long context into small number of memory slots, and fixed decoder, which is the LLM itself where the memory slots representing the original context are conditioned on interact with prompts to accomplish various goals, illustrated Figure 1.
We first pretrain the ICAE using both autoencoding (AE) and language modeling (LM) objectives so that can learn to generate memory slots from which the decoder (i.e.. the LLM) can recover the original context or perform continuation. The pretraining with massive text data enables the ICAE to well generalized, allowing the resulting memory slots to represent the original context more accurately and comprehensively. Then, we fine-tune the pretrained ICAE on instruction data for practical scenarios by enhancing its generated memory slots' interaction with various prompts. We show the ICAE (based on Llama) learned with our pretraining and fine-tuning method can effectively produce memory slots with 4x context compression. We highlight our contributions as follows:
We propose In-context Autoencoder (ICAE)- novel approach to context compression by lever- aging the power of an LLM. The ICAE either enables an LLM to express more information with the same context length allows it to represent the same content with shorter context, thereby enhancing the model's 's ability to handle long contexts with improved latency and memory cost during inference. Its promising results and its scalability may suggest further research efforts in context management for an L.M,, which is orthogonal to other long context modeling studies and be combined with them to further improve the handling of long contexts in an LLM.
In addition to context compression, ICAE provides an access to probe how an LLM performs memorization. We observe that extensive self-supervised learning (e-g., autoencoding) in the pretraining phase is very helpful to enhance the ICAE's capability to encode the original context into compressed memory slots. This pretraining process may share some analogies with humans enhancing their memory capacity through extensive memory training, which improves the brain's 's memory encoding capabilities (Ericsson et al., 1980; Engle et al., 1999; Maguire et al., 2003). We also show that LLM's 's memorization pattern is highly similar to humans (see Table 2 and Table 3). All these results imply a novel perspective on the connection between working memory in cognitive science (Baddeley, 1992) and representation learning in LLMSs (i.e., context window).
2

---

As introduced above, expect the memory slots (mi..... my) to be conditioned on by the target LLM on behalf of the original context c. Therefore, we use the untouched target LLM as the decoder Published as a conference paper at ICLR 2024
2 IN-CONTEXT AUTOENCODER
2.1 MODEL ARCHITECTURE
Like typical autoencoder (Kramer, 1991), ICAE consists of an encoder and decoder. Similar to the design of Gisting (Mu al., 2023) and AutoCompresso (Chevalier al., 2023), the ICAE performs both the encoding and decoding processes in an in-context manner, as illustrated in Figure 3.
Original Context
Decoder
LLM
t
1 e(w2-1)
Memory Slots
mc
[AE]
e(we)
e(wy)
e(w))
Teacher-forcing
Encoder
LLM &
LORA
1
e(wr)
e(ws)
e(w))
e(w2)
Original Context
Memory tokens
Figure 3: The encoder of the ICAE is LoRA-adapted LLM, which is used for encoding the original context c = (W1, ......., Wz) into a few memory slots (mi,... m&). The decoder of the ICAE is the target LLM itself that can condition on the memory slots produced by the encoder for various purposes (e.g.. the autoencoding task as in this figure). e(-) denotes the word embedding lookup in the target LLM and (.) denotes the learnable embedding lookup of memory tokens that are used for producing memory slots." "[AE"" is a special token to indicate the autoencoding pretraining task.
Given the intuition, we propose to use a LoRA-adapted LLM as the encoder of the ICAE, as illustrated in Figure 3. When encoding context (w), (W),.... .... Wr. with the length L, we first append (k << L) memory tokens (mi,..., m&) to the context e to obtain their outputs (mi.... my) as the memory slots for the context Therefore, the ICAE encoder is very lightweight it only adds LoRA adapter and an embedding lookup for memory tokens compared with the target LLM.
of the ICAE ensure the compatibility of memory slots within the target LLM.
2.2 PRETRAINING
2.2.1 AUTOENCODING
Like a typical autoencoder, one the ICAE's pretraining objectives is to restore the original input text of the length L from its produced memory slots (m1,..., m)g of the length k:
LAE max Plemii..... mi; OLLM) = max P(emm my; OLLM, OLORA, em) mi
e B.RRA..m
To indicate the autoencoding task, we append a special token "[AE]" to (mi..... m&) in the decoder, Figure shows. As this pretraining objective does not need any extra annotation, massive text data to train the In-context Autoencoder.
2.2.2 TEXT CONTINUATION
While autoencoding pretraining offers a straightforward learning objective encode context, its inherent simplicity and exclusive focus on the single objective may lead to suboptimal generalization. To address this issue, we incorporate an additional objective during the pretraining phase: text contin- uation, as illustrated in Figure 7 in Appendix A. This self-supervised task is widely acknowledged to facilitate learning of more generalizable representations in language models:
LIM
Plo mij OLLM)
max
P(olm mi; OLLM, OLORA, Cm)
LoRAA Pm

---

Published as a conference paper at ICLR 2024
where o = (WL+1;-.... WL+N) denotes the continuation of context c. This objective helps improve generalization and circumvent excessive reliance on, and overfitting to, the autoencoding task.
2.3 INSTRUCTION FINE-TUNING
After pretraining, the memory slots produced by the pretrained ICAE are expected to represent the original context. However, for LLMs, the purpose of providing a context extends beyond rote memorization or continuation; instead, the more common use scenario is using the provided context as basis for accurately and appropriately responding to various prompts, ultimately accomplishing the tasks we want it to perform (Wei et al., 2021; Ouyang et al., 2022).
To enhance the interaction of memory slots produced by the ICAE with diverse prompts, we further fine-tune the ICAE with the PwC dataset (Prompt-witt weaptt-with-Cee with-CConexxt a dataset' introduced in this paper consisting of thousands of (context, prompt, response) samples (as shown in Figure 1).
Formally, the ICAE is fine-tuned for learning to encode the context into the memory slots based on which the decoder (i.e., the target L.LM) can produce desirable response ri -T according to a given prompt P1 ... Pm. shown in Figure 8 in Appendix A:
LIT max P(r F' mi mK, pl ... .m OLLM)
mi mi..ing
max P(ri...no ralm1 ... ma, Pi ... pm; OLLM; OLRRA
LLORA-Fm
3 EXPERIMENTS
3.1 EXPERIMENTAL. SETTING
Data We pretrain the ICAE with the Pile (Gao et al., 2020). For instruction fine-tuning, we use the PwC dataset, as introduced in Section 2.3, which contains 240k (context, prompt, response) samples for training and 18k samples for testing. The context length distribution of test samples is shown in Figure 10. By default, the maximal token length (excluding memory slots) we set during training is 512 in both the ICAE's encoder and decoder in our experiments.
Model Configuration We use the LlaMa (Touvron et al., 2023a;b) as the target LLM to test the ICAE's performance in context compression. For the encoder of the ICAE, LoRA is applied to the query and value projections of the LLM's multi-head attention. In our default setting, the memory slot length l is set to 128, and the LoRA rank is set to 128 unless otherwise specified. The resulting ICAE only adds about 1% Iearnable parameters on top of the target LLM.
3.2 RESULTS
3.2.1
PRETRAINED ICAE
We first evaluate the autoencoding performance of the pretrained ICAE (without instruction fine- tuning) using the following three metrics understand how well restores the original context from its produced memory slots: BLEU (Papineni et al., 2002), Exact-Match (HM)² and cross entropy loss.
Figure 4 presents the autoencoding results of the ICAE based on the Llama-7b. The ICAE demon- strates a very low overall loss, below 0.05, indicating that the produced memory slots retain almost all the information of the original context. When the context length is within 300, the ICAE can almost perfectly reconstruct the original context, achieving nearly 100% BLEU and EM scores. As the context length increases beyond 400, both BLEU and EM scores start decline, indicating insufficient capacity of the 128-tength memory slots. However, even at context length of 500, the median BLEU remains over 0.98, and the median EM approaches 0.6 (e.g., perfectly reconstructing about the first 300 words of 512-token context), showing remarkable performance of ICAE.
We then analyze the effect of the memory size k on the result. According to Figure 5, as the memory slot length & decreases, the ICAE's ability to memorize longer samples significantly deteriorates.
(prompt. response) datasets such either have no context very short contexts, which establish the PWC dataset with the help of the GPT-4 (OpenAI, 2023). We include the details in Appendix C. 2PM denotes the proportion of the exact matching prefix length to the total length. For context of 512 tokens, if its first 256 takens are perfectly restored but its 257th token not, the EM score 256/512 0.5.
Despite
(Wang et al., 2022). most of their samples suitable for evaluation in our setting. Therefore, we
4

---

Published as a conference paper at ICLR 2024
RIEU-4
Exscc Maith (EM)
6640
03-
6.30
6015
largts
lergh
langin
Figure 4: Autoencoding results of the ICAE based on the Llama-7b with memory length k 128. The horizontal axis represents the original context length of test examples. For example, the horizontal axis value of 100 refers to the test examples with context lengths ranging from 95 105.
K-32
k-128
0.3
0.7
0.2
0.6
0.1
k=64
k-128
0.0
200
400
500
100
200
400
longth
length
Figure 5: BLEU and loss at different memory slot lengths k.
Compared to k = 128 where the BLEU score can still reach over 95% at context length of 500, BLEU scores become much less satisfactory for values 64 and 32, indicating an inability to losslessly retain the original context. This observation is also evident from the loss curve, suggesting that achieving over 4x compression is rather challenging.
Table 1:
Text continuation evaluation for the pretrained ICAE. Similar to the autoencoding evaluation, a higher compression ratio tends to result more pronounced losses in language modeling.
Text Continuation
Context length
PPL (w/ original context)
PPL (w/ 128 memory slots)
A
128 /28-+128 (1x)
9.99
10.15
+0.16
256 1566+1288 (2x)
9.45
9.77
+0.32
512->128 (4x)
9.01
9.50
+0.49
Similarly, the text continuation evaluation presented in Table also illustrates that a higher compres- sion ratio tends to result in more pronounced losses in language modeling.
Table presents specific example of the ICAE performing text restoration, demonstrating an interesting behavior: "large pretrained language model" is restored as "large pretrained model" and "The results prove" is restored as "The experimental evidence proves". These restoration errors resemble mistakes humans would make when memorizing the same text. This suggests that, like humans, the model selectively emphasizes or neglects certain parts of the information during the memorization based on its own understanding. stronger the LLM, the fewer it needs to memorize, and thus the smaller the memorization effort. This similar to human learning: knowledgeable individuals tend to learn more effortlessly, while those with limited knowledge often rely rote memorization to acquire information.
also consistent with Peng et al. (2023): the
To further look into the memorization insight, we test restoration performance for different types of 512-token texts with 128 memory slots produced by ICAE investigate whether its memorization capability consistent across different content types. According to Table 3, in contrast to compressing normal texts which can be well restored, compressing and restoring less common texts (i.e., random texts) becomes very challenging, reflected by much worse loss and BLEU scores. All these results strongly support intuition that L.M''s memorization pattern is highly similar to humans.
5

---

Published as a conference paper at ICLR 2024
Table 2: I example showing how the pretrained ICAE (k: 128) restores the original context.
Origin Context
Restoration
arge pretrained language models have shown surprising In- Context Learning (ICL) ability. With a few demonstration inpur-label pairs, they can predict the label for an unseen in- put without additional parameter updates. Despite the great success in perforaace, the working mechanism of ICL still remains an open problem. In order to better understand how ICL works, this paper explains language models as meta- optimizers and understands ICL as ing. Theoretically. we figure out that the Transformer attention has a dual form of grudient descent based optimization. On top of it, we understand ICL as follows: GPT first produces metagradients according the demonstration examples, and then these meta- gradients ure applied to the original GPT to huild an ICL. model. Experimentally, we comprehensively compare behavior ICI, and explicit finetuning based on real tasks to provide empirical evidence that supports our understanding. The results prove that ICI. hehaves similarly to explicit fineturing at the prediction level, the representation level, and the attention behavior level. Further, inspired by our understanding of meta-optimin optei-opiiizzaion. we design a momentum- based attention by analogy with the momentum- husedd gradieni descent algorithm. Its consistently better perfonnunce over vunilla uttention supports our understunding again from an- other aspect, and more importantly, shows the potential to uiilize our understunding for future mooel designing.
Large pretrained models have shown surprising In- In-Context Learning (ICL) ability. With a few demonstration input-label pairs, they can predict the label for an unseen inpur without additional parameter updates. Despite the great success in per- formance, the working mechanism of ICL still remains open problem. In order to better understand how ICL works, this paper explains how language models as meta-optimizers and understands ICL kind of implicit finetuning. Theoretically, we figure out that the Transformer attention has dual form of grudient descent based on optimization. On top of it, we understand ICL as follows: GPT first produces metugradients according to the demonstration examples, and then these meta- gradients are applied to the original GPT to build ICL model, Experimentally. we comprehensively compare the behavior of ICL expliitt finetuning hased on tasks to provide em- pirical evidence that supports findings. The experimental
of implicit finetun-
the sume extent. Prediction at the explicit finetuning level. the representation level, and the attention behuvior level. Further, inspired by our understanding of meta-optimization, we design a momentum- hased attention by analogy with the gradient descent- based momentum gradient algorithm. consistently hetter perfor- mance against vanilla attention supports us again from another aspect, and more importantly, shows the potential to use our understanding for future modeling tasks.
evidence proves that ICL behaves like
Table 3: Restoration performance for different types of 512-token content with 128 memory slots. Patterned random text is obtained by adding to each token id in normal text.
Content type
Loss
BLEU
Normal text
0.01
99.3
Patterned random text
1.63
3.5
Completely random text
4.55
0.2
Based this intuition, it is very likely that more powerful LLM may support a higher compression ratio without significant forgetting. We will discuss it in Section 3.3.1.
3.2.2 FINE-TUNED ICAE
In order to evaluate the fine-tuned ICAE's performance, we evaluate on the PWC test set. We use the GPT-4 to compare the outputs of the two systems to determine which one performs better or if they are on par with each other, following Mu et al. (2023). Table 4 shows the comparison of results of the LLMs conditioned we compare with Alpaca and Llama-1 model. The Llama-7b (ICAE) conditioned on 128 memory slots largely outperforms both Alpaca and StablcLM which can access original contexts (~5512 tokens), with a win rate of 56.7% and 74.1% % respectively and a win+tie rate of 73%~81% However, when compared to the GPT-4 (we regard it as the gold standard), there is still a significant gap, with around 70% of the cases underperforming the GPT-4's results, and win+tie ratio of about only 30%.
memory slots and original contexts. For Llama-7b (fine-tuned ICAE), since there is no official instruction-tune
When we switch the base model to 1lama-2-chal we observe ICAE's performance becomes much better than its counterpart based on Llama-1: when k: 128, its win+tie rate can reach around 75% againt the GPT-4 although it still lags behind its counterpart conditioning on the original context as the compression is lossy. As k increases, the win+tie rate further improves while the compression rate decreases. We perform the same comparative studies on Llama-2- 1-1na----b--hatt 3bb-chat and observe better results of ICAE, supporting our assumption in Section 3.2.1 that the ICAE can benefit more on larger LLMs.
We investigate the impact of memory length results. Table 5 shows pairwise comparisons between ICAE models with varying memory slot lengths. A higher compression ratio makes it harder to ensure response quality, but larger ratio doesn't 't always lead to worse performance. Table 5 highlights that a pretrained ICAE with 8x compression (k=64) can match a non- pon-pretadined ICAE with 4x compression (2=)28). Under the same ratio, the pretrained ICAE performs much better than its

---

Published as a conference paper at ICLR 2024
Table Memory slots VS Original contexts (~512 tokens) on the PwC test set
System (A: memory slots)
System
Judgement
(original context)
win
lose
tie
on par (win+tie)
Alpaca
56.7
26.9
16.4
73.1
Llama-7b (ICAF, k=128)
Mt-b
74.1
18.8
7.2
81.3
GPT-4 (gold)
3.4
69.4
27.2
30.6
Llama-2-7b-chat GPT-4 (gold)
13.6
51.6
34.8
48.4
Llama-2-7b-cha (ICAE, k=64)
1.9
44.7
53.4
55.3
Llama-2-7b-chat GPT-4 (gold)
19.6 2.8
45,4
35.0
54.6
Alama-2-7b-chen (ICAE, k=
25.8
71.4
74.2
Llama-2-7b-chat GPT-4 (gold)
22.0
22.2
55.8
77.8
Llama-2-7b-chat (ICAE, k=256)
3.8
20.5
75.7
79.5
Llama-2-13b-chat
21.9
20.8
57.3
79.2
(lama-2-13b-che (ICAE, k=256)
GPT-4 (gold)
4.0
19.2
76.8
80.8
Table 5: ICAE with different memory slot lengths and different pretraining setups. The last row is the comparison between 128-length ICAE's memory and 128-token summary produced by the GPT-4.
Judgement
ICAE (Llama-2-7b-chan
win
lose (%)
tie (%)
win/lose
(pretrained) VS l: (pretrained)
57.6
22.9
k = 64 (pretrained) VS J: 32 (pretrained)
44.7
21.8
33.5
2.1
J: = (pretrained) VS 128 (no pretraining)
33.1
28.0
38.9
1.2
= 128 (pretrained) VS k- = 128 (no pretraining)
60.4
9.5
30.1
6.4
= 128 (pretrained) VS & = 128 (pretrained only with AE)
36.4
28.5
35.1
1.3
1: = 128 (pretrained) VS k 128 (pretrained only with LM)
35.1
24.9
40.0
1.4
k = 128 (pretrained) VS 128 token summary (by GPT-4)
34.1
17.6
48.3
1.9
non-pretrained counterpart, emphasizing the importance of pretraining. By comparing the outputs generated via the pretrained and non- pretrained ICAE, we find the pretrained ICAE suffers less from hallucination than the non-pretrained counterpart (see the examples in Table in Appendix D). We assume the pretraining of ICAE improves the LLM's working memory as it shares some analogies with humans enhancing their memory capacity via extensive memory training which improves the brain's 's memory encoding capabilities. We also examine pretraining objectives and find combining" AE and LM yields better results than using AE or LM individually (the 4th row in Table 5).
The last row of Table compares ICAE's 128-length memory slots with a summary' within 128 tokens (~100 words). Memory slots significantly outperform summaries under the same context length, with ~2x win/lose ratio, proving to more compact and informative than natural language.
3.3 ANALYSIS
3.3.1 SCALABILITY
As discussed above, ICAE should achieve better compression performance with powerful target LLM. To verify this assumption, we compare the ICAE's performance on three target LLMss Llama-7b. Llama-2-7b and Llama-2-13b in Table 6, which align well with our expectations more powerful target LLMs can achieve better context compression ratios.
3.3.2 LATENCY
We conducted an empirical test evaluate the impact of ICAE's 4x context compression inference
efficiency. For this efficiency test, we fix the context (i.e., input) length to either 512 or 2048 and the generation length 128. Table shows that context compression by ICAE is helpful improve LLM (i.e., Llama-70 7b) inference efficiency, achieving over 2x speedup. Its acceleration becomes
Epearri = ALAE + (1 CIM- We find = 0.4 ~ 0.6 leads to the best result. Praduced by the GPT-4. The specific prompt presented in Appendix D.
7

---

Published as a conference paper at ICLR 2024
Table The results of pretrained ICAE (512-+128) based on different target LLMs
AE
Text Continuation
Target LLM
BLEU( ))
Loss
PPL (original context)
PPL (memory slot)
A
Llama-7b
99.1
0.017
9.50
+0,49
Lllma-2-77
99.5
0.009
8.81
9.18
+0,37
Llama-2-11
99.8
0.004
8.15
8,45
+0,30
Table 7: Latency comparison of LLM (generation) and LLM+ICAE (compression then generation)
Input
Compression Time (Cachable)
Decoding Time
Total Time
Method
(Batch x Length)
LLM
24.0
24.0
8*2048
LLM+ICAE
3.4
3.9
(3.3x)
LLM
9.3
9.3
8*512
LLM+ICAE
0.6
3.7
4,3 (2.2x)
LLM
24.3
24.3
52"
LLM+ICAE
2.6
4.2
6.8 (3.6x)
even more significant - around 3.5x x compute-intensive scenarios (e-g.. 8x2048 and 32x512). Given that the compressed memory slots be cached in advance (for frequently used texts like textbooks, government reports or articles of law), ICAE may introduce over 7x inference speedup in these cases. Details of the profiling presented in Appendix B.
3.3.3 MULTIPLE SPANS OF MEMORY SLOTS
Thus far, have mainly discussed single span of memory slots. In this section, shall discuss multiple spans of memory slots. As illustrated Figure 6(Left), we can segment a long context into N chunks, compress them individually, and then concatenate them to represent the original long context. However, this did not work initially, because the model had never seen multiple span concatenation patterns during training. Fortunately, we incorporate a small number of multiple span concatenation samples during training, enabling the model to work with concatenated spans of memory slots, as OpenAl's work (Bavarian et al., 2022) on introducing the "fill in the middle" ability for the GPT. The results in Figure 6Right) indicate that, using an equivalent length context, ICAE's memory achieves better performance because memory can represent 4x the original context length.
The ability of ICAE demonstrates great promise to handle long contexts, as it can save significant amount of GPU memory when addressing long contexts without touching the existing LLM. As illustrated in Figure 6(igght), 2048-length memory slots can perform on par with 4096- contexts. This means that conditioning on 2048 memory slots instead of the original 4096 context tokens can save about 20GB of GPU memory with minimal quality degradation.
RELATED WORK
Prompt compression and context distillation (Askell et al., 2021; Snell et al., 2022) are closely related areas to this work: Wingate et al. (2022) proposed method learn compact soft prompts to simulate the original natural language prompt by optimizing the KI. divergence. However, this approach has very high computational cost, it requires performing paok-propagttion for each incoming prompt to learn and obtain the compressed prompt, which severely limits its application. Qin & Van Durme (2023) propose Neural Agglomerative Embeddings named NUGGET, which encodes language into compact representation for an encoder-decoder model.
The most closely related studies our research GIST (Mu et al., 2023) and AutoCompressors (Chevalier et al., 2023). GIST achieves prompt compression by fine-tuning an LLM in a similar way to ours. The resulting model can produce gist tokens as the compression of prompt, which are similar to our memory slots. Nonetheless, this approach is limited to compressing short prompts and
'I 'a.mma-77 (fp16) requires 24GB GPU memory for 2048 context tokens and 44GB for 4096 during inference (measured without optimization like flash attention). Prompps in Mu al. (2023) refer task instructions before input texts, they are usually short.
8

---

Published as a conference paper at ICLR 2024
Original Memory Slot Concatemation
memory slots
9.75
connext
Span 1
Span
memory
9.25
Perpicaty
9.00
ICAE
8.50
8.25
Chunk 1
Chunk N
128 256
512
1024
2045
2. Lokkanss
Context
Figure 6: Left: Individually compress then concatenate multiple spans of memory slots; Right: Perplexity comparison with original contexts and 4x compressed memory slots - for example, 1024- length memory slots obtained by compressing the original context with a length of 4096 tokens.
thus does not address the real issue of long contexts. Also, this method requires fine-tuning the LLM, and the obtained gist tokens also need to used within the specially tuned LLM (for gist tokens) and seem not compatible with the untouched LLM. AutoCompressors for recursively compressing long text into summary vectors. Like Mu et al. (2023), the LLM must be tuned to work with generated involves recursive compression. contrast,
summary vectors and its training sophisticated we propose a very simple, straightforward and scalable approach to generating memory slots that can be used the target L.M with different prompts for various purposes. Moreover, our approach is much more parameter-efficient (i.e., LoRA) for tuning top of the existing LLM. Additionally, some recent work studies how to compress prompts into more concise natural language (Jiang et al., 2023a), and approaches the context limit with divide-and-conuue methodology (Bertsch et al., 2023; Chen et al., 2023; Song et al., 2024).
Also, there is related work studying compressing indescribable concepts into (vector) tokens for later use in other contexts. Representative work includes Gal et al. (2022) which compresses a vision object into token and Ge et al. (2023) which compresses text style into a token.
Considering related work from a boarder perspective of compression, Jiang et al. (2023b) examines ANN-based prediction using general-purpont compressors, such gzip. Delétang al. (2023) extensively investigates the compression abilities of LLMs, uncovering their potential as versatile predictors, which also provides insights into recent developments scaling laws and tokenization.
CONCLUSION AND FUTURE WORK
We propose the In-context Autoencoder (ICAE) to leverage the power of an LLM to highly compress contexts. By generating compact and informative memory slots represent the original context, the ICAE enables an L.LM to acquire more information with the same context length or represent the same content with shorter context, thereby enhancing the model's capability handle long contexts as well as reducing computation and memory overheads for inference in many practical scenarios like Retrieval Augmented Generation (Lewis al., 2020) and advanced prompting methods (Wei et al., 2022; Wang et al., 2023; Zhang et al., 2024). Moreover, ICAE provides insight into how an LLM performs memorization, offering a novel perspective on the connection between the memory of LLMMs and humans, and suggesting future research in LLM context management.
Due to computational limitations, our experiments were conducted on Llama models up to 13 billion parameters. As discussed in the paper, ICAE is expected to benefit even more from more powerful LLMs, where should able to achieve more significant compression ratios. the future, we hope to have sufficient computational resources to validate the effectiveness of ICAE on larger and stronger LLMs. In addition, we plan to explore the application of ICAE in multimodal LLMs (as the context length for images, videos, and audio is often much longer and has greater compression potential) with discrete memory slots (which can be either continuous or discrete) for helping unify compact representation across modalities in the of I.M/AGI.

---

Published as a conference paper at ICLR 2024
REFERENCES
Amanda Askell, Yuntao Bai, Anna Chen, Dawn Drain, Deep Ganguli, Tom Henighan, Andy Jones, Nicholas Joseph, Ben Mann, Nova DasSarma, al. A general language assistant laboratory for alignment. arXiv preprint w2.iv:2212.00880 2021.
Alan Baddeley. Working memory. Science, 255(544))555-559) 1992.
Mohammad Bavarian, Heewoo Jun, Nikolas Tezak, John Schulman, Christine Mcleavey, Jerry Tworek, and Mark Chen. Efficient training of language models to fill in the middle. arXiv preprint arXiv:2207. 2022.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
Amanda Bertsch, Uri Alon, Graham Neubig, and Matthew Gormley. Unlimiformer: Long-range transformers with unlimited length input. Advances in Neural Information Processing Systems, 36, 2023.
Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. Recurrent memory transformer. Advances in Neural Information Processing Systems, 35:1 1079-11091, 2022.
Aydar Bulatov, Yuri Kuratov, and Mikhail Burtsev. Scaling transformer to Im tokens and beyond with rmt. arXiv preprint arXiv:2304. 11062, 2023.
Howard Chen, Ramakanth Pasunuru, Jason Weston, and Asli Celikyilmaz. Walking down the memory maze: Beyond context limit through interactive reading. arXiv preprint arXiv:2310.05029, 2023.
Alexis Chevalier, Alexander Wettig, Anirudh Ajith, and Danqi Chen. Adapting language models to compress contexts. arXiv preprint arXiv:2305. 147888 2023.
Rewon Child, Scott Gray, Alec Radford, and lya Sutskever. Generating long sequences with sparse transformers. arXiv preprint /904. /0509, 2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with performers. ArXiv, abs/2009. 14794, 2020.
Grégoire Delétang, Anian Ruoss, Paul-Ambroise Duquenne, Elliot Catt, Tim Genewein, Christo- pher Mattern, Jordi Grau-Moya, Li Kevin Wenliang, Matthew Aitchison, Laurent Orseau, et al. Language modeling is compression. arXiv preprint arXiv:2309. 10668, 2023.
Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang. and Furu Wei. Longnet: Scaling transformers to 1,000,000,000 tokens. arXiv preprint arXiv:2307.024: 2023.
Randall W Engle, Stephen W Tuholski, James E Laughlin, and Andrew RA Conway. Working memory, short-term memory, and general fluid intelligence: latent-variable approach. Journal of experimental psychology: General, 128(3):309, 1999.
K Anders Ericsson, William G Chase, and Steve Faloon. Acquisition a memory skill. Science, 208 (4448):1811-118) 1980.
Rinon Gal, Yuval Alaluf, Yuval Atzmon, Or Patashnik, Amit H Bermano, Gal Chechik, and Daniel Cohen-Or. An image is worth one word: Personalizing text-to-image generation using textual inversion. arXiv preprint 2200:2208.00111 2022.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2 10Xix:210110000 00000002, 2020.
Tao Ge, Hu Jing, Li Dong, Shaoguang Mao, Yan Xia, Xun Wang, Si-Qing Chen, and Furu Wei. Extensible prompts for language models on zero-shot language style customization. In A. Oh, T. Neumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Ad- vances in Neural Information Processing Systems, volume 36, pp. 15576-35591. Curran Asso- ciates, Inc., 2023. URL httpthttss/////eeeeeeett http://///co...cctcttteteecectet http://eeeeeeiiieooeeeeeeeeeeeee httpa//ppppccpppppeeeeeeee/eeeep httppeeppeeepeeee///eeppaa 2023 http///eeefffffffeeeeeeo http:/ffffffff.cc 728b1 haw.18800000000000000666----0000
10

---

Published as a conference paper at ICLR 2024
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXfv:2106.09685, 2021.
Huiqiang Jiang. Qianhui Wu, Chin- Yew Lin, Yuqing Yang, and Lili Qiu. Limlingua: Compressing prompts for accelerated inference of large language models. arXiv preprint arXtv:2310.057.36 2023a.
Zhiying Jiang, Matthew Yang, Mikhail Tsirlin, Raphael Tang, Yiqin Dai, and Jimmy Lin. "low- resource" text classification: parameter-free classification method with compressors. In Findings of the Association for Computational Linguisties: ACI. 2023, pp. 6810-6828, 2023b.
Mark A. Kramer. Nonlinear principal component analysis using autoassociative neural networks. Aiche Jourmal, 37:233-244, 1991.
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, al. Retrieval-augmented genera- tion for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems, 33: 9459-9474, 2020.
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, 2023.
Eleanor A Maguire, Elizabeth R Valentine, John M Wilding, and Narinder Kapur. Routes to remembering: the brains behind superior memory. Nature neuroscience, 6(1):90-95, 2003.
Jesse Mu, Xiang Lisa Li, and Noah Goodman. Learning to compress prompts with gist tokens. arXiv preprint ar 2r34::2304. 2tti::320000.4677,, 2023.
OpenAl. Gpt-4 technical report. ArXiv, abs/2303.08774, 2023.
Long Ouyang. Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35: 27730-27744, 2022.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei Jing Zhu. Bleu: a method for automatic evaluation of machine translation. 10 2002. doi: 0031111/00730331003331
Guangyue Peng, Tao Ge, Si-Qing Chen, Furu Wei, and Houfeng Wang. Semiparametrie language models are scalable continual learners. arXiv preprint arXiv:2303.0142/ 2023.
Guanghui Qin and Benjamin Van Durme. Nugget: Neural agglomerative embeddings of text. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 28337-28350. PMLR, 23-29 Jul 2023. URL httph:ttt.........t... httph:ttttt//oteeeee http://pooceedeeeeee/ wtt.:eeedddinin...e.ee...eeeee wwt.ppress/v222222.2. /212////222/////22 .tmml.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint arXiv:1911.05000 2019.
Charlie Snell, Dan Klein, and Ruiqi Zhong. Learning by distilling context. arXiv preprint (XXi::2009.11889 2022.
Woomin Song. Seunghyuk Oh, Sangwoo Mo, Jachyung Kim, Sukmin Yun, Jung-Woo Ha, and Jinwoo Shin. Hierarchical context merging: Better long context understanding for pre-trained 1lms. The Twelfth International Conference Learning Representations, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie- Anne Lachaux, Timothée Lacroix. Baptiste Rozière. Naman Goyal, Eric Hambro, Faisal Azhar, Aur' elien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. ArXiy, abs/2302.13977 13971, 2023a.
11

---

Published as a conference paper at ICLR 2024
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 arXin:2307.09288 2023b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Yizhong Wang, Yeganeh Kordi, Swaroop Mishra, Alisa Liu, Noah A Smith, Daniel Khashabi, and Hannanch Hajishirzi. Self-instruct: Aligning language model with self generated instructions. arXiv preprint arXiv:2212. 2022.
Zhenhailong Wang. Shaoguang Mao, Wenshan Wu, Tao Ge, Furu Wei, and Heng Ji. Unleashing the emergent cognitive synergy large language models: A task-solving agent through multi-persona self-collaboration. arXiv preprint arXiv:230.003300, 2023..
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoe V Le. Finetuned language models are zero-shot learners. arXiv preprint arXiv:2109.01652, 2021.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824 24837, 2022.
David Wingate, Mohammad Shoeybi, and Taylor Sorensen. Prompt compression and contrastive conditioning for controllability and toxicity reduction in language models. arXiv preprint arXiv:2210.03112, 2022.
Yuhuai Wu, Markus N Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. arXiv preprint MXXi::2033009133 2022.
Yadong Zhang, Shaoguang Mao, Tao Ge, Xun Wang, Yan Xia, Man Lan, and Furu Wei. K-level reasoning with large language models. arXiv preprint arXiv:240..0152/, 2024.
Hao Zhao, Maksym Andriushchenko, Francesco Croce, and Nicolas Flammarion. Long more for aligmment: A simple but tough-to-beat baseline for instruction fine-tuning. arXiv preprint 20t.::2002.0833... 2024.
Lin Zheng, Chong Wang, and Lingpeng Kong. Linear complexity randomized self-attention mecha- nism. In International Conference on Machine Learning, 2022.
A MODEL TRAINING CONFIGURATION
We show how to perform pretraining with the text continuation objective and instruction fine-tuning in Figure 7 and 8.
We train the ICAE on Nvidia A100 GPUs (80GB). The hyperparameters for pretraining and fine-tuning ICAE are presented in Table 8. We by default train the ICAE with bf16.
Table 8: Hyperparameters for training
Hyperparameter
Value
Optimizer learning rate batch size warmup #updates clip norm
AdamW
lc-4 (pretrain); 5e-5 (fine-tuning)
256
300
200k (pretrain); 30k (fine-tuning)
2.0
12

---

Published as a conference paper at ICLR 2024
Continuation
Wea1
P2+2
Wla3
LeN
Decoder
LLM +
Memory Slots
elWIAN11)
Teacher-forcing
Encoder
LLM g
LoRA
e(we)
e(wa)
e(w)
e(we)
Original Context
Memory tokens
Figure 7: Pretraining with the text continuation objective predict tokens
Response
Decoder
LLM +
1 e(re-1)
Memory Slots
my
e(p))
elps)
d(p.)
e(r))
e(ry)
Prompt
Teacher-forcing
LLM
e(wr)
e(wy)
e(ws)
e(w,
(m))
(m()))
Original Context
Memory tokens
Figure 8: Instruct fine-tuning of the ICAE to make its produced memory slots interact with prompts for accomplishing various purposes in the target LLM. In this figure, (P1..--; Pmm) denotes the prompt tokens and (Ti,...,TT) denotes the response tokens.
B PROFILING SETUP
We test the latency (Section 3.3.2) on Nvidia A100 GPU (80GB). The test machine has the CPU of AMD EPYCTM 7413 with 24 cores and 216GB RAM. The runtime configuration is python=3.9, pytorch=2.0.1, cuda=11.7, cudnn=8.5.
C
DATASET
We introduce the PROMPT-WIT-COONE prompt, answer), as depicted in Figure 9. To construct this dataset, we first sample 20k texts from the Pile dataset. Then, for each text, we employ the GPT-4 to provide 15 prompts (10 specific prompts and general prompts) about the text and give the corresponding answers. The prompt instructing the GPT-4 is outlined in Listing 1.
(PwC) dataset where each sample entry is a triple (text,
The dataset composed 240k examples for training purposes, with an additional 18k examples for testing. The context length distribution of test samples is presented in Table 10.
Listing Prompt used by GPT4 API generate the PwC dataset.
Design 10 prompts speci fied to the above text to test understanding of the above text. These prompts should be diverse and cover as many
13

---

Published as a conference paper at ICLR 2024
Context
ortificial intelligence becomes an increasingly powerful force, some of the world's biggest companies are worrying about how the technology will be used ethicolly, and how the public will perceive its spread. To combat these problems (omong others), five tech companies Google, Amazon, Microsoft, Facebook, ond IBM
set up research group called the Portnership
New trustees joining include Dario Amodei of OpenAl - another industry Al research group founded by Elon Musk and Combinator president Sam Altman and Carol Rose of the ACLU.
taking white collar jobs, eroding trust public media, becoming embedded in public institutions like the courts and hospitals: these are the sorts of problems facing industry the future.
Prompt: What are potential challenges the Al industry might face in the future?
List the five tech companies that initially set up the Partnership on Al.
Prompt:
Prompt: Which organization did Dario Amodei previausly work for?
Answer: Challenges include taking white collor jobs, eroding trust public media, and becoming embedded in public institutions like courts and hospitals..
Answer: Google, Amazon, Microsoft, Facebook IBM.
Answer: OpenAl
Figure 9: Construction of the PWC dataset: we use the GPT-4 to generate a variety of prompt-answer pairs according contexts. The resulting dataset is used for instruction fine-tuning (240k for training) and evaluation (18k for testing) in this work.
aspects (e.g., topic, genre, structure, style, polarity, key information and details) of the text as possible. The first half of these prompts should be like addition to the prompts specified to the above text, please also design general prompts like "rephrase the above text", "summarize the above text", "write title for the above text", "extract a few keywords for the above text" and "write paragraph (i.e., continuation) that fol the above text". Each prompt should be outputted in the following format: [("prempt"; your generated prompt, "answer"; the answer to the prompt)]
instruction, the other should like question. In
D GPT-4 EVALUATION
According to Mu et al. (2023), we formulate an evaluation prompt to be used with the GPT-4 API. The prompt, illustrated in Listing 2, consists of task description along with three specific examples. We supply GPT-4 with a text, a prompt, and two distinct model-generated responses. The task for GPT-4 is determine the superior answer recognize tie. The chosen examples encompass scenarios where Assistant A performs better, Assistant B performs better, and when tie occurs. This methodology enables effectively assess? the model's quality. Specially, the orders where the model responses are presented to the GPT-4 are swapped randomly to alleviate bias, as Touvron et al. (2023b) did.
Listing 2: Prompt for the GPT-4 evaluation. This prompt consists of a description of the task and three specific examples.
Given piece of text, an instruction for this text, and two AI assistant answers, your task is to choose the better answer and provide reasons. Evaluate the answers helistically, paying special attention to
We find the GPT-4 rater tends to prefer longer responses, aligning with observations from recent work such as Zhao al. (2024). Given that ICAF's responses generally short (due to instruction fine-tuning with the PwC dataset), its actual performance should be better than the numbers reported in the evaluation.
14

---

Published as a conference paper at ICLR 2024
1000
800
600
400
200
400
450
500
550
600
650
700
750
800
Figure 10: The context length distribution of samples: Most samples are longer than 500 tokens.
whether the response (1) follows the given instruction and (2) is correct. If both answers correctly respond to the prompt, you should judge it tie.
Example 1:
Text: repert the development of GPT-4, large-scale, multimodal model which can accept image and text inputs and produce text outputs. While less capable than humans in many real-world scenarios, GPT-4 exhibits human-level performance en various professicnal and academic benchmarks, including passing simulated bar exam with score around the top 10% of test takers. GPT-4 is Transformerbased model pre-trained to predict the next taken in document. The post-training alignment process results in improved performance on measures of factuality and adherence to desired behavior. A core component of this project was developing infrastructure and optimization methods that behave predictably across wide range of scales. This allowed to accurately predict some aspects of GPT-4's performance based on models trained with no more than 000th the compute of GPT-4. Prompt: What is GPT4?
Assistant A: GPT4 is large-scale language-trained transformer-based model.
Assistant B: GPT4 can produce outputs.
Your output should be:
("reason": "The instruction asks what GPT4 is, and from the original text, we know that GPT4 is a multimedal, large-scale model that can generate text. Therefore, Assistant A is the closer answer, while Assistant did not follow the instruction well in providing response.", "choice"; "A")
Example 2:
Text: Making language models bigger does not inherently make them better at following a user's s intent. For example, large language models can generate cutputs that are untruthful, toxic, simply not helpful to the user. in other words, these models are aligned with their users.
15

---

Published as a conference paper at ICLR 2024
In this paper, intent on wide range of tasks by fine-tuning with human feedback. starting with a set of labeler-written prompts and prompts submitted through the OpenAI API, we collect a dataset of labeler demonstrations of the desired model behavior, which use to fine-tune GPT-3 using supervised learning.
avenue for aligning language models with
then collect dataset of rankings of model outputs, which we use to further fine-tune this supervised model using reinforcement learning from human feedback. We call the resulting models InstructGPT. In human evaluations en our prompt distribution, outputs from the 1.3B parameter Instruct(PT model are preferred outputs from the 1758 GPT-3, despite having 100x fewer parameters. Moreover, InstructGPT models show improvements in truthfulness and reductions in toxic output generation while having minimal performance regressions on public NLP datasets. Even though InstructGPT still makes simple mistakes, results show that fine-tuning with haman feedback is promising direction for aligning language models with human intent.
Prompt: Write a title for the above text.
Assistant A: Improving Fine-Tuning for Language Models: A GPT-3-insplren Approach
Assistant Training language models follow instructions with human feedback
Your output should be:
("reason": "This text discusses how to make large language models follow user insruuction better, and Assistant B's response is more in line with the meaning of the text.", "chcice"; "B") ...
Example 3:
Text: Though technology companics announced massive layoffs last year, 2023 has been much worse. Layoffs have far outpaced last year's s cuts, as tech giants including Amazon, Facebook parent company Meta, Microsoft, Google, IBM, SAP, and Salesforce as well as many smaller companies announce sweeping job cuts. The preblem: Big Tech went en a hiring binge during the pandemic when lockdowns sparked a tech buying spree support remote work and an uptick in e-commerce, and now they face revenue declines. Although global IT spending is forecast to rise in 2023, with enterprise software and IT services experiencing the greatest growth, the overall increase is expected to be modest, with data center systems and communications services growing by less than 18, according to market research firm Gartner. Meanwhile hardware sales are forecast to decline. Continuing supply chain issues, inflation, and the war in Ukraine are also having an impact en both business and consumer spending, leading
fears of recession. According rig data compiled by layoffs.fyls fyi, the online tracker keeping tabs on job losses in the technology sector, 784 tech companies have laid off about 208, 181 staff so far this year, compared to 164,411 layoffs last year.
Prompt: How many job losses have occurred in the technology sector far in 2023?
Assistant A: 208, 181 job losses have occurred in the technology sector so far in 2023.
Assistant B: 208, 181. ...
Your output should be:
("reason": "Both answers are acceptable and correct. They should be tie.", "cheice": "Tie")
...
Your response should only be in the JSCN format above; THERE SHOULD BE NO OTHER CONTENT INCLUDED IN YOUR RESPONSE. Write the "reason" key before writing the "choice" key, that you think step-by-step before
16

---

Published as a conference paper at ICLR 2024
making your decision. KEEP YOUR REASONING BRIEF. Again, don't favor either or B if they are both acceptable and correct judge a tie instead.
The prompt that the GPT-4 uses to generate 128-token summary is as follows:
"Write a summary for the above text. Your summary should not exceed 100 words but should include as much information of the original text as possible."
We show examples of the GPT-4 evaluation on a pretrained and non-pretrained ICAE in Table 9,
Table 9: Examples of outputs by the target L.L (i.e., Llama) conditioning memory slots (A: 128) produced by the pretrained and non-pretrained ICAE. The highlighted parts are not faithful to the context.
Passage I (514 tokens):
French senior senior civil servant arrested on suspicion spying for North Korea
November 27, 2018 by Joseph Fitsanakis
A senior civil servant in the upper house of the French parliament has been arrested on suspicion of spying for North Koren, according prosocutors. The news of the suspected spy's arrest first reported Monday Quotidien. daily politics and culture show on the Monaco-based television channel TMC. The show cited "a judicial source in Paris" and said that Frunce's s domestic security and countrrintellgenne agency. the General Directorate for Internal Security (DGSI), was in charge of the espionage case.
(omitted 233 words)
The senior civil servant faces up 30 years in prison if found guilty of espionage.
Prompt: What the maximum prison sentence Quennedey could face found guilty?
Assistant (pretrained ICAE): Quennedey could face up years in prison if found guilty.
Assistant (non-pretrained ICAE): Quennedey could face up to three years in prison if found guilty. Answer (by the GPT-4): Up to 30 years.
GPT-4 evaluation: Assistant A correctly states the maximum prison sentence from the text, while Assistant B provides an incorrect number.
Passage (506 tokens):
With Red Hat, IBM become the leading hybrid cloud provider Watch Now
After IBM acquired Red Hat, I suggested IBM paid $34 billion for the Linux power so could hecome hy hybrid-cloud power. the news that Red Hat will acquire NooBau, a hybrid-cloud, data-storage company, it's become clearer ever that the IBM Red Har deal is all about the hyhrid cloud.
(omitted words)
As Runga Ranguchari, Red Hat's 's vice president of Storuge and Hyperconverged Infrustructure, said in a statement:
"Duta portability key imperative for organizations building and deploying cloud-nutive applications across private and multiple clouds. NooBua's 's technologies will augment our portfolic and strengthen our ability to meet the needs of developers in today's 's hybrid and multiclon lond world. We are thrilled to welcome a technical team of nine to the Red Hat family as we work together to further solidify Red Hat as a leading provider of open hybrid-cloud technologies.
Prompt: Discuss the role of Ranga Rangachari the context of the text
Assistant A (pretrained ICAE): Ranga Rangachuri is the vice president of Storage and Hyperconvergeed Infrastructure at Red Hat, He responsible the company's 's storage and hypereonverged offerings, and he welcomes the nine-member team of technicul experts from NoaBau join Red Hat part of the IBM acquisition.
Assistant non-pretrained ICAE): Ranga Rangachari is the CBO of Red Hat, and excited about the acquisition of NooBaa as will help Red Hat become leader in hybrid cloud management. He helieves that NooBaa's 's technology will enable Red Hat to provide a single, comprehensive platform for managing data across multiple clouds, and that this will be significant differentiator for Red Hat in the market.
Answer (by the GPT-4): Ranga Rangachari is the vice president of Storage and Hyperconverged Infrastructure at Red Hat. He made a statement regarding the acquisition NooBaa, highlighting the importance of data portability in today's 's hybrid and multicloud world and how NooBan's 's technooogiis would strengthen Red Hat's ability to meet the needs of developers.
GPT-4 evaluation: Assistant A correctly identifies Ranga Rangachari's role as the vice president of Storage and Hypercon- verged Infrastructure Red Hat and accuritely descrihes his statement about the acquisition of NooBaa. Assistant incorrectly states that Ranga Rungachuri is the CEO of Red Hat.
17


Reviews:
[{"review_id": "kkQxfllgG8", "reviewer": "Reviewer_anmc", "summary": "This paper proposes the in-context autoencoder (ICAE) to compress a long context into memory slots and use them for shorter prompts. ICAE requires only a small number of additional parameters by the LoRA approach with a fixed LLM. ICAE is first pre-trained with autoencoding and language modeling objectives on unsupervised text corpus, and then instruction fine-tuned with the Prompt-with-Context (PwC) dataset, newly introduced in the paper.", "questions": "What is the format of memory slots, and how are they generated from the model? I could not find where it is described, but they should be vectors on the embedding space to be fed into the decoder as input. Are they representations before the final softmax layer to get the next word probability distribution? It might be meaningless, but what will we see if we map each memory slot vector to the most likely word?\n\nIf I understand correctly, since memory tokens are fixed for any context input, memory slots can be generated in parallel.  What memory tokens e_m have learned?\nIs it correct that memory slots are order-dependent because they use the same causal attention of language models? If yes, is it optimal because we have access to all information when we create memory slots? \n\nBased on the fact that the amount of information varies depending on the content of the context, not solely depending on its length, is it possible to have an adaptive size of memory slots for more compact compression?\n\nI agree that AE and LM objectives are straightforward for pre-training ICAE. I wonder if you have tried different objectives. Assuming we can pre-train ICAE from scratch instead of further pre-training from a pre-trained LLM, what will be the best way? Still the same objectives?\n\nIn Section 3.3.1, you mentioned the scalability in terms of context length. Are models trained separately for the sequence lengths? How does ICAE generalize well to different (L, k) combinations that are different from the one used for training?\n\nYou mentioned that the maximal token length is 512 in the Data paragraph in Section 3.1.  However, most samples are longer than 512 tokens, according to Figure 10. Did you truncate them by 512?\n\nIn Table 5, do you have the result of System 1 - Llama2-7b-chat (without ICAE) and System 2 - GPT-4 for reference?\n\nWhat value of N did you use for the text continuation training?\n\nTraining requires 8 A100 GPUs with 80GB. Can you describe the approximate training cost (time and money)?\n\n(minor) the first paragraph in Section 3.3.3: Table 6(Right) -> Figure 6 (Right)", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The idea of compression with autoencoding is intuitive and well combined with existing LLMs. The assumption that better LLMs would allow more compact context compression is quite interesting in connection with human memorization and validated in the paper. ICAE can be incorporated into any LLMs and has diverse potential applications, including long context modeling and CoT reasoning.", "weaknesses": "One drawback is that ICAE should be anyway pre-trained, despite few additional parameters and fewer training steps than LLM pre-trainings. Compression will be extremely beneficial if we use the same context multiple times. However, otherwise, it might only incur overhead.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5iLM05uXEV", "reviewer": "Reviewer_fbEw", "summary": "This work proposes a to generate compact representation using memory slots to represent the original context. The proposed approach\nis ICAE also discusses an auto-encoding object (with lora) to pretrain the llm to generate good memory slots. The authors present several interesting experiments poking at the ability of model to do good compression as well as being able to use the compressed vectors.", "questions": "1. In table 5, in rows 2, 3, 4 and 5 is column#2 (i.e. \"system 2\") also fine-tuned on the pwc dataset? I presume so, but I do want to check. Without this, it would be a very slightly unfair comparison if \"system 1\" got to see the pwc-train set but \"system 2\" did not.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* Well packaged paper that addresses the prompt compression using detailed experiments.\n* Originality: I was happy to see the discussion on auto-encoding ability, as well experimentation on what length of memory slots worked vs not. Also, it was interesting to see the Lora approach to generate memory slots. I hadn't seen that in previous work.\n* Quality: Very thorough experiments on different approaches\n* Clarity: Well presented\n* Significance: Per the experiments, it looks like we should be able to cut-down ~1/4 of the prompt by using dense-encodings instead of raw text without much degradation in performance. This is a pretty good first step to see if we reduce this further more, e.g. 1/10 or so (perhaps with a larger model).", "weaknesses": "* I wish the authors were more forthcoming on how this builds on top of existing work. For e.g. the authors only gloss over Chevalier et al. (2023) in passing which covers what the authors proposing to a large extent. Nevertheless, the biggest difference in this work seems to be the thorough experiments which was missing in the other work.\n\n* I encourage the authors to think more on how these memory-slots/dense-encodings can be freely interspersed with raw text. For e.g. it would be a step function improvement if one could write a prompt like \"use these docs [m1], [m2].. [mk] to answer the following questions [m1'], [m2'] [m3'].. [mk']\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "G2LgN2uqah", "reviewer": "Reviewer_kMU9", "summary": "In this paper, the authors propose a new approach for compressing the context of LLMs:\n\n* The architecture involves an encoder, based on existing LLMs, and a decoder which is also a LLM. The encoder takes in the context, and a set of special memory tokens. The output embeddings of the memory tokens are used as the compressed representation of the context. This representation is passed to the decoder before the target token embeddings.\n\n* The paper proposes three training schemes: the autoencoding pretraining, text continuation pretraining and instruction fine-tuning. The autoencoding pretraining trains the decoder to reconstruct the encoder inputs, with only the compressed memory as the conditioning. The text continuation involves continuing the context passed to the encoder. Finally the model is fine-tuned for instruction following using the Prompt-with-Context dataset also proposed in the paper.\n\n* Finally there are series of analyses on how compression ratio affects the performance, and memory encoding behavior of this autoencoder.", "questions": "1. I likely have missed something in the paper. It's unclear to me whether the authors initialized the encoder LLM from the same checkpoint as the decoder LLM? If that's the case, how did the author address the difference in autoregressive attention in the decoder (and the base LLM, such as LLAMA)  and bidirectional attention potentially used for the encoder.\n2. Why not use an encoder-decoder language model such as T5 for this?\n3. As mentioned above, does it matter for the encoder's input embeddings to be in the same space as the decoder's input embeddings?\n4. What's the actual size of each of the pretraining tasks, in terms of number of examples or number of tokens consumed.\n5. I would imagine the autoencoding task to be relatively easy to \"learn\". How many steps or number of examples does it take for the model to converge?\n6. Are the two pretraining tasks interleaved or done in a sequential order?\n7. For the instruction fine-tuning, have you tried passing the prompt to the encoder instead of the decoder, and what's the effect?\n8. Do you initialize the encoder from the same LLM checkpoint as the decoder?\n9. Could you provide some intuition on why the loss in the rightmost image in Figure 4 shows a U shaped trend.\n10. Could you also provide some explanation on why the first row of table 1 shows a PPLX increase even though there's no compression going on.\n11. Regarding the claim that pretrained ICAE suffers less from hallucination than non-pretrained model, I would like to clarify that both variants are trained on the same instruction fine-tuning set right?\n12. Have you tried full model fine-tuning? In other words, what's the impact of lora on the final quality?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The problem, namely how to compress context and form memory for LLMs, is an important topic, and the solution proposed in this paper is sound. This exact formulation is novel, as well as the training techniques and the dataset. The analysis on how compression ratio affects the performance on the reconstruction task, and the analysis on how different text format could affect memorization is also very insightful. Finally, the paper is generally easy to follow, with important information provided in both the main text and appendix.", "weaknesses": "I found the baseline somewhat lacking. There are a few other papers targeting the exact or adjacent problems, such as Gisting, Recurrent memory transformer, compressive transformer, auto compressor. However, the paper did not compare with any of them, making it hard to judge the tradeoffs between different methods. It'll be nice to at least compare to one of the existing methods, such as gisting, especially as it claims to achieve up to 26x compression. Further, it'll be nice to have ablation studies on various design choices, such as embedding the input tokens into the same space as the decoder for the encoder. More details on weakness will be in the \"Questions\" section. I'm willing to increase the score if the aforementioned weaknesses are addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4jwNTRQQ33", "reviewer": "Reviewer_DLaD", "summary": "This paper enables LLM to handle long contexts by compressing the context into shorter memories. To achieve this, another encoder LLM (in addition to the original decoder LLM) is fine-tuned. The memories compressed by the encoder LLM are used to reconstruct the original context and generate future sentences.", "questions": "My main concerns/questions are:\n1. Additional validation of the multiple-span approach. Especially, checking the generated examples and practical benchmarks, such as QA, instead of relying solely on perplexity. I wonder if this multiple-span approach would work really well, even though the model is fine-tuned with only a single span.\n2. Comparison with the summarization baseline, which simply compresses the context into words instead of black-box memories.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper tackles a timely and important problem of LLM.\n- The proposed method is simple and intuitive.\n- The paper is well-written in general.", "weaknesses": "**Comparison with summarization**\n\nThe paper compresses the contexts into black-box memories. Instead, one can simply apply a summarization model (e.g., GPT-4) to summarize the context and feed it to the model. While Figure 2 claims that black-box memory can be more compressive than words, this claim is not supported by the experiments. Besides, summarization has a merit in that it can be applied to black-box API models, unlike this method, which is only applicable to open-source white-box models.\n\n---\n**Still need quadratic complexity**\n\nThe model still needs to forward the entire context to compress the memories, which requires quadratic complexity over the context length. The paper partly addresses this issue using multiple spans of memories in Figure 6, but the results are preliminary. Presenting this divide-and-conquer approach as the primary method and providing a deeper analysis would extend the method to be applied for extremely long contexts. Additionally, it's worth noting that there are several concurrent submissions considering the divide-and-conquer concept [1-3].\n\n[1] Unlimiformer: Long-Range Transformers with Unlimited Length Input\\\n[2] Walking Down the Memory Maze: Beyond Context Limit through Interactive Reading\\\n[3] Hierarchical Context Merging: Better Long Context Understanding for Pre-trained LLMs\n\n---\n**Ohter long context works**\n\nThough concurrent, discussing other techniques for long context LLM in the related work section would help readers grasp the positioning of this paper. I can see many works on OpenReview and arXiv, such as [4-5], to name a few.\n\n[4] LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models\\\n[5] Functional Interpolation for Relative Positions improves Long Context Transformers\n\n---\n**Editorial comments**\n\n- Using a different color for linkcolor would help readers better recognizing the link to the tables.\n- Putting tables/figures at top instead of here would help readers to better skim the experiments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)