Focused Transformer: Contrastive Training for Context Scaling

Summary

This paper proposes an improved training approach for memory-augmented Transformers on language modeling tasks. In particular, the authors identify the distraction issue in memory-augmented models, whereby the attention mechanism tends to focus on irrelevant contexts in the regime of long sequences. To address this concern, the authors propose cross-batch training, which is inspired by contrastive learning and exposes the attention mechanism to process both relevant and irrelevant documents. Extensive experiments are conducted to validate that the proposed cross-batch training can effectively shape the key-value space, mitigate the distraction issue, and successfully extend the model context size.

Strengths

- The paper is clearly written and well structured. - The identified distraction issue is compelling and makes a valuable contribution to the long-context research community, highlighting the need for more effective attention mechanisms or retrieving techniques. - The method is straightforward to use and demonstrates its significant effectiveness in identifying relevant information.

Weaknesses

My main concern is about the increased training computational costs. The computational complexity for the memory-augmented attention layer increases from $O(bn^2)$ to $O(b d n^2)$, where $b$ represents the batch size, $d$ is the number of cross documents, and $n$ is the sequence length. Employing a large $d$ might introduce a significant computation overhead.

Questions

1. How does the proposed technique compare to conventional packing techniques commonly used during pretraining (e.g., concatenating multiple documents together, possibly interleaved with [EOS] tokens)? This technique is mostly used to reduce padding but also generates a sequence with context from diverse documents. 2. How well does the proposed model converge in comparison to vanilla transformers? Intuitively, as the attention mechanism in FoT is exposed to more irrelevant information, the learning process might slow down. 3. How can one determine the optimal transformer layer for memory augmentation? Does allowing more transformer layers access to the memory aid in the retrieval of relevant information?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

The authors have adequately addressed the limitations.

Summary

This paper proposes the Focused Transformer (FoT), which modifies one layer of the Transformer model's Attention layer to Memory Attention, thus enabling the model to learn almost infinite length context without being limited by the constraints of local attention. Specifically, the paper also proposes a cross-batch training method that can adapt the model to any length context at a relatively low fine-tuning cost. Experimental results demonstrate the superiority of FoT on Long Context tasks.

Strengths

1. By only retaining one layer of Global Attention (Memory Attention), both training and inference are efficient; 2. During the inference stage, it can flexibly follow the results of kNN retrieval; 3. The cross-batch training method is very effective, and the training cost is also relatively low; 4. The increase factor of the Context length for LLM reached 2^15!

Weaknesses

1. With all due respect, I believe that the main text of the paper does not clearly introduce FoT and cross-batch training. I didn't understand it until I read the code in the appendix; 2. The inference stage needs to be combined with kNN, which is inconsistent with the training mode and may affect the upper limit of the model's ability; 3. There is a lack of comparison with other similar schemes for increasing Context length (such as Parallel Context Window).

Questions

1. I don't quite understand why, during the cross-batch training stage, we need to distinguish between C_prev and C_curr? Isn't it better to treat all tokens across batches as negative learning directly? 2. In the inference stage, how are the K and V to be input into Memory Attention calculated after kNN? 3. After cross-batch training, will the model overfit long Context? Is its performance on short Context still good?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

The authors have clearly and sufficiently explained the limitations of their research

Summary

This work presents a modified method to train Transformers with memory. The memory is based on k-Nearest Neighbors (with exact match). The work is based on the Memorizing Transformer paper from Y. Wu et al. 2022. The model assumes a transformer model with a memory attention (an attention layer that has inputs from both current sequence and memory output). The (k-NN) memory returns the top-k matches for a query, and these are included in the input to the memory attention. One difference is the removal of a gating mechanism that mixes between memory and local context, becoming all input to an attention layer. In addition, contrastive learning is proposed as a method to enhance the model capability to differentiate relevant keys from irrelevant ones in the memory layer. This is achieved by including negative samples from other documents, applying a forward pass up to the memory attention layer, and using these key,value pairs as simulated responses from the k-NN memory. These samples are mixed with the previous context block from the current sample. This method is differentiable and allows to backprop through these positive and negative samples, requiring only a modification to the data pipeline. The authors suggest that the contrastive learning is needed do reduce a "distracting issue" caused in a multi-document setup. Eventually, the authors show the value of their method with a set of experiments that compare mainly against the memorizing transformer. The results show that memory can extend the effective context length for language modeling at inference time, they can solve a copy task using memory, a pre-trained model can be fine-tuned to introduce the memory attention layer, improves performance in single and multiple document scenarios, and an ablation study with the importance of the negative samples and the differentiable keys.

Strengths

The paper mixes the idea of negative sampling and a twist to the memorizing transformer model with a k-NN memory. The topic is of significance to this community. * At inference time this model can increase the context length beyond the training length. Inheriting the original memorizing transformer property. * It is possible to fine-tune from a pre-trained model with no memory to a model with external memory. * The experiments show the value of the negative samples. The training methodology behind the idea is relatively simple.

Weaknesses

The contribution of this work has limited novelty: both negative sampling (and contrastive learning [2]) and memorizing transformers have been proposed in the past for language modeling. The clarity and explanations in the work could be improved. The work continuously omits important details or descriptions, sending the reader in most cases to the appendix. The references could make a better exploration of contrastive learning methods applied to language modeling (e.g., [2]). Please, see questions for additional details that may need clarifications. Retrieval-augmented language models are able to utilize memory to gain information from multiple documents. The "distracting issue" described in the paper suggests that the memory attention mechanism gets distracted by keys from different documents. However, this seems to contradict the findings in previous work (retrieval methods) that earn performance using multiple documents. The experiments are mainly limited to the memorizing transformer and the vanilla transformer. However, retrieval language models solve the same multi-document problem. Also the dataset chosen follow the memorizing transformers, failing to compare with existing benchmarks in the long context use case, like SCROLLS [1]. [1] Shaham et al., SCROLLS: Standardized CompaRison Over Long Language Sequences, 2022 [2] Jain et al., CONTRACLM: Contrastive Learning For Causal Language Model, 2022

Questions

I would appreciate it if the authors could answer the following questions: * Line 188: Why is layer 8 picked as the memory layer? How would someone choose the "right" layer for connecting the memory? * What is the improvement of an existing language model (like OPT, LLaMA, etc.) when fine-tuned to introduce the memory attention layer? * Why is your model not reaching perfect accuracy in the synthetic task? (Shallow) transformers are able to obtain and copy information. Is the same key related to multiple values in different documents in the dataset? * What is stored in memory during the evaluation in Section 4.3? What is the data in the memory? * What is the distance measure for k-NN? * How is the number of neighbors $k$ chosen for the memory retrieval?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

The authors address limitations of their work in Section 5. Scaling (approximate) k-NNs and memory has been addressed for quite some time. See the NeurIPS 2021 competition on Billion Scale Approximate Nearest Neighbor Search.

Summary

The paper proposes a key challenge (distraction issue) when extending the attention layer to external (key, value) pairs, either from previous context of the same document or from other documents. This paper proposes the Focused Transformer (FoT), a model that can utilize a long context by retrieving kNN (key, value) pairs from memory. The architecture of FoT is similar to Memorizing Transformer (how retrieved items are integrated is different), but the training technique is different. This paper proposes to use crossbatch to train the model, which simply includes both previous context of the same document and also contexts from different documents. So the model can learn how to distinguish between useful information and distraction.

Strengths

* The paper focuses on an important and potentially impactful problem that extends the context window size of transformer models. * The proposed model and the training method is straightforward and easy to implement. * The experiments are fairly solid and have supported the points that the paper has raised. * I also like the synthetic dictionary task as an evaluation task to test if a language model has the ability to attend to desirable context and gather necessary information.

Weaknesses

1. One main weakness of this paper is that it is not clear to me if the proposed model is designed to incorporate long context or to incorporate external memory which may come from a large corpus. As the paper includes both single-doc and multi-doc experiments, I assume it is the latter case. Based on this assumption, the paper identifies “distraction issue”, which essentially means the model attends to more other documents when considering more documents during inference. I did not get the point here. If we allow the memory to contain items from other documents, we actually expect the model to extract useful information instead of treating them as “distraction”. When showing the problem of “distraction issue” (Fig 3), we only consider 64 documents, which is far from the real case where we want to use external corpus. Indeed, prior works have shown that using more external information can help the model to achieve better results instead hurting it, e.g., (Khandelwal et al., 2019), (Zhong et al., 2022). While if the paper is mainly considering only the long-context cases, we also should question if the “distraction issue” exists during inference or not, because we can always control the model to only attend to a single document. So, I don’t think the paper is well-motivated or well-positioned. The authors are encouraged to clearly state what testing situations they are addressing and re-evaluate the proposed issue in that situation. 2. The paper can be presented better. For example, it is never clear to me what “external memory” exactly means in the paper. Why “previous local context of the given document” is called positive and “contexts from unrelated documents” is called negative, given there are really no distinctions in the training objective (correct me if it is not the case)? 3. The proposed CrossBatch technique is based on including both previous contexts from the same document and contexts from different documents in the same training batch. This method is very similar to the data batching method proposed in (Zhong et al., 2022). The authors are encouraged to discuss the differences.

Questions

* L144: why is there a positive and negative? Did you actually distinguish positive/negative in the contrastive loss? * Table 2: why do you use token-level accuracy here but use perplexity elsewhere. * How do you define external memory? Anything that is out of the original input can be called external.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The authors have well discussed the limitations of the research.

discussion

We hope that our explanations and changes are clear enough. If there is something else, we would be happy to address it in the remaining time of the rebuttal period.