Paper: ReFusion: Improving Natural Language Understanding with Computation-Efficient Retrieval Representation Fusion (ICLR-2024)

Reviewer comment (Novelty):
The argument for keeping prompts short is growing less compelling as models can handle more and more context. The other savings from the current approach seem to be small, since we still need to create representations for every contextual element.

Author response:
**W1**: Thanks for your suggestions. Although existing modern language models can handle quite a long input sequence (e.g., 128k for GPT-4-turbo), the resource overhead for self-attention is still non-negligible, i.e., larger memory requirements as well as computational costs. When the long input sequence is due to retrievals, our ReFusion can significantly reduce the resource overhead, maintaining almost the same FLOPs (from 188.6541 GFLOPs to 188.6548 GFLOPs shown in Figure 1). For other savings, the retrieval representations are encoded by a frozen pre-trained language model and stored in a task-agnostic key-value database. These processes are done offline before the training/inference. During the fusion process, our ReFusion only retrieves the pre-built representations and directly fuses them into the language models, thus introducing few overheads.

Reviewer comment (Writing):
In terms of clarity: overall, the paper is good, and the methods are mostly easy to follow, but I am not sure how the architecture search works. Section 3.3.1 creates one set of expectations, and section 3.3.2 does not seem to meet those expectations, or even respond to them directly. In particular, we are led to believe that the search module will decide what method to use on a per-layer basis, but then it seems to turn out to be a global choice. But I may well simply not be understanding the description and its connection to the experimental results.

Author response:
**W2**: Thanks for your suggestions. We recognize that there are some unclear descriptions in Section 3.3.2, especially the presentation of notations. We will carefully revise the Section 3.3. The hierarchy of architecture search can be classified into three levels, i.e., network level (global), layer level, module level. We are working on the finest granularity optimization (module level). Specifically, we replace several original modules with our search modules per layer. 

To make it clear, we reformulate the notations used in Section 3.3.2. 
* The superscript $l$ indicates the index of the layer.
* The subscript $\text{key/value/query}$ indicates the module our search module will replace.
* The subscript $m$ indicates the number of candidates used in each search module.

Taking the key module of the $l$-th layer as an example, our search module contains three ranking scheme candidates (i.e., $m=3$) and a vector of architecture parameters, i.e., $\alpha\_{\text{key}}^l=\left( \alpha\_{\text{key},1}^l,\alpha\_{\text{key},2}^l,\alpha\_{\text{key},3}^l \right)$. Then, we use a softmax operation over all candidates within the search module to relax the categorical choice,

$$
\hat{o}\_{\text{key}}^l (h)=\sum\_i \frac{\exp⁡{\alpha\_{\text{key},i}^l}}{\sum\_j \exp⁡{\alpha\_{\text{key},j}^l}} o\_{\text{key},i}^l (h)
$$

where $o\_{\text{key},i}^l (h)$ represents the output of the $i$-th candidate within the search module which replaces the key module in the $l$-th layer, $\hat{o}\_{\text{key}}^l (h)$ represents the final output of the search module that replaces the key module in the $l$-th layer. Finally, we use bilevel optimization to jointly optimize the architecture parameters and model weights.

Reviewer comment (Experiments):
Could you provide the details on how many train examples the various methods represented in Table 1 need, either as supervised examples or as demonstrations? I think this is a significant factor in how models perform on the chosen tasks. If all of them use all the train examples, that suffices as an indication that the playing field is level. (I am not asking for an efficiency analysis or anything here.)

Author response:
**Q1**: Thanks for your comments. The number of training examples used in Table 1 for all baselines as well as our ReFusion is 16 per class, which is the experimental setting used in LM-BFF. For example, the task SST-2 has two classes, thus we use 32 training examples. For fair comparisons, we don’t add any examples as demonstrations for either LM-BFF or our ReFusion.

Reviewer comment (Writing):
Some terms are confusing. The authors use the term "neural architecture search" for the search module, but rather, the search module simply performs hard gate operation among the fusion modules.

Author response:
**W1**: Thanks for your suggestions. We acknowledge the concern regarding the use of the term “neural architecture search” and recognize that using this term may imply a broader scope of architecture search than what is used in our search module. We will carefully revise the manuscript and use the term “Adaptive Retrieval Integrator (ARI)” to specify our methods in the search module.

For the architecture search module, we first define the search space and then describe the details about how to find the optimal architecture using bilevel optimization. The search space is basically determined by two factors, i.e., the number of ranking scheme candidates and the number of original modules that will be replaced by our search module. Within our search module, there are multiple ranking scheme candidates with a learnable architecture weight. We use bilevel optimization to train the search process. On the upper level, we aim to find the optimal combination of ranking schemes that maximizes performance on a validation set. During this level, we only update the architecture weights. On the lower level, we aim to fine-tune the model given the current selection of ranking schemes from the upper level. This involves minimizing the loss of downstream tasks. For the lower level, we only update model’s weights on a training set.

Reviewer comment (Writing):
Notations can be improved. When the authors write "h_y_{<cls>}", this seems like the output of the final layer. However, the authors perform fusion at every layer. Therefore, I suggest the author to denote layer with superscript l.

Author response:
**W2**: Thanks for your suggestions. We revised the notations with the superscript that denotes the layer index.

Reviewer comment (Evaluation):
The authors stress "computation-efficient" part, but it would be interesting to see the actual throughput and speed compared to other baselines methods as the proposed method brings in multiple modules.

Author response:
**W3**: Thanks for your comments. We evaluate the inference latency per sentence of our ReFusion, the baseline LM-BFF, and the retrieval-augmented baseline CA-512 on 15 NKI tasks. Our ReFusion can significantly accelerate the inference efficiency compared to CA-512, and makes a better trade-off between performance and efficiency compared to LM-BFF. Please refer to the Table 1 in the response to the **Reviewer 9abu** for more details.

Reviewer comment (Experiments):
The authors should add information on the number of parameters. It is not clear whether the empirical gains are obtained by simply adding more learnable parameters compared to that of LM-BFF.

Author response:
**W4**: Thanks for your suggestions. We counted the number of parameters of LM-BFF and our ReFusion during the training stage in Table 1. Compared to LM-BFF, our ReFusion only increases the number of parameters by 0.88%, which is quite trivial. However, our ReFusion can improve the performance by about 2.5 on average.

Table 1. The number of parameters of ReFusion and LM-BFF.
| **Methods** | **# Params (M)** |
| --- | --- |
| LM-BFF | 357.51 |
| ReFusion | 360.66 |

Reviewer comment (Writing):
I don't get how section 3.2.1 works. The authors states that "reranker is a 1D learnable vectors of k dimensions". If so, each rank (decided by retriever) is assigned with a learned scalar. This means that the learned 1D vector is shared among all the retrievals. Could you clarify on this point?

Author response:
**Q1**: Thanks for your comments. The reranker is a $k$-dim weight vector, where the $i$-th value in vector indicates the task-related importance of the $i$-th retrieval. The reranker is designed to learn the mappings between task-related importance and task-agnostic similarity, rather than the specific importance of retrievals. For example, ReFusion uses L2 distance to retrieve the top-$k$ nearest neighbors. But these retrievals might not be semantically suitable for downstream tasks. Thus, ReFusion adopts a simple $k$-dim vector to re-weight the importance of retrievals for the downstream tasks.

Reviewer comment (Experiments):
Scope of Experiments: The research confines its experimentation to masked language models, leaving the effectiveness of the proposed ReFusion method on populr autoregressive language models unexplored. Clarification on its adaptability to such models would be beneficial for broader application.

Author response:
**W1**: Thanks for your suggestions. We applied our ReFusion to the popular autoregressive language models, such as T5 [1]. As shown in Table 1, we evaluate ReFusion (T5-base) on 15 NKI tasks with 5 random seeds. Experimental results show that our ReFusion can be adapted to autoregressive models, improving the model performance by 1.4 on average. 

Table 1. Results with T5-Base.
| **Methods**        | **SST-2** | **SST-5** | **MR**   | **CR**   | **MPQA** | **Subj** | **TREC** | **CoLA** | **Avg-S** |
|--------------------|-----------|-----------|----------|----------|----------|----------|----------|----------|-----------|
| T5-Base            | ${91.9}_{0.3}$      | ${46.6}_{1.7}$      | ${87.6}_{1.0}$     | ${91.4}_{0.2}$     | ${85.9}_{1.2}$     | ${84.5}_{3.7}$     | ${91.8}_{5.0}$     | ${51.1}_{3.3}$     | ${78.8}_{0.7}$      |
| ReFusion (T5-Base) | $\boldsymbol{92.4}_\boldsymbol{0.3}$  | $\boldsymbol{47.2}_\boldsymbol{0.9}$  | $\boldsymbol{88.3}_{0.9}$ | $\boldsymbol{91.4}_\boldsymbol{0.4}$ | $\boldsymbol{86.1}_\boldsymbol{1.5}$ | $\boldsymbol{87.4}_\boldsymbol{1.3}$ | $\boldsymbol{94.9}_\boldsymbol{0.5}$ | $\boldsymbol{53.7}_\boldsymbol{4.1}$ | $\boldsymbol{80.2}_\boldsymbol{0.6}$  |

Table 1 (cond.). Results with T5-Base.
| **Methods**        | **MNLI** | **MNLI-mm** | **SNLI** | **QNLI** | **RTE**  | **MRPC** | **QQP**  | **Avg-P** | **Avg-all** |
|--------------------|----------|-------------|----------|----------|----------|----------|----------|-----------|-------------|
| T5-Base            | ${73.9}_{0.8}$     | ${75.8}_{0.6}$        | ${79.9}_{2.1}$     | ${82.1}_{1.9}$     | ${79.0}_{1.9}$     | ${75.5}_{3.6}$     | $\boldsymbol{76.2}_\boldsymbol{0.5}$ | ${77.5}_{0.3}$      | ${78.2}_{0.4}$        |
| ReFusion (T5-Base) | $\boldsymbol{75.3}_\boldsymbol{1.1}$ | $\boldsymbol{77.3}_\boldsymbol{1.0}$    | $\boldsymbol{83.0}_\boldsymbol{0.8}$ | $\boldsymbol{83.6}_\boldsymbol{0.5}$ | $\boldsymbol{79.1}_\boldsymbol{1.6}$ | $\boldsymbol{77.8}_\boldsymbol{2.7}$ | ${76.1}_{0.8}$     | $\boldsymbol{78.9}_\boldsymbol{0.5}$  | $\boldsymbol{79.6}_\boldsymbol{0.3}$    |

Reviewer comment (Experiments):
Performance on Knowledge-Intensive Tasks: While the paper presents improved results for non-knowledge-intensive tasks using ReFusion, it remains unclear how this method stacks up against others in knowledge-intensive scenarios. Given that retrieval-augmented language models are often specifically leveraged for their prowess in knowledge-intensive tasks, a direct comparison in this primary use case would be valuable.

Author response:
**W2**: Thanks for your suggestions. We evaluate our ReFusion on typical knowledge-intensive tasks, such as SQuAD. We also follow the 16-shot setting and evaluate our ReFusion with 5 random seeds. We compare our ReFusion with the baseline method and other SOTA methods [2,3] with the same settings. As shown in Table 2, our ReFusion can also achieve a significant improvement.

| **Dataset** | **T5-Base** | **ReFusion (T5-Base)** | **Splinter [2]** | **FewshotBART [3]** |
|-------------|-------------|------------------------|------------------|---------------------|
| SQuAD       | ${46.6}_{3.9}$        | $\boldsymbol{58.6}_\boldsymbol{0.9}$               | ${54.6}$             | ${55.5}_{2.0}$                |

[1] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J. Liu. Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. J. Mach. Learn. Res. 21: 140:1-140:67 (2020).

[2] Ori Ram, Yuval Kirstain, Jonathan Berant, Amir Globerson, Omer Levy. Few-Shot Question Answering by Pretraining Span Selection. ACL/IJCNLP (1) 2021: 3066-3079.

[3] Rakesh Chada, Pradeep Natarajan. FewshotQA: A simple framework for few-shot learning of question answering tasks using pre-trained text-to-text models. EMNLP (1) 2021: 6081-6090.

Reviewer comment (Presentation):
The caption for Figure 2 is not clear.

Author response:
**Q1**: Thanks for your suggestions. We have revised the caption of Figure 2.

Reviewer comment (Novelty):
The discussion on related works in the paper does not appear to be exhaustive. For example, Important references, such as the context-augmentation method and retrieval representation fusion, are absent.

Author response:
**Q2**: Thanks for your suggestions. We have revised the Related Work section and added the descriptions of context-augmentation methods [8, 9, 10] and retrieval fusion methods [11, 12, 13]. Please check page 9 for the new related works.

[8] Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Leyton-Brown, Yoav Shoham. In-Context Retrieval-Augmented Language Models. CoRR abs/2302.00083 (2023).

[9] Gautier Izacard, Edouard Grave. Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering. EACL 2021: 874-880.

[10] Zhicheng Guo, Sijie Cheng, Yile Wang, Peng Li, Yang Liu. Prompt-Guided Retrieval Augmentation for Non-Knowledge-Intensive Tasks. ACL (Findings) 2023: 10896-10912.

[11] Patrick S. H. Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel, Douwe Kiela. Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks. NeurIPS 2020.

[12] Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George van den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, Diego de Las Casas, Aurelia Guy, Jacob Menick, Roman Ring, Tom Hennigan, Saffron Huang, Loren Maggiore, Chris Jones, Albin Cassirer, Andy Brock, Michela Paganini, Geoffrey Irving, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack W. Rae, Erich Elsen, Laurent Sifre. Improving Language Models by Retrieving from Trillions of Tokens. ICML 2022: 2206-2240.

[13] Zonglin Li, Ruiqi Guo, Sanjiv Kumar. Decoupled Context Processing for Context Augmented Language Modeling. NeurIPS 2022.

Reviewer comment (Novelty):
Have authors explored the work in ''Unlimiformer: Long-Range Transformers with Unlimited Length Input (Bertsch et al., 2023)'' where authors claimed that their LM can handle unlimited input sequences?

Author response:
**Q1**: Thanks for your comments. There are several key differences between our ReFusion and Unlimiformer. Firstly, our ReFusion can be applied to any transformer-based model while Unlimiformer is only designed for encoder-decoder models. Secondly, the retrieval module of Unlimiformer only resides in memory while our ReFusion has an on-disk retrieval module, which is not limited by the memory capacity and has a larger knowledge database. Thirdly, Unlimiformer uses a cross-attention module to fuse retrievals while our ReFusion uses a more computation-efficient module. Fourthly, Unlimiformer prefers to retrieve from long-range inputs, while our ReFusion prefers to retrieve from external knowledge databases. Finally, the goal of Unlimiformer, making transformers process unlimited length input, is different from that of our ReFusion, efficiently introducing retrievals. Actually, our ReFusion is orthogonal to Unlimiformer.

Reviewer comment (Experiments):
Authors should consider extending this study to other LMs except Roberta.

Author response:
**Q2**: Thanks for your suggestions. We applied our ReFusion to the popular autoregressive language models, such as T5 [1]. Please refer to Table 1 in the response to the **Reviewer aM36**.  Experimental results show that our ReFusion can be adapted to autoregressive models, improving the model performance by 1.4 on average. 

[1] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J. Liu. Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. J. Mach. Learn. Res. 21: 140:1-140:67 (2020)