Better Prompt Compression Without Multi-Layer Perceptrons

Large language models (LLMs) display incredible usefulness across many natural language tasks, and generally have increased utility with increasingly long and complex prompts (Agarwal et al., 2024; Bertsch et al., 2024). The downside of lengthier prompts is increased computational load and response time, motivating research into compressing prompts into a smaller number of tokens, known as prompt compression.

While some methods focus on compressing prompts by pruning information in the prompt/text space (Li et al., 2023; Jiang et al., 2023a, b), one can also consider compressing prompts into a lower dimensional latent space (Wingate et al., 2022; Mu et al., 2023; Chevalier et al., 2023). The In-context Autoencoder (ICAE) (Ge et al., 2024) exemplifies this approach by training a LLM encoder to compress prompts into a shorter sequence of learned memory tokens and uses a learned [AE] autoencoder token for decoding the original prompt. This latent representation retains the information of the prompt and is used with the original frozen (meaning not further trained) LLM decoder to reduce the number of tokens at inference time. 500xCompressor (Li et al., 2024) works similarly, but compresses prompts into neural attention (Vaswani et al., 2017) key-value pairs instead of explicit tokens, and uses a pretrained [BOS] token instead of a learned [AE] token. Notably, both ICAE and 500xCompresson use Low-Rank Adaptation (LoRA) (Hu et al., 2021) to train encoders from the frozen decoder LLM used for inference, which requires more computational resources to perform compression than may be necessary.

We demonstrate that using the entire decoder LLM as an encoder is unnecessary and introduce an alternative in the Attention-Only Compressor (AOC). Instead of learning the encoder as a LoRA of the decoder LLM, we first remove the multi-layer perceptron (MLP) layers before training the entire encoder. By removing MLPs, AOC’s prompt compression encoder has roughly 67% less parameters compared to previous methods’ encoders, while improving or maintaining similar compression ability. These results emphasize that prompt compression encoders do not need identical architecture to their decoders and that there exist compression models that with higher performance and lower inference-time computational requirements compared to recent approaches using frozen-LLM-based compressors.

Our contributions can be summarized as follows:

• •

  We introduce the Attention-Only Compressor (AOC), a novel prompt compression encoder that removes the MLP layers from a LLM, resulting in an encoder that performs comparably to baseline compression encoders that are roughly three times larger.

• •

  Preliminary experimental results on regeneration demonstrate that compression encoders do not need architecture identical to their decoders, which motivates further research into more efficient compressors.

• •

  To further study compression encoders, we present examples of interpolating between the embeddings of two compressed prompts, showcasing a novel classifier-free approach to merging separate prompts and understanding the latent space of compressed prompts.

Model. Our proposed model consists of a learned prompt compression encoder $\mathbf{E}$ and a pretrained LLM decoder $\mathbf{D}$ that is always frozen throughout training and inference. The encoder is architecturally identical to the decoder as in 500xCompressor and ICAE, with the key exception that the MLP layers have been replaced with the identity operation within each block of the Transformer (Vaswani et al., 2017):

	$\displaystyle h_{\ell}$	$\displaystyle={\rm LN_{pre}}(h_{\ell-1})$	$\displaystyle h_{\ell}$	$\displaystyle={\rm LN_{pre}}(h_{\ell-1})$		(1)
	$\displaystyle h_{\ell}$	$\displaystyle={\rm MHA}(h_{\ell})+h_{\ell}$	$\displaystyle h_{\ell}$	$\displaystyle={\rm MHA}(h_{\ell})+h_{\ell}$		(2)
	$\displaystyle h_{\ell}$	$\displaystyle={\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{\rm MLP}({\rm LN_{post}}(h_{\ell}))}+h_{\ell}$	$\displaystyle h_{\ell}$	$\displaystyle={\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\rm LN_{post}}(h_{\ell})}+h_{\ell}$		(3)

$h_{\ell-1}$ denotes the input hidden state to the $\ell$th Transformer block, $\rm LN_{pre}$ and $\rm LN_{post}$ are layer norms (Jimmy Lei Ba and Hinton, 2016), and $\rm MHA$ denotes multi-headed attention (Vaswani et al., 2017).

Let the input for the encoder be represented by the concatenation of $n$ prompt tokens $\mathbf{X}_{n}=(x_{1},\dots,x_{n})$ with the encoder’s $m$ learned memory tokens $\mathbf{Y}_{m}=(y_{1},\dots,y_{m})$. $\mathbf{Z}=\mathbf{E}([\mathbf{X}_{n},\mathbf{Y}_{m}])$ is the latent representation from the encoder output. For 500xCompressor, $\mathbf{Z}=\{\mathbf{KV}(h_{\ell}^{\mathbf{Y}_{m}})\forall\ell\}$: the encoder’s per-layer attention key-value pairs corresponding to $\mathbf{Y}_{m}$. The input to the decoder is $\mathbf{Z}$ concatenated with a regeneration token [REGEN], which is used to regenerate $\mathbf{X}$ using the latent information from $\mathbf{E}$. For both 500xCompressor and AOC [REGEN] is the [BOS] token. Therefore, the regeneration of $\mathbf{X}_{n}$ from the latent representation $\mathbf{Z}$ is given by

$$\hat{\mathbf{X}}_{n}=\mathbf{D}\left(\left[\mathbf{Z},\mathbf{[REGEN]}\right]\right)=\mathbf{D}\left(\left[\mathbf{E}([\mathbf{X}_{n},\mathbf{Y}_{m}]),\mathbf{[BOS]}\right]\right)$$

(4)

The standard cross-entropy loss between the decoder logits and the input $\mathbf{X}$ is used to train the encoder via backpropagation (LeCun et al., 1989). For all experiments, we use Llama 3.2 1B Instruct (Llama Team, 2024) as the pretrained LLM in bfloat16 (Wang and Kanwar, 2019) precision, AdamW (Kingma and Ba, 2015) with a 300-step warmup to a learning rate $2\times 10^{-4}$ as the optimizer in PyTorch (Paszke et al., 2019) conducting training using Transformers (Wolf et al., 2020) on a single NVIDIA A6000 GPU. LoRAs are trained on the queries, keys, values, and output projections in the multi-headed attention components for 500xCompressor and LoRA ablations on AOC.

LoRA Ablations. Due to the lower number of total parameters in the encoder for AOC, we perform full training instead of LoRA to learn a strong prompt compressor. However, this causes the total number of parameters in memory at both training and inference time to be slightly larger with AOC compared to the baseline 500xCompressor which use a LoRA of the decoder LLM. This trade-off of increased memory for decreased compression time motivates ablations on learning the AOC encoder using LoRA (LoRA-AOC) instead of training the entire encoder.

Compressed Prompt Interpolation. The latent information $\mathbf{Z}$ from compressing a prompt has not been extensively studied beyond classifier-guided generation by Wingate et al. (2022). As an initial step toward better understanding the compressed forms of prompts, we conduct linear interpolations between compressed prompts and qualitatively inspect the intermediary output. The interpolation between $\mathbf{Z}_{0}$ and $\mathbf{Z}_{1}$ with a given weight $w$ is given by:

$$\mathbf{Z}_{\rm interp}=\mathbf{Z}_{0}+w(\mathbf{Z}_{1}-\mathbf{Z}_{0})$$

(5)

Data. Experiments are performed using random samples from the arXiv dataset (arXiv.org submitters, 2024). AOC is trained on 300,000 abstracts from the arXiv dataset first submitted before July 1, 2023 and validated on 3,000 abstracts first submitted after January 4, 2024. Final evaluations were conducted on a held-out test set of 3,000 abstracts from after January 4, 2024. These dates were chosen based on the Llama 3.2 training cutoff of December 2023, and are identical to the cutoffs presented in (Li et al., 2024). The amount of training data was determined while accounting for limited computational resources.

Metrics. We evaluate AOC on text regeneration as performed using Equation 4. Following (Ge et al., 2024), we report the Bilingual Evaluation Understudy (BLEU) (Papineni et al., 2002) and Exact-Match (EM) scores. Notably, the EM metric defined by (Ge et al., 2024) is the proportion of identical prefix length to total target length. Given a regenerated sequence of length $n^{\prime}$, this proportional EM metric is defined as:

$${\rm EM}(\mathbf{X}_{n},\hat{\mathbf{X}}_{n^{\prime}})=\frac{1}{n}\sum_{i=1}^{n}\mathbf{1}_{\mathbf{X}_{i}=\hat{\mathbf{X}}_{i}}(\mathbf{X}_{i},\hat{\mathbf{X}}_{i})$$

(6)

In contrast, the EM metric defined by (Li et al., 2024) is a binary metric equal to $1$ when the regeneration $\hat{\mathbf{X}}_{n^{\prime}}$ is identical to $\mathbf{X}_{n}$ and $0$ otherwise, introducing a discrepancy in notation. We report the EM metric as defined in Equation 6 since it is more informative. Additionally, we report the Recall-Oriented Understudy for Gisting Evaluation Longest Common Subsequence (ROUGE-L) (Lin and Och, 2004) F1 scores which evaluate overall sequence similarity, following (Li et al., 2024).

Baseline Comparison. To demonstrate the benefits of AOC, we compare to 500xCompressor with a variety of input prompt lengths $n\in\{96,192,288,384,480\}$ and number of memory tokens $m\in\{1,4,16\}$.

As seen in Table 1 and Table 2, AOC outperforms 500xCompressor across all prompt lengths with 4 or 16 memory tokens despite having 67% less encoder parameters. Based on the results in Table 1 we find that AOC and 500xCompressor only perform similarly when restricted to a single memory token. The large variance in EM between models can be attributed to differences in early parts of the regeneration, as the EM metric is based on the proportion of identical prefix matching. Interestingly, LoRA-AOC tends to perform worse than AOC and the baseline 500xCompressor across all metrics, which suggests that the effectiveness of LoRA in Transformers relies in part on the frozen MLPs, in line with prior work on freezing Transformer components Lu et al. (2022).

It can be seen in Table 3 that for $m=1$, AOC performs on-par with 500xCompressor, although both display poor regeneration abilities for some of the largest compression ratios in our experiments. Upon inspection of the loss curves from training the $m=1$ models in Table 3, we discover that they are likely under-trained due to computational budget constraints. Based on these results, it appears that increasing the amount of memory tokens $m$ may allow for a smaller training data set.

Latent Space Inspection. In Table 4 we show the result of linearly interpolating between the compressed information $\mathbf{Z}$ from the prompt $p_{0}=$"We present an awesome new idea." and the prompt $p_{1}=$"Large planets may have many moons." for AOC, color-coding by similarity to $p_{0}$ or $p_{1}$. As can be observed, the interpolation of the two latent representations results in a regenerated mixture of prompts, such as when the interpolation weight $w=0.5$ planet which is more closely related to planets from $p_{1}$ than idea from $p_{0}$. For $w=0.53$, many moons from $p_{1}$ appears in a regeneration that shares the same prefix as $p_{0}$. Similarly, for interpolation weights $w=0.55$ and $w=0.6$, amazing and wonderful, which are more closely related to awesome from $p_{0}$, appear in a regeneration almost identical to $p_{1}$ with the same two-word prefix. We also note that Table 4 shows both $p_{0}$ and $p_{1}$ were perfectly regenerated from their unaltered compressed states with zero information loss.

We introduce AOC, a prompt compression encoder using only attention layers from a decoder LLM that demonstrably achieves comparable or better compression to LoRA baselines with identical architecture to the decoder LLM. Experiments show that the memory tokens learned with AOC can encode similar amounts of information to baselines with $3\times$ the amount of parameters. In future work, we hope to further explore encoder architectures, as our results indicate that a prompt compression encoder need not have the same architecture as the decoder LLM. Additionally, we seek to better understand the latent space formed by compressed prompts and extend the use of compressed prompts beyond the interpolation example presented in this work. While this work was performed with limited computational resources, we aim to study more diverse and larger datasets, model architectures, and compression ratios in the future.