ReFusion: Improving Natural Language Understanding with Computation-Efficient Retrieval Representation Fusion
Authors: Shangyu Wu
Ying Xiong
Yufei CUI
Xue Liu
Buzhou Tang
Tei-Wei Kuo
Chun Jason Xue
Venue: ICLR 2024 poster (2024)
Keywords: Retrieval Augmentation; Non-Knowledge-Intensive Task; Natural Language Understanding;
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: Retrieval-based augmentations (RA) incorporating knowledge from an external database into language models have greatly succeeded in various knowledge-intensive (KI) tasks. However, integrating retrievals in non-knowledge-intensive (NKI) tasks is still challenging.
Existing works focus on concatenating retrievals with inputs to improve model performance. Unfortunately, the use of retrieval concatenation-based augmentations causes an increase in the input length, substantially raising the computational demands of attention mechanisms.
This paper proposes a new paradigm of RA named \textbf{ReFusion}, a computation-efficient \textbf{Re}trieval representation \textbf{Fusion} with bi-level optimization. Unlike previous works, ReFusion directly fuses the retrieval representations into the hidden states of models.
Specifically, ReFusion leverages an adaptive retrieval integrator to seek the optimal combination of the proposed ranking schemes across different model layers. Experimental results demonstrate that the proposed ReFusion can achieve superior and robust performance in various NKI tasks.

Published as a conference paper at ICLR 2024
REFUSION: IMPROVING NATURAL LANGUAGE UN- DERSTANDING WITH COMPUTATION-EFFCCIENT RE-
TRIEVAL REPRESENTATION FUSION
Shangyu Wu'; Ying Xiong" Yufei Cui
Xue Liu Buzhou Tang². Tei-Wei Kuo*5 Chun Jason Xue
City University of Hong Kong 2 Harbin Institute of Technology, Shenzhen
MILA, MeGill University National Taiwan University
5 Mohamed bin Zayed University of Artificial Intelligence
ABSTRACT
Retrieval-based augmentations (RA) incorporating knowledge from external database into language models have greatly succeeded in various knowledge- intensive (KI) tasks. However, integrating retrievals in non-knowledge-intensiww (NKI) tasks is still challenging. Existing works focus on concatenating retrievals with inputs to improve model performance. Unfortunately, the use of retrieval concatenation- based augmentations causes an increase in the input length, sub- stantially raising the computational demands of attention mechanisms. This pa- per proposes new paradigm of RA named ReFusion, computation-efficien Retrieval representation Fusion with bi-level optimization. Unlike previous works, ReFusion directly fuses the retrieval representations into the hidden states models. Specifically, ReFusion leverages an adaptive retrieval integrator to seek the optimal combination of the proposed ranking schemes across different model layers. Experimental results demonstrate that the proposed ReFusion can achieve superior and robust performance in various NKI tasks.
INTRODUCTION
Recent advances in language models (Khandelwal et al. 2020 Borgeaud et al. 2022b Guu et al., 2020 Lewis et al. 2020 Li et al. 2022) have demonstrated that retrieval-based augmentations (RA) can achieve remarkable performance on variety of knowledge-intensive (KI) tasks. The basic idea of RA is first leveraging
neaaatteeenaretteieehhte sewooor-eercc--abaeeaaaa indexing to retrieve the top-k related knowledge from external key-value store database, then incorporating the retrieved knowledge into language models with different fusion methods. KI tasks such as question-anse answering and text generation have an inherent retrieval-based property (Chen et al. 2017 Karpukhin al. 2020 as answers can be sourced or deduced from external knowledge databases.
(NKI) tasks, such text classification, still chal- lenging. Unlike KI tasks, NKI tasks often require understanding and categorizing given only one sentence rather than generating answers with contexts (Wang al., 2019). Previous works (Guo et al.) 2023b; Izacard & Grave 2021) have proposed to leverage plain texts such as Wikipedia to build the retrieval database. These methods concatenate retrievals with inputs the context for the model to make predictions. However, these methods significantly increase the input length, impos- ing considerable computational overheads in the attention module. Besides, the constraints on the model's 's max input length would also limit the concatenation of retrievals, which may result in truncated context and lead to semantically incomplete information, thus degrading performance.
However, RA in non-knowleege-intensii
This paper introduces new paradigm of retrieval-based augmentation named ReFusion, computation-efficien Retrieval representation Fusion framework with bi-level optimization. Dif- ferent from previous retrieval-basce augmentations (Borgeaud et al. 2022b Guo et al. 2023b),
AAtthors contributed equally to this research "orresponniig author

---

Published as a conference paper at ICLR 2024
RF-add (acc)
RC (acc)
Single (SST-2)
Pair (STS-B)
RF-add (fops))
RC (flops)
900
800
700
400
Segence Lengen
92
8600
Max sequence length (512) of RoBERTa
Accuracy
Perf. drops
5000
due to truncation
300
300
250
88
200
100
150
16
Retrieval (a) Input Sequence length
Retrieval (b) Accuracy and FLOPs
Figure 1: Impact the number concatenated retrievals on inputs and its effect on model's aceu- racy and FLOPs. RC (acc) and RF-add (acc) refer to the accuracy of retrieval-comcatenatioo -based augmentation (RC) and retrieval representation fusion with addition (RF-add). RC (flops) and RF- add (flops) refer to the FLOPs of RC and RF-add.
ReFusion directly fuses the representations of retrievals into the hidden states of models. ReFu- sion consists of three types of modules: the retriever module, the retrieval fusion module, and the adaptive retrieval integrator. Specifically, the retriever module includes query encoder to encode query texts and an efficient retriever for retrieving the representations of similar information. The retrieval fusion module includes two effective ranking schemes, the reranker-based scheme and the ordered-mask--asee scheme, to refine the representation of retrievals. The refined representations are subsequently fused into the hidden states. The adaptive retrieval integrator aims to find the optimal fusion structure with bi-level optimization.
Finally, we conducted comprehensive experiments on 15 different NKI tasks. These tasks vary from sentiment analysis, opinion polarity, natural language inference, etc. The experimental setting follows Gao et al. (Gao et al., 2021). Experimental results show that the ReFusion outperforms other comparisons and achieves superior results various tasks. Codes available at
The main contributions of this paper are:
This paper demonstrates the bottleneck of existing retrieval concatenation based augmen- tations, ie., significant computational overheads and limited performance improvements.
This paper proposes a new paradigm of retrieval-based augmentation that fuses the retrieval representations into the hidden states of models, making a better trade-off between perfor- mance and efficiency.
This paper designs two ranking schemes for retrieval fusion and an adaptive retrieval inte- grator for searching the best combination of different ranking schemes. Experimental results demonstrate that our ReFusion framework can significantly improve models' understanding capability, achieving superior and robust performance.
BACKGROUND AND MOTIVATION
Retrieval-augmented Prompt-based Fine-tuning The common NKI tasks involves inputting a sentence as Xsingle
[CLS] [SEP] or two sentences Z1, as Xpair [CLS]æ] [SEP]<> [SEP], and outputting a categorical label y, where [CLS] is classifica- tion token and [SEP] is sentence separate token in language models. In traditional methods, the [CLS] token is utilized represent the overall contextual information of the input and to facilitate classification tasks, i.., Viogits softmax(W. hrcLs) ). where h poIs; is the final hidden states of [CLS] token. To harness the potential of transformer- models in terms of mask predictre capabilities, which serve the pre-training objective, recent studies (Gao al. 2021 Zhang al.
2

---

Published as a conference paper at ICLR 2024
[8]22 by
Truncato
retricval texts
Large computations due longer embeddings
XL
Truncate retrievals due max sequence length
MLLM Head
[c]x,
[812g watch.
EMB
ATTA
RKA
great (0.1)
temible (0.9) v
Label space
Retrievals
Input texts
Template
Retrieval-augmenttd prompt-based inputs
Embeddings
(a) Retrieval concatenation-base augmentation (RC) via prompts.
Retrieval-sugnaellll fusion
XL
retrieval representations
MILM Head
.n P
EMB
ATTN
KEN
great (0.1)
[C] no resson to watch, was [M] 18)
fembble (0.9)
space
Input texts
Templati
Prompt-based inputs
Embeddings
(b) Retrieval representation fusion (RF).
Figure 2: Comparisons between retrieval concatenation- based augmentation (RC) and retrieval rep-
resentation fusion (RF).
2022
Guo et al. 2023b) suggest employing prompt-based fine-tuning. These works involve trans-
forming the original inputs into prompt-based inputs incorporating mask token for prediction. A simple example of prompt-based inputs single sentence or pair of sentences shown below,
X Xpromps single [CLS] ø It was [MASK] [SEP]
xprompt-pair [CLS] xi [MASK] x2 [SEP]
(2)
The categorical label is transformed into semantic token You, e.g., using 'positive' to replace '1' and 'negative' to replace '0', Then, the objective is to maximize the probability of the label word corresponding to the [MASK] token,
Loss log p(u Xpromp=) log p( [MASK] Vw X Xpromp=)
To further improve prompt-based fine-tuning, recent works (Guo et al. 2023b Chen et al. 2022) concatenate retrieval information as contexts or demonstrations of the given input, thus hel helping language models have better semantic understanding. For example, as shown in Figure 2 a). they first retrieve top-k similar sentences = { 2 from an external key-value database for the given input æ. Then, they concatenate all retrievals Z with prompt-based inputs X Xprompt. The retrieval- auggeentd prompt-based input of a single sentence is:
Xretrievel single [CLS] ] 21 [SEP] [SEP] = [SEP] x was [MASK] [SEP]
(4)
The objective is then optimized based on the petrieval-augmenter prompt-based inputs.
cooppt-based fine-tuning trades off the amount of contextual in- formation against the computational efficiency. Concatenating retrievals into the input enables lan- guage models to gather more contextual information, thus further improving language models' per- formance. However, due to the quadratic computational complexity of the attention mechanism, feeding longer input would introduce significant computational overhead. Besides, the limited length of the input is typically constrained by the hyper-parameter max.length, a longer input would be truncated before being fed in models [
Motivations Retrieval-augmeente
To quantify the impact of retrieval- based augmentations, this paper measures the performance of two baseline methods, i.e., Retrieval Concatenation-basee augmentation (RC) and Retrieval- representation Fusion with addition (RF-add), in terms the accuracy and the number of floating- point operations (FLOPs) in Figure (b). Specifically, RC concatenates retrievals with the given
3

---

Published as a conference paper at ICLR 2024
Retriever Module
Indexing One
no to [8]
Rerrieval we...
representations mo... P 1itt
ATTN
xl.
ARI
MLM Head
EMB
Linear
great (0.1)
MatMal
terrible (0.9)
FFN
Label space
Reranker
(3)
Order Mask
Prompt-based imputs
Emboddings
Linear
Figure 3: The architecture of the proposed ReFusion and the detailed structure of proposed modules.
input, while RF-add adds the retrieval representations the hidden states of the [CLS] ] token Fig- ure [b) demonstrates that concatenating more retrievals can slightly improve RC's accuracy while producing a significant amount of FLOPs. The degradation in RC's performance with 16 retrievals is due to the truncation of inputs, which results in a loss of semantic completeness. Conversely, when fusing more retrieval representations, the FLOPs RF-add remains almost unchanged, yet its accuracy continues to exhibit an upward trend. Notably, RF-add achieves even higher accuracy than RC.
Consequently, this paper aims to propose a new paradigm of retrieval-based augmentation that di- rectly integrates the retrieval representations into language models as shown Figure (2b)). The intuition is to embed the knowledge directly into the thought process of models by merging mod- els' hidden states with external information in an efficient fusion way. The preliminary experiments demonstrate the potential of the proposed paradigm. However, there still some drawbacks to the baseline method (RF-add),
1. The external knowledge is obtained by task-agnostic index based simple similarity metric, such as L2-Norm, rather than a more sophisticated semantic similarity metric;
2. RF-add pays the same attention to all retrievals on each representation dimension;
3. Not every layer of the model necessarily requires augmentation with retrievals.
These observations motivate propose computation-efficient retrieval representation fusion with bi-level optimization.
REFUSION: A COMPUTATION-EFICIA RETRIEVAL REPRESENTATION FUSION WITH BILEVEL OPTIMIZATION
In this section, we introduce the details of the proposed ReFusion framework. The ReFusion frame- work can be adapted to any transformer based architecture (Vaswani et al. 2017). or any architecture that contains the attention module. As shown in Figure in the framework, which retrieves the representations of top-k similar sentences. Then, we present the retrieval fusion module containing -adered-mask scheme. Finally, we propose the adaptive retrieval integrator, which learns to choose the best combination different ranking schemes across layers.
first present the retriever module used
ranking schemes, i.e., the reranker-based scheme and the
3.1 THE RETRIEVER MODULE
The retriever module includes a query encoder for encoding query texts, a task-agnostic index such as FAISS (Johnson et al. 2019) or ScaNN (Guo et al., 2020) built offline over billions of dense vectors, and a key-value store vector database for storing all text representations. The retrieving
Thh term of the retrieval representations in this paper refers to the final layers' output vectors of retrievals. The term of the hidden states of token refers the intermediate vectors of the token.
4

---

Published as a conference paper at ICLR 2024
Sign.
ARI
Rzs
Reran
Ordered Mask
Bzg
Rgg
hag
Ranked by learnable part.
Ranked by L2-norm
Prefer2
Reranker
Ranked by ordered mask
Ordered Mask
hzs
Rag
hzg
Search to select the ranking scheme
Sign.
Mire
Figure 4: Two different ranking schemes used in the fusion module.
process in the framework is online performed, which means that for every forward, the query en- coder first passes the representation fg of the imput text a to the index, then the index retrieves the representations Hz ( h22 of top-k similar sentences Z {=1...., 2k to the fusion module.
3.2 THE RETRIEVAL FUSION MODULE
The retrieval fusion module first ranks the representations of top-k similar sentences, then add them into the hidden states of existing modules. Specifically, this paper introduces two effective ranking schemes shown in Figure 4
3.2.1 RERANKING THE RETRIEVALS
In the retriever module, the retrievals are selected by task-agnostic similarity metric, e.g., L2 norm. Directly adding the representations to the hidden states can only bring limited improvements. That because 1) The ranking of retrievals should be different for each module layer in the model, each module layer has different functionality (Jawahar et al. 2019); 2) The models should pay different attention to those retrievals in case of overemphasizing less relevant information. Therefore, this paper proposes learnable reranker to learn the ranking distribution tailored to each module in the model. As shown in the top of Figure the importance of retrievals is re-assigned after reranking.
Specifically, the reranker is k-dimensional learnable vector, i.e., R (ri....,,s* where r E I.. The reranker is first normalized and then multiplied by the retrievals. Finally, the averaged representation of all reranked retrievals is added to the hidden states of the model, e.g., [CLS] token BERT-like models (Liu al. 2019b Devlin et al. 2019). The formal steps are as follows, exp(r;)
Ti = Σ, exp(r,)
(5)
h 9icu:
Σ ri. r
(6)
Picsss
where b M the input and output hidden states of the [CLS] token in the I-th layer.
3.2.2 ORDERED MASK OVER RETRIEVAL REPRESENTATIONS
Rippel et al. (Rippel et al. 2014) proposed nested dropout that directly drops the representation units from the sampled index thus yielding an inherent importance ranking on the representation dimensions. Based on the nested dropout, recent works (Cui et al. [023 2020 2021; Mao et al. 2022) proposed the ordered mask that modeled the dropping process with a chain Bernoulli variables and made it differentiable using the ce-parameterization trick. Such methods rank the representation units for a better feature representation.
As shown in the bottom of Figure this paper follows the main idea of the ordered mask but ranks the representation units over & retrievals on each representation dimension. This is a more fine- grained ranking scheme compared to the previous ranking scheme. Specifically, let hap hz, be
5

---

Published as a conference paper at ICLR 2024
the top-k D D-dimensional retrieval representations. For d-th dimension of retrieval representation, the ordered mask is modeled by a chain of Bermoulli variables V = foy..... " where of ~ Bernoulli<x;) indicates whether drop the d-th representation unit of the i-th retrieval. Following the property of nested dropout, the variable of is conditioned on uf thus, we can obtain the marginal distribution p(v
After that, we use the re-parameterization trick, e.g., choosing the Gumbel Softmax distribu- tion (Jang et al. 2017) as the tractable variational distribution g(v With Gumbel Softmax distri- bution, if ca Gumbel(3, r), then of = 1 cumsum, where cd is sample choice of the dropped index over k retrievals on the dimension d, and cumsum,(c") =Ei= Sj. In the Gumbel Softmax distribution, the temperature variable that controls the smoothness of the step at the dropped index.
a learnable parameter in the differentiable function of g(e;; 3) and is
Finally, we obtain D different ordered masks v'...., V over the representation dimension. We use them to mask the retrievals in a fine-grained grained way. Then, the masked retrievals would be fused into the hidden states similarly. The formal steps are as follows,
d Gumbel(3,7 r)
(7)
of = - cumsumy(c" ir
(8)
o h
h
h (
Σ h
(10)
k
where id. is the d-th masked representation unit of i-th retrieval.
3.3 THE ADAPTIVE RETRIEVAL INTEGRATOR
As shown in Figure 4 it is difficult to tell which ranking scheme is better on each module layer in the model. Therefore, this paper proposes the adaptive retrieval integrator (ARI) to find the optimal combination of ranking schemes across layers.
3.3. SEARCH SPACE
This paper proposes to search for the combination of ranking schemes from two aspects, i.e., module level (whether it contains a retrieval fusion module in the module layer) and ranking scheme level (which ranking scheme is used for the module). For the module level, this paper replaces the linear module, e.g.. query/key/value module in the attention, with the proposed adaptive retrieval integra- tor. For the ranking scheme level, the proposed adaptive retrieval integrator consists of two parallel retrieval fusion modules with different ranking schemes (e.g.. the fusion module with the reranking scheme) and the original module without any ranking scheme.
Although the total number of candidates in the adaptive retrieval integrator small (3 most in this work), the whole search space is still quite large. Given a transformer- based language model with N hidden layers, if only the key and value modules in every attention module are replaced, there are at least (3 x 3) = & combinations of different ranking schemes. Taking the RoBERTa-large as an example (24 layers), the search space can be septillion level large.
3.3.2 SEARCHING WITH BI-LEVEL OPTIMIZATION
The searching challenge can be addressed by the bi-level optimization (Liu et al. 2019a). For the lower level, this paper aims to fine-tune the backbone model with the selected combination of ranking schemes on the training set. For the upper level, this paper aims to find the optimal combination of ranking schemes that maximizes performance on the validation set. Specifically, this paper creates m architecture parameters in each adaptive retrieval integrator. For example, this paper uses aley
f I koyym }o to learn the choice for the key module in the (-th layer. To make the search space continuous, this paper also relaxes the categorical choice of a particular candidate to the softmax value over all possible candidates within each adaptive retrieval integrator, exp(atey.))
&ley (h) = Σ E, exp((hey.)) (h)
(11)

---

Published as a conference paper at ICLR 2024
Table 1: Main results with RoBERTa-laree.
Methods
SST-2
SST-5
MR
CR
MPQA
SUBJ
TREC
CoLA
Avg-S
LM-BFF
92.7a.9
47.42.
87.0
90.31.0
84.72.
91.21.1
84.85.1
9.373
73.4
DART
93.5p.5
88.21.
91.8p.5
90.71.
87.138
KPT
90.31.6
86.818
88.83.7
CA-512
91.31.4
46.7
1..
88.3
76.928
88.0
82.24.4
7.433
70.7
ReFusion
93.40.6
49.81.4
87.91.1
91.70.3
86.71.1
92.5p.8
90.33.7
11.44.1
75.5
Methods
MNLI
MNLI-m
SNLI
QNLI
RTE
MRPC
QQP
Avg-P
Avg-all
LM-BFF
68.32.3
70.5
77.23.7
64.54.2
69.1
74.55.3
65.55.3
69.9
71.8
DART
67.52.6
75.81.6
66.73.7
78.34.5
67 67.83.2
KPT
61.42.1
61.528
71.62.7
CA-512
66.21.0
67.81.3
71.62.2
66.93.2
66.63.1
73.56.9
64.01.9 68.93.3
68.1
69.5
ReFusion
69.31.5
10.91.5
80.61,4
73.01.1
70.92.3
77.03.6
72.9
74.3
The results of LM-BFF; DART refer to their original paper. The results of KPT refer to (Chen
al. 2022).
where key, (h) represents the output hidden states of the i-th candidate retrieval fusion module o;(-),
6(-) indicates the overall output hidden states.
EXPERIMENT
In this section, we first introduce the dataset and the experimental setting. Then, we present main results 15 NKI tasks. We also give the ablation study to demonstrate the efficacy of each proposed module. Finally, we analyze the efficiency of the proposed ReFusion.
4.1 DATASETS AND EXPERIMENTAL SETTING
Datasets We conduct comprchensive experiments across 15 NKI tasks, including 8 tasks from GLUE benchmark (Wang al. 2019), SNLI, SST-5, MR, CR, MNLI, MNLI-mm, Subj and TREC. There are 8 single-sentence tasks and 7 sentence- pair tasks. These tasks cover sentiment analysis, opinion polarity analysis, grammatical judgment, natural language inference, paraphrasing, etc. All these datasets' configurations are identical to that of LM-BFF (Gao et al.) (021).
Experimental Settings
The proposed method was implemented using PyTorch framework, utilizing the computational power of two NVIDIA V100 GPUs. The experiments were conducted with the same settings as LM-BFF, which measures the average performance of five different sampled Dtrain for each task with fixed set of seed Saeed {13. 21, 42, 87, 100). In the dataset, there are 16 samples per class. The hyperparameters are listed as follows: the learning rate is le-5, the batch size is 32, the maximum sequence length is 128, the maximum steps 1000, the number & of similar sentences retrieved is set to 64 and we save the last checkpoint. We use Adam as the optimizer. The models based RoBERTa-large for fair comparison with LM-BFF.
To validate the effectiveness, we compared ReFusion with several other models: (1) LM-BFR: a prompt-based fine-tuning approach; (2) DART et al.) 2022): differentiable prompt-based model, that can automatically search for the optimal prompt; (3) (P et al. (022b): prompt- based approach incorporating knowledge into the prompt verbalizer; and (4) CA-512: a retrieval- augmented prompt-based method concatenating retrievals with inputs. Unlike LM-BFF, DART, and KPT, we did not employ a grid search for the best parameters but instead chose default template and parameters based on LFF. Consequently, there is potential for the ReFusion to improve further through grid search. The templates are listed Appendix A
4.2
MAIN RESULTS
Table presents the main experimental results of the ReFusion and comparisons on 15 NKI tasks. The results shown in the form of means and variances, with the variance denoted by subseript.

---

Published as a conference paper at ICLR 2024
Table 2: Ablation studies on different modules.
Methods
MPQA
SUBJ
TREC
SNLI
QNLI
RTE
Avg
Roberta-L
83.62.5
90.32.
83.85.2
73.55.2
65.0s.0
64.120
76.7
Reranker
84.22.2
91.31.3
85.04.2
74.34.6
68.81.4
65.63.1
78.2
Ordered Mask
83.31.9
90.81.4
83.0s.8
74.94.0
68.31.4
65.83.1
77.7
ReFusion with Reranker
86.91.3
92.41.3
90.82.5
80.31.9
73.518
69.22.
82.2
ReFusion with Ordered Mask
87.01.s
92.40.7
90.73.0
80.31.3
73.01.o
70.42.5
82.3
ReFusion with All Rankings
86.71.1
92.50.8
90.33.7
80.61.4
73.01.1
70.9
82.3
For tasks with single sentences (S-Task), ReFusion consistently demonstrates superior performance across almost all benchmarks. ReFusion achieves state-of-the-art performance on tasks over 8 tasks. Besides, ReFusion improves the average performance the S-Task benchmark by about 2.1% than LM-BFF. Specifically, on the TREC task, ReFusion (90.3%) exhibits the maximum im- provements over LM-BFF (84.8%). For tasks consisting of pair sentences (P-Task), ReFusion con- tinues to demonstrate strong performance. ReFusion also achieves state-of-the-ar results on 5 tasks over 7 tasks. ReFusion can improve the average performance on the P-Task benchmark by about 3.0% than LM-BFF. For instance. the QNLI and SNLI benchmark, ReFusion (73% for QNLI, 80.6% for SNLI) significantly exceeds LM-BFF (64.5% for QNLI, 77.2% for SNLI).
The Avg-all represents the average performance of all 15 NKI tasks. For overall average perfor- mance, ReFusion achieves a score of 74.3%, marginally surpassing LM-BFF's 71.8%. This further highlights ReFusion's 's consistent and superior performance. Besides, ReFusion surpasses other mod- els like DART, CA-512, and KPT, delivering superior or comparable results. Notably, the standard deviation of ReFusion considerably smaller than that of other models, indicating that ReFusion produces stable results and offers superior robustness.
4.3 ABLATION STUDY
We conduct ablation experiments six representative tasks to show the contributions of each mod- ule to the overall performance. As shown Table compared to the baseline method (Roberta- large), adding ranked retrieval representations into hidden states (lines 2-3) can slightly improve the model's performance. However, the improvement is limited as the ranking schemes may not al- ways be suitable for every module in the model. With the adaptive retrieval integrator searching for the best combination of different ranking schemes, those methods using the ReFusion framework (lines 4-6) achieve significant improvements compared to those baselines (lines 1-3). The ReFusion searching among all ranking schemes (ReFusion with All Rankings) achieves the best performance. This indicates that the adaptive retrieval integrator can learn the functionality of each module and choose the most suitable ranking schemes for each module.
4.4 ANALYSIS
In this section, we analyze the ReFusion in terms of efficiency. As shown in the Table 3 we re- port the model performance and the inference latency based on the different queries, e.g., 'Input' which retrieves based on the input texts and 'Hidden' which retrieves based on the hidden states. Experimental results show that retrieving based on the hidden states can further improve the model performance but at the expense of increasing computational times. The benefit of 'Hidden' that the retrieval representations can dynamically change hidden states change. However, 'Hidden' requires retrieving at each retrieval fusion module, which may be only suitable for scenarios where adaptability is more important than time constraints, particularly in offline contexts.
Table 4 shows the latency breakdown of the model forward process, including the baseline model (Roberta-large) and the retrieval-augmented model (CA-512 and ReFusion). The inference latency of the baseline model is about 7.38 seconds per sentence. For retrieval concatenation- mentation (CA-512), the time cost of the retrieving process is about 6.63 seconds. However, due to the long input length and the quadratic complexity of attention, the time cost of the model forward significantly increases (31.93 seconds), thus leading longer inference latency. In the proposed
aug-
8

---

Published as a conference paper at ICLR 2024
Table 3: The trade-off between model performance and efficiency. 'Input' and 'Hidden' refer to re- trieving knowledge based on input texts and hidden states, respectively. 'Ace' and 'Latcy' refer to the accuracy of the task and the inference latency per
Table 4: Latency breakdown of model forward. 'F' 'R' and 'All' refer to the latency the model for- ward, retrieving, and overall latency, respectively. All results are shown seconds. More results can be found in the Appendix.
sentence in seconds, respectively.
Queries
MPQA
SUBJ
TREC
Methods
MPQA
SUBJ
TREC
Acc
87.9
92.1
84.0
Roberta-L.
F
7.35
7.41
7.37
Input
Laicy
3.7
3.8
11.1
R
5.71
8.46
5.72
Acc
88.7
91.7
85.6
CA-512
F
31.69
32.24
31.86
Hidden
Latcy
108.2
109.4
294.9
All
37.40
40.70
37.58
R
11.27
13.56
11.43
ReFusion
F 8.95
9.03
9.02
AII
20.22
22.59
20.45
ReFusion, the number of retrievals is not limited by the hyperparameter of max length, the triever module can retrieve more related knowledge at the cost of a bit more time (12.09 seconds). As expected, the inference latency of the model forward in the ReFusion is only bit longer than the baseline due to the time cost of retrieval fusion modules. Overall, the proposed ReFusion makes better trade-off between the model performance and the model efficiency for retrieval-basee aug- mentation.
RELATED WORK
Retrieval-Augmentee Transformers Retrieval-augmennte transformers [Borgeaud et al. 2022b
Lewis al. 2020, Shi al. 2022 Zhou al., 2023 Frisoni al. (022) leverage the information retrieval techniques to augment modern transformer- based language models with external knowl- edge databases. These retrieval-based augmentations can be classified into two types of work, i.e.,
retrieval concatenation-base augmentation (RC) and retrieval representation fusion (RF). The RC generally concatenates retrieval texts or representations. Ori Ram et al. (Ram et al. 2023) concate- information. Different from this, other works (Izacard & Grave 202I Guo et al., 2023a) first leveraged encoders to encode retrievals, then fed the feature representation concatenated by the retrieval representations and the in- put representation into a decoder. Those techniques will result in long sequence and then introduce large computations in the attention module. The RF involves introducing retrieval representations inside models. Some works (Borgeaud et al., [202aa] Li et al. (022) proposed to fuse the retrievals using a newly added cross-atted auss-attention module. However, those works require pre-traine pre-training the newly added cross-attentin module and the backbone model. Besides, the cross cross-attention module results in large computational overheads. The ReFusion in this paper introduces a computation-efift efficiien retrieval fusion module that can be fine-tuned with a few-shot dataset.
nated retrievals the context of inputs, thus providing models with
CONCLUSION
This paper proposes a new paradigm of retrieval-based augmentations, solving the bottleneck of retrieval concatenation-base augmentations. This paper proposes a computation-efffiien retrieval representation fusion framework with bi-level optimization named ReFusion. Experimental results demonstrate that ReFusion achieves superior performance while saving considerable computational ranking schemes. Besides, it is also
resources. For the future work, we plan to investigate worthwhile to explore the effect different byper-parameters. parameters in the framework, such as the module to fuse retrievals and the number of retrieved representations.

---

Published as a conference paper at ICLR 2024
7 ACKNOWLEDGEMENT
The work described in this paper was supported by grant from the Research Grants Council of the Hong Kong Special Administrative Region, China (Project No. CityU 11209122).
REFERENCES
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Mil- lican, George van den Driessche, Jean Saptistt 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, and Laurent Sifre. RETRO-Improving Language Models by Retrieving from Trillions of Tokens. In International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, vol- ume 162 of Proceedings of Machine Learning Research, pp. 2206 PMLR, 2022a.
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katic Mil- lican, George van den Driessche, Jean- Bann-stpe 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, and Laurent Sifre. Im- proving language models by retrieving from trillions of tokens. In International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 2206-2240. PMLR, 2022b.
Rakesh Chada and Pradeep Natarajan. Fewshotqa: A simple framework for few-shot learning of question answering tasks using pre-trained text-to-text models. In Marie-Francine Moens, Xuan- jing Huang, Lucia Specia, and Scott Wen-tau Yih (eds.), Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021, pp. 6081-6090. Association for Computational Lin- guistics, 2021.
Danqi Chen, Adam Fisch, Jason Weston, and Antoine Bordes. Reading wikipedia to answer open- domain questions. In Proceedings of the 55th Annual Meeting of the Association for Computa- tional Linguistics, ACL 2017, Vancouver, Canada, July 30 August 4, Volume 1: Long Papers, pp. 1870-1879, 2017.
Xiang Chen, Lei Li, Ningyu Zhang, Xiaozhuan Liang, Shumin Deng, Chuanqi Tan, Fei Huang, Luo Si, and Huajun Chen. EETROPPROMPT-Dcouplic augmented Prompt Learning. In Advances in Neural Information Processing Systems, volume 35, pp. 23908-23922. Curran Associates, Inc., 2022.
Knowledge from Memorization: Retrieval-
Yufei Cui, Ziquan Liu, Wuguannan Yao, Qiao Li, Antoni B. Chan, Tei-Wei Kuo, and Chun Jason Xue. Fully nested neural network for adaptive compression and quantization. In Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI 2020, pp. 2020.
Yufei Cui, Ziquan Liu, Qiao Li, Antoni B. Chan, and Chun Jason Xue. Bayesian nested neural networks for uncertainty calibration and adaptive compression. IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2021, virtual, June 19-25, 2021, pp. 2392-2401, 2021.
Yufei Cui, Yu Mao, Ziquan Liu, Qiao Li, Antoni B. Chan, Xue (Steve) Liu, Tei-Wei Kuo, and Chun Jason Xue. Variational nested dropout. IEEE Trans. Pattern Anal. Mach. Intell., 45(8); 10519-10534, 2023.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter the Association for Computational Linguistics: Human Language Technologies. NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume I (Long and Short Papers), pp. 4171-4186 Association for Computational Linguistics, 2019.
10

---

Published as a conference paper at ICLR 2024
Giacomo Frisoni, Miki Mizutani, Gianluca Moro, and Lorenzo Valgimigli. BioReader: a Retrieval- Enhanced Text-to-Teat Text Transformer for Biomedical Literature. In Yoav Goldberg, Zomnisaa Kozareva, and Yue Zhang (eds.). Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, EMNLP 2022, Abu Dhabi, United Arah Emirates, December 7-11, 2022, pp. 5770-5793. Association for Computational Linguisties, 2022.
Tianyu Gao, Adam Fisch, and Danqi Chen. Making pre-trained language models better few- shot learners. In Proceedings of the 59th Annual Meeting of the Association for Computa- lional Linguistics and the 11th International Joint Conference Natural Lunguage Processing, ACIAJCCLPP 2021, (Volume I: Long Papers), Virtual Event, August 1-6, 2021, pp. 3816-3830. Association for Computational Linguisties, 2021.
Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simcha, Felix Chern, and Sanjiv Kumar. Accelerating large-scale inference with anisotropic vector quantization. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 3887-3896. PMLR, 2020.
Zhicheng Guo, Sijie Cheng, Yile Wang, Peng Li, and Yang Liu. Prompt-Guider Retrieval Aug- mentation for Non- lee---ooowlle--eeeeeeeee Intensvv Tasks. In Findings of the Association for Computational Linguistics: ACL 2023, Toronto, Canada, July 9-14, 2023, pp- 10896-10912 Association for Computational Linguistics, 2023a.
Zhicheng Guo, Sijie Cheng, Yile Wang, Peng Li, and Yang Liu. Prompt-guided retrieval augmen- tation for non non- boon--oowlege--ntensii
tasks. In Anna Rogers, Jordan L. Grabbe,,, and Naoaki Okazaki (eds.). Findings of the Association for Computational Linguisties: ACI. 2023, Toronto, Canada, July 9-14, 2023, pp. 10996-10912. Association for Computational Linguistics, 2023b.
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming- Wei Chang. REALM: retrieval- augmented language model pre-training. CoRR, abs/2002.08909, 2020.
King-Wei
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. The Tenth Inter- national Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022a.
Shengding Hu, Ning Ding, Huadong Wang, Zhiyuan Liu, Jingang Wang, Juanzi Li, Wei Wu, and Maosong Sun. Knowledgeable prompt-tuning: Incorporating knowledge into prompt verbalizer for text classification. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022, pp. 2225-2240. Association for Computational Linguistics, 2022b.
Gautier Izacard and Edouard Grave. Leveraging passage retrieval with generative models for open domain question answering. In Paola Merlo, Jörg Tiedemann, and Reut Tsarfaty (eds.), Pro- ceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume, EACL 2021, Online, April 19 23, 2021, pp. 874-880. Association for Computational Linguistics, 2021.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizztioo with gumbel-softuax, In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Truck Proceedings. OpenReview.net, 2017.
Ganesh Jawahar, Benoît Sagot, and Djamé Seddah. What does BERT learn about the structure of language? In Proceedings of the 57th Conference of the Association for Computational Linguis- fics, ACL 2019, Florence, Italy, July 28- August 2019, Volume 1: Long Papers, pp. 3651-3657, 2019.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data, 7(3):535-55) -473555-77,-5 2019.
Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick S. H. Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020, pp. 6769-6781. 2020.
11

---

Published as a conference paper at ICLR 2024
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. In Sth International Conference on Learning Representations, ICLR 2020. Addis Ababa, Ethiopia, April 26-30, 2020, 2020.
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, and Douwe Kiela. augmenee generation for showledge-intensing NLP tasks. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Pro- cessing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
Zonglin Li, Ruiqi Guo, and Sanjiv Kumar. Decoupled context processing for context augmented language modeling. In NewrIPS, 2022.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: differentiable architecture search. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans. LA, USA, May 6-9, 2019. OpenReview.net, 2019a.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized BERT pretraining approach. CoRR, abs/1907.1 1692. 2019b.
Yu Mao, Yufei Cui, Wei-Wei Kuo, and Chun Jason Xue. Accelerating general-purpose lossless compression via simple and scalable parameterization. In João Magalhaes, Alberto Del Bimbo, Shin'ich Satoh, Nicu Sebe, Xavier Alameda-Pineda. Qin Jin, Vincent Oria, and Laura Toni (eds.), MM '22: The 30th ACM International Conference Multimedia, Lisboa, Portugal, October 10 14, 2022, pp. 3205-3213, ACM, 2022.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to- text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020. URL http:
77 7/jmlr. wwtp.orpppppp.///.1... 200 074 html
Ori Ram, Yuval Kirstain, Jonathan Berant, Amir Globerson, and Omer Levy. Few-shot question an- swering by pretraining span selection. In Chengqing Zong, Fei Xia, Wenjie Li, and Roberto Nav- igli (eds.). Proceedings of the 59th Annual Meeting of the Association for Computational Linguis- fics and the 11th International Joint Conference on Natural Language Processing, ACL/IJCNLP 2021, ( (Volume 1: Long Papers). Virtual Event, August 1-6, 2021, pp. 3066-3079. Association for Computational Linguistiss 2021.
Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Leyton- Leyton-Brown, and Yoav Shoham. In-context -uggreet-- language models. CORR, abs/2302.00008 000833, 2023. doi: 2044850/ARXXV223322000 00083
URL tittps://don httg//11..4550//rxiii 2302.
Oren Rippel, Michael A. Gelbart, and Ryan P. Adams. Learning ordered representations with nested dropout. In Proceedings of the 3/th International Conference on Machine Learning, ICML 2014, Beijing, China, 21-26 June 2014, volume 32 of JMLR Workshop and Conference Proceedings, pp. 1746-1754. JMLR.org 2014.
Ensheng Shi, Yanlin Wang, Wei Tao, Lun Du, Hongyu Zhang, Shi Han, Dongmei Zhang, and Hong- bin Sun. RACE: auumened Commit Message Generation. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, EMNLP 2022, Abu Dhabi, United Arab Emirates, December 7-11, 2022, pp. 5520-5530. Association for Computational Linguistics, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 1998-6008. 2017.
Alex Wang, Amanpreet Singh. Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding, In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
12

---

Published as a conference paper at ICLR 2024
Ningyu Zhang, Luoqiu Li, Xiang Chen, Shumin Deng. Zhen Bi, Chuanqi Tan, Fei Huang, and Huajun Chen. Differentiable prompt makes pre-trained language models better shot learners. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022, 2022.
Shuyan Zhou, Uri Alon, Frank F. Xu, Zhengbao Jiang, and Graham Neubig, DocPrompting: Gen- erating Code by Retrieving the Docs. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023, 2023.
13

---

Published as a conference paper at ICLR 2024
A TEMPLATES ON ALL TASKS
T an overview of the manual templates and selected label words used for each dataset in this paper. These templates and label words were created following LM-BFF (Gao et al. (001).
Table 5: Templates and label words used in this paper.
Task
Prompts
Label word
SST-2
[CLS] x It was [MASK]. [SEP]
"0":"terriblle, "I"."great"""
SST-5
[CLS] It [MASK]. [SEP]
"0"""terrible"""" "bad",
"2" "oka",, "3"" "gooo""""""""""""... "4"" "great"
MR
[CLS] It was [MASK]. [SEP]
"O":"terrible", "I"."great"
CR
[CLS] It was [MASK]. [SEP]
"O"""terrible". "I":"great"
MPQA
[CLS] was [MASK] [SEP]
"O":"terrible", "1":"great"
SUBJ
[CLS] This is [MASK] [SEP]
"O":"subjecttvve" "b"etojecctv"""""""""
TREC
[CLS] [MASK] æ [SEP]
"I.:"Deccrittion"""""""tiiit" """""i""""""""""esiisii"i"""""""
" "S":b"u
CoLA
[CLS] x It was [MASK]. [SEP]
MNLI
[CLS] I1 [MASK], T2 [SEP]
"contradiction". "entailmmmmmmmt
"neutral": "Maybe"
MNLI-m
[CLS] J7 [MASK], 32 [SEP]
"contradiction" "No","entailmmem""""We
"neutral") "Maybe"
SNLI
[CLS] T1 [MASK] 42 [SEP]
"contradiction" "ettimmena "Yes""
"neutral"" "Maybe"
QNLI
[CLS] 01 [MASK] 1 T2 [SEP]
"not entailment" "No "No "entailment""""Ve" "Ye"" "not entailment", " ""entaimmnn"" "Ye""
RTE
[CLS] 01 [MASK], 12 [SEP]
MRPC
[CLS] 27 [MASK]. 32 [SEP]
"O":"No", "e":"Yes"
QQP
[CLS] X1 [MASK] [SEP]
"O":"No", "es""Yes"
B RESULTS ON AUTOREGRESSIVE MODELS
This paper also evaluates the ReFusion with the classical autoregressive language model T5 Raf- fel et al. (2020). As shown in Table 6 ReFusion outperforms the T5-base model on most tasks, raising the overall average score from 78.2 to 79.6. The results are consistent with that of Roberta- large, demonstrating the robustness of ReFusion and the effectiveness of applying ReFusion transformer-based models.
Table 6: Results with T5-base.
Methods
SST-2
SST-5
MR
CR
MPQA
SUBJ
TREC
CoLA
Avg-S
T--base
91.0p.3
46.61.7
87.61. 88.3p,9
91.4a.2
85.9
84.51.7
91.8s.
51.13,3
78.8
ReFusion (T5)
92.40.3
47.2n.9
91.4g.4
86.1
87.41.3
94.90.5
53.74.1
80.2
Methods
MNLI
MNLI-m
SNLI
QNLI
RTE
MRPC
QQP
Avg-P
Avg-all
T5-base 73.90.8
75.80.6
79.92.
82.
79.01.9
75.5
76.20.s
77.5
78.2
ReFusion (T5)
75.31.1
77.3
83.0g.s
83.6a.
79.1
77.827
76.
78.9
79.6
C RESULTS ON FULL TRAINING SET
This paper compares the proposed ReFusion with LM-BFF (Gao et al. 2021) 1) on several tasks un- der the prompt-based setting with the full training set. As shown in Table 7 ReFusion generally demonstrates superior performance compared to LM-BFF. The average performance of ReFusion surpasses that of LM-BFF by 0.9%. The results suggest that ReFusion's performance superiority is consistent and not dependent on the size of the dataset. This also implics that ReFusion is robust and can generalize well across varying amounts of data.
14

---

Published as a conference paper at ICLR 2024
Table 7: Full training set results compared with LM-BFF
Methods
SST-2
SST-5
MR
CR
MPQA
SUBJ
TREC
CoLA
RTE
Avg
LM-BFF
95.0
58.7
90.8
89,4
87.8
97.0
97.4
62.6
80.9
84.6
ReFusion
95.6
61.0
92.3
91.4
84.4
97.1
97.6
62.8
85.2
85.3
D THE MODULE CHOICE WHERE THE RETRIEVALS ARE FUSED
Table presents the results of fusing the retrievals into different modules, i.e., key, value, query, and fin (feed-forward network) modules. Experimental results show that fusing retrievals into only key or query, key, and value modules can achieve the best performance. However, except for the module, fusing retrievals into the modules in the attention modules yields competitive performance. This indicates that the proposed ReFusion can improve the model's 's focus.
Table 8: Different module choices where the retrievals are fused.
Module
MPQA
SUBJ
TREC
SNLI
QNLI
RTE
Avg
key
86.7
93.5
94.2
79.2
76.6
71.1
83.6
value
86.6
93.4
94.0
78.5
76.3
71.8
83.4
query
87.3
93.6
93.6
78.6
76.3
71.8
83.5
key and value
86.3
93.3
93.0
81.0
72.1
73.3
83.2
query and key
86.6
93.6
94.2
78.6
74.8
72.6
83.4
query and value query and key and value ffn
86.7
93.9
93.0
80.1
74.8
72.6
83.5
86.3
92.9
94.2
79.2
76.4
72.6
83.6
84.0
93.7
94.2
78.0
75.5
70.4
82.6
E IMPACT OF RETRIEVER METRICS
The retriever metrics measure the similarity different data the retrieval database. Table shows the results of the ReFusion retrieving based on two similarity metries. Experimental results show that using the inner product has a slightly higher overall performance than that of the L2 norm. This observation suggests the potential for further enhancements in ReFusion through the adoption of inner product metrics.
Table 9: The impact of retriever metrics. 'L2' and 'IP' refer the metric of L2-norm and inner- product.
Metries
MPQA
SUBJ
TREC
SNLI
QNLI
RTE
Avg
L2
86.7
92.5
90.3
80.6
73.0
70.9
82.3
IP
86.4
92.1
91.2
80.9
73.3
71.8
82.6
F IMPACT OF RETRIEVAL NUMBER k
Table 10 illustrates the impact of varying the number of retrievals. As k increases from 1 to 8, the proposed ReFusion improves performance on most tasks. This suggests that fusing more retrievals can add valuable information to enhance the model's performance. As & continues to increase, there is a slight fluctuation performance. For example, the TREC shows a peak performance at A = 16, and the SNLI peaks at & = 32. The average performance remains relatively stable, around 83.5 to 83.8. When is larger than 64, there a general trend of either plateauing or slightly declining performance across most tasks. This could imply that there may be diminishing returns with too many retrievals, with additional information possibly introducing noise or irrelevant data that complicating the model's filtering process.
15

---

Published as a conference paper at ICLR 2024
Table 10: The impact of retricval number k.
k
MPQA
SUBJ
TREC
SNLI
QNLI
RTE
Avg
86.7
92.6
91.0
78.2
74.3
72.9
82.6
2
85.9
93.6
90.2
79.4
74.8
71.5
82.6
86.1
93.4
93.8
79.6
74.6
72.2
83.3
8
87.8
93.8
93.0
79.1
74.9
72.6
83.5
16
87.0
93.8
94.4
79.8
75.3
71.5
83.6
32
87.8
92.9
94.4
81.1
74.6
72.2
83.8
64
86.3
93.3
93.0
79.6
74.7
71.8
83.1
128
86.6
93.0
93.6
79.4
74.9
72.2
83.3
256
85.1
92.4
93.4
80.4
74.5
72.2
83.0
512
87.0
93.8
93.2
80.2
74.9
72.2
83.6
G RESULTS ON FEW-SHOT KNOWLEDGE-INTENSE TASKS.
This paper also evaluates the ReFusion knowledge-intensst tasks, such the SQuAd dataset, in the few-shot setting (16-shot). Table IT shows that the ReFusion significantly enhances the baseline model's performance, outperforming other comparison methods. This suggests the ReFusion sub- stantially contributes to the model's understanding and reasoning abilities knowledge-intensinn tasks.
Table 11: Results on few-shot knowledge-intensive tasks.
Methods
SQuAD
T5-base
46.63.9
Splinter
|Ram et 2021
54.6
FewShotBART (Chada & Natarajan 2021
55.52.0
ReFusion (T5)
58.6n.9
H RESULTS WITH LORA.
This paper also evaluates the ReFusion with different fine-tuning techniques, such as LoRA (Hu et al. 2022a). Table 12 shows the results of using LoRA to fine-tune the backbone model with/without the ReFusion. Generally, the ReFusion can also outperform the baseline from 63.8 to 64.7. However, compared to full parameters fine-tuning, the degree of improvements in Table T2 not competitive. This may lie in the fact that in the few-shot setting, fine-tuning the model with LoRA and RA NKI tasks is quite difficult task for deep learning models, requiring stronger learning and alignment capability, which is still a promising domain to be explored.
Table 12: Results with Roberat-large tuned by LoRA (Hu al.
2022ap.
Methods
SST-2
SST-5
MR
CR
MPQA
SUBJ
TREC
CoLA
Avg-S
Roberta-L
91.40.1
45.24.3
87.5p.
90.2pg
78.92.3
85.3u.7
41.40.
4.018
65.5
ReFusion
91.1a7
44.8n.
88.5n.4
90.8p.1
81.10.
89.1
46.549
2.502
66.8
Methods
MNLI
MNLI-m
SNLI
QNLI
RTE
MRPC
QQP
Avg-P
Avg-all
Roberta-L
57.0n.1 56.80.1
57.40.6
68.95.2
62.4p.5
56.7 56.5
72.36.4
59.41.3
62.0
63.8
ReFusion
57.308
67.8p66
64.2
71.23.7
62.43,7
62.3
64.7
16


Reviews:
[{"review_id": "2DMn7Rp9AD", "reviewer": "Reviewer_9abu", "summary": "With an aim to inject external knowledge into language models (LMs) for performing non-knowledge-intensive (NKI) tasks, including sentiment analysis, opinion polarity analysis, grammatical judgment, natural language inference, etc., towards which retrieval-based augmentation has been struggling as a solution due to the incapability of LMs to handle long sequence during the process, this paper proposes \"ReFusion\" a computation-efficient retrieval representation fusion framework as a solution to handle the same. Towards this, at first, top-k similar sentences are retrieved using an online retrieval module; thereafter, based on the neural architecture search (NAS) optimal ranking scheme, several modules are replaced with either fusion module with the reranker-based scheme, the fusion module with the ordered-mask-based scheme, or the original module. Authors have conducted comprehensive experiments across 15 NKI tasks comprising 8 single-sentence tasks and 7 sentence-pair tasks. Experimentations are performed using the RoBERTa LMs where ReFusion achieves state-of-the-art performance on 5 tasks over 8 single-sentence tasks and also achieves state-of-the-art performance on 5 tasks over 7 sentence-pair tasks. The study of the averaged performance reflects ReFusion is unbeatable in any of the cases over all 15 NKI tasks. Authors have also performed an ablation study to understand the importance of the sub-components, and the results reported validate that combining different ranking schemes on different tasks is necessary. Moreover, ranking schemes are not always suitable for every layer in LMs. Therefore, ReFusion disables the fusion module at some layers, thus integrating all effective candidate fusion modules.", "questions": "1. 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?\n2. Authors should consider extending this study to other LMs except Roberta.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The authors claim to be the first to propose fusing the representations of retrievals directly into models to solve the performance and efficiency bottleneck of prompt-based techniques.\n2. Ranking-based weighted representations are fused to LMs to enhance their performance over NKI tasks, which helps to overcome the issue of LMs' inability to handle long sequences during the retrieval-based augmentation processes.", "weaknesses": "1. The fusion of representation directly to LMs' layers is based on top-k sentence retrieval based on similarity. However, the authors need to study the impact of some more metrics while retrieving top-k similar sentences.\n2. This approach is suitable for situations where adaptability is more important than time constraints.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pSbEjvG1Xc", "reviewer": "Reviewer_aM36", "summary": "The paper discusses an approach to enhance prompt-based fine-tuning techniques for language models by integrating retrieval-augmented methods. The main contribution is the introduction of the ReFusion framework, which fuses representations of retrieved similar texts directly into the model to address performance and efficiency bottlenecks in existing techniques.", "questions": "1. The caption for Figure 2 is not clear. \n2. 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.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. ReFusion introduces a novel approach to combining retrieval information with model representations, which is a step forward in prompt-based learning.\n2. The experimental results indicate that ReFusion achieves superior performance over other models, suggesting a better understanding capability of the language models.\n3. Evaluation comprehensivenss: The paper conducts a series of experiments on 15 different natural language understanding (NKI) tasks, which is a broad evaluation of the ReFusion method.", "weaknesses": "1. 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.\n\n2. 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.\n\n3. Compatibility with Current Language Models: Modern language models, such as GPT-4 and LongChat, are designed to handle extended sequences, which facilitates the use of context augmentation methods. However, the representation fusion approach proposed by ReFusion necessitates fine-tuning of language models. This requirement may not be compatible with the current trend of employing language models as 'black-box' functions, where fine-tuning is either not possible or practical due to access or resource constraints.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "p7Anttsz1i", "reviewer": "Reviewer_UsSB", "summary": "Retrieval-augmented models have gained much attention for knowledge-intensive tasks, yet the core drawback is the increased computation brought by the augmented context. The paper proposes retrieval-augmented natural language understanding model that is computation-efficient. Specifically, the authors introduces three modules: 1) retrieval,  2) reranker, and 3) architecture search module. The retrieval module retrieves relevant sentences, then reranker reweighs retrieved representation, and finally, search module finds suitable fusion module for each layer. At inference, the search module simply picks a module with the highest score and utilizes the module for inference.", "questions": "Q1. 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The authors illustrate the effectiveness of the proposed model on 15 NLU tasks. \n- The proposed model is technically sounding. Injecting retrievals via representation fusion is reasonable.\n- Experimental settings are well listed out for reproducibility\n- Appendix covers meaningful discussions, such as the effect of retrieval with different k", "weaknesses": "- 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. \n- 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. \n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LWpiXLdFJW", "reviewer": "Reviewer_UAxr", "summary": "This paper develops a new method, ReFusion, for combining retrieved contextual representations with example representations to improve models of non-knowledge intensive tasks like sentiment and NLI. The approach adds a weighted average of retrieved passage representations into the CLS token representations. These representations can be ranked by a learned weighting scheme or a learned mask over dimensions. For ReFusion, this choice is made by a search module. Overall, ReFusion does well in evaluations, and the paper includes valuable ablation studies that help us understand which pieces contribute most. Those studies suggest that the search component is especially important.", "questions": "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 favor 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.)", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "This is a thoughtful paper that poses a well-defined, tightly circumscribed problem and addresses it in a creative and successful way. Assuming we are working with models with small window sizes, this seems to be a good way to handle lots of context. It's also valuable to have additional evidence that retrieval (in the form of additional similar examples) helps even for tasks like the ones addressed here.", "weaknesses": "1. 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.\n\n2. 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 on 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 many well simply not be understanding the description and its connection to the experimental results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)