Paper: Nugget 2D: Dynamic Contextual Compression for Scaling Decoder-only Language Models (ICLR-2024)

Reviewer comment (Presentation):
The results of section 4.2 are not provided in a table, but depicted in a figure: Figure 4. Figure 4 is hard to interpret and not clear at all. It is unclear what the bars represent there, and the ICAE bar is missing for x-axis values 100, 200, 300.

Author response:
We did not provide a table to present the result because the baseline ICAE uses boxplots, and we are unable to access their raw results, as we said in our footnote.  The bars in fig4 are the results of ICAE for sequence lengths of 100, 200, and 300. They are nearly 100%, so they are at the top of the graph. We added comments to the caption to clarify the plots.

Reviewer comment (Experiments):
For results in 5.2, it is unclear whether the number of trainable parameters for Compressive and nugget2D are comparable. Nugget2D employs an encoder and an LM, whose parameters are fine-tuned separately. How does this compare to Compressive?

Author response:
We admit that Nugget2D contains more trainable parameters than Compressive Transformers. However, it is hard for us to change the number of parameters of Compressive Transformers without changing its design. On the other hand, the parameters of both models are much fewer than those of the LLaMA because of LoRA (fewer than 1%). Therefore, the difference between their sizes is nearly negligible. We provide a breakdown of the number of parameters in Appendix D.3, which may alleviate your concerns.

Reviewer comment (Experiments):
In section 2.5, Choice of k, it is mentioned that the scorer is not trained on the autoregressive text continuation task, but taken from the autoencoding experiments. The reason for this choice is not explained.

Author response:
We did not train the scorer for autoregressive LM because there is a chicken-and-egg situation. Instead of taking a certain ratio of tokens as nuggets, autoregressive LM uses a threshold $\overline{s}$ on scores to select nuggets, which is decided by running scorer on all documents such that $\overline{s}$ selects out a certain ratio of tokens as nuggets *on average*. However, training the scorer needs a specified $\overline{s}$ so that the training pipeline can be started. Given the features fed into the scorer do not depend on $\theta$ or $\phi$, we find that re-using the scorer in section 2.4 is convenient and works well. We have updated the description in section 2.5 to make this point clearer.

Reviewer comment (Writing):
Typo: Before Eq (8) I think you mean "\textit{recent} tokens $w_{\tau+1:t}$" -- $t$ instead of $i$.

Author response:
We have adopted other minor suggestions you pointed out, including using $\mathbf{z}$ to consistently denote the hidden states selected by $\phi$, fixing the typos you pointed out, and other minor issues.

Reviewer comment (Writing):
Please define the NUGGET2D function in Sec. 2.2 NUGGET2D, including what is $\phi$, instead of letting readers figure out themselves. Later you use "NUGGET2D" as a function, but in Sec. 2.2 you only use it as a method name.

Author response:
## The definition of the Nugget2D function

$\mathtt{Nugget2D}$ takes a sequence of texts as inputs. It first encodes the text with the encoder ($\phi$), and then uses the scorer ($\varphi$) to select nuggets. The output of $\mathtt{Nugget2D}$ is the hidden states of the nugget tokens.  We added eq 4 to the main paper to define the $\mathtt{Nugget2D}$ function. Thanks for your suggestion!

Reviewer comment (Writing):
Perhaps you could consistently use $\mathbf{z}$ to denote hidden states from $\phi$. In Eq. (2) (3) (4), are $\mathbf{x}_i^l$ computed using $\phi$ but not $\theta$?

Author response:
## Is $\mathbf{x}_i^l$ computed with $\phi$ not $\theta$?

The hidden states $\mathbf{x}_i^l$ in eq 2 and 3 are computed with $\phi$ as they are on the $\mathtt{Nugget2D}$ side. In eq 4 (now eq 5), the $\mathbf{z}_j^l$  is computed with $\phi$ because it is a nugget, but $\mathbf{x}_i^l$ is computed with $\theta$ because it is on the $\mathtt{LM}$ side. We added subscript $\phi$ to eq 2 and 4 to clarify their parameters.

Reviewer comment (Reproducibility):
You mention in footnote 4 that Scorer quickly converges, and then $s_j$ can be removed from Eq. (5). How do you decide when to remove it? Do you first train all trainable parameters (based on the task) with $s_j$, and then just continue to still train all trainable parameters without $s_j$? Removing $s_j$ will cause all activations to change suddenly; not sure if it's a problem.

Author response:
## Removing the residual connection

We empirically found that the nugget selection converges after 3000 steps of training. At this stage, removing the residual connection does not cause much performance drop, as the model can quickly adapt to the new attention weights. We did not find any significant performance difference after dropping the residual connection, therefore we did not detach the residual connection for the experiments in the main paper.

Your concern is valid: Removing the residual connection causes the activation to suddenly change and affect the forward propagation. In Appendix B, we discussed an alternative solution to train the Scorer based on the *straight-through estimator*. The basic idea is to subtract  $s_j$ from the attention weights to cancel its effect on the forward pass. However, it does not show any improvement in the experiments of autoencoding. We reckon that LLaMA is so large that it can flexibly adapt to any new attention patterns.

Reviewer comment (Experiments):
In Sec. 2.5, can you clarify how you pick $\tau$?

Author response:
## How do we pick $\tau$? Do we ignore the beginning tokens? How about gist tokens?

Suppose Nugget2D is decoding the $t$-th token, then $\tau$ is set as $t-\omega_r$, which means we compress texts that are more than $\omega_r$ tokens away. 

We indeed ignore tokens that are $\omega_r+\omega_d$ tokens away to make the inference time finite for arbitrarily long sequences. However, we argue that the information is not truncated at $\omega_r+\omega_d$: A token can access the information that is $L\cdot(\omega_r+\omega_d)$ tokens earlier, where $L$ is the number of transformer layers, because the information is propagated not only from left to right but also from bottom to top. For example, let $\omega = \omega_d + \omega_r$, then $x_\omega^{l=1}$ can access $x_0^{l=0}$, $x_{2\omega}^{l=2}$ can access $x_\omega^{l=1}$, …, and $x_{L\cdot\omega}^{l=L}$ can access $x_{(L-1)\cdot\omega}^{l=(L-1)}$. It is similar to the idea of TransfomerXL, which is a recurrence-based transformer variant.

One could use other compression methods like gist tokens to incorporate more distant information. However, our ultimate goal is to process infinitely long sequences with finite resources, and methods like gist tokens assume a finite sequence length. Moreover, the function of gist tokens overlaps with Nugget2D. For these reasons, we did not introduce a separate compressing method like gist tokens.
We have included the above discussion in the revised paper.

Reviewer comment (Experiments):
Can you provide analysis about how much space are used by the extra model parameters when you introduce NUGGET2D? I notice that $\phi$ and $\theta$ are not tied.

Author response:
## Analysis of additional model space

It is true that $\phi$ and $\theta$ are not tied. Because they are both initialized from the same LLaMA checkpoint, we implement $\phi$ and $\theta$ with LoRA, which means we only add a marginal number of parameters to the LLaMA model. With a LoRA rank of 32, the number of all trainable parameters is fewer than 1% of LLaMA, and LLaMA itself is kept frozen. We added a parameter breakdown in Appendix D.3.

Reviewer comment (Theory):
The nugget tokens are picked based on $\phi$ but not $\theta$. Is it possible that some tokens' contextualized representations according to $\theta$ has much information but is not selected based on $\phi$?

Author response:
## Does $\theta$ prefer different nuggets as $\varphi$ does?

We argue that $\varphi$ tends to select tokens that $\theta$ prefers, because the parameters of Scorer, $\varphi$, are trained with a residual connection from the $\mathtt{Nugget2D}\_\phi$ side to the $\mathtt{LM}\_\theta$ side. 
In section 2.3, we conclude that this training objective will assign higher scores to the nuggets that are most attended to by the target tokens (i.e., tokens encoded with $\theta$). 

To fully answer this question, we conducted an experiment in Appendix A, where we designed a greedy algorithm to approximate the *optimal nugget selection* that achieves the best downstream performance. We run experiments on language modeling and compare the performance of optimal nugget selection to the nugget Scorer. In conclusion, we found that the nugget selection with Scorer is only 24% different from the optimal selection, and their performance gap is less than 8%, meaning the selection done by Scorer is very close to the oracle.

Reviewer comment (Writing):
First paragraph on page 5 is confusing, e.g., you said $\theta$ in eq (11) but there's no $\theta$ in eq (11). Please clarify what re-assignment means with more understandable notations.

Author response:
## Parameter re-assignment

We have re-organized this section to make the narrative more readable. To answer your question: In the scenario of section 2.3, texts are split into 2 segments (e.g. document and prompt + response), where $\mathtt{Nugget2D}\_{\phi,\varphi}$ is used to compress one part of text while the $\mathtt{LM}_\theta$ encodes the other part. In the scenario of autoregressive Nugget2D in section 2.5, the text is not split but some tokens are selected nuggets. Therefore, $\phi$ is used to encode the nugget tokens, while $\theta$ is used to encode the non-nugget tokens. With this re-assignment, Nugget2D is able to decode and compress new tokens in a streaming way.

Reviewer comment (Experiments):
Is it true that the model only processes $w_r+w_d$ tokens and ignore the beginning tokens if the context is long? Would it be helpful if, like in many long-context processing papers, you still have some "gist" tokens to summarize the beginning tokens?

Author response:
## How do we pick $\tau$? Do we ignore the beginning tokens? How about gist tokens?

Suppose Nugget2D is decoding the $t$-th token, then $\tau$ is set as $t-\omega_r$, which means we compress texts that are more than $\omega_r$ tokens away. 

We indeed ignore tokens that are $\omega_r+\omega_d$ tokens away to make the inference time finite for arbitrarily long sequences. However, we argue that the information is not truncated at $\omega_r+\omega_d$: A token can access the information that is $L\cdot(\omega_r+\omega_d)$ tokens earlier, where $L$ is the number of transformer layers, because the information is propagated not only from left to right but also from bottom to top. For example, let $\omega = \omega_d + \omega_r$, then $x_\omega^{l=1}$ can access $x_0^{l=0}$, $x_{2\omega}^{l=2}$ can access $x_\omega^{l=1}$, …, and $x_{L\cdot\omega}^{l=L}$ can access $x_{(L-1)\cdot\omega}^{l=(L-1)}$. It is similar to the idea of TransfomerXL, which is a recurrence-based transformer variant.

One could use other compression methods like gist tokens to incorporate more distant information. However, our ultimate goal is to process infinitely long sequences with finite resources, and methods like gist tokens assume a finite sequence length. Moreover, the function of gist tokens overlaps with Nugget2D. For these reasons, we did not introduce a separate compressing method like gist tokens.
We have included the above discussion in the revised paper.

Reviewer comment (Experiments):
The experiments are not sufficient to support the main contribution of scaling autoregressive LLMs to long contexts, comparing with recent long-context LLM studies (e.g. RMT, LongLLAMA). See details in Questions.

Author response:
Thanks for pointing out the recently published works. However, there are technical difficulties for us to conduct experiments to establish new baselines. The paper LongLLaMA released their checkpoints based on OpenLLaMA 3B and 7B, while our experiments are based on LLaMA-2 7B. We are unable to expediently train another LongLLaMA on LLaMA-2 7B during the rebuttal period. Moreover, the LongLLaMA paper was released on July 6, 2023, 3 months within the submission deadline of ICLR 2024, thus we reckon Nugget2D is a contemporary work to LongLLaMA.

For similar reasons, the experiments with RMT cannot be shown now. However, we want to emphasize that the idea of Nugget2D is orthogonal to these works, as integrating kNN or recurrent Transformer to Nugget2D is straightforward. We will bring discussion of their relationship in the next version of the Nugget2D paper.

Reviewer comment (Experiments):
The baseline is relatively limited. It would be good to see other long-context LLMs (e.g. LongLLAMA[1]) as baselines.

Author response:
Thanks for pointing out the recently published works. However, there are technical difficulties for us to conduct experiments to establish new baselines. The paper LongLLaMA released their checkpoints based on OpenLLaMA 3B and 7B, while our experiments are based on LLaMA-2 7B. We are unable to expediently train another LongLLaMA on LLaMA-2 7B during the rebuttal period. Moreover, the LongLLaMA paper was released on July 6, 2023, 3 months within the submission deadline of ICLR 2024, thus we reckon Nugget2D is a contemporary work to LongLLaMA.

For similar reasons, the experiments with RMT cannot be shown now. However, we want to emphasize that the idea of Nugget2D is orthogonal to these works, as integrating kNN or recurrent Transformer to Nugget2D is straightforward. We will bring discussion of their relationship in the next version of the Nugget2D paper.

Reviewer comment (Experiments):
I wonder how the memory grows as the context length increases against other context scaling methods. It seems the memory complexity in training stage is $O(n^2 * r)$. What is the limit of the context length of your model?

Author response:
If Nugget2D is used as a context compressor, then the space complexity on the encoder side is still $O(n^2)$, and the complexity on the decoder side is $O(m^2 + m\cdot n)$, where $n$ ($m$) is the number of tokens in the encoder (decoder) side. Note that the nugget ratio $r$ is a constant throughout experiments so is only a coefficient. 

In the settings of autoregressive transformers, the space complexity of Nugget2D depends on the local window size $w\_r$ and distant windows size $w\_r$ and is $O((\omega\_r+\omega\_d)^2)$. If they are set as constant, the complexity is $O(1)$, though the actual computational overhead is still 7B-level.

We conduct experiments with a single V100 GPU card with 32GiB memory, and the maximum context length for encoding is ~1200 tokens. 
We argue that Nugget2D is not designed to lower the theoretical space complexity but the coefficient $r$. Even with the same complexity, a compression ratio of 10x can still greatly increase the context lengths of LLMs in practice.

Reviewer comment (Experiments):
The context length of downstream tasks is really short. Do you ever try other dataset for long-context downstream task evaluation, such as Multi-News[2], Narrative QA[3] and CUAD[4].

Author response:
We thank the reviewer for pointing out several alternative datasets. We did consider running Nugget2D on datasets with longer sequence lengths, however, they are relatively small, and training LLaMA on them would cause significant overfitting. CNN DailyMail is the best trade-off we found, which contains 287k training examples with an average length of ~1000 tokens. Among the datasets you pointed out, the largest one is MultiNews, which only contains 44k examples, which we reckon is not sufficient for a 7B model.

We also have technical difficulty running very long documents. We note that even with Nugget2D, it is still tough to conduct experiments with contexts longer than 2000 tokens due to our limited resources. We argue that Nugget2D can be viewed as a proof of concept. With sufficient resources and training on longer documents, it may achieve better performance on much longer documents.