SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yingjun Du<sup>1</sup>, Zehao Xiao<sup>1</sup>, Shencai Liao<sup>2</sup>, Cees G. M. Snoek<sup>1</sup>  
<sup>1</sup>AIM Lab, University of Amsterdam <sup>2</sup>Inception Institute of Artificial Intelligence
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.
This paper considers prototype-based meta-learning, where models are trained to swiftly adapt to current tasks and perform classification through metric-based comparisons between examples and newly introduced variables - the prototypes of the classes. This approach, rooted in works by Reed [34] and further developed by Snell et al.[42], generalizes deep learning models to scenarios where labeled data is scarce. The fundamental idea is to compute the distances between queried examples and class prototypes and perform classification based on these distances to predict classes for queried examples. Derived from the prototypical network [42], several prototype-based methods have demonstrated their effectiveness in few-shot learning [2, 8, 11, 61]. However, as prototypes are estimated from a limited number of sampled examples, they may not accurately capture the overall distribution [54, 61]. Such biased distributions could lead to biased prototypes and subsequent classification errors, suggesting that the current prototype modeling approach might lack the ability to represent universal class-level information effectively. Rather than relying on points in the prototype embedding space by using a simple average function, we propose to generate the distribution of prototypes for each task.
To overcome the challenges of estimating deterministic prototypes from limited examples in few-shot learning, we propose ProtoDiff, an innovative framework that leverages a task-guided diffusion model during the meta-training phase. ProtoDiff introduces a three-step process to address these limitations comprehensively. Firstly, we optimize a set of prototypes to achieve per-task prototype overfitting,
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  $S^i$ . Subsequently, through the process of diffusion sampling, the diffused prototype  $\mathbf{z}_{t-1}$  is derived from the combination of  $\mathbf{z}_t$  and  $\tilde{\mathbf{z}}$ . Finally, the prediction  $\mathbf{y}$  for the query set  $\mathcal{Q}^i$  is generated by utilizing the diffused prototype  $\mathbf{z}_0$  in conjunction with the query set  $\mathbf{x}$ . The diffusion forward and sampling processes are indicated by dashed and solid arrows within the gray rectangular area.
ensuring accurate retrieval of overfitted prototypes specific to individual tasks. These overfitted prototypes serve as the ground truth representations for each task. Secondly, a task-guided diffusion process is implemented within the prototype space, enabling meta-learning of a generative process that smoothly transitions prototypes from their vanilla form to the overfitted state. ProtoDiff generates task-specific prototypes during the meta-test stage by conditioning random noise on the limited samples available for the given new task. Finally, we propose residual prototype learning, which significantly accelerates training and further enhances the performance of ProtoDiff by leveraging the sparsity of residual prototypes. The computational graph of ProtoDiff, illustrated in Figure 1, showcases the sequential steps involved in the forward diffusion process on the prototype and the generative diffusion process. The resulting generated prototype, denoted as  $\mathbf{z}_0$ , is utilized for query prediction. By incorporating probabilistic and task-guided prototypes, ProtoDiff balances adaptability and informativeness, positioning it as a promising approach for augmenting few-shot learning in prototype-based meta-learning models.
To validate the effectiveness of ProtoDiff, we conduct comprehensive experiments on three distinct few-shot learning scenarios: within-domain, cross-domain, and few-task few-shot learning. Our findings reveal that ProtoDiff significantly outperforms state-of-the-art prototype-based meta-learning models, underscoring the potential of task-guided diffusion to boost few-shot learning performance. Furthermore, we provide a detailed analysis of the diffusion mechanism employed in ProtoDiff, showcasing its ability to capture the underlying data structure better and improve generalization. This thorough investigation highlights the strengths of our approach, demonstrating its potential to offer a more effective solution for few-shot learning tasks in various applications and settings.
Before detailing our ProtoDiff methodology, we first present the relevant background on few-shot classification, the prototypical network, and diffusion models.
Few-shot classification. We define the N-way K-shot classification problem, which consists of support sets  $S$  and a query set  $\mathcal{Q}$ . Each task  $\mathcal{T}^i$ , also known as an episode, represents a classification problem sampled from a task distribution  $p(\mathcal{T})$ . The way of an episode denotes the number of classes within the support sets, while the shot refers to the number of examples per class. Tasks are created from a dataset by randomly selecting a subset of classes, sampling points from these classes, and subsequently dividing the points into support and query sets. The episodic optimization approach [50] trains the model iteratively, performing one episode update at a time.
Prototypical network. We develop our method based on the prototypical network (ProtoNet) by Snell et al. [42]. Specifically, the ProtoNet leverages a non-parametric classifier that assigns a query point to the class having the nearest prototype in the learned embedding space. The prototype  $\mathbf{z}^c$  of an object class  $c$  is obtained by:  $\mathbf{z}^c = \frac{1}{K}\sum_k f_\phi(\mathbf{x}^{c,k})$ , where  $f_\phi(\mathbf{x}^{c,k})$  is the feature embedding of the support sample  $\mathbf{x}^{c,k}$ , which is usually obtained by a convolutional neural network. For each query sample  $\mathbf{x}^q$ , the distribution over classes is calculated based on the softmax over distances to the prototypes of all classes in the embedding space:
$$
p \left(\mathbf {y} _ {n} ^ {q} = c \mid \mathbf {x} ^ {q}\right) = \frac {\exp \left(- d \left(f _ {\phi} \left(\mathbf {x} ^ {q}\right) , \mathbf {z} ^ {c}\right)\right)}{\sum_ {c ^ {\prime}} \exp \left(- d \left(f _ {\phi} \left(\mathbf {x} ^ {q}\right) , \mathbf {z} ^ {c ^ {\prime}}\right)\right)}, \tag {1}
$$
where  $\mathbf{y}^q$  denotes a random one-hot vector, with  $\mathbf{y}_n^q$  indicating its  $n$ -th element, and  $d(\cdot, \cdot)$  is some (Euclidean) distance function. Due to its non-parametric nature, the ProtoNet enjoys high flexibility and efficiency, achieving considerable success in few-shot learning. To avoid confusion, we omit the superscript  $c$  for the prototype  $\mathbf{z}$  in this subsection.
Diffusion model. In denoising diffusion probabilistic models [16], a forward diffusion process,  $q(\pmb{x}_t|\pmb{x}_{t-1})$ , is characterized as a Markov chain that progressively introduces Gaussian noise at each time step  $t$ , beginning with a clean image  $\pmb{x}_0 \sim q(\pmb{x}_0)$ . The subsequent forward diffusion process is formulated as follows:
$$
q \left(\boldsymbol {x} _ {T} \mid \boldsymbol {x} _ {0}\right) := \prod_ {t = 1} ^ {T} q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t - 1}\right), \quad \text {w h e r e} \quad q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t - 1}\right) := \mathcal {N} \left(\boldsymbol {x} _ {t}; \sqrt {1 - \beta_ {t}} \boldsymbol {x} _ {t - 1}, \beta_ {t} \boldsymbol {I}\right), \tag {2}
$$
where  $\{\beta\}_{t=0}^{T}$  is a variance schedule. By defining  $\alpha_{t} := 1 - \beta_{t}$  and  $\bar{\alpha}_{t} := \prod_{s=1}^{t} \alpha_{s}$ , the forward diffused sample at time step  $t$ , denoted as  $\boldsymbol{x}_{t}$ , can be generated in a single step as follows:
$$
\boldsymbol {x} _ {t} = \sqrt {\bar {\alpha} _ {t}} \boldsymbol {x} _ {0} + \sqrt {1 - \bar {\alpha} _ {t}} \boldsymbol {\epsilon}, \quad \text {w h e r e} \quad \boldsymbol {\epsilon} \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I}). \tag {3}
$$
Since the reverse of the forward step,  $q(\pmb{x}_{t-1}|\pmb{x}_t)$ , is computationally infeasible, the model learns to maximize the variational lower bound using parameterized Gaussian transitions,  $p_{\theta}(\pmb{x}_{t-1}|\pmb{x}_t)$ , where the parameter is denoted as  $\theta$ . Consequently, the reverse process is approximated as a Markov chain with the learned mean and fixed variance, starting from a random noise  $\pmb{x}_T \sim \mathcal{N}(\pmb{x}_T; \pmb{0}, \pmb{I})$ :
$$
p _ {\theta} \left(\boldsymbol {x} _ {0: T}\right) := p _ {\theta} \left(\boldsymbol {x} _ {T}\right) \prod_ {t = 1} ^ {T} p _ {\theta} \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} _ {t}\right), \tag {4}
$$
$$
p _ {\theta} \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} _ {t}\right) := \mathcal {N} \left(\boldsymbol {x} _ {t - 1}; \boldsymbol {\mu} _ {\theta} \left(\boldsymbol {x} _ {t}, t\right), \sigma_ {t} ^ {2} \boldsymbol {I}\right), \quad \boldsymbol {\mu} _ {\theta} \left(\boldsymbol {x} _ {t}, t\right) := \frac {1}{\sqrt {\alpha} _ {t}} \left(\boldsymbol {x} _ {t} - \frac {1 - \alpha_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \boldsymbol {\epsilon} _ {\theta} \left(\boldsymbol {x} _ {t}, t\right)\right). \tag {5}
$$
Here,  $\epsilon_{\theta}(\pmb{x}_t,t)$  is the diffusion model trained by optimizing the following objective function:
$$
\mathcal {L} _ {\theta} = \mathbb {E} _ {t, \boldsymbol {x} _ {0}, \epsilon} \left[ \| \epsilon - \epsilon_ {\theta} \left(\sqrt {\bar {\alpha} _ {t}} \boldsymbol {x} _ {0} + \sqrt {1 - \bar {\alpha} _ {t}} \epsilon , t\right) \| ^ {2} \right]. \tag {6}
$$
Upon completing the optimization, the learned score function is integrated into the generative (or reverse) diffusion process. To sample from  $p_{\theta}(\pmb{x}_{t-1}|\pmb{x}_t)$ , one can perform the following:
$$
\boldsymbol {x} _ {t - 1} = \boldsymbol {\mu} _ {\theta} (\boldsymbol {x} _ {t}, t) + \sigma_ {t} \boldsymbol {\epsilon} = \frac {1}{\sqrt {\alpha_ {t}}} \left(\boldsymbol {x} _ {t} - \frac {1 - \alpha_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \boldsymbol {\epsilon} _ {\theta} (\boldsymbol {x} _ {t}, t)\right) + \sigma_ {t} \boldsymbol {\epsilon}. \tag {7}
$$
In the case of conditional diffusion models [22, 38, 39], the diffusion model  $\epsilon_{\theta}$  in equations (6) and (7) is substituted with  $\epsilon_{\theta}(\mathbf{c}, \sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon, t)$ , where  $\mathbf{c}$  represents the corresponding conditions, e.g., other images, languages, and sounds, etc. Consequently, the matched conditions strictly regulate the sample generation in a supervised manner, ensuring minimal changes to the image content. Dhariwal and Nichol [7] suggested classifier-guided image translation to generate the images of the specific classes along with a pre-trained classifier. Taking inspiration from conditional diffusion models that excel at generating specific images using additional information, we introduce the task-guided diffusion model that generates prototypes of specific classes by taking into account contextual information from various few-shot tasks. Note that our task-guided diffusion model operates not on an image  $\mathbf{x}$ , but on the prototype  $\mathbf{z}$ .
This section outlines our approach to training prototypical networks via task-guided diffusion. We begin by explaining how to obtain task-specific overfitted prototypes in Section 3.1. Next, we introduce the task-guided diffusion method for obtaining diffused prototypes in Section 3.2. In the same section, we also introduce residual prototype learning, which accelerates training. Figure 2 visualizes the diffusion process of diffused prototypes by our ProtoDiff.
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$ .
Figure 3: ProtoDiff illustration. We first obtain an overfitted prototype by the per-task prototype overfitting. We then add noise to the overfitted prototype, which inputs the diffusion forward process. The input of task-guided diffusion includes the vanilla prototype  $\tilde{\mathbf{z}}$  and a random time step  $t$ . The resulting output is the diffused prototype  $\mathbf{z}_t$ .
3.1 Per-task prototype overfitting
Our approach utilizes a modified diffusion process as its foundation. Firstly, a meta-learner, denoted by  $f_{\phi}(\mathcal{T}^{i}) = f(\mathcal{T}^{i},\phi)$ , is trained on the complete set of tasks,  $\mathcal{T}$ . Here,  $\phi$  corresponds to the learned weights of the model over the entire task set. Subsequently, fine-tuning is performed on each individual task  $\mathcal{T}^i$  to obtain task-specific overfitted prototypes denoted by  $\mathbf{z}^{i,*} = f\phi^i (S^i,\mathcal{Q}^i)$ . This involves running the meta-learner  $f_{\phi}$  on the support set  $S^i$  and a query set  $\mathcal{Q}^i$  of the task  $\mathcal{T}^i$ .
$$
\phi^ {i} = \phi - \eta \sum_ {(\mathbf {x}, \mathbf {y}) \sim \mathcal {T} _ {i}} ^ {\mathcal {T} _ {i}} \mathcal {L} _ {\mathrm {C E}} \left(f \left(\mathcal {S} ^ {i}, \mathbf {x} ^ {q ^ {i}}, \phi\right), \mathbf {y} ^ {q ^ {i}}\right), \tag {8}
$$
where  $\mathcal{L}_{\mathrm{CE}}$  is the cross-entropy loss minimized in the meta-learner training. The support set  $S^i$  and query set  $\mathcal{Q}^i = \{\mathbf{x}^{q^i},\mathbf{y}^{q^i}\}$  correspond to the data used for fine-tuning task  $\mathcal{T}^i$ . During fine-tuning, we obtain the  $\phi^i$  through several iterations of gradient descent. We illustrate the overall framework of the per-task prototype learning in the appendix.
Our proposed method, ProtoDiff, for few-shot learning, relies on acquiring task-specific overfitted prototypes that can be considered as the "optima" prototypes due to their high confidence in final predictions, approaching a value of 1. To achieve this, we employ fine-tuning of the meta-learner and extract task-specific information that surpasses the accuracy and reliability of generic prototypes used in the meta-training stage. However, accessing the query set in the meta-test stage is not feasible. Hence, we need to rely on the vanilla prototype to meta-learn the process of generating the overfitted prototype during the meta-training stage. In the forthcoming section, we will introduce the task-guided diffusion model, which facilitates the learning process transitioning from the vanilla prototype to the overfitted prototype.
3.2 Task-guided diffusion
Our generative overfitted prototype is based on diffusion [43], a robust framework for modeling the prototypes instead of images. The length of the diffusion process  $T$  determines the number of forward passes needed to generate a new prototype, rather than the dimensionality of the data. Our model uses diffusion to progressively denoise the overfitted prototype  $\mathbf{z}^*$ .
Meta-training phase. Diffusion models can be configured to predict either the signal or the noise when presented with a noisy input [16, 31]. Previous research in the image domain has suggested that noise prediction is superior to signal prediction. However, we discovered empirically that signal prediction performs better than noise prediction in our experiments, so we parameterized the diffusion model to output the prototype.
Figure 4: Diffusion sampling during meta-test. Sampling starts from a random noise  $\mathbf{z}_{\mathrm{T}}$  and gradually denoises it to the diffused prototypes  $\mathbf{z}_{\mathrm{t}}$ . At each time step  $t$  is sampled by taking the  $\mathbf{z}_{\mathrm{t} + 1}$ , vanilla prototypes  $\tilde{\mathbf{z}}$ , and time step  $t$  as inputs. The diffused prototypes  $\mathbf{z}_0$  are used to predict the query set.
The generative diffusion process is devised to reconstruct the overfitted prototype  $\mathbf{z}^*$  by iteratively operating on a random noise vector  $\mathbf{z}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ , which possesses the same dimensions as  $\mathbf{z}^*$ . This process continues with  $\mathbf{z}_t$ , where  $t$  denotes the number of diffusion steps. Consequently, the reconstructed  $\mathbf{z}_0$  should exhibit proximity to  $\mathbf{z}^*$  for a given task.
Specifically, during the forward diffusion process at time step  $t$ , we obtain the noised overfitted prototype  $\hat{\mathbf{z}}_t$ . Subsequently, we input the noised prototype  $\hat{\mathbf{z}}_t$ , the vanilla prototype  $\tilde{\mathbf{z}}$ , and the time step  $t$  into the task-guided diffusion. This yields the diffused prototype  $\mathbf{z}_t$ , which then allows us to predict the final results of the query set using Equation (1). For each task, our task-guided diffusion entails two components: the variational lower bound  $\mathcal{L}_{\mathrm{diff}}$  for the diffused prototype, and the cross-entropy loss  $\mathcal{L}_{\mathrm{CE}}$ .
The objective is to minimize the simplified variational lower bound, which involves predicting the denoised overfitted prototype:
$$
\mathcal {L} _ {\text {d i f f}} = \left| \mathbf {z} ^ {*} - \mathbf {z} _ {\theta} \left(\sqrt {\bar {\alpha} _ {t}} \mathbf {z} ^ {*} + \sqrt {1 - \bar {\alpha} _ {t}} \boldsymbol {\epsilon}, \tilde {\mathbf {z}}, t\right) \right| ^ {2}, \tag {9}
$$
Here,  $\mathbf{z}_{\theta}(\cdot ,\cdot)$  is implemented by the transformer model [49] operating on prototype tokens from  $\mathbf{z}^*$ ,  $\tilde{\mathbf{z}}$ , and the time step.
By utilizing equation (1), we derive the final prediction  $\hat{\mathbf{y}}^q$  using the diffused prototype  $\mathbf{z}_t$ . The ultimate objective is expressed as follows:
$$
\mathcal {L} = \sum_ {(\mathbf {x}, \mathbf {y}) \sim \mathcal {Q}} ^ {| Q |} \left[ - \mathbb {E} q \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {t + 1}, \tilde {\mathbf {z}}\right) \left[ \log p \left(\mathbf {y} ^ {q} \mid \mathbf {x} ^ {q}, \mathbf {z} _ {t}\right) \right] + \beta \left| \mathbf {z} ^ {*} - \mathbf {z} _ {\theta} \left(\sqrt {\bar {\alpha} _ {t}} \mathbf {z} ^ {*} + \sqrt {1 - \bar {\alpha} _ {t}} \boldsymbol {\epsilon}, \tilde {\mathbf {z}}, t\right) \right| ^ {2} \right], \tag {10}
$$
where  $\beta$  represents a hyperparameter, and  $|Q|$  denotes the query size.
To prepare the two input prototypes  $\mathbf{z}^*$  and  $\tilde{\mathbf{z}}$  for processing by the transformer, we assign each position token inspired by [10]. We also provide the scalar input diffusion timestep  $t$  as individual tokens to the transformer. To represent each scalar as a vector, we use a frequency-based encoding scheme [28]. Our transformer architecture is based on GPT-2 [33]. The decoder in the final layer of our ProtoDiff maps the transformer's output to the diffused prototype. Note that only the output tokens for the noised overfitted prototypes  $\hat{\mathbf{z}}_t$  are decoded to predictions. The overall meta-training phase of ProtoDiff is shown in Figure 3.
Meta-test phase. We can not obtain the overfitted prototype during the meta-test phase since we can not access the query set. Thus, diffusion sampling begins by feeding-in Gaussian noise  $\mathbf{z}_T$  as the  $\hat{\mathbf{z}}$  input and gradually denoising it. Specifically, to handle a new task  $\tau = \{S, Q\}$ , we first compute the vanilla prototype  $\tilde{\mathbf{z}}$  using the support set  $S$ . We then randomly sample noise  $\epsilon$  from  $\mathcal{N}(\mathbf{0}, I)$ . We input both  $\tilde{\mathbf{z}}$  and  $\epsilon$  to the learned task-guided diffusion model to obtain the diffused prototype  $\mathbf{z}_{T-1} = \mathbf{z}_{\theta}(\mathbf{z}_T, \tilde{\mathbf{z}}, T)$ . After  $T$  iterations, the final diffused prototype can be obtained as  $\mathbf{z}_0 = \mathbf{z}_{\theta}(\mathbf{z}_1, \mathbf{z}, t_0)$ . Once the final diffused prototype  $\mathbf{z}_0$  is obtained, we calculate the final prediction  $\hat{\mathbf{y}}^q$  using equation (1). Figure 4 illustrates sampling in the meta-test stage. We also provide the detailed algorithm of the meta-training and meta-test phase in the appendix.
Residual prototype learning. To further enhance the performance of ProtoDiff, we introduce a residual prototype learning mechanism in our framework. We observe that the differences between the overfitted prototype  $\mathbf{z}^*$  and its vanilla prototype  $\tilde{\mathbf{z}}$  are not significant, as the vector  $\mathbf{z}^* -\tilde{\mathbf{z}}$  contains many zeros. Therefore, we propose to predict the prototype update  $\mathbf{z}^* -\tilde{\mathbf{z}}$  instead of directly predicting
the overfitted prototype  $\mathbf{z}^*$  itself. This approach also enables us to initialize ProtoDiff to perform the identity function by setting the decoder weights to zero. Moreover, we find that the global residual connection, combined with the identity initialization, significantly speeds up training. By utilizing this mechanism, we improve the performance of ProtoDiff in few-shot learning tasks.