[Paper Review] ProtoDiff: Learning to Learn Prototypical Networks by Task-Guided Diffusion

Yingjun Du, Zehao Xiao+2 more|arXiv (Cornell University)|Jun 26, 2023

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

ProtoDiff introduces a task-guided diffusion model to learn probabilistic, task-specific prototypes in few-shot learning, improving upon deterministic prototypical networks by generating refined prototypes from noise during meta-testing. It achieves state-of-the-art performance across within-domain, cross-domain, and few-task few-shot benchmarks by meta-learning a diffusion process that transitions from vanilla to overfitted prototypes, with residual prototype learning accelerating training and boosting accuracy.

ABSTRACT

Prototype-based meta-learning has emerged as a powerful technique for addressing few-shot learning challenges. However, estimating a deterministic prototype using a simple average function from a limited number of examples remains a fragile process. To overcome this limitation, we introduce ProtoDiff, a novel framework that leverages a task-guided diffusion model during the meta-training phase to gradually generate prototypes, thereby providing efficient class representations. Specifically, a set of prototypes is optimized to achieve per-task prototype overfitting, enabling accurately obtaining the overfitted prototypes for individual tasks. Furthermore, we introduce a task-guided diffusion process within the prototype space, enabling the meta-learning of a generative process that transitions from a vanilla prototype to an overfitted prototype. ProtoDiff gradually generates task-specific prototypes from random noise during the meta-test stage, conditioned on the limited samples available for the new task. Furthermore, to expedite training and enhance ProtoDiff's performance, we propose the utilization of residual prototype learning, which leverages the sparsity of the residual prototype. We conduct thorough ablation studies to demonstrate its ability to accurately capture the underlying prototype distribution and enhance generalization. The new state-of-the-art performance on within-domain, cross-domain, and few-task few-shot classification further substantiates the benefit of ProtoDiff.

Motivation & Objective

• To address the fragility of deterministic prototype estimation in few-shot learning due to limited support samples.
• To model the underlying class distribution more robustly by learning a generative process for prototypes rather than relying on simple averaging.
• To improve generalization and adaptability in few-shot classification by learning a diffusion-based transition from random noise to task-specific prototypes.
• To accelerate training and enhance performance through residual prototype learning, exploiting sparsity in prototype updates.
• To establish a new state-of-the-art in few-shot learning across diverse benchmarks, including cross-domain and few-task settings.

Proposed method

• Optimize a set of prototypes to achieve per-task prototype overfitting, using fine-tuning to obtain accurate, task-specific representations as ground truth.
• Introduce a task-guided diffusion process within the prototype space, learning a denoising process that transitions from random noise to overfitted prototypes.
• Condition the diffusion sampling at meta-test time on limited support set examples to generate task-specific prototypes from noise.
• Utilize a residual prototype learning scheme that models only the difference between the vanilla and overfitted prototypes, improving training efficiency and performance.
• Train the diffusion model using a reverse process that learns to denoise the prototype space, guided by the overfitted prototypes as targets.
• Combine the generated prototype with query features via metric learning (e.g., cosine similarity) for final classification.

Figure 1 : Computational graph of ProtoDiff. In the initial stage, the computation of the vanilla prototype $\tilde{\mathbf{z}}$ is performed based on the support set $\mathcal{S}^{i}$ . Subsequently, through the process of diffusion sampling, the diffused prototype $\mathbf{z}_{t-1}$ is derived fro

Experimental results

Research questions

• RQ1 Can a diffusion model be effectively used to generate task-specific prototypes in few-shot learning, improving upon deterministic averaging?
• RQ2 How does task-guided diffusion in the prototype space enhance generalization across few-shot benchmarks?
• RQ3 To what extent does residual prototype learning accelerate training and improve performance in prototype-based meta-learning?
• RQ4 Does ProtoDiff generalize effectively to cross-domain few-shot learning, where domain shift exists between support and query sets?
• RQ5 Can ProtoDiff achieve state-of-the-art performance in few-task few-shot learning, where only a small number of tasks are available for meta-training?

Key findings

• ProtoDiff achieves 44.75% accuracy on mini-ImageNet-S in the 5-way 1-shot setting, surpassing MetaModulation by 1.54% and setting a new SOTA.
• On the challenging DermNet-S dataset, ProtoDiff achieves 51.53% accuracy in 5-way 1-shot, outperforming the previous best method by 1.08%.
• In cross-domain few-shot learning on ChestX, ProtoDiff achieves 28.54% accuracy in 5-way 5-shot, outperforming HVM by 1.39%.
• The model consistently improves over baselines across all benchmarks and settings, including within-domain, cross-domain, and few-task few-shot learning.
• Residual prototype learning reduces training time and improves performance, demonstrating the effectiveness of sparsity-aware prototype updates.
• Despite higher inference cost due to diffusion sampling, ProtoDiff maintains strong performance and generalization, validating the trade-off between accuracy and efficiency.

Figure 2 : Visualization of the diffusion process. ProtoDiff randomly selects certain areas to predict during the diffusion process, with the lowest probability at the beginning time step. As time progresses, the prototype gradually aggregates towards the dog , with the highest probability at t=0.