Source-linked AI summary

Lift Yourself Up: Retrieval-augmented Text Generation with Self Memory

Xin Cheng, Di Luo, Xiuying Chen, Lemao Liu, Dongyan Zhao, Rui Yan

arXiv:2305.02437v3cs.CLcs.AI

TL;DR

Selfmem addresses the limitation of fixed-corpus memory retrieval by using a retrieval-augmented generator to create an unbounded memory pool and selecting self-generated outputs as memory. Across multiple text-generation tasks, it surpasses strong baselines and improves state-of-the-art performance, while requiring additional computational resources for memory selection.

  • Problem

    Fixed-corpus retrieval constrains memory-augmented generation because the finite retrieval space limits memory quality and potential.

  • Method

    Selfmem iteratively generates an unbounded candidate memory pool and uses a memory selector to choose one self-generated output as memory for subsequent generation.

  • Results

    Selfmem surpasses strong baselines and improves state-of-the-art performance across multiple text-generation tasks and generation paradigms, including 50.3 ROUGE-1 on XSum and 62.9 ROUGE-1 on BigPatent.

  • Takeaways & Limitations

    Self-generated outputs can serve as effective memory, supporting retrieval-augmented generation across translation, summarization, and dialogue tasks.

  • Takeaways & Limitations

    Selfmem requires additional computational resources for memory selection, a more serious concern for large datasets with long contexts because transformer computation has quadratic time complexity.

Abstract

from arXiv · show

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, XSum (50.3 ROUGE-1), and BigPatent (62.9 ROUGE-1), 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 bottlenecks and provide insights for future research.

1 Introduction

Retrieval-augmented generation traditionally relies on similar memories from a fixed corpus, but Selfmem instead uses generated outputs as recursively selected memory. The framework evaluates this approach across multiple generation tasks and reports strong benchmark results.

  • Retrieval-augmented generation combines input text with memories retrieved from an external database, typically the training corpus.
  • The primal problem assumes that retrieving memory most similar to the input improves generation quality.The paper motivates this principle through a reported correlation between memory similarity and translation quality.
  • Selfmem addresses fixed-corpus limitations by generating an unbounded memory pool and selecting one output for the next generation round.It combines the primal principle with the dual premise that better generation can produce better memory.
  • Selfmem uses a retrieval-augmented generator under either fine-tuned small-model or few-shot LLM paradigms, followed by memory selection.The generator is trained with labeled data in the small-model setting and uses retrieved in-context samples with a fixed LLM.
  • Selfmem achieves state-of-the-art results in four JRC-Acquis directions, XSum with 50.3 ROUGE-1, and BigPatent with 62.9 ROUGE-1.

2 Related Work

Related work has used external memory across many text-generation applications and has developed increasingly sophisticated retrieval and reranking methods. Selfmem differs by feeding selected candidates back as memory for another generation round rather than making them the final output.

  • External memory and retrieval have been applied to machine translation, dialogue, summarization, code generation, language modeling, and other NLP tasks.
  • Neural text generation research has explored generative and discriminative reranking to reduce exposure bias and optimize desired metrics.
  • Selfmem uses selected candidates as memory for subsequent retrieval-augmented generation, enabling later rounds to produce candidates with enhanced memory.

3 Methods

Selfmem combines a retrieval-augmented generator with a memory selector to iteratively generate candidates and reuse selected outputs as memory. Its framework supports fine-tuned small models and few-shot LLMs, with candidate generation and metric-based memory selection forming the core pipeline.

  • Generation as memory: 38.89 BLEU for training data versus 58.58 BLEU for the model’s own output motivates using generated outputs as memory.The motivating experiment compares training-data memory with the model’s output within the unbounded generation space.
  • Overall framework: Selfmem uses a retrieval-augmented generator and memory selector in complementary primal and dual stages.The generator takes source and memory to produce candidates, while the selector uses the source and candidates to select memory for later generation.
  • Retrieval-augmented generator: The generator supports fine-tuned small models and few-shot LLMs, using retrieved memory during training or inference.For small models, the generator is trained with labeled data and retrieved memory; for LLMs, a fixed black-box model uses retrieved in-context examples.
  • Memory selector: The memory selector chooses one candidate from the generator’s pool using model-free metrics such as BLEU or ROUGE rather than generation probability.The selected candidate becomes memory for the subsequent generation round, and the selector is trained to match metric-derived scores.
  • Combine generator and selector: Selfmem separates hypothesis mode from candidate mode: one output supports evaluation, whereas N outputs train the selector and support memory selection.The two modes are integrated into the complete Selfmem framework.

4 Experimental Setup

Selfmem is evaluated across translation, summarization, and dialogue generation using seven datasets and both trainable and LLM-based generation settings. The experiments use task-specific generators, selectors, retrieval, candidate generation, and validation-based iteration.

  • Datasets: Selfmem is evaluated on seven datasets spanning neural machine translation, abstractive summarization, and dialogue generation.The translation experiments use four JRC-Acquis directions; summarization uses XSum and BigPatent; dialogue uses DailyDialog.
  • Retrieval and generation: The framework uses BM25 retrieval and beam-search candidate generation with beam width 50 across tasks.The number of Selfmem iterations is determined by validation-set performance.
  • Translation: Translation uses a randomly initialized Transformerbase generator for the trainable-small-model setting and XGLM for few-shot LLM inference.Translation evaluation includes BLEU, TER, and chrF++; the selector uses XLM-Rbase with BLEU as its selection metric.
  • Summarization and dialogue: BARTbase is used for BigPatent and dialogue generation, while BRIO is used for XSum summarization.Summarization uses ROUGE (R1/2/L), and dialogue uses BLEU (B-1/2) and Distinct (D-1/2).
  • Memory selection: Memory selectors use RoBERTabase for summarization and dialogue, with task-specific combinations of BLEU or ROUGE metrics as selection objectives.Dialogue uses a linear combination of B-1/2, while summarization uses R-1/2/L.

5 Experimental Results

Selfmem improves retrieval-augmented generation across translation, in-context learning, summarization, and dialogue by replacing retrieved memory with selected self-generated memory. Results also show that the selector’s optimization target can shape dialogue attributes.

  • Machine Translation: Selfmem evaluates trainable and few-shot translation systems against vanilla and retrieval-augmented baselines on four JRC-Acquis directions.The translation evaluation uses BLEU, with TER and chrF++ reported additionally because BLEU alone may not guarantee superior translation quality.
  • Machine Translation: Self-memory is more similar to the reference than retrieval memory and serves as more effective memory for generating the final translation.Unlike reranking, Selfmem selects candidates as memory, so they need not exceed beam hypotheses in quality.
  • Machine Translation: Few-shot translation improves with larger XGLM models and more similar demonstrations, although the cross-lingual LLM remains below supervised baselines.Performance improves from random to kNN to Self demonstrations across model sizes.
  • Summarization: Self-memory substantially improves summarization, yielding +1.2 R1 for BRIO and +18.5 R1 for BART while achieving state-of-the-art results on XSum and BigPatent.The gain is significantly larger on BigPatent, whose patent documents exhibit considerable similarity.
  • Dialogue Generation: 3.08 B-1 and 0.6 B-2 are gained over retrieved memory on BARTjoint in dialogue generation, while optimizing D-1,2 targets more diverse and informative responses.BLEU-based optimization improves B-1/2 but falls behind on D-1/2, motivating direct optimization of the diversity metrics.

6 Further Analysis

Further analyses attribute Selfmem’s iterative gains primarily to memory selection rather than generator tuning. They also connect generated-token frequency with accuracy and show that the candidate pool improves across iterations.

  • Selector Analysis: Improving the selector significantly outperforms the 60.11 BLEU state-of-the-art result, while oracle candidate-pool scores improve during iteration.The analysis examines both generation quality and candidate-pool quality under different selector settings.
  • Selector Analysis: Oracle, quartile, average, and minimum candidate-pool scores all improve during iterative self-memory generation.The generator benefits from superior memory selected from its own unbounded output, then produces a higher-quality pool for the next round.
  • Generator Analysis: A fixed trained generator can distinguish good from bad memory, indicating that the generator is not Selfmem’s current bottleneck.This finding supports keeping the generator parameters fixed while focusing on memory selection.
  • Frequency Analysis: Tokens encountered more frequently during training receive higher generated-output 1-gram F1 scores.The frequency analysis groups generated translation tokens by training-corpus frequency and evaluates their 1-gram F1.

7 Conclusion

The paper introduces Selfmem to address bounded memory in retrieval-augmented generation by combining better generation with better memory. Experiments across tasks and generation paradigms show improved state-of-the-art performance and identify system bottlenecks.

  • Conclusion: Selfmem addresses bounded memory by uplifting retrieval-augmented generation with the model’s own output across tasks and generation paradigms.The framework is evaluated with trainable small models and few-shot prompted LLMs.
  • Conclusion: Selfmem surpasses strong baselines and improves state-of-the-art performance on several datasets.The paper also analyzes crucial components and identifies an existing system bottleneck for future research.

Limitations

Selfmem improves generation quality but has higher computational costs and uses a uniform transformer-based design across tasks, leaving room for task-specific architectures and methods.

  • Computational cost: Selfmem requires more computational resources for memory selection, especially on large datasets with long contexts such as BigPatent.The limitation is more pronounced because transformer architectures have quadratic time complexity.
  • Architectural scope: The framework is evaluated only with transformer-based generators and memory selectors.The same architecture is retained across all text generation tasks.
  • Future improvements: Task-specific architectures, training objectives, and generation methods may further improve performance across different text generation scenarios.

A Dataset Details

The paper reports dataset statistics and task-specific training configurations for machine translation, summarization, and dialogue generation.

  • Dataset Details: Dataset statistics for the three tasks are provided in Table 7.
  • Machine Translation: Machine translation uses a randomly initialized Transformerbase generator with joint BPE and shared memory-source encoder parameters.Training uses dropout 0.1, label smoothing 0.1, gradient clipping 1.0, Adafactor, 4,000 warm-up steps, a 4.4e-2 maximum learning rate, and 30 epochs.
  • Summarization: Summarization initializes the generator with BARTbase for BigPatent and BRIO for XSum.The maximum input lengths are 512 for XSum and 1,024 for BigPatent.
  • Dialogue Generation: Dialogue generation uses BARTbase on DailyDialog with tuned learning rates, dropout 0.1, batch size 64, label smoothing 0.1, and a maximum input length of 120.RoBERTabase is used as the memory-selector backbone for both summarization and dialogue generation.

C Evaluation Details

Evaluation uses task-specific standard metrics and reports additional baselines for XSum and BigPatent alongside detailed JRC-Acquis translation results.

  • Machine Translation: Machine translation is evaluated with BLEU, TER, and chrF++ using SACREBLEU3 signatures.The signatures are listed in Table 8.
  • Summarization: Summarization uses the standard ROUGE Perl package with R-1, R-2, and R-L settings and PTB tokenization.
  • Dialogue Generation: Dialogue generation is evaluated with NLTK BLEU using space tokenization and smoothing method 1, plus Distinct scores.
  • Reported Results: JRC-Acquis En→De results are reported for BLEU, TER, and chrF++, while XSum and BigPatent include additional baselines.Confidence regions for the SOTA models on XSum and BigPatent are also reported.

F Empirical analysis of latency

The paper analyzes Selfmem latency against a retrieval-augmented baseline across datasets and hardware, and separately compares generated outputs using human and GPT-4 pairwise judgments.

  • Latency Measurement: Selfmem latency is measured in seconds against a retrieval-augmented baseline across datasets using CPU and CUDA.The latency experiments use one Selfmem iteration on the same device with one NVIDIA A100 GPU and one AMD EPYC 7V13 64-Core Processor.
  • Human and GPT-4 Evaluation: Human and GPT-4 annotators perform pairwise rankings of Selfmem and baseline outputs.The evaluation samples 50 translation examples and 20 examples each for summarization and dialogue.
  • Evaluation Results: Figure 5 presents human and GPT-4 evaluation results.
Loading 2305.02437v3…