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.

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.

Other minor questions

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.

Thank you for your careful reading and all the useful suggestions! We have updated the paper according to your comments, with some major revisions in section 2.5. Because the main paper is already full, we put some discussions and a new experiment in the appendices. We hope our replies below can address your concerns:

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!

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.

How do Qin and Van Durme (2023) train the Scorer?

They build a residual connection between the encoder and decoder to back-propagate the gradients to the scorer. In the revised paper, we added a description of the residual connection in Qin and Van Durme (2023) in section 2.1. Moreover, we discussed its architecture and compared it with the scorer training in Nugget2D in Appendix C.

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.

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.

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.

Dear reviewer Hb1S,

Thank you again for your valuable time and insightful comments.

Given that the author-reviewer discussion period is coming to a close soon, would you please let us know if our responses have resolved their concerns and if there are any other questions we can address to help recommend our paper?

Best regards!

Authors

Thank you for a detailed summary of the paper and for acknowledging the contribution of our paper! We made revisions to our paper according to your comments. Due to the space limit, we put some of the new contents in the appendix. Below are our answers to your questions:

Novelty of the key idea

We argue that Nugget2D is a non-trivial extension of the Nugget, not only changing the form of nuggets and applying to decoder-only transformers:

• It has a redesigned Scorer module, which does not need to freeze transformer encoder parameters as Nugget does. Also, to train the scorer, Nugget requires a dual decoder, which will be discarded. In contrast, Nugget2D builds residual connections between decoder tokens directly and does not need an auxiliary module.
• Nugget2D is much more flexible than Nugget: It can flexibly compress any part of a text, achieving different levels of granularity. E.g., in section 6, we compress the document into nuggets but keep the prompt as tokens. Moreover, Nugget2D can do streaming language modeling, autoregressively compressing generated tokens.

We provide a detailed comparison between the 2 models in Appendix C, covering technical details and their applications.

Interpretation of autoencoding results

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.

Fairness of comparison to Compressive Transformers

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.

Does Nugget2D use informed nugget encoding?

No. We empirically found it is not helpful for downstream tasks. We have 2 speculations (a more thorough discussion can be found in Appendix C):

• A decoder-only transformer uses the hidden state of the last token to predict the next token, thus it naturally gathers context information. Informing the model about the nugget selection might be unnecessary.
• In the scenario of autoregressive Nugget2D, nugget tokens are encoded with separate parameters, thus a type embedding is obviously redundant.

Scorer is not trained on the autoregressive text continuation task

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.

Attention pattern within Nugget2D

The attention pattern within Nugget2D is a standard causal masking for decoder-only transformers. We have updated Fig 2 to make it clear.

Indices in eq 4 (now eq 5)

The $j$ indices in eq 5 are the indices of nuggets. The i indices in eq 5 are all tokens that attend to the nuggets. We updated the paper to make this point clearer.

Dear reviewer TZ9E,

Thank you again for your valuable time and insightful comments.

Given that the author-reviewer discussion period is coming to a close soon, would you please let us know if our responses have resolved their concerns and if there are any other questions we can address to help recommend our paper?

Best regards!

Authors

We thank the reviewer WuCW for your recognition of our contributions. Below are our responses to your questions:

Compare to other baselines

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.

Novelty of the idea

We argue that the novelty of Nugget2D is not simply splitting contexts and compressing the distant ones. Instead, the model can learn how to compress the distant tokens by selecting certain hidden states as representations of the past texts, while this discrete step can be learned in an end-to-end manner with a residual connection.

Space complexity

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.

Experiments on other datasets

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.

The sampling method for autoregressive LM experiments

100k indicates the total number of tokens for evaluation, summed over all evaluation documents. The context length is shown in Table 1.

We held out a part of the Pile for evaluation. During the evaluation, we iterate over documents in the evaluation split, filtering out documents shorter than the required context lengths. We keep collecting documents until 100k tokens are used. We will clarify the sampling process in the paper and release the codes upon paper acceptance.