SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xin Cheng $^{1}$  Di Luo $^{2}$  Xiuying Chen $^{3}$  Lemao Liu $^{4}$  Dongyan Zhao $^{1}$  Rui Yan $^{2}$
1 Peking University 2 Remin University of China
3 KAUST 4 Tencent AI Lab
chengxin1998@stu.pku.edu.cn
With direct access to human-written reference as memory, retrieval-augmented generation has achieved much progress in a wide range of text generation tasks. Since better memory would typically prompt better generation (we define this as primal problem). The traditional approach for memory retrieval involves selecting memory that exhibits the highest similarity to the input. However, this method is constrained by the quality of the fixed corpus from which memory is retrieved. In this paper, by exploring the duality of the primal problem: better generation also prompts better memory, we propose a novel framework, Selfmem, which addresses this limitation by iteratively employing a retrieval-augmented generator to create an unbounded memory pool and using a memory selector to choose one output as memory for the subsequent generation round. This enables the model to leverage its own output, referred to as self-memory, for improved generation. We evaluate the effectiveness of Selfmem on three distinct text generation tasks: neural machine translation, abstractive text summarization, and dialogue generation, under two generation paradigms: fine-tuned small model and few-shot LLM. Our approach achieves state-of-the-art results in four directions in JRC-Acquis translation dataset, 50.3 ROUGE-1 in XSum, and 62.9 ROUGE-1 in BigPatent, demonstrating the potential of self-memory in enhancing retrieval-augmented generation models. Furthermore, we conduct thorough analyses of each component in the Selfmem framework to identify current system bottlenecks and provide insights for future research<sup>1</sup>.
In recent years, retrieval-augmented text generation has attracted growing interest across various fields, including neural machine translation[28, 17, 2], dialogue response generation[81, 6, 46], and language modeling[36, 77, 19]. This innovative generation paradigm initially equips a fine-tuned small model or a large language model (LLM) with access to an external database (typically the training corpus) using information retrieval techniques. Subsequently, the generation process is conducted based on both the input text and the retrieved memory.
In this paradigm, the guiding principle for memory retrieval is to find the memory that exhibits the highest similarity to the current input [36, 96, 49]. This aligns with the human intuition that a more similar demonstration sample typically offers more hints. As demonstrated in Figure 1, for a retrieval-augmented translation model, the memory similarity alone exhibits a strong correlation with the final translation quality, regardless of other factors that may influence translation quality (e.g.,
$^{1}$ Code and data available at: https://github.com/Hannibal046/SelfMemory
polysemy, morphology, and coreference). We define this as the primal problem: better memory prompts better generation. Consequently, numerous studies have focused on how to retrieve better memory, ranging from sparse retrieval to dense retrieval [10, 63], from a fixed retriever to a learnable retriever [41, 8], and from sentence-level memory to more fine-grained token-level memory [36, 35].
However, a fundamental limitation exists in all previous works: the memory is retrieved from a fixed corpus and is constrained by the corpus's quality. Due to the finite retrieval space, bounded memory significantly restricts the potential of memory-augmented generation models [97]. In this paper, we explore the duality of the primal problem, which posits that better generation also prompts better memory. We propose a novel framework called Selfmem, which iteratively employs a retrieval-augmented generator to create an unbounded memory pool and uses a memory selector to choose one output as memory for the subsequent generation round. By combining the primal and dual problem, a retrieval-augmented generation model can elevate itself using its own output, referred to as self-memory. The key insight behind Selfmem is that the text more closely resembling the data distribution during inference is not the training data [87], but the model's own output.
Figure 1: Relation between memory and hypothesis on JRC-Acquis  $\mathrm{En} \rightarrow \mathrm{De}$  dataset. The hypothesis is generated by a retrieval-augmented translator whose memory is retrieved from the training set. The X-axis represents the similarity between memory and the reference.
Selfmem consists of two complementary components:
a retrieval-augmented generator and a memory selector. The generator operates under two distinct paradigms: fine-tuning a small model or few-shot prompting an LLM. For the former, we train the generator with labeled data and retrieved memory, while for the latter, we employ a fixed black-box LLM exclusively for inference alongside retrieved in-context learning samples. We then use the generator's output to train a memory selector based on a specific performance metric. By simply replacing the retrieved memory with unbounded generated memory, we achieve higher-quality generation output (primal problem), which subsequently serves as memory for the next round after being refined by the memory selector (dual problem).
To evaluate the efficacy of the Selfmem, we carry out comprehensive experiments in three distinct text generation tasks: neural machine translation, abstractive text summarization, and dialogue generation. We witness substantial enhancements over robust baselines, attaining state-of-the-art outcomes in JRC-Acquis (four directions), XSum (50.3 ROUGE-1), and BigPatent (62.9 ROUGE-1). To gain deeper insights into the Selfmem, we meticulously investigate each crucial component and pinpoint the existing system bottleneck to guide future research endeavors.
2.1 Retrieval-augmented Text Generation
Since the world is not a snapshot once the training corpus is collected, we can never expect an ever-large model to capture everything in its parameters, even for LLMs like GPT-4 [62]. Therefore, it is crucial to equip these models with an external memory bank to store additional knowledge or useful demonstration examples for solving various NLP tasks[41, 78, 95].
In the translation domain, retrieval techniques have long been employed by the localization industry to enhance human translators' productivity and consistency even before the advent of machine translation [94]. Early works on machine translation primarily focused on utilizing memory for statistical machine translation (SMT) systems [80, 50]. For neural machine translation (NMT), [28] were the first to use search engines to retrieve memory from the training set and incorporate it with an external memory network. Subsequent research explored various aspects of retrieval-augmented NMT, such as memory encoding methods [92, 93, 31], joint training of retrievers and generators with monolingual data [8], memory granularity [35], and memory diversity [17]. For few-shot LLM generation, strategies for in-context example selection have been proposed to improve translation quality [2]. Furthermore, in-context machine translation has been shown to be effective for on-the-fly adaptation [79]. For dialogue response generation tasks, employing exemplar/template
retrieval as an intermediate step has proven advantageous for generating informative responses [89, 91, 6, 7]. In-context learning example retrieval also aids in controllable dialogue [46]. Other applications include abstractive summarization [64, 14, 18, 15], code generation [30], paraphrase generation [34, 83], language modeling [36, 105], counterfactual data generation [24], open domain question answering [12, 33] and semantic parsing [99].
2.2 Neural Text Reranking
By alleviating the discrepancy between training and inference (i.e., exposure bias) and directly optimizing desired metrics, two-stage reranking methods have facilitated significant progress in various text generation tasks. In machine translation, pioneering works by [75] and [61] introduced and popularized discriminative reranking for SMT. In the context of NMT, research has focused on two primary reranking approaches: generative reranking [56, 32, 88] and discriminative reranking [39, 71, 23]. For syntactic parsing, [21] were the first to employ a two-stage reranking method to select outputs from a base parser, while [11] introduced a maximum entropy reranker. In text summarization, RefSum [53] proposed a second-stage summarization framework to address train-test distribution mismatches. SimCLS [54] used pairwise Learning To Rank (LTR) to select candidates with the highest matching scores. SummaReranker [68] adopted a multi-task mixture-of-experts framework to leverage different metrics capturing various aspects of generated candidates. BRIO [55] reused the base model for a second round of fine-tuning with both cross-entropy loss and a candidate-level ranking loss. JGR [76] employed an alternate training paradigm to train the generator and reranker.
A key limitation of these reranking methods is that they only represent a one-way process, wherein the selected candidates become the system's final output. In contrast, our framework innovatively utilizes the chosen candidates as memory for the subsequent generation round of a retrieval-augmented generator, which can produce better candidates with enhanced memory.
In this section, we begin with a motivating experiment on generation as memory (§ 3.1). Then, we introduce Selfmem, a framework comprising a retrieval-augmented generator (§ 3.2) and a memory selector (§ 3.3). The complete framework and algorithm are illustrated in Figure 2 and Algorithm 1.
3.1 Generation as Memory
The primary motivation behind our framework stems from the observation that the memory, which is more similar in distribution to the data during inference, is not the training data (38.89 BLEU, as shown in the first row of Table 1). Instead, it is the model's own output (58.58 BLEU) within the unbounded generation space. One interesting exploration involves directly utilizing the generated output as memory in relation to the primal problem: better memory prompts better generation.
We conduct experiments on the JRC-Acquis En  $\rightarrow$  De dataset. The first row in Table 1 represents conventional retrieval-augmented training with retrieved memory and achieves a 58.58 BLEU score. However, directly incorporating beam output of this trained model as memory (Beam) back into the generation model does not yield any improvements (row 2), despite its higher similarity to the reference compared to the retrieved ones. We hypothesize two potential reasons for this: (1) the retrieval-augmented generator may not generalize effectively in this context due to the
Table 1: Experiments on the relation between memory quality and the final hypothesis quality, measured by the BLEU score with ground truth translation. The retrieval-augmented translator keeps fixed while the memory is obtained from different sources.
Memory Source	Memory Quality	Hypothesis Quality
Retrieval	38.89	58.58
Beam	58.58	58.43
Reference	100	90.43
Random	1.14	49.08
memory distribution shift (from 38.89 to 58.58), and (2) the beam memory does not offer any information gain compared to the retrieved one, even it exhibits more overlap with the references.
To investigate the first hypothesis, we conduct experiments under the oracle and random scenarios by using the reference as memory (Reference) and randomly sampled sentences as memory (Random). The result is shown in Table 1 and it illustrates that a retrieval-augmented generator (trained with
Figure 2: Overall framework. There are two components in Selfmem, a retrieval-augmented generator (a) and a memory selector (b). For the primal problem, (a) takes source and memory as input to generate candidates for (b). For the dual problem, (b) takes as input source and generated candidates to select memory for (a).
retrieved memory) has already learned to discriminate between different memories in both oracle and random scenarios, without updating the model weights.
To evaluate the second conjecture, we first define the token sets of the reference, retrieved memory, and beam memory as  $\mathcal{R},\mathcal{M}$ , and  $\mathcal{B}$ , respectively. The overlap token set, denoted by  $\mathcal{O}$ , is defined as the tokens that overlap with the references in the beam memory but not in the retrieved memory, which is represented as  $\mathcal{R} \cap \mathcal{B} - \mathcal{R} \cap \mathcal{M}$ .  $\mathcal{O}$  is considered as the additional information provided by the beam memory. Inspired by the confidence analysis of NMT model [58], we compute the set confidence score,  $\psi(\cdot)$ , as follows:
$$
\psi (\cdot) = \frac {1}{| \cdot |} \sum_ {y ^ {i} \in \cdot} p \left(y _ {i} \mid x, y _ {<   i}\right) \tag {1}
$$
where  $p(y_{i}|x,y_{<i})$  is defined by the generation model.  $\psi(\cdot)$  measures the confidence with which the generation model generates the tokens. The value of  $\psi(\mathcal{R})$  is 0.58, while that of  $\mathcal{O}$  is 0.76, indicating that the generator is relatively confident in generating tokens in  $\mathcal{O}$ , and therefore does not need to resort to external memory [38]. Beam search ranks generated candidates based on  $p(y|x)$ , where the selected memory falls within the confidence region of the generator and consequently provides no information gain. This observation motivates us to select memory according to metrics other than  $p(y|x)$  in the memory selector (§3.3).
3.2 Retrieval-augmented Generator
Given a text pair  $(x,y)$ , where  $x = \{\mathbf{x}_1,\dots,\mathbf{x}_{|x|}\}$  is the source,  $y = \{\mathbf{y}_1,\dots,\mathbf{y}_{|y|}\}$  is the target. They could be (document, summary) in summarization, (context, response) in dialogue generation or (source, target) in machine translation. The retrieval-augmented generation would first use  $x$  to retrieve memory  $m$  from datastore  $\mathbb{D}$ . Then the generator  $G_{\xi}(x,m)$ , parameterized by  $\xi$ , would take both  $x$  and  $m$  as input to generate the target sentence  $y$ . In this paper, following standard practice, we choose the training set as  $\mathbb{D} = \{(x^i,y^i)\}_{i=1}^{\mathbb{D}|}$ . For LLM as  $G_{\xi}$ , we use the standard in-context learning format to give  $(x,y)$  as demonstration example. For tunable generator  $G_{\xi}$ , we only keep the target side of top-1 retrieval results as memory and we consider two commonly used architectures: Joint-Encoder [29, 87, 41] and Dual-Encoder [92, 8, 17].
Joint-Encoder This architecture is the standard encoder-decoder-based model [3, 84]. The input is the concatenation of  $x$  and  $m$ . The encoder would first map the input into the hidden states  $H$ :
$$
H = \operatorname {E n c o d e r} (x [ \text {S E P} ] m) \tag {2}
$$
And the decoder would incorporate  $H$  by attention mechanism and generate tokens in an auto-regressive manner:
$$
h ^ {i} = \operatorname {D e c o d e r} (\operatorname {C r o s s A t t n} (H), y _ {<   i}) \quad P _ {G _ {\xi}} (\cdot | x, y _ {<   i}) = \operatorname {S o f t m a x} \left(h ^ {i}\right) \tag {3}
$$
Dual-Encoder Instead of treating  $x$  and  $m$  as a long sequence, this architecture has two encoders, one for  $x$  and the other for  $m$ . Their outputs are sequentially attended by the decoder with dual cross attention as in [17]:
$$
H _ {x} = \operatorname {S o u r c e E n c o d e r} (x) \quad H _ {m} = \operatorname {M e m o r y E n c o d e r} (m) \tag {4}
$$
$$
h ^ {i} = \operatorname {D e c o d e r} \left(\operatorname {C r o s s A t t n} \left(H _ {x}, H _ {m}\right), y _ {<   i}\right) \tag {5}
$$
We use Transformer [84] as the building block for both architectures and optimize  $G_{\xi}$  with NLL loss:
$$
\mathcal {L} _ {\mathrm {n l l}} = - \sum_ {t = 1} ^ {| y |} \log P _ {G _ {\xi}} \left(y _ {t} \mid x, m, y _ {<   t}\right) \tag {6}
$$
3.3 Memory Selector
The role of memory selector  $S_{\theta}(x,c)$ , parameterized by  $\theta$ , is to select one candidate  $c$  from the candidate pool  $\mathbb{C}$  generated by  $G_{\xi}$  based on a specific metric  $\Delta(\cdot, \cdot)$ . The chosen candidate  $c$  is then utilized as memory  $m$  for the subsequent generation round of  $G_{\xi}$ . As discussed in §3.1, using  $p_{G_{\xi}}(y|x)$  as the metric  $\Delta(\cdot, \cdot)$  would result in falling into the confidence region of  $G_{\xi}$ , leading to no information gain. Moreover, a larger value of  $p_{G_{\xi}}(y|x)$  does not necessarily guarantee improved generation quality [59]. Consequently, we define  $\Delta(\cdot, \cdot)$  as model-free metrics that are widely employed for assessing generation quality, such as BLEU for Neural Machine Translation (NMT) and ROUGE for Summarization. Our memory selector takes the concatenation of the source  $x$  and candidate  $c_i$  as input, and produces a multinomial distribution  $p_{S_{\theta}}(\cdot|x)$  over  $\mathbb{C}$ .
In this paper, we focus on the role of the memory selector,  $S_{\theta}(x,c)$ , which is parameterized by  $\theta$ . The objective of this selector is to choose a single candidate  $c$  from the candidate pool  $\mathbb{C}$ , generated by  $G_{\xi}$ , based on a specific metric,  $\Delta(\cdot,\cdot)$ .
$$
p _ {S _ {\theta}} \left(c _ {i} \mid x\right) = \frac {\exp \left(S _ {\theta} \left(x [ \mathrm {S E P} ] c _ {i}\right)\right)}{\sum_ {j = 1} ^ {| \mathrm {C} |} \exp \left(S _ {\theta} \left(x [ \mathrm {S E P} ] c _ {j}\right)\right)} \tag {7}
$$
In accordance with [39], the training goal for  $S_{\theta}$  is to minimize the discrepancy between the  $S_{\theta}$ 's predictions and the scores determined by  $\Delta(\cdot, \cdot)$ . This divergence is quantified using the Kullback-Leibler (KL) divergence.
$$
\mathcal {L} _ {\mathrm {k l}} = - \sum_ {i = 1} ^ {| \mathbb {C} |} p _ {M} \left(c _ {i}\right) \log p _ {S _ {\theta}} \left(c _ {i} \mid x\right) \quad \text {w h e r e} \quad p _ {M} \left(c _ {i}\right) = \frac {\exp \left(\Delta \left(c _ {i} , y\right) / \tau\right)}{\sum_ {j = 1} ^ {| \mathbb {C} |} \exp \left(\Delta \left(c _ {j} , y\right) / \tau\right)} \tag {8}
$$
$\tau$  is the temperature to control the smoothness of the distribution. At inference, the output of the  $S_{\theta}$  is  $\arg \max_{c_i\in \mathbb{C}}p_{S_\theta}(c_i|x)$ .
3.4 Combine Generator and Selector
We define two generation modes for  $G_{\xi}$ . The first mode, referred to as the hypothesis mode, generates a single output for each input, which is utilized for system evaluation. The second mode, known as the candidate mode, produces N outputs for a given input, and is employed for training  $S_{\theta}$  as well as memory selection. By integrating two modes together, we present the complete framework of our proposed model, Selfmem, as illustrated in Algorithm 1.
We assess the performance of Selfmem on three generation tasks, utilizing a total of seven datasets. Translation. We evaluate our framework on JRC-Acquis datasets [82], a collection of parallel
Algorithm 1 Selfmem Framework
Require: a dataset  $\mathbb{D}$ , a retriever  $R$ , a memory selection metric  $\Delta(\cdot, \cdot)$ , a retrieval-augmented generator  $G_{\xi}$ , and a memory selector  $S_{\theta}$
1: retrieve memory  $\mathbb{M}$  in  $\mathbb{D}$  with  $R$
2: train  $G_{\xi}$  with  $\mathbb{D}$  and  $\mathbb{M}$  (if not LLM)
3: use  $G_{\xi}$  to generate candidate pool  $\mathbb{C}$  with  $\mathbb{M}$  in candidate mode
4: train  $S_{\theta}$  on  $\mathbb{C}$  with  $\Delta (\cdot ,\cdot)$
5: while not converged in the validation set do
6:  $S_{\theta}$  selects memory from  $\mathbb{C}$  as  $\mathbb{M}$
7:  $G_{\xi}$  generates candidate pool  $\mathbb{C}$  with  $\mathbb{M}$  in candidate mode
8: end while
9:  $G_{\xi}$  generates the final hypothesis with  $\mathbb{M}$  in hypothesis mode
legislative text of European Union Law. It is the benchmark dataset used in translation memory-augmented NMT task [28, 92, 8, 17]. We choose 4 translation directions, namely, Spanish  $\leftrightarrow$  English (Es  $\leftrightarrow$  En), German  $\leftrightarrow$  English (De  $\leftrightarrow$  En). Summarization. We evaluate on 2 summarization datasets: 1) XSum [60], extreme summarization, a single-document summarization dataset with highly abstractive articles from British Broadcasting Corporation. 2) BigPatent [73], consisting of 1.3 million records of U.S. patent documents along with human-written abstractive summaries. Dialogue. We experiment on DailyDialog [44], which contains multi-turn dialogs on daily life topics and is used by [13, 4, 103]. The detailed statistics for these datasets can be found in the Appendix A.
We utilize the BM25 algorithm [70] for retrieval purposes. For all tasks, the candidate generation method consists of beam search with a beam width of 50. The number of iterations is determined by the performance on the validation set. For translation, we follow the approach of [93, 8, 17], employing a randomly initialized Transformer<sub>base</sub> architecture as  $G_{\xi}$  for trainable small model and XGLM [48] for LLM in-context learning. Evaluation metrics include BLEU, TER, and  $\mathrm{chrF++}$  obtained from SACREBLEU[66]. The memory selector  $S_{\theta}$  utilizes an XLM- $R_{base}$ [22] as backbone, with BLEU serving as  $\Delta(\cdot, \cdot)$ . For summarization, we initialize  $G_{\xi}$  with  $\mathrm{BART}_{\mathrm{base}}$ [40] for BigPatent and employ BRIO [55] for XSum. The evaluation metric comprises ROUGE (R-1/2/L) [47]. For dialogue generation,  $\mathrm{BART}_{\mathrm{base}}$  serves as the backbone for  $G_{\xi}$ . Our dialogue system is evaluated using BLEU (B-1/2) and Distinct (D-1/2) scores [43]. For both dialogue and summarization tasks, we adhere to the methods of [54, 26], adopting  $\mathrm{RoBERTa}_{\mathrm{base}}$ [52] as the backbone for  $S_{\theta}$ . The linear combination of B-1/2 is chosen as  $\Delta(\cdot, \cdot)$  for Dialogue Generation, while R-1/2/L is used for Summarization, following [76]. For further implementation details, please refer to the Appendix B and Appendix C for evaluation metrics.
5.1 Machine Translation
We select four translation directions and experiment with two generation paradigms: trainable small models and few-shot prompted LLMs [85, 20]. For trainable models, we explore two architectures (joint and dual, as detailed in §3.2). The baselines comprise two types of translation systems: one being the vanilla sequence-to-sequence model [3, 84] without memory augmentation, and the other consisting of retrieval-augmented translation models focusing on memory encoding [28, 92], memory construction [101], memory retrieval [8], and memory diversity [17]. Based on the experimental results<sup>2</sup> shown in Table 2, Selfmem significantly enhances the performance of  $G_{\xi}$  across four translation datasets and two different architectures. This is noteworthy, given that the parameters of the  $G_{\xi}$  remain fixed, with the only variable being the input memory. This finding is consistent with the primal problem which posits that improved memory typically leads to better generation results.
As higher BLEU scores in this range do not necessarily guarantee a superior translation system [9], we also evaluate our system using TER and  $\mathrm{chrF + + }$  . The results can be found in the Appendix D.