QUITO-X: An Information Bottleneck-based Compression Algorithm with Cross-Attention

The Attention Mechanism (Vaswani et al. 2017) is crucial in today’s machine learning. It is widely used in many fields including image generation (Rombach et al. 2021), image recognition (Caron et al. 2021), and language processing (Brown et al. 2020; Radford et al. 2019). It starts with an input sequence transformed into query (Q), key (K), and value (V) matrices through linear projections. And the attention matrix is computed as follows:

Among these, cross attention is often used as a way of exchanging information between two different types of information, such as between image and text modalities (Rombach et al. 2021; Cho et al. 2021) , or for translation between different languages. In practice, Q or query is generally designed based on the type of output that is desired from the model to extract the most important information needed from K and V. In ViT (Dosovitskiy et al. 2020), it is proposed that by observing the attention heatmap, one can determine which NLP token is most important for which part. (Maharjan et al. 2018) uses a constant query ”which representation is most important” to weight different representations. In our work, we use the T5 model (Chung et al. 2022), with the prompt + query as the KV pair, and the first token of the answer as Q (acting as ”which tokens are most important” because in normal training, it is necessary to score the token weights to generate answers), to perform cross attention and determine which tokens in the prompt are most important.

Generative LLMs have achieved strong performance across various tasks, but they encounter computational challenges when processing long documents and extended conversations due to increased token counts and context truncation. ICL (Brown et al. 2020) helps mitigate some of these issues by providing task-relevant context directly, reducing the need for specific task Supervised Fine-Tuning (SFT) and lowering costs. However, ICL also increases token numbers and inference costs.

The Selective Context method was developed to address these challenges by compressing input context through the removal of redundant information. Li et al. evaluates tokens based on self-information, retaining only the most informative content. This approach reduces token numbers and inference costs while mitigating the ”lost in the middle” problem (Tay et al. 2021) associated with long contexts. Experiments demonstrate that Selective Context reduces memory usage and generation latency while maintaining comparable performance, especially in tasks involving lengthy documents.

The LLMLingua (Jiang et al. 2023b) series, including LongLLMLingua (Jiang et al. 2023a) and LLMLingua2 (Pan et al. 2024), enhance context compression by optimizing context selection strategies and introducing advanced memory network structures. These models aim to improve efficiency and accuracy in processing longer documents and complex dialogues, optimizing context usage without compromising performance.

QUITO (Wang et al. 2024) introduces self-attention as a metric for context compression, diverging from traditional methods that rely on perplexity or self-information. By concatenating context and query, QUITO uses a small transformer decoder-only model to compute attention, retaining tokens with high attention values.

However, in the self-attention module of QUITO, the query, key, and value are trained together, which primarily aims to obtain a comprehensive global representation, and may lead to information mixing between the input and output. In contrast, we employ cross-attention in an encoder-decoder architecture, where the input and output are separated. We keep the key-value pairs fixed and only train the query representations, which forces the query to learn to distinguish which token representations are most crucial for generating the output.

The information bottle-neck(IB) (Tishby, Pereira, and Bialek 1999; Fischer 2020) is a simplistic concept: when facing a task, one should try to accomplish it using the minimal information. In our case, we want to compress the context while retaining the accuray of output answer, which makes it well-suited for modeling using IB theory. So we can use the IB score to formulate our objective as follows:

where $\bar{X}$ stands for compressed context, $Q$ stands for query, while $Y$ stands for output. $X=x_{1}x_{2}x_{3}...x_{n}$, in which $x_{i}$ is the $i^{th}$ token, and $\beta$ is related to compressed ratio. The first item serves to enhance efficiency while the second term serves to retain as much useful information as possible to enable the LLM generate target outputs.

Notice that since our compress method is to delete some token from $X$ to get $\bar{X}$, so for a fixed compressed radio, the first item $I(\bar{X};X|Q)$ for efficiency can be ignored. Then our objective is to maximize $I(\bar{X};Y|Q)$

Using the chain rule of Mutual Information and r-Markov properties, we have

(where $I_{Q}$ stands for conditioning on Q). So, for a fixed compressed radio, we can calculate the exact number k such that we need to delete k tokens while maximizing the mutual Information. We define $I(x_{i};Y|x_{1},...,x_{i-1},Q)$ as the conditional mutual information of the ith term, noting that the ith token, as well as the $r-1$ tokens preceding it, are related to the ith term. Therefore, to delete k tokens, we can consider finding the k terms with the smallest conditional mutual information.

According to the above analysis, we need a metric to reflect the relative size of $I_{Q}(x_{i};Y|x_{i-1},...x_{i-r+1})$. It should satisfy the following properties:

1. 1.

   Conditioning on $Q$, and it can capture the information from previous $x_{i}$

2. 2.

   Under the above condition, it can also reflect the importance or similarity of $x_{i}$ to the output $Y$

We found that cross-attention in the encoder-decoder architecture can effectively satisfy the above properties. Specifically, we put the context (corresponding to $X$) and the query (corresponding to $Q$) together in the encoder. Through self-attention, we obtain a representation of the context that can capture the previous information and query information, which we use as the key-value (KV) pair for cross-attention (corresponding to Property 1). Then, we use the first token in the decoder marked as the start as the query ( q), performing cross-attention with the KV pair. At this time, the first token in the decoder as q can be considered a constant query: ’which representations are most important to generate output’ (corresponding to Property 2). Therefore, the attention scores obtained from the cross-attention between q and the KV pair can be used to represent the relative size of the conditional mutual information in Equation (2).

From Property (1), we need to ensure that after removing some tokens, we can still capture the information preceding $x_{i}$. According to the r-Markov assumption, we should retain the $r$ tokens preceding those important tokens. In practice, we found that adding a Gaussian filter, which is a softer implementation, works better and is easier to implement than the hard constraints that directly retaining the preceding $r$ tokens. The detailed process is presented in the following Algorithm section.

In practice, we use the word as the smallest unit of compression, following the work of previous researchers (Li et al. 2023). We utilize the T5 model’s cross-attention mechanism to evaluate the importance of each word in the context relative to a given query. The procedure is as follows:

1. 1.

   Concatenation of Context and Query: The context and query are concatenated to form a single input sequence. This sequence is fed into the T5 model’s encoder to generate a feature representation.

2. 2.

   Encoding and Decoding: Leveraging the T5 model’s encoder-decoder architecture, the concatenated input is encoded, and the model begins decoding from a start token [start]. During this decoding process, the model generates output tokens while computing cross-attention scores that highlight the importance of each token in the context with respect to the generated output.

3. 3.

   Get the Cross-Attention Scores: Cross-attention weights from the final layer are extracted, specifically focusing on tokens corresponding to the context. The weights from each attention head are averaged to obtain a more robust representation. These averaged weights are then normalized using a softmax function to produce attention-based importance scores for each token.

4. 4.

   Smoothing with Gaussian Filter: To ensure that tokens adjacent to those with high attention scores also receive appropriate attention, we apply a Gaussian filter to the attention scores. This smoothing process helps to distribute the attention more evenly across nearby tokens, enhancing the model’s ability to capture relevant context. The Gaussian smoothing process is mathematically defined as:

   $\displaystyle\tilde{s}_{t}=\frac{1}{\sqrt{2\pi\sigma^{2}}}\sum_{k=-K}^{K}s_{t+k}\cdot\exp\left(-\frac{k^{2}}{2\sigma^{2}}\right)$

   (4)

   where $\tilde{s}_{t}$ is the smoothed attention score for token $t$, $s_{t+k}$ represents the attention score for neighboring tokens, $\sigma$ is the standard deviation that controls the smoothness, and $K$ is the window size for smoothing.

5. 5.

   Reconstruction Based on Word Importance: Based on our practice that words are the smallest semantic units, the reconstructed context is derived by considering the importance of each word. The ‘reconstruct‘ function groups tokens into words and then re-evaluates their importance using the smoothed attention scores. This approach ensures that the semantic integrity of the text is preserved during compression.

6. 6.

   Context Compression: Finally, the context is compressed by selectively retaining words with the highest importance scores. The compression ratio can be adjusted to control the level of detail retained in the compressed context, balancing between context size and information retention.

Formally, after obtaining the attention scores for each token, we filter words as follows:

Here, $\text{score}_{w}(\text{word}_{i})$ is the cumulative attention score for each word, calculated by summing the attention scores of all tokens within the word. The words are then filtered by selecting the top-ranked ones based on their cumulative scores, ensuring that the selected words collectively satisfy the desired ratio of the total token count. Figure 1 2provides an overview of our algorithm.

This approach, by focusing on word-level importance via cross-attention scores, allows for effective compression of long contexts while maintaining essential information. This is particularly useful for tasks that require handling extensive text, where preserving critical details is crucial for maintaining model performance.

By focusing on the cross-attention mechanism of the T5 model, our method effectively identifies the most relevant information in a context, enabling efficient compression without significant loss of critical information. This approach is particularly beneficial in scenarios involving long contexts, where the preservation of key information is paramount for maintaining high accuracy in downstream tasks.