Retrieval-augmented Prompt Learning for Pre-trained Foundation Models

TL;DR

This work addresses the instability and memorization tendencies of prompt learning in pre-trained foundation models by introducing RetroPrompt, a retrieval-augmented framework that decouples knowledge from memorization through an open-book knowledge-store and dense retrieval. By integrating $k$NN-guided training and $k$NN-based probability for cloze-style inference, RetroPrompt reinforces contextual cues from the training data during input, training, and inference, applicable to both language and vision tasks. Across zero-shot, few-shot, and fully-supervised regimes,RetroPrompt demonstrates superior generalization, reduced memorization, and robust performance, including cross-domain transfer. The approach offers a practical, scalable path to enhance prompt-based learning in PFMs without extensive parameter updates, with potential extensions to generative and multilingual scenarios.

Abstract

The pre-trained foundation models (PFMs) have become essential for facilitating large-scale multimodal learning. Researchers have effectively employed the ``pre-train, prompt, and predict'' paradigm through prompt learning to induce improved few-shot performance. However, prompt learning approaches for PFMs still follow a parametric learning paradigm. As such, the stability of generalization in memorization and rote learning can be compromised. More specifically, conventional prompt learning might face difficulties in fully utilizing atypical instances and avoiding overfitting to shallow patterns with limited data during the process of fully-supervised training. To overcome these constraints, we present our approach, named RetroPrompt, which aims to achieve a balance between memorization and generalization by decoupling knowledge from mere memorization. Unlike traditional prompting methods, RetroPrompt leverages a publicly accessible knowledge base generated from the training data and incorporates a retrieval mechanism throughout the input, training, and inference stages. This enables the model to actively retrieve relevant contextual information from the corpus, thereby enhancing the available cues. We conduct comprehensive experiments on a variety of datasets across natural language processing and computer vision tasks to demonstrate the superior performance of our proposed approach, RetroPrompt, in both zero-shot and few-shot scenarios. Through detailed analysis of memorization patterns, we observe that RetroPrompt effectively reduces the reliance on rote memorization, leading to enhanced generalization.

Figures (5)

• Figure 1: Decoupling knowledge from memorization. To achieve a harmonious balance between generalization and memorization in prompt learning, we put forward a method that separates knowledge from mere memorization. Our approach involves creating a knowledge-store accessible for reference and retrieval throughout the training and inference phases.
• Figure 2: Illustration of RetroPrompt.
• Figure 3: We present the results on 9 image classification datasets in both the zero-shot and few-shot settings. For the few-shot setting, we employ $k$NN-train, which involves retrieving $k$NN to guide the training process. On the other hand, for the zero-shot setting, we utilize $k$NN-test, where $k$NN is retrieved to interpolate predictions.
• Figure 4: Fully-supervised performances.
• Figure 5: Case examples of Top-5 neighbors from trainset of ImageNet. The numbers before and after the arrow in "Predict Prob" represent the values before and after retrieval.