\section{Method}
\label{sec:method}

In this section, we introduce the general problem setting for few-shot classification and its related formulations. Following, we describe the overall method as presented in \Fig{overview}, explain the attention mechanism and it's components, including the learning objective.    

\subsection{Problem Setting}
\label{method:setting}
% 
\paragraph{Overview} We consider the standard few-shot learning problem in which we are given a labeled train set $\mathcal{D}_{t}$, an unlabeled query set $\mathcal{D}_{q}$, and a few labeled examples from a support set $\mathcal{S}$ sharing the same label space with the query set. The label space of the training and query set are non-overlapping \ie, $\{Y_t \cap Y_q\} = \emptyset$, where $\{Y_t\}$ and $\{Y_q\}$ denote the training and query labels, respectively. Thus, given a set of labeled tasks from $\mathcal{D}_{t}$, and a few examples from the support set; the goal is to train a model that can generalize to novel tasks in the query set $\mathcal{D}_{q}$. 

Towards this goal, meta-learning with episodic training can be employed. This scheme aims to improve generalization by mimicking the low-data setting encountered in inference by creating balanced episodes. An episode is formed by sampling disjoint data points from the support set consisting a few labeled points, and the query set where labels are used to calculate prediction errors per episode. Moreover, each episode defines an "$N$-way,$k$-shot" task, where $N$ indicates the number of classes per episode, and $k$ is the number of support examples per class. 

\paragraph{Prototypical Baselines} In the context of Prototypical networks~\cite{Snell2017PrototypicalNF}, one aims to construct an embedding space in which points cluster around a single \textit{prototype} $\mathbf{p}_k \in \mathbb{R}^Z$  representation of each class $k$. Here, the objective is to learn an embedding function $\mathbf{F}_{\theta}(\text{x}): \text{x} \in \mathbb{R}^{H \times W \times C} \rightarrow \mathbb{R}^Z$ to transform an input into a $\mathbf{z}$-dimensional vector. To convey the description of the class as meta-data, the prototype of each class $\mathbf{p}_k$ is defined as the average of all embedded samples in that class:

\begin{equation}
	\mathbf{p}_k = \frac{1}{|S_k|} \sum_{(\text{x}_i,y_i) \sim S_k}^{} \mathbf{F}_{\theta}(\text{x}_i), \label{eq:1}
\end{equation} 

\noindent where $S_k$ denotes samples from class $k$. To obtain the probability distribution over classes for a query sample $\text{x}_q$, the softmax function on distances of the query to the prototypes is employed:
\begin{equation}
	p_{\theta}(y_i = k|\text{x}_q) = \frac{\text{exp}(-d(\mathbf{F}_{\theta}(\text{x}_q),\mathbf{p}_k))}{\sum_{k^\prime}^{} \text{exp}(-d(\mathbf{F}_{\theta}(\text{x}_q),\mathbf{p}_{k^\prime}))}, \label{eq:2}
\end{equation}

\noindent where $d(\cdot{})$ is a distance function between query and the prototype. Consequently,  $p_{\theta}(y_i = k|\text{x}_q)$ is used to assign the correct class per sampled task, and the model learns to minimize the predictive error on query samples per task. 

Though effective in practice, we argue that the prototype embeddings are not ideal to learn a discriminative model. Since every element in $S_k$ is processed independently in the pooling operation (\ie via mean), some information regarding interactions between them may be discarded making it difficult for the model to learn robust embeddings.  In what follows, we describe how to co-adapt the embeddings via a transformer style set-to-prototype function using self-attention.  

\subsection{Instance Embedding Adaptation}
\label{method:coadapt}

In order to co-adapt support and query embeddings $\text{x}_s$ and $\text{x}_q$, we introduce an additional step where the initial embeddings are further transformed using self-attention. Here, attention is employed to find the relevant features to attend to in $\text{x}_s$ and $\text{x}_q$ by considering discriminative feature similarity only, and thus enables us to re-weight the features. Formally, let $z_q = \mathbf{F}_{\theta}(\text{x}_q)$ and $z_s = \mathbf{F}_{\theta}(\text{x}_s)$ be the initial query and support embeddings after the penultimate layer, each having vector dimension $m$ \ie, $z \in \mathbb{R}^m $. Our goal is to obtain adapted embeddings $z^{\prime}_q$ and $z^{\prime}_s$. 

\paragraph{Self-Attention} Following the work of Vaswani \etal\cite{vaswani2017attention} and its related extensions presented by Lee \etal\cite{lee2019set} for set-based tasks, we employ an attention function $\varphi(\mathbf{Q},\mathbf{K},\mathbf{V})$ that measures similarity between a query vector $\mathbf{Q}$ with key-value pairs $\mathbf{K} \in \mathbb{R}^{n \times m}$ and $\mathbf{V} \in \mathbb{R}^{n\times m}$ via: 

\begin{equation}
	\varphi(\mathbf{Q},\mathbf{K},\mathbf{V}) = \text{softmax}(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{m}})\mathbf{V}. \label{eq:3}
\end{equation}

Concretely, the formulation introduced by Vaswani \etal\cite{vaswani2017attention} was extended to a multi-head attention block (MAB) where vectors $\mathbf{Q},\mathbf{K},\mathbf{V}$ are first projected onto $h$ different dimensional vectors. Here, $\varphi(\mathbf{Q},\mathbf{K},\mathbf{V})$ then becomes $\varphi(\mathbf{Q}W_j^{\mathbf{Q}},\mathbf{K}W_j^{\mathbf{K}},\mathbf{V}W_j^{\mathbf{V}})$, with each transformation having its own learnable parameters $W$. In this setting, each $W$ can be modeled by functions $\mathbf{f}^{q}_{\theta}$, $\mathbf{f}^{k}_{\theta}$ and $\mathbf{f}^{v}_{\theta}$, respectively. Also, an additional function $\mathbf{f}^{o}_{\theta}$ takes the output of $\varphi$ and adds the residual followed by a layer normalization operation. Mathematically,

\begin{equation}
	\varphi^o = \varphi(\mathbf{Q},\mathbf{K},\mathbf{V}) + \mathbf{Q}, \label{eq:4}
\end{equation}

\begin{equation}
	\text{MAB}_h(\mathbf{Q},\mathbf{K},\mathbf{V}) = \Phi(\varphi^o + \gamma(\mathbf{f}^{o}_{\theta}(\varphi^o))) \label{eq:5}
\end{equation}

\noindent where $\gamma$ is the ReLU function and $\varphi^o$ is the output of applying attention on query key-value pairs after individual feature projection using functions $\mathbf{f}^{q,k,v}_{\theta}$ along with the concatenation of $\mathbf{Q}$. $\text{MAB}_h$ defines a single attention block ($h=1$) with an optional normalization layer $\Phi$. It is worth noting that the $\mathbf{Q}$ notation used here is specifically related to features in the $\text{MAB}_h$ module, and should be considered separate from the notion of queries in the few-shot setting. 

\input{fig/self_att}

\paragraph{Cross-Adaptation with MAB} In the context of this work, to adapt features relative to the other, we consider two aspects: (i) \textit{support}-\textit{query}, and (ii) \textit{query}-\textit{support}, as presented in \Fig{selfatt}. For scenario (i), we employ the attention module as $\text{MAB}_h(z_s,z_q,z_q)$ (\Fig{selfatt}(a)), where support features $z_s$ are considered as the query ($\mathbf{Q}$), and the initial query embeddings $z_q$ are key-value pairs. For (ii), the opposite is applied, \ie $\text{MAB}_h(z_q,z_s,z_s)$ (\Fig{selfatt}(b)). Intuitively, this should produce embeddings that are weighted by attention reflecting feature similarity to obtain $z^{\prime}_s = \text{MAB}_h(z_q,z_s,z_s)$, and $z^{\prime}_q = \text{MAB}_h(z_s,z_q,z_q)$, respectively.

To this end, we wish to note that Lee \etal\cite{lee2019set} employ \Eq{4} \& \Eq{5} by making $\mathbf{Q}$ a learnable parameter in order to perform attention-based pooling in set-based tasks \ie $(\mathbf{K} \in \mathbb{R}^{n\times k\times m},\mathbf{V}\in \mathbb{R}^{n\times k\times m})$ are a set of features and $\mathbf{Q} = W \in \mathbb{R}^{n\times m}$ is a tensor initialized from a normal distribution to produce a pooled tensor $\mathbf{O}\in \mathbb{R}^{n\times m}$, with $n$ as the batch dimension, and $k$ the number of samples in the set, respectively. This approach has been shown to produce better pooled embeddings over simply employing the mean directly. In this work, employing \Eq{4} reveals a tensor shape mismatch since few-shot inputs consist $q$-queries and $k$-shot samples for the support set in a meta-setup. To address this, we modify \Eq{4} and \Eq{5} as: 

\begin{equation}
	\hat{\varphi}^o = \varphi(\mathbf{Q},\mathbf{K},\mathbf{V}) + \mathbf{K}, \label{eq:6}
\end{equation}

\begin{equation}
	\text{MAB}_h(\mathbf{Q},\mathbf{K},\mathbf{V}) = \Phi(\hat{\varphi}^o + \gamma(\mathbf{f}^{o}_{\theta}(\hat{\varphi}^o))). \label{eq:7}
\end{equation}

In addition, note that $\varphi(\cdot)$ will result in different attention scores for \textit{query}-\textit{support} and \textit{support}-\textit{query} \ie, $\mathbf{A}_s \in \mathbb{R}^{n\times q\times k}$ and $\mathbf{A}_q \in \mathbb{R}^{n\times k\times q}$ for $k$-shot with $q$ query examples. Thus, in each setting we take the average over $k$ or $q$ to obtain the final scores. For example, in support adaptation, this implies we obtain $\mathbf{A}_s \in \mathbb{R}^{n\times 1\times k}$ and taking the dot product with $\mathbf{V}$ produces a prototype embedding, vice versa for queries. The support/query prototypes reflect the relevant features later added to the initial $\mathbf{K}$ to obtain the final embeddings. Consequently, the prototype of the adapted support embeddings are now obtained via: 

\begin{equation}
	 \mathbf{\hat{p}}_k = \frac{1}{|S_k|} \sum_{(\text{x}_i,y_i) \sim S_k}^{} z^{\prime}_s,
	 \label{eq:8}
\end{equation} 

\noindent and the distribution over classes is:

\begin{equation}
	p_{\theta}(y_i = k|\text{x}_q) = \frac{\text{exp}(-d(z^{\prime}_q,\mathbf{\hat{p}}_k))}{\sum_{k^\prime}^{} \text{exp}(-d(z^{\prime}_q,\mathbf{\hat{p}}_{k^\prime}))}. \label{eq:9}
\end{equation} 

To optimize the backbone and proposed modules, we train the model following the standard setting proposed in Prototypical networks (inductive only) using a cross-entropy loss on the negative distances to minimize the query prediction error per task. During inference, all modules are used as is, maintaining the design of the training phase without omission of any components.

