Prompt Diffusion Robustifies Any-Modality Prompt Learning

Yingjun Du,Gaowen Liu,Yuzhang Shang,Yuguang Yao,Ramana Kompella,Cees G. M. Snoek

2024-10-26

Abstract:Foundation models enable prompt-based classifiers for zero-shot and few-shot learning. Nonetheless, the conventional method of employing fixed prompts suffers from distributional shifts that negatively impact generalizability to unseen samples. This paper introduces prompt diffusion, which uses a diffusion model to gradually refine the prompts to obtain a customized prompt for each sample. Specifically, we first optimize a collection of prompts to obtain over-fitted prompts per sample. Then, we propose a prompt diffusion model within the prompt space, enabling the training of a generative transition process from a random prompt to its overfitted prompt. As we cannot access the label of a test image during inference, our model gradually generates customized prompts solely from random prompts using our trained, prompt diffusion. Our prompt diffusion is generic, flexible, and modality-agnostic, making it a simple plug-and-play module seamlessly embedded into existing prompt learning methods for textual, visual, or multi-modal prompt learning. Our diffusion model uses a fast ODE-based sampling strategy to optimize test sample prompts in just five steps, offering a good trade-off between performance improvement and computational efficiency. For all prompt learning methods tested, adding prompt diffusion yields more robust results for base-to-new generalization, cross-dataset generalization, and domain generalization in classification tasks tested over 15 diverse datasets.

Machine Learning,Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The problem this paper attempts to address is the lack of generalization ability in existing prompt-based classification methods when dealing with unseen samples due to distributional shifts. Specifically, traditional fixed prompt methods perform poorly when faced with new data because these prompts are usually optimized for specific training sets and cannot adapt well to new data distributions.

To solve this problem, the authors propose a method called "prompt diffusion." This method uses a diffusion model to gradually optimize the prompt for each sample, thereby generating customized prompts. In this way, prompt diffusion can better adapt to the characteristics of different samples, improving the model's generalization ability in zero-shot and few-shot learning tasks.

The main contributions include:
1. Proposing a Transformer-based prompt diffusion method that can learn the generation path from random prompts to personalized prompts in the prompt space.
2. Introducing a per-sample overfitting strategy to obtain the optimal prompt for each data sample, allowing the diffusion Transformer to effectively transition from general prompts to highly personalized prompts.
3. The prompt diffusion method is general, flexible, and modality-agnostic, and can be seamlessly integrated into existing text, visual, or multimodal prompt learning methods.

Experimental results show that adding the prompt diffusion module significantly improves the performance of various prompt learning methods across multiple datasets, especially in terms of generalization from base to new tasks, cross-dataset generalization, and domain generalization.