Efficient Long-range Language Modeling with Self-supervised Causal Retrieval

Authors: Xiang Hu, Zhihao Teng, Wei Wu, Kewei Tu
Primary area: foundation or frontier models, including LLMs
Keywords: long-range language modeling, Retrieval-based LM, self-supervised learning

Abstract:
Recently, retrieval-based language models (RLMs) have received much attention. However, most of them leverage a pre-trained retriever with fixed parameters, which may not adapt well to causal language models. In this work, we propose Grouped Cross-Attention, a novel module enabling joint pre-training of the retriever and causal LM, and apply it to long-context modeling. For a given input sequence, we split it into chunks and use the current chunk to retrieve past chunks for subsequent text generation. 
Our innovation allows the retriever to learn how to retrieve past chunks that better minimize the auto-regressive loss of subsequent tokens in an end-to-end manner.
By integrating top-$k$ retrieval, our model can be pre-trained efficiently from scratch with context lengths up to 64K tokens. 
Our experiments demonstrate that our model achieves superior performance in various tasks against strong baselines, and 100\% accuracy in the needle-in-a-haystack (NIAH) test with a 16M context length.

Decision: reject
Reviewer scores: 6, 6, 5, 6

Reviews and discussion:
<thread>
creation date:
1729512534956

thread_type:
official_review

<thread_object>
creator: reviewer_uh71

<content>
summary:
The paper proposes a novel long-range language model called Differentiable Retrieval-based Transformer (DRT), which introduces a Grouped Cross-Attention (GCA) module to enable joint pre-training of the retriever and causal language model. The main innovation lies in splitting input sequences into chunks and dynamically retrieving relevant past chunks to enhance the prediction of subsequent tokens. The retriever learns how to minimize the auto-regressive loss by retrieving past chunks in a differentiable manner, allowing more efficient training with context lengths up to 64K tokens.
With GCA, the model maintains random-access flexibility during auto-regressive generation while reducing memory and computational costs. Experiments show that DRT outperforms baseline long-range language models by achieving lower perplexity, faster inference speeds, and more efficient memory usage.

strengths:
1. The introduction of Grouped Cross-Attention enables joint optimization of the retriever and language model, solving the issue of fixed retriever parameters in traditional models and making the retriever more adaptive to causal language modeling tasks.
2. The model minimizes memory consumption by performing chunk-based retrieval and leverages hardware optimization to accelerate both training and inference.
3. By offloading historical chunk representations to CPU memory and reloading them as needed, the model significantly reduces GPU memory usage and achieves up to 100 times faster inference compared to existing baselines.

weaknesses:
The method requires pre-training from scratch, as it cannot simply add modules to fine-tune existing models, leading to a substantially higher computational cost.

questions:
1. In Figure 2, the tokens in chunk c7 rely on the relevance score from the LMK in chunk c6 regarding the previous chunks instead of from the token in c7 regarding the previous chunks. Is there any inconsistency?
2. As noted in previous work, there are some problems associated with using PPL as a metric. Do you have any other tasks or metrics that could demonstrate the effectiveness of your method?
3. Could you provide the performance results of your model on long-context benchmarks such as LongBench?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730116060753

thread_type:
official_review

<thread_object>
creator: reviewer_xkmg

<content>
summary:
This paper introduces a new approach to enhancing the efficiency and effectiveness of retrieval-based language models through a mechanism called Grouped Cross-Attention (GCA). The authors address the limitations of existing models, which typically rely on separately pre-trained retrievers with fixed parameters that may not optimally align with causal language modeling objectives. The chunk selection part of GCA can be optimized with language modeling loss, and this method can be implemented based on Flash-attention 2, which remarkably reduces the training and inference cost.

strengths:
1. the  proposed GCA method  facilitates the joint pre-training of the retriever and the causal language model in an end-to-end manner. This allows the retriever to dynamically learn to select past chunks.
2. The integration of top-k retrieval and hardware-aware implementations (based on FlashAttention-2) ensures that DRT achieves lower perplexity with reduced training time, outperforming  baselines like Retrieval-Pretrained Transformer and Landmark Attention.
3. During inference, DRT employs memory offloading techniques to manage GPU memory usage efficiently, substantially reducing the memory footprint while maintaining high performance.

weaknesses:
1. The proposed GCA module is difficult to understand because the authors do not provide sufficient motivation for the design of each step. Additionally, there are alternative methods to implement the differentiable operation, but the paper does not include any analysis or comparison of these alternatives. This lack of clarity and justification makes it challenging for the community to follow and reproduce the work.
2. My primary concern is the lack of scalability. Most of the baselines use models with fewer than 200M parameters. The authors mention having access to 8 A100 GPUs, which should enable the training of larger-scale models. I believe this method could benefit from being adopted with a fine-tuning approach, such as CEPE[1], in combination with existing  LLMs like Llama to demonstrate its effectiveness at a larger scale.
3. The authors only provide experiments on language modeling. Evaluating the method across a diverse range of downstream tasks would allow the community to better assess the overall effectiveness and applicability of this work.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730682055224

thread_type:
official_review

<thread_object>
creator: reviewer_uyq7

<content>
summary:
This paper proposes Differentiable Retrieval-based Transformers (DRT). Following a previous line of work on retrieval-augmented LMs (e.g. RETRO and RPT), the key novelty of this paper is Grouped Cross Attention – instead of concatenating top k retrieved units, the authors proposed to use retrieval score to fuse outputs derived independently from each unit in the top K. The authors pre-train a DRT model and several baselines (RPT, Landmark Attention, etc.) from scratch and conduct perplexity evaluation on long-range language modeling tasks. Results show that DRT obtains roughly on-par or better performance compared to the baseline, with better performance as the context length increases.

strengths:
* The paper proposes a method to train retrieval-augmented LMs, aiming to improve LM’s capability to handle long context, which is a practical use-case and active research area.
* The paper contains comprehensive experiments against several baselines as well as ablation studies.

weaknesses:
While I appreciate the inclusions of different baselines (RPT, Landmark attention) and the motivation of training RALM in an end-to-end manner,  I am finding it a little bit hard to situate this paper against previous work  and the major benefit of the proposed method:
* It seems like the authors are advocating for the use of GCA, which enables training time efficiency due to end-to-end training. Yet, if I am reading table 1 correctly, it seems like baseLM (+2 layers) achieve lower perplexity on PG19 consistently compared to DRT and also with lower training wall clock time. 
 
* It is also unclear how to isolate the benefit of GCA, given that the baselines have several other different design choices (RPT might be the closest, but it doesn’t use the landmark tokens as this paper, nor grouping the layers) and there is no ablation of using GCA v.s. a different method to train the retrieval module (e.g. the method employed by the RPT model which involves using a reference LM to score the retrieved chunks). It would be helpful to include an ablation study on GCA, besides those that are already included in Table 1.

questions:
* The original RPT paper uses representation from the lower layer to retrieve chunks for upper layers, and train the retriever using a reference model. This paper implements RPT using a contreiver to retrieve chunks, how is that exactly done, i.e. what is the representation for the query and the chunks? And is a fixed contriever used? If so, this seems like a weaker baseline than the originally-proposed RPT, it would be helpful to include a discussion on this.

* The paper mentioned that the proposed method enables “multi-hop retrieval” several times in the paper (e.g. second paragraph in section 3.1), what does this mean and what’s the experiment that demonstrates the ability?

* It seems like the proposed method perform better on the Arxiv-math dataset compared to the PG-19 dataset and it'd be helpful to include analysis to understand the reason. It'd also be helpful to experiment on other domains that are tested in previous work, such as Code [0] and books [1].

[0] https://arxiv.org/pdf/2306.13421
[1] https://arxiv.org/pdf/2402.16617
</content>
</thread_object>

</thread>
<thread>
creation date:
1730705736661

thread_type:
official_review

<thread_object>
creator: reviewer_ddqm

<content>
summary:
This paper aims to build a differentiable retrieval-based transformers (DRT) that optimizes the retriever together with generation in an end-to-end manner. It proposes to replace the Chunked Cross Attention with Group Cross Attention that allows the gradient to back-propagate to the retriever so that retriever can be trained as well where in previous work, the retriever is not trained. The evaluation on PG19 and ArXiv-math shows the effectiveness of the proposed method while being much faster than landmark attention.

strengths:
The authors successfully trained a differentiable retrieval-based transformers where adding grouped cross attention shows better perplexity compared to the baseline model.

weaknesses:
1. All the models are trained using very small model size (i.e. 133M), compare to the popular LLMs, e.g. 1B or 7B model size. The improvement could diminish when scaling to larger model sizes and a larger model size should be reported to make the claim more convincing.
2. It is not surprising that the model gets better perplexity by using more compute. Existing SOTA models mostly choose to scale the data size to better utilize the compute. Thus, a more interesting research question could be that given fixed amount of compute, should we allocate the extra compute to train more data or change the architecture to a differentiable retrieval-based transformer.

questions:
1. Will the improvement over baseline model diminish when scaling up the model size?
2. Will the authors open source the codebase? Making it open source might benefit more researchers.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731809406898

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response (part. 1)

comment:
We sincerely thank all reviewers for their valuable comments.

**1. Evaluation of downstream tasks**

We conduct evaluations on several downstream tasks and have some exciting findings. All models are pre-trained from scratch, then fine-tuned and evaluated on long-context tasks, including needle-in-a-haystack (NIAH) tests like RULER (https://arxiv.org/abs/2404.06654) and long text summarization (https://aclanthology.org/D18-1206/ , https://aclanthology.org/K16-1028/). Notably, DRT excels in summarization tasks and maintains high accuracy in the NIAH tests, even with context lengths exceeding 1 million tokens. We believe GCA significantly contributes to long-range language modeling, as it not only can extrapolate to contexts much longer than those seen during pre-training, but effectively utilizes distant information as well (single NIAH test), and even can perform multi-hop reasoning (2-hop NIAH test).

Summarization:
| Models | xsum:|R-1 | R-2 | R-L | CNN: |R-1 | R-2 | R-L(Cnn)|
|-|-|-|-|-|-|-|-|-|
| BaseLM | | 29.43| 8.04 | 23.26 | | 32.38|12.57|22.61|
| +2layers| |29.74| 8.26 | 23.48 | | 34.14 | 14.09 | 23.60 |
| Landmark Attn. | | 27.98 | 6.96 | 21.99 | | 34.06 | 13.80 | 23.77  |
| DRT$_{\text{retrieval} \times 1}$ | | 30.30 | 8.59 | 23.92 | | 36.27 | 15.88 | 25.08 |
| DRT$_{\text{retrieval} \times 2}$ | | $\textbf{30.39}$ | $\textbf{8.64}$ | $\textbf{23.98}$ | | $\textbf{36.39}$ | $\textbf{15.96}$ | $\textbf{25.15}$ |


Single NIAH test results:
| Models | Retrieval| 1K | 4K | 16K | 64K | 16M |
|-|-|-|-|-|-|-|
| BaseLM(+2 layers)   | -   | 60.43   | 15.37 | 3.89 | 0.0 | -
| Block Recurrent TFM   | -   | 66.80   | 13.96 | 6.01 | 4.29 | -
| RPT$_\text{contriever}$  | fixed   | 42.13   | 11.66 | 4.24 | 2.86 | -
|Landmark Attn.  | adaptive   | $\textbf{99.98}$  | $\textbf{99.82}$ | 97.88 | 0.0 | -
|DRT$_{\text{contriever}\times1}$  | adaptive   | 98.23   | 98.50 | 98.59 | 100.00 | 100.00
|DRT$_{\text{contriever}\times2}$  | adaptive   | 99.69  | 99.65 | $\textbf{99.65}$ | $\textbf{100.00}$ | $\textbf{100.00}$

An example of the single NIAH: 
...(essays) ... The passkey is: YOBDC. ...(essays)...
What is the passkey? The passkey is: (Answer: YOBDC)

2-hop NIAH with noises:
| Models | 1K | 4K | 16K | 64K | 4M
|----------|----------|----------|----------|----------|----------|
| BaseLM(+2 layers)   | 16.60   | 5.83   | 1.06 | 0.0 | -
|Landmark Attn.  | $\textbf{90.82}$   | $\textbf{88.35}$ | $\textbf{86.41}$ | 0.0 | -
|DRT$_{\text{contriever}\times1}$  | 41.07  | 33.39 | 39.93 | 38.57 | 34.28
|DRT$_{\text{contriever}\times2}$  | 88.52  | 84.45 | 86.21 | $\textbf{81.43}$ | $\textbf{79.41}$ |

An example of the 2-hop NIAH with noises:
...(essay)... DEF A->K , ..., DEF G->P, ..., DEF M->A, ..., DEF Z->L ..
Q: The path from M is: (Answer: A, K)

To the best of our knowledge, **DRT is the first model to achieve such high accuracy in the single-NIAH test over 16M and the 2-hop NIAH test over 4M, just with an average attention field of 768 tokens.** Such results fully demonstrate the self-supervised causal retrieval could be a promising direction for long-range language modeling.
A detailed experimental analysis is provided in the updated manuscript.

**2. Advantages of the DRT**

We highlight that the benefits of DRT extend beyond efficient training, as it not only achieves lower perplexity but also exhibits several distinctive properties. Below is a comparison table of various long-context models with similar perplexities under similar pre-training quotas:

| Model | Training Complexity | Inference Time Complexity | Inference Space Complexity | Effective Context Length  |
|-|-|-|-|-|
| Transformer (full attention)| O($NL^2$)  | O($NL^2$) | O($NL$) | $1 \times L$ |
| Transformer (sliding window)| O($NL$)  | O($NL$) | O($C$) | $1 \times W$ |
| Recurrent Models    | O($NL$) | O($NL$) | O($C$) | Limited  |
| Landmark Attention  | O($NL^2$) | O($NL+ \frac{NL^2}{S}$) | O($\frac{NL}{S}$) | $\sim 64 \times L$  |
| DRT | O($NL + \frac{L^2}{S^2}$)| O($NL+ \frac{L^2}{S^2}$) | O($\frac{L}{S}$) | $>256 \times L$  |
| DRT$_\text{offload landmarks}$ | O($NL + \frac{L^2}{S^2}$)| O($NL$) | O($C$) | $>256 \times L$  |

N=number of layers, S = chunk size, L = sequence length, W = sliding window size, C = constant.
The effective context length is measured using the NIAH tests, similar to RULER. DRT$_\text{offload landmarks}$ denotes offloading landmark representations to a faiss datastore, thus reducing time and space complexity of inference to O($NL$) and $O(C)$, respectively.

While performance improvements in language modeling might become marginal with growing parameters, DRT retains advantages in inference speed, low KV cache costs, and ultra-long effective context length. These benefits remain consistent even when scaling up.
</content>
</thread_object>

</thread>


Metareview:
metareview:
While DRT presents an interesting approach to long-context modeling through Grouped Cross-Attention and joint training of retriever and language model, the paper's experimental evaluation falls significantly short of the standards. The most critical weakness is the limited scope of experiments: models are only trained at a relative small scale, and evaluations are restricted to basic language modeling metrics without meaningful real-world benchmarks or comprehensive comparisons to state-of-the-art long-context models. The lack of experiments on larger model scales makes it difficult to assess whether the benefits would translate to practically relevant scenarios. Additionally, the baseline comparisons are incomplete. The absence of downstream task evaluations and limited ablation studies further weaken the empirical validation. While the technical approach shows promise, the insufficient experimental rigor and lack of comprehensive benchmarking prevent a clear assessment of the method's practical value, making it fall below the bar for ICLR acceptance.