IDEAL: Leveraging Infinite and Dynamic Characterizations of Large Language Models for Query-focused Summarization

With the growth in model sizes, fine-tuning methods have advanced significantly, modifying only a small number of parameters or adding new ones to a frozen language model for specific tasks (Li and Liang, 2021; Lester et al., 2021; Hu et al., 2021; He et al., 2022; Zhang et al., 2023;Zhang et al., 2024b;). These methods often achieve performance comparable to full model fine-tuning. In this paper, we use three types of PEFT methods, including prompt tuning, parallel adapter, and LoRA, as baselines to investigate our approach.

As shown in Figure 1(a), Prompt tuning can add soft prompts to the hidden states in attention layers to guide model learning and adapt to new tasks, where only the soft prompts are updated during training. LLaMA-Adapter-v1 (Zhang et al., 2023) introduces a zero-initialized attention mechanism into prompt tuning, which adaptively incorporates the knowledge from soft prompts. We use this LLaMA-Adapter-v1 as our prompt tuning baseline.

Parallel adapters (He et al., 2022) aim to incorporate additional learnable networks in parallel with distinct sublayers within the backbone model. To reduce the number of parameters, small bottleneck networks are used as parallel adapters. In transformer-based LLMs, parallel adapters can be applied to both the feedforward and self-attention modules in each transformer block. Hu et al. (2023) conducted experiments showing that applying parallel adapters only to the feedforward module achieves the best results on math reasoning datasets. As shown in Figure 1(c), we also apply parallel adapters only to feedforward module in LLaMA’s transformer block.

LoRA (Hu et al., 2021) adds trainable low-rank decomposition matrices in parallel to existing weight matrices (Figure 1(b)). For a pre-trained weight matrix $\boldsymbol{W}\in\mathbbm{R}^{d\times k}$, LoRA constrains its update by adding low-rank matrix pairs, resulting in $\boldsymbol{W}+\Delta\boldsymbol{W}=\boldsymbol{W}+\boldsymbol{BA}$, where $\boldsymbol{B}\in\mathbbm{R}^{d\times r}$, $\boldsymbol{A}\in\mathbbm{R}^{r\times k}$, and the rank $r\ll\min(d,k)$. During training, $\boldsymbol{W}$ is frozen while $\boldsymbol{B}$ and $\boldsymbol{A}$ are trainable. LoRA initializes $\boldsymbol{A}$ randomly and $\boldsymbol{B}$ to zero, ensuring that $\Delta\boldsymbol{W}=\boldsymbol{BA}$ starts from zero at the beginning of training, thereby preserving the pre-trained knowledge as much as possible.

Previous works (Ivison and Peters, 2022; Zhao et al., 2024) indicate that hypernetworks can learn the parameter information of the main neural network under different input scenarios and efficiently adjust the target network’s parameters to adapt to this information. We propose generating query-focused adapters conditioned on the query instruction using a hypernetwork.

Our hypernetwork is a bottleneck network that consists of an encoder to transform the mean-pooling of the query representation $\boldsymbol{H}_{query}$ into a low-dimensional representation $\boldsymbol{h}$, and a decoder to convert $\boldsymbol{h}$ into the parameters of the target adapters. For example, the computation of IDEAL_LoRA is as follows:

$$\boldsymbol{h}=dropout(ReLU(\boldsymbol{W}_{0}mean(\boldsymbol{H}_{query})+\boldsymbol{b}_{0}))$$

(1)

$$\hat{\boldsymbol{A}}_{q}=\boldsymbol{W}_{1}\boldsymbol{h}+\boldsymbol{b}_{1}$$

(2)

$$\hat{\boldsymbol{A}}_{k}=\boldsymbol{W}_{2}\boldsymbol{h}+\boldsymbol{b}_{2}$$

(3)

where $\hat{\boldsymbol{A}}_{q}$ and $\hat{\boldsymbol{A}}_{k}$ correspond to $\boldsymbol{W}_{q}$ and $\boldsymbol{W}_{k}$ in self-attention, respectively. We only generate the $\boldsymbol{A}$ matrix in the LoRA module, initializing $\boldsymbol{B}$ to zero and updating it during training as in the original LoRA. This ensures that $\Delta\boldsymbol{W}=\boldsymbol{B\hat{A}}$ starts from zero at the beginning of training. Unlike IDEAL_LoRA, IDEAL_Prompt and IDEAL_PAdapter generate all the parameters of the target adapters in the required layers.

In addition, each layer that needs to generate the target adapters has its own encoder, as shown in Equation 1, and shares a single decoder. This allows for generating different parameters for each layer and reduces the number of trainable parameters.

A key-value (KV) cache is typically used in LLMs for fast and efficient inference. To maintain fine-grained local attention, for each segment, multi-head self-attention $\boldsymbol{\mathcal{A}}_{local}\in\mathbbm{R}^{L\times d_{value}}$ is computed with a fixed KV length $L$ in both the training and inference stages using the KV cache. In detail, when the last segment length is less than $L$, we use the KV cache to extend the length of the current KV states to $L$ for computing the local attention and compress the remaining KV cache into the memory.

For the $s$-th segment with length $L$, before computing the local attention, we update the full context memory $\boldsymbol{M}_{s-1}^{all}\in\mathbbm{R}^{d_{key}\times d_{value}}$ and the query-focused memory $\boldsymbol{M}_{s-1}^{query}\in\mathbbm{R}^{d_{key}\times d_{value}}$, and a normalization term $\boldsymbol{z}_{s-1}\in\mathbbm{R}^{d_{key}}$ is then used for memory retrieval as follows:

$$\boldsymbol{M}_{s-1}^{all}\leftarrow\boldsymbol{M}_{s-2}^{all}+\sigma(\boldsymbol{K}_{cache})^{T}\boldsymbol{V}_{cache}$$

(4)

$$\boldsymbol{M}_{s-1}^{query}\leftarrow\boldsymbol{M}_{s-2}^{query}+\sigma(\boldsymbol{K}_{cache})^{T}\hat{\boldsymbol{V}}_{cache}$$

(5)

$$\boldsymbol{z}_{s-1}\leftarrow\boldsymbol{z}_{s-2}+\sum_{t=1}^{L}\sigma(\boldsymbol{K}^{t}_{cache})$$

(6)

where $\sigma$ is a nonlinear activation function. Following the work of Katharopoulos et al. (2020) and Munkhdalai et al. (2024), we employ element-wise ELU+1 as the activation function (Clevert et al., 2015). The term $\sigma(\boldsymbol{K})^{T}\boldsymbol{V}$ on the right side of Equation 4 and 5 is referred to as an associative binding operator (Schlag et al., 2020). The query-focused memory $\boldsymbol{M}_{s-1}^{query}$ differs from the full context memory only in the value states $\hat{\boldsymbol{V}}_{cache}$ used within the associative binding operator. We ultilize the query states $\boldsymbol{Q}_{query}$ of query instruction to scale the value states, and keep only query-related information $\hat{\boldsymbol{V}}_{cache}$ as

$$\alpha_{i}=sigmoid\left(\frac{mean(\boldsymbol{Q}_{query})(\boldsymbol{K}_{cache}^{i})^{T}}{\sqrt{d_{model}}}\right)$$

(7)

$$\hat{\boldsymbol{V}}_{cache}=\boldsymbol{\alpha}\odot\boldsymbol{V}_{cache}.$$

(8)

Here, we use the mean pooling of $\boldsymbol{Q}_{query}$ and the key states to compute a related score for each representation.

After updating the memory, we retrieve new content $\boldsymbol{\mathcal{A}}_{all}\in\mathbbm{R}^{L\times d_{value}}$ and $\boldsymbol{\mathcal{A}}_{query}\in\mathbbm{R}^{L\times d_{value}}$ from the full context memory $\boldsymbol{M}_{s-1}^{all}$ and the query-focused memory $\boldsymbol{M}_{s-1}^{query}$, respectively. This retrieval is performed using the query states $\boldsymbol{Q}\in\mathbbm{R}^{L\times d_{key}}$ as follows:

$$\boldsymbol{\mathcal{A}}_{all}=\frac{\sigma(\boldsymbol{Q})\boldsymbol{M}_{s-1}^{all}}{\sigma(\boldsymbol{Q})\boldsymbol{z}_{s-1}}$$

(9)

$$\boldsymbol{\mathcal{A}}_{query}=\frac{\sigma(\boldsymbol{Q})\boldsymbol{M}_{s-1}^{query}}{\sigma(\boldsymbol{Q})\boldsymbol{z}_{s-1}}$$

(10)

First, we apply a linear layer to aggregate $\boldsymbol{\mathcal{A}}_{all}$ and $\boldsymbol{\mathcal{A}}_{query}$. Then, we aggregate the retrieved content and the local attention $\boldsymbol{\mathcal{A}}_{local}$ using a learned gating scalar $\boldsymbol{\beta}$:

$$\boldsymbol{\gamma}=sigmoid(\boldsymbol{W}_{g}\boldsymbol{\mathcal{A}}_{query})$$

(11)

$$\boldsymbol{\mathcal{A}}_{ret}=\boldsymbol{\gamma}\odot\boldsymbol{\mathcal{A}}_{query}+(1-\boldsymbol{\gamma})\odot\boldsymbol{\mathcal{A}}_{all}$$

(12)

$$\boldsymbol{\mathcal{A}}=sigmoid(\boldsymbol{\beta})\odot\boldsymbol{\mathcal{A}}_{ret}+\\ (1-sigmoid(\boldsymbol{\beta}))\odot\boldsymbol{\mathcal{A}}_{local}$$

(13)

where $\boldsymbol{W}_{g}\in\mathbbm{R}^{1\times d_{value}}$ is a trainable weight that dynamicly merges the two retieved contents. $\boldsymbol{\beta}$ contains a single scalar value per head as training parameter, enabling a learnable trade-off between the long-term and local information flows in the model.

To incorporate query instructions into the model, we concatenate the query instruction with the document as the input of model. During local attention, the query states $\boldsymbol{Q}_{query}$ of the query instruction are utilized to compute query-focused memory within each segment. However, when generating summaries, the retrieved information from memory fails to effectively guide the model in producing summaries that adhere to the query instructions. To address this issue, we employ a straightforward approach: we replicate the query instruction at the end of the document. This ensures that the query instruction is within the window of the local attention computation when generating summaries, enabling the model to accurately generate query-relevant summaries.

As shown in Table 1, we compare the performance of IDEAL on different Adapter with same input size. On the CovidET dataset, the performance differences among the three adapters on IDEAL were minimal. However, on the QMSum(Golden) dataset, IDEAL_LoRA outperformed IDEAL_PAdapter by 1.48 ROUGE-L points under the same input length of 768. Overall, IDEAL_LoRA achieves the best results on four datasets.

In Table 4, we evaluated the effectiveness of Query-focused Infini-attention through comparative testing. First, we implemented Infini-attention based on LoRA as Lora+Inf and observed significant improvements compared to LoRA alone under the same GPU memory constraints, with increases of 1.55 and 1.33 points in ROUGE-L and ROUGE-Lsum on QMSum dataset, respectively. These results indicate that compressing the key-value states of historical segments enables summarization of long documents within limited GPU memory. Furthermore, we enhanced IDEAL_LoRA with Infini-attention, achieving better results than Lora+Inf in ROUGE-L. The IDEAL_LoRA method integrated with Query-focused Infini-attention as IDEAL${}_{LoRA}^{QF\_Inf}$ outperformed both IDEAL_LoRA+Inf and Lora+Inf in all metrics, demonstrating that our proposed Query-focused Infini-attention effectively compresses query-related information. For the IDEAL_LoRA+Inf method, we observed a significant decline in all metrics after removing the repeated query instruction at the end of the input document, demonstrating the necessity of repeating the query instruction.

IDEAL_LoRA consistently demonstrates improved performance as input length increases. However, this comes at the cost of increased GPU memory consumption. Table 4 illustrates this trade-off, showcasing IDEAL_LoRA performance on input lengths of 1.6K, 3.8K, and 8K, requiring 24G, 40G, and 80G of memory, respectively. In contrast to IDEAL_LoRA, our proposed IDEAL${}_{LoRA}^{QF\_Inf}$ exhibits memory efficiency when handling long inputs. IDEAL${}_{LoRA}^{QF\_Inf}$ maintains a consistent memory footprint 24G regardless of input length. Notably, on the QMsum dataset, IDEAL${}_{LoRA}^{QF\_Inf}$ outperforms IDEAL_LoRA with an input length of 1.6K on all metrics within a same 24GB memory constraint. Moreover, it surpasses IDEAL_LoRA with an input length of 3.8K in 40GB memory on the ROUGE-L metric and achieves performance close to IDEAL_LoRA with an input length of 8K in 80GB memory.

In Table 3, we compare the performance of different IDEAL HyperExperts under the same parameter count. The Prompt-tuning method can adjusts parameter count only by controlling prompt length, with experiments from Hu et al. (2023) indicating optimal performance at a prompt length of 10. Despite having the fewest trainable parameters, its performance on the QMSum(Golden) dataset is the lowest. With the same parameter count, LoRA with a rank of 16 still significantly underperforms compared to IDEAL_LoRA, highlighting the effectiveness of HyperExpert. IDEAL_PAdapter can improve performance by increasing the bottleneck size, but even with 89.5M parameters, it is still inferior to IDEAL_LoRA with 24.5M parameters. Overall, IDEAL_LoRA achieves the best performance and parameter efficiency.

Figure 3 presents the performance of IDEAL${}_{LoRA}^{QF_{I}nf}$ under varying local context sizes (LC). On the QMSum dataset, the model exhibits stable performance when LC is beyond 400, achieving nearly the best overall performance at LC=800. Similarly, on the SQuALITY dataset, the optimal LC is observed at 1.6K. These findings indicate that IDEAL${}_{LoRA}^{QF_{I}nf}$ differs from IDEAL_LoRA, the limited memory for the former is enough to handle extremely long inputs.

Table 4 presents the optimal input length for IDEAL${}_{LoRA}^{QF\_Inf}$ on the QMsum and SQuALITY datasets. The results suggest that information relevant to the query in the QMsum dataset is primarily concentrated within the first 6000 tokens, while in the SQuALITY dataset, the relevant information is more evenly distributed throughout the document.