Matching Networks with Attention

Matching Networks offer a distinct approach within metric-based meta-learning, framing few-shot prediction as a form of weighted nearest neighbors in a learned embedding space. Unlike methods that compute fixed class prototypes, Matching Networks directly compare a query (test) sample $x'$ to each available support sample $(x_i, y_i)$ in the support set $S = \{(x_j, y_j)\}_{j=1}^{k \times N}$ for an N-way, k-shot task. The prediction for $x'$ is a weighted sum of the support set labels $y_i$, where the weights reflect the similarity between the query and each support example.

The original formulation often relied on a simple cosine similarity between the embeddings of the query and support samples, produced by an embedding function $f_\phi$. While effective, this assumes a fixed similarity suffices across all comparisons. However, for complex tasks and high-dimensional embeddings derived from foundation models, the relevance of a support example $x_i$ to a query $x'$ might depend heavily on the context provided by other support examples or specific features of the query itself.

This is where attention mechanisms provide a significant enhancement. Instead of using a static similarity function, we can learn an attention mechanism $a(\cdot, \cdot)$ that dynamically computes the importance (weight) $\alpha_i$ of each support example $x_i$ relative to the query $x'$. The prediction $\hat{y}'$ for the query sample $x'$ becomes:

$$\hat{y}' = \sum_{i=1}^{k \times N} \alpha_i y_i$$

Where the attention weights $\alpha_i$ are typically computed via a softmax over similarity scores between the query and support embeddings:

$$\alpha_i = \frac{\exp(a(f_\phi(x'), g_\phi(x_i)))}{\sum_{j=1}^{k \times N} \exp(a(f_\phi(x'), g_\phi(x_j)))}$$

Here, $f_\phi$ embeds the query sample and $g_\phi$ embeds the support samples. These embedding functions can be the same network or different ones. The attention function $a(\cdot, \cdot)$ itself can range from a simple cosine similarity (recovering the original Matching Network) to more sophisticated learned functions, such as a scaled dot-product or a small neural network that takes the pair of embeddings as input.

Full Contextual Embeddings (FCE)

A powerful concept introduced with Matching Networks is the use of Full Contextual Embeddings (FCE). The idea is to make the embedding of each sample dependent on the entire support set context. This allows the model to capture richer relationships and dependencies within the task definition provided by the support set.

Typically, FCE is implemented using bidirectional LSTMs or similar recurrent architectures.

1. Support Set Embedding $g_\phi(x_i, S)$: An LSTM processes the sequence of embeddings of the support samples $\{f_\phi(x_j)\}_{j=1}^{k \times N}$. The hidden state corresponding to $x_i$ after processing the full sequence (often in both forward and backward directions) becomes its context-aware embedding $g_\phi(x_i, S)$. This allows the representation of $x_i$ to be influenced by all other $x_j$ in $S$.
2. Query Embedding $f_\phi(x', S)$: Similarly, the query embedding $f_\phi(x')$ can also be processed through an LSTM that takes the support set $S$ (or its processed states) as context. This results in a query embedding $f_\phi(x', S)$ that is aware of the specific task defined by $S$.

The attention weights are then computed using these context-aware embeddings:

$$\alpha_i = \frac{\exp(\text{cosine}(f_\phi(x', S), g_\phi(x_i, S)))}{\sum_{j=1}^{k \times N} \exp(\text{cosine}(f_\phi(x', S), g_\phi(x_j, S)))}$$

While FCE significantly increases the representational power by incorporating the full support set context into each embedding, it also introduces substantial computational overhead due to the recurrent processing involved, especially as the support set size ($k \times N$) grows.

Flow of information in a Matching Network with Attention. Support and query samples are embedded (potentially using FCE), attention weights are computed based on query-support similarity, and the final prediction is a weighted sum of support labels.

Advantages and Considerations

Advantages:

• Dynamic Weighting: Attention allows the network to assign importance scores dynamically based on the specific query-support pair, potentially capturing finer-grained relationships than fixed prototypes or simple cosine similarity.
• Contextualization (with FCE): FCE enables the model to reason about the support set as a whole, potentially leading to more informed embeddings and comparisons.
• End-to-End Training: The embedding functions ($f_\phi, g_\phi$) and the attention mechanism ($a$) are typically trained end-to-end within the meta-learning loop.

Considerations:

• Computational Complexity: Calculating pairwise attention scores across the support set for each query is more compute-intensive than methods like Prototypical Networks. FCE adds further overhead with its recurrent components.
• Scalability to Large Support Sets: The $O(k \times N)$ complexity per query for attention calculation can become a bottleneck for tasks with many shots or many ways.
• Architecture Choices: The performance can depend significantly on the choice of embedding architecture (e.g., CNN, Transformer), the attention function $a(\cdot, \cdot)$, and whether FCE is used.

Integration with Foundation Models

Matching Networks with attention can effectively utilize embeddings from large foundation models. The foundation model can serve as the primary feature extractor ($f_\phi$, $g_\phi$), potentially frozen or minimally adapted during meta-training. The attention mechanism then operates on these rich, high-dimensional embeddings. Techniques like applying projection layers before the attention calculation might be necessary to manage dimensionality or adapt the embeddings to the specific needs of the attention function. Using pre-trained embeddings reduces the burden of learning complex feature extractors from scratch during meta-training, allowing the focus to be on learning the metric and attention space for rapid adaptation. The choice between using simple attention on static embeddings versus implementing FCE on top of foundation model features involves a trade-off between performance gains and computational cost.