# FaceCoresetNet: Differentiable Coresets for Face Set Recognition

Gil Shapira1,2 and Yosi Keller 1

###### Abstract

In set-based face recognition, we aim to compute the most discriminative
descriptor from an unbounded set of images and videos showing a single person.
A discriminative descriptor balances two policies when aggregating information
from a given set. The first is a quality-based policy: emphasizing high-
quality and down-weighting low-quality images. The second is a diversity-based
policy: emphasizing unique images in the set and down-weighting multiple
occurrences of similar images as found in video clips which can overwhelm the
set representation. This work frames face-set representation as a
differentiable coreset selection problem. Our model learns how to select a
small coreset of the input set that balances quality and diversity policies
using a learned metric parameterized by the face quality, optimized end-to-
end. The selection process is a differentiable farthest-point sampling (FPS)
realized by approximating the non-differentiable Argmax operation with
differentiable sampling from the Gumbel-Softmax distribution of distances. The
small coreset is later used as queries in a self and cross-attention
architecture to enrich the descriptor with information from the whole set. Our
model is order-invariant and linear in the input set size. We set a new SOTA
to set face verification on the IJB-B and IJB-C datasets. Our code is publicly
available 111https://github.com/ligaripash/FaceCoresetNet.

Figure 1:  Template coreset selection. in the round blue rectangle, an input
face set with 12 images containing a video clip sample ($I_{1}$ to $I_{6}$)
and six photographs of the same person ($I_{7}$ to $I_{12}$). The coreset
selection process balances quality and diversity. The final coreset in the red
round rectangle contains four high-quality images with diverse poses and
expressions.

## Introduction

Face Set Recognition (FSR) is a variant of Face Recognition (FR) that involves
recognizing an individual’s identity based on a template, that is an unbounded
and unordered set of images and video clips depicting a particular face
instead of a single photograph. A template is shown in the Fig. 1. Given the
added data, FSR is more accurate compared to FR. The increased accuracy and
proliferation of face media available on the Internet and public cameras make
FSR an attractive solution for ’in the wild’ situations, such as in
unconstrained surveillance and Homeland Security scenarios (Kalka et al. 2018)
where image quality is lacking due to low resolution, extreme pose,
illumination, occlusions, and recognition of a single image is difficult.

FSR consists of verification, which means confirming the identity of an input
query template (also referred to as a probe) and identification, which entails
searching a database of registered subjects (known as a gallery) to identify
an unknown probe person. FSR schemes encode the images in a template into a
fixed-size vector descriptor to allow computationally efficient similarity
search. As the templates consist of an arbitrary number of images and video
frames, the image aggregation process has to be efficient with respect to the
template’s size. Furthermore, since the template is an unordered set of
images, the aggregation should be permutation invariant. Additionally, the
aggregation should prioritize informative images while downplaying
uninformative, low-quality (blurry, poorly lighted) images. Finally, the
process of feature fusion should factor in ”face burstiness”. If a template
includes a burst of similar frames as in video clips, the burst should not
overpower the template descriptor.

Some aggregation schemes apply average pooling (Chen et al. 2015; Parkhi,
Vedaldi, and Zisserman 2015) to the set of features. However, this method
presents a significant drawback. All images within the set are weighted
equally, regardless of their quality. To address this issue, (Kim, Jain, and
Liu 2022; Meng et al. 2021) suggest facial features whose magnitude is
proportional to image quality. To improve the aggregation, other works model
intra-set relationships using reinforcement learning (Liu et al. 2018),
recurrent models (Gong, Shi, and Jain 2019), or self-attention (Gong et al.
2019; Liu et al. 2019; Liu, Yan, and Ouyang 2017; Xie and Zisserman 2018). In
contrast to intra-set modeling, face burstiness received only limited
attention in the literature (Wang, Zhao, and Wu 2022).

While Recurrent methods (Graves and Graves 2012; Gong, Shi, and Jain 2019) are
effective in sequence modeling, they are unsuitable for modeling orderless
sets as they are not permutation invariant. On the other hand, Multihead Self
Attention (MSA) (Vaswani et al. 2017; Liu et al. 2018) is an effective
permutation-invariant approach to model self-attention among template
features, but it becomes impractical for large templates since it is quadratic
in the template size.

In this work, we propose a feature aggregation approach, denoted
FaceCoresetNet, that selects a small coreset we dub Core-Template of the face
template with balanced quality and diversity. Given $N$ data points, a coreset
is a smaller subset of $K<<N$ sampled points that approximate the original
dataset with respect to a given metric. It is often used to speed up training
or reduce the memory and storage requirements of large datasets
(Mirzasoleiman, Bilmes, and Leskovec 2020; Campbell and Broderick 2018;
Pooladzandi, Davini, and Mirzasoleiman 2022).

We select the Core-Template using a greedy differentiable farthest point
sampling (FPS) process with a learned metric parameterized by the quality of
the face images. As computing the farthest point requires an Argmax operation,
which is not differentiable, we suggest a novel differential FPS by replacing
the Argmax operation with sampling from the differentiable Gumbel-Softmax
(Jang, Gu, and Poole 2016) distribution of quality-aware distances. Unlike
conventional FPS which starts with a point selected at random, our selection
process starts with highest quality feature, hence the selection is
permutation invariant. An example of a face template and the selected Core-
Template is depicted in Fig. 1. As far as we know, this is the first
differential coreset approach that can be fully optimized by backpropagation.
Importantly, our differentiable Core-Template selection reduces the unbounded
template size $N$ to a fixed size $K$, facilitating the overall linear
algorithm complexity.

To enhance the Core-Template representation, we compute self-attention amongst
the Core-Template features. To enrich the representation with information from
the full template missing from the Core-Template, we compute cross-attention
between the Core-Template and the full template, using the Core-Template
features as queries and the full template as Key-Values. By using the Core-
Template features as queries for self-attention and cross-attention, we
compute attention in linear time, unlike the typical quadratic time complexity
in other works modeling intra-template relations, such as in the well-written
paper by (Kim et al. 2022).

As a final step, we sum and normalize the enriched Core-Template features to
get the final feature representation for the template. The final template
representative feature is used to compute the vanilla FR adaptive margin loss
(Kim, Jain, and Liu 2022) to optimize the model, which is simple and
effective. The Core-Template selection process and attention computation are
linear in the template size, facilitating an overall linear time complexity of
FaceCoresetNet. The proposed scheme is shown to achieve SOTA accuracy on the
IJB-B (Whitelam et al. 2017), IJB-C (Maze et al. 2018) FSR benchmarks.

In particular, we propose the following contributions:

  * •

We propose a novel differential, permutation invariant, coreset selection for
SFR, with a differentiable FPS strategy based on a quality-aware metric to
balance quality and diversity. The coreset selection and metric parameters are
part of the trained model and optimized end-to-end. As far as we know, this is
the first differential coreset selection method with potential value to other
fields.

  * •

We model intra-template relations using the selected coreset as queries in
self and cross-attention architecture to extract information from the whole
set.

  * •

As far as we know, our algorithm is the first to compute intra-template
relations in linear time.

  * •

We use a vanilla FR AdaFace loss with no additional loss terms for a simple
design.

  * •

We establish a new SoTA for face verification on IJB-B and IJB-C datasets

Figure 2: FaceCoresetNet Architecture

## Related Work

A simple method to combine features involves computing the mean of a
collection of features (Chen et al. 2015; Parkhi, Vedaldi, and Zisserman
2015). A limitation of this elementary technique is that inferior-quality
images can dominate the combined descriptor, diminishing its discriminatory
capacity. During training with margin-based softmax, superior-quality faces
indirectly produce features with greater magnitude (Meng et al. 2021; Kim,
Jain, and Liu 2022). By summing the features, high-quality images are
inherently emphasized without needing explicit image quality evaluation. Both
CFAN (Deng et al. 2019) and QAN (Liu, Yan, and Ouyang 2017) employ a self-
attention mechanism to learn image weighting through an explicit quality
assessment (Gong et al. 2019; Liu, Yan, and Ouyang 2017). Nevertheless, these
approaches share a common drawback: they do not consider the relationships
within the set during the weight computation stage.

To model intra-set relationships, the nonlocal neural network and RSA (Liu et
al. 2019; Wang et al. 2018) utilize intermediate feature maps $U_{i}$ of size
${C_{m}\times H\times W}$, during aggregation. This approach leverages the
rich and complementary information feature maps provide and considers spatial
relationships to refine them. However, the main limitation is the intensive
computation required for attention calculation. Specifically, for a set of N
feature maps, an attention module requires an affinity map of size ${(N\times
H\times W)}^{2}$. CAFace, cluster and Aggregate network (Kim et al. 2022)
generates an affinity map of size $N^{2}$, resulting in improved computation
efficiency for attention calculation. Using our differential selection of
coresets, this work reduces the feature fusion complexity to $O(N)$.

Set burstiness suppression network (Wang, Zhao, and Wu 2022), tackles the face
burstiness, a problem mostly overlooked by FSR research, by quantizing the
face set feature Gram matrix, which has a size of $N^{2}$. In the main FSR
datasets IJB-B, IJB-C, each image in each template is labeled by its media
source (either a video clip or a singular image). MagFace and AdaFace (Kim,
Jain, and Liu 2022; Meng et al. 2021) workaround the face burstiness problem
by using these labels to fuse each media source separately and then sum the
per-media features to produce the final fused feature. In this method, each
media source contributes equally to the final descriptor, regardless of size,
mitigating the burstiness problem. In our work, we solve face burstiness in
linear time without resorting to media labels, which may not be available, and
still achieve superior performance.

DAC (Liu et al. 2018) proposes an RL-based quality estimator, while MARN
(Gong, Shi, and Jain 2019) is an RNN-based quality estimator. However, these
methods are not agnostic to input order, making them unsuitable for modeling
orderless sets. In contrast, our method is permutation invariant, and well
suited for set modeling.

A long-standing practice in fundamental approaches such as kNN and kMeans
(Har-Peled and Kushal 2005) or mixture models (Feldman, Faulkner, and Krause
2011) is the use of subsets of the core to approximate the structure of an
available set. These subsets enable the finding of approximate solutions with
considerably reduced costs (Agarwal et al. 2005). Recently, coreset-based
methods have made their way into Deep Learning techniques, such as network
pruning (Mussay et al. 2021), active learning (Kim and Shin 2022), and
increasing effective data coverage of mini-batches for improved GAN training
(Sinha et al. 2020), or representation learning (Roth et al. 2020). The latter
two have achieved success through a greedy coreset selection mechanism.

As far as we know, all previous coreset selection methods are non-
differentiable. In contrast, in our proposed method, the coreset selection is
integrated into our PyTorch model and trained end-to-end. Quality-based
distance metric parameters are optimized during training, allowing an optimal
balance between quality and diversity in the selected coreset.

To implement our differentiable Core-Template selection, we use the Gumbel-
Softmax (Jang, Gu, and Poole 2016) technique, which enables the smoothing of
discrete categorical distributions, making them suitable for backpropagation.
It employs the Gumbel-Max reparametrization trick, which efficiently draws
samples from a Categorical distribution. Let $z$ be the categorical variable
with class probabilities $\pi_{1},\pi_{2}\dots\pi_{k}$. The trick asserts that
sampling from the categorical distribution is equivalent to sampling from the
following expression:

$z=one\\_hot(\underset{i}{\mathrm{argmax}}[g_{i}+log\pi_{i}])$

Here, $log\pi_{i}$ denotes the logits of the categorical distribution, and
$g_{i}$ follows a Gumbel distribution with parameters (0,1). To avoid non-
differentiability introduced by the Argmax operation, it is replaced with a
Softmax function:

$y_{i}=\dfrac{exp((log(\pi_{i})+g_{i})/\tau)}{\Sigma_{j=1}^{k}exp((log(\pi_{j})+g_{j})/\tau)}$
for $i=1,...,k$.

The temperature parameter ($\tau$) controls the spreading of the distribution.
As $\tau$ approaches zero, the Gumbel-Softmax method converges to the discrete
categorical distribution. During training, a $\tau$ value of 1 is typically
used, while during inference, $\tau$ approaches zero ($\tau\rightarrow
0^{+}$). (Yang et al. 2019) employ the Gumbel-Softmax for a differentiable
sampling of point clouds, facilitating tasks such as point cloud segmentation
and classification.

Multi-head self-attention (MSA) (Vaswani et al. 2017) is a commonly used
function representing intra-set and inter-set relationships through an
affinity map. It is a critical component of transformer architectures, which
have surpassed CNNs in various visual tasks (Dosovitskiy et al. 2021; Zhang et
al. 2022; Zheng et al. 2021). In principle, each MSA output value is a
weighted average of other values controlled by the queries and keys affinity.
The MSA time complexity is $O(N^{2})$; hence computing intra-set attention on
a large template is prohibitive. To reduce the computation time, we compute
intra-set attention on a fixed-size coreset, and cross-attention between the
coreset and the template in $O(N)$ time.

## Method

Let $\bm{T}=\\{\bm{x}_{1},\bm{x}_{2},\ldots,\bm{x}_{N}\\}$ be a template of
$N$ facial images of the same identity. The goal is to produce a
discriminative single compact feature vector $\bm{f}$ from $\bm{T}$. To
compute $\bm{f}$, we employ an Extract, Select, Attend and Aggregate paradigm.
To Extract (fig. 2 left) compact feature $\bm{f_{i}}$ from each image
$\bm{x}_{i}$, we use a pre-trained frozen single image FR model
$E:\bm{x}_{i}\rightarrow\bm{f}_{i}$ following (Meng et al. 2021; Kim, Jain,
and Liu 2022). The set of features produced is
$\bm{F}=\\{\bm{f}_{1},\bm{f}_{2},\ldots,\bm{f}_{N}\\}$ or in a matrix notation
$\bm{F}\in\mathbb{R}^{N\times N_{c}}$ where $N_{c}$ is the number of channels
in the compact feature. During training, we create a batch of $B$ templates of
$B$ different individuals. Each template of size $N$ where $N_{min}\leq N\leq
N_{max}$ , i.e. $\bm{F}\in\mathbb{R}^{B\times N\times N_{c}}$. The template
size $N$ is randomly selected for each mini-batch.

### Core Template Selection

To reduce the template feature size $\bm{F}$ from unbounded size $N$, we
Select (Fig. 2 green block) a core template denoted by $\bm{CT}$ of fixed size
$K$. Our core template selection process strives to select high-quality images
with diverse appearances to address the burstiness problem. As established in
(Kim, Jain, and Liu 2022), $||\bm{f}_{i}||$ is a proxy for $\bm{x}_{i}$’s
image quality. Our core template selection starts by selecting the feature
$\bm{f}_{i}$ with the largest norm in $\bm{F}$, the template’s highest quality
feature. Following the farthest point sampling (Eldar et al. 1997) paradigm,
in each selection iteration, we choose the farthest feature in
$\bm{F}\setminus\bm{CT}$ from the set of features already selected in
$\bm{CT}$. The farthest point in our setting should correlate to a high-
quality image with a diverse appearance relative to the features already
selected for the core template. Cosine distance between two features
$\bm{f}_{i}$ and $\bm{f}_{j}$ is defined as:

$\displaystyle d_{c}(\bm{f}_{i},\bm{f}_{j})\coloneqq
1-\frac{\bm{f}_{i}\cdot\bm{f}_{j}}{||\bm{f}_{i}||||\bm{f}_{j}||}$ (1)

Large cosine distance between features correlates with large appearance
variations (face pose, expression, age, etc.) between the corresponding
images. Selecting the farthest feature using this metric encourages core
template appearance diversification, but ignores image quality, resulting in
the selection of diverse but low-quality images. To account for image quality,
we scale the cosine distance by the feature norm exponentiated by a learned
parameter $\gamma$ to balance between quality and diversity:

$\displaystyle d_{q}(\bm{f}_{i},\bm{f}_{j};\gamma)\coloneqq
d_{c}(\bm{f}_{i},\bm{f}_{j})||\bm{f}_{j}||^{\gamma}$ (2)

Here, $\bm{f}_{i}\in\bm{CT}$ and $\bm{f}_{j}\in\bm{F}\setminus\bm{CT}$ and
$d_{q}$ is a non-symmetric distance function modulated by image $\bm{x}_{j}$’s
quality. $\gamma$ is a trained parameter that reflects the optimal importance
of the image quality. If $\gamma$ is 0, there is no weight to quality and
$d_{q}$ is reduced to $d_{c}$. If $\gamma>>1$, the feature quality dominates,
and we have
$d_{q}(\bm{f}_{i},\bm{f}_{j};\gamma)\approx||\bm{f}_{j}||^{\gamma}$ as
illustrated in Fig. 4. The effect of using our quality aware distance function
vs. regular cosine distance is illustrated in Fig. 3. In Fig. 3A, we have the
template $\bm{F}=\\{\bm{f}_{1},\bm{f}_{2},\bm{f}_{3}\\}$ and $\bm{f}_{1}$
already in the Core Template ($\bm{CT}=\\{\bm{f}_{1}\\}$), marked by the red
dashed ellipse, and we want to choose the next feature for the core template
out of $\bm{F}\setminus\\{\bm{f}_{1}\\}=\\{\bm{f}_{2},\bm{f}_{3}\\}$ (marked
by the blue dashed ellipse). Using the regular cosine distance, the feature
$\bm{f}_{3}$ of low-quality image 3 is farther from $\bm{f}_{1}$ than high-
quality image 2 due to its larger pose difference. As a result, the low
quality $\bm{f}_{3}$ feature is selected for CT
($\bm{CT}=\\{\bm{f}_{1},\bm{f}_{3}\\}$), reducing the discriminative power of
$\bm{CT}$. In Fig. 3B, using our quality aware metric Eq. 2, the distance
$d_{c}(\bm{f}_{1},\bm{f}_{2})$ is increased by $\bm{f}_{2}$’s larger quality
measure, and the distance $d_{c}(\bm{f}_{1},\bm{f}_{3})$ is reduced by
$\bm{f}_{3}$ lower quality, reversing the distance order and selecting the
high quality $\bm{f}_{2}$ for the core template
($\bm{CT}=\\{\bm{f}_{1},\bm{f}_{2}\\}$)

Figure 3: Cosine distance (A) vs. quality-aware distance (B) for core template
selection. The distance from $f_{1}$ to $f_{2}$ increases in B relative to A,
as $f_{2}$ has high quality. The distance from $f_{1}$ to $f_{3}$ is decreased
in B, as $f_{3}$ has low quality.

Core Template Selection

Figure 4: The core template selection is influenced by the value of $\gamma$.
In this example, the template size $N=5$ and $K=3$. In the top row, when
$\gamma$ is large, the selection process prioritizes image quality over
diversity. As a result, three similar high-quality images are chosen, marked
with red rectangles. Conversely, in the bottom row, with $\gamma=0$, diversity
precedes image quality, leading to a selection that emphasizes a diverse range
of images.

The Core-Template selection algorithm extends the FPS algorithm with important
modifications. Instead of relying on the non-differentiable Argmax operation
to choose the feature in $\bm{T}$ farthest from $\bm{CT}$, we adopt a
differentiable approach by sampling from the Gumbel-Softmax distribution of
distances (Jang, Gu, and Poole 2016). This enables us to incorporate
differentiability and leverage backpropagation during training. Additionally,
we introduce a quality-based distance function that incorporates a learned
parameter $\gamma$, allowing the algorithm to achieve an optimal balance
between quality and diversity in the selected Core-Template features.
Furthermore, unlike the conventional FPS method that randomly selects the
first feature, we prioritize the feature with the highest norm (highest
quality) to be inserted first into the Core-Template $\bm{CT}$. This selection
criterion ensures permutation invariance and enhances the algorithm’s overall
performance. The pseudocode for our core template selection algorithm is in
Algorithm 1.

Algorithm 1: Differential Core-Template Selection

⬇

1

2def Core-Template_selection(F, K, F_norms):

3 """

4 Input:

5 F: normalized features data, [B, N, C]

6 K: The intended coreTemaplate (CT) size (scalar)

7 F_norms: F norms, [B, N]

8 gamma: learned quality vs. diversity balance parameter (scalar)

9 Return:

10 CoreTemplate CT with size K

11 """

12 temperture = 1.0 if train else 1e-10

13 #Soft select the feature with highest quality (max norm)

14 1_hot_max = torch.gumbel_softmax(F_norms, tau=temperture, hard=True)

15 #Extract the feature with maximum norm

16 CT = F @ 1_hot_max

17 #After extraction compute the distance from CT to F

18 d_CT_to_F = quality_aware_dist(CT, F_norms, F)

19 for i in range(K-1):

20 1_hot_max = torch.gumbel_softmax( d_CT_to_F, tau=temperture, hard=True)

21 # extract next feature to add to CT

22 new_f = F @ 1_hot_max

23 d_new_f_to_CT = quality_aware_dist(new_f, norms, F, gamma)

24 d_CT_to_F = torch.min(d_CT_to_F, d_new_f_to_CT)

25 CT = torch.cat([CT, new_f])

26 return CT

27

28def quality_aware_dist(candidate_point, F_norms, F, gamma):

29 inner_product = torch.bmm(F, candidate_point)

30 cosine_dist = (1 - inner_product)

31 quality_dist = torch.pow(norms, gamma) * cosine_dist

32 return quality_dist

The algorithm starts by sampling from the Gumbel-Softmax distribution of the
template’s norms $\bm{F}$ (line 14). The sampling operator returns a 1-hot
vector that is used in line 16 to insert the highest quality feature to the
Core-Template $\bm{CT}$. In line 18, we calculate the distance from each
feature in $\bm{F}$ to the newly selected feature in $\bm{CT}$. In each
iteration of the loop starting at line 19 we add one new member to $\bm{CT}$
that is farthest from it. We use the quality aware distance function in line
28 for all the distance calculations. An example of one iteration in this
process is depicted in Figure 5. Consider the feature set
$\bm{F}=\\{\bm{f}_{1},\bm{f}_{2},\bm{f}_{3}\\}$ and the initial Core-Template
$\bm{CT}=\\{\bm{f}_{1}\\}$. The quality aware distances from $\bm{F}$ to
$\bm{CT}$, serve as input to the selection process. In the next stage, Gumbel
noise is added to each distance value. These values are then multiplied by the
temperature parameter $\tau$, and the softmax operation is applied to obtain
the final Gumbel-Softmax distribution. A mask is generated to identify the
maximal value in the distribution, which is then utilized to determine the
next feature to add to $\bm{CT}$.

Figure 5: Illustration of a single iteration of the differential core-template
selection (lines 19 to 25) in Alg. 1

Importantly, our proposed core template selection process time complexity is
$O(NK)$ where $N$ is the template size, and $K$ is the core template size with
$K<<N$. Assuming $K$ is a small constant (we found $K=3$ optimal in our
experiments), we have a selection process with linear complexity. In addition,
the core template selection process is permutation invariant, as it depends
solely on the geometrical attributes of each feature in $\bm{F}$ and not on
their order.

### Attend

Our core template selection process performs an important data reduction from
unbounded template size $N$ to a small fixed-size core template. This data
reduction enables the application of attention mechanisms like multi-head-
self-attention (Vaswani et al. 2017), which have quadratic time complexity and
are prohibitively slow for large templates.

We use Multihead Self Attention (MSA) as an encoder to enhance the selected
core template representation power and exchange information among its features
(Fig. 2).

Inspired by recent developments in detection transformers such as Efficient-
DETR (Yao et al. 2021), which demonstrate the superiority of leveraging image
features as queries instead of relying on learned queries for decoding as in
the original transformer (Vaswani et al. 2017) and DETR (Carion et al. 2020),
we leverage the core template features as transformer decoder queries that
interact with key values taken from the full template features in $\bm{F}$.
This allows the core template features to extract additional information from
the full template that is missing from the core template.

Norm encoding. It appeals to our intuition that explicit image quality data
can aid information extraction in multi-head attention settings, facilitating
the suppression of low-quality features and highlighting high-quality
features. To incorporate the quality information, we encode the feature norms
as a quality proxy, using the standard Sinusidial conversion established in
(Vaswani et al. 2017), and add them to the original features. Our experiments
validate the positive impact of Norm encoding on the model’s overall accuracy.

Importantly, our attention process interpolates between template features
using self and cross-attention realized with MHA. We avoid the Feed-Forward
block in the transformer encoder and decoder blocks, which extrapolates the
template features and deteriorates performance.

The complexity of the encoder is $O(K^{2})$, which is constant. The complexity
of the decoder is $O(NK)$. The total complexity of the attention block is
$O(N)$.

Aggregate. To represent the entire template with one feature $\bm{f}$, we
combine the $K$ features in the enhanced core template by adding them up and
normalizing the result to have a unit length.

### Training and Loss

After the aggregation phase, we get a single vector $\bm{f}$ representing the
whole template. From this point on, the problem is reduced from SFR to FR,
facilitating the use of SoTA FR loss functions. In this work, we choose
AdaFace (Kim, Jain, and Liu 2022) loss for its superior accuracy. AdaFace
computes a class margin adaptively based on the image quality of the face
approximated by the face feature norm. In our case, the margin is computed by
the template quality instead of the image (See Fig. 2 right-hand side). Using
a simple margin-based loss for template feature fusion facilitates a
simplistic implementation compared with elaborate loss terms found in (Kim,
Jain, and Liu 2022) and others.

Large modern FR datasets like WebFace (Zhu et al. 2021) contain millions of
identities but only a limited number of images per identity (WebFace4M
contains about 20 images per identity on average) with no video clips. On the
other hand, SFR benchmarks like IJB-B, IJB-C contain rich templates with a
diverse combination of still images and video clips for each identity. To
bridge this gap, we online augment the training set and simulate IJB-B, IJB-C
style templates. To compensate for the lack of video clips, we simulate a
video clip by randomly sampling an image and applying various geometrical
transformations (translations, rotations, etc.) on it to create multiple
”video frames”. Each training template comprises a random collection of still
photos and simulated video clips with random addition of photometric noise to
each image.

## Experiments

To allow for a valid comparison with the current SOTA, we strictly follow the
experiments protocol of CAFace (Kim et al. 2022) (NeurIPS’22). Our training
dataset is WebFace4M (Zhu et al. 2021), which contains 4.2 million facial
images from 205,990 identities. We use CAFaces’s pre-trained face recognition
model, $E$, an IResnet-101, trained with AdaFace loss on the entire WebFace4M
dataset . To train the FaceCoresetNet, we use CAFace’s subset of 813,482
images from 10,000 identities. As VGG-2 (Cao et al. 2018) and MS1MV2 (Guo et
al. 2016; Deng et al. 2019) creators have withheld them due to privacy and
other issues, we do not use them in our experiments. We train FaceCoresetNet
for $\bf{2}$ epochs till convergence, compared with CAFace’s $\bf{10}$ epochs
training schedule.

We test on IJB-B (Whitelam et al. 2017) and IJB-C (Maze et al. 2018) datasets
as they are intended for template-based face recognition. IJB-B comprises 1845
individuals and includes 11,754 images, 55,025 frames, and 7,011 videos. Each
template in the dataset consists of a diverse collection of still images and
video frames from various sources. These images and videos were gathered from
the Internet and are considered unconstrained, as they exhibit considerable
variation in pose, lighting conditions, and image quality. Following CAFace,
we focus our experiments on face verification, where IJB-B offers 10,270
genuine template pair comparison and 8,000,000 impostor comparisons,
facilitating valid measurements of very small FAR values on the TAR@FAR ROC
curve. IJB-C extends IJB-B by adding 1,661 new subjects, with increased
emphasis on occlusion and diversity.

# |  | Differential
---
Core-Template
Selection
Self Att | Cross Att |  | Norm
---
Encoding
| IJB-B:TAR@
---
FAR=1e-6
1 | $\times$ | $\times$ | $\times$ | $\times$ | 38.45
2 | ✓ | $\times$ | $\times$ | $\times$ | 48.01
3 | ✓ | ✓ | $\times$ | $\times$ | 49.24
4 | ✓ | ✓ | ✓ | $\times$ | 51.71
5 | ✓ | ✓ | ✓ | ✓ | 52.56

Table 1: FaceCoresetNet Ablations

|  |  | IJB-B 1:1 Verification TAR | IJB-C 1:1 Verification TAR
---|---|---|---|---
Method |  Complexity ($N$: Template Size) | FAR=1e-4 | FAR=1e-5 | FAR=1e-6 | FAR=1e-4 | FAR=1e-5 | FAR=1e-6
AdaFace (Kim, Jain, and Liu 2022) | NA | 94.84 | 90.86 | 38.45 | 96.42 | 94.47 | 87.90
RSA (Liu et al. 2019) | $O(N^{2})$ | 95.00 | 91.22 | - | 96.49 | 94.58 | -
SD + VBA (Wang, Zhao, and Wu 2022) | $O(N^{2})$ | 95.38 | 90.89 | 48.96 | 96.65 | 94.85 | 90.02
CAFace (Kim et al. 2022) | $O(N^{2})$ | 95.78 | 92.78 | 47.21 | 97.30 | 95.96 | 90.56
FaceCoresetNet | $O(N)$ | 95.28 | 91.44 | 52.56 | 96.96 | 95.21 | 91.12

Table 2: A performance comparison of recent methods on IJB-B (Whitelam et al.
2017) and IJB-C (Maze et al. 2018) datasets. The results are color-coded, with
green representing the best, blue the second best, and red the third-best
outcomes. FaceCoresetNet achieves the best results on TPR@FPR=1e-6 for both
IJB-B and IJB-C while being the most efficient.

### Ablation and Analysis

To demonstrate the effectiveness of each component in FaceCoresetnNet, we
conduct a series of experiments by ablating individual components. Initially,
we turned off all components, resulting in a configuration where the template
feature was simply the average feature computed from all the templates’
features. The obtained TAR@FPR=1e-6 for the IJB-B dataset was 38.45 (Table 1).
Subsequently, we enabled the differential core-template selection. We computed
the final template descriptor by averaging the pooled features within the core
template. Surprisingly, by effectively selecting 3 features, the accuracy
increased significantly to 48.01. We hypothesize that our core template’s high
quality and diversity reduce the amount of noise present in the original
template, which may explain the observed improvement in accuracy. Furthermore,
we integrated self-attention into the core-template features, boosting
accuracy to 49.24. Additionally, by employing cross-attention between the core
and full templates, the accuracy increased to 51.71. Finally, with the
incorporation of Norm encoding, the accuracy reached 52.56.

$\gamma$ policy |  | IJB-C:
---
TAR@FAR=1e-6
Fixed 0 (Diversity priority) | 90.84
Fixed 10 (Quality priority) | 89.69
Trained | 91.12
Table 3: Training $\gamma$ for optimal balance between quality and diversity
yields best results, compared with fixed $\gamma$ values

Effect of Core-Template size. To determine the optimal core-template size, we
searched a range of sizes from 1 to 6. The outcomes of this search are
summarized in Table 1 of the supplementary materials. Notably, the ideal
coreset size was identified as 3, and this optimal value remained constant
throughout all our subsequent experiments.

$\gamma$ computation policy. Our model optimizes the $\gamma$ value in the
quality-aware distance. To showcase the effectiveness of this optimization, we
conducted 3 experiments. One setting $\gamma$ to the constant 10 value,
prioritizing quality over diversity. Conversely, we set $\gamma$ to 0, which
reduces the quality aware distance to the vanilla cosine distance,
prioritizing diversity over quality. In the last experiment, we optimize
$\gamma$ end-to-end. The results in Table 3 demonstrates the effectiveness of
our approach.

Comparison with SoTA methods. In Table. 2, we compare our results with the
state-of-the-art (SOTA) in face verification, focusing on the IJB-B and IJB-C
datasets. Our method achieves the highest accuracy in the low false acceptance
rate regime, TAR@FAR=10e-6, for both IJB-B and IJB-C, outperforming others
significantly. Additionally, we obtain the second-best results for
TAR@FAR=10e-5 while maintaining the advantage of utilizing the lowest linear
time compute.

Figure 6: Computation complexity comparison with SoTA. FaceCoresetNet is
linear in the template size, while CAFace is quadratic.

Computation Complexity. To validate our algorithm’s claimed linear time
complexity, we employ (fvcore 2023) to calculate the FLOPS in our
FaceCoresetNet and compare it against CAFace (Kim et al. 2022) using the
authors’ code. The FLOPS measurements for different template sizes are
illustrated in Fig. 6. The results demonstrate that our algorithm indeed
operates in linear time, while CAFace exhibits quadratic time complexity.

## Conclusion

In this work, we employ a method to fuse face template features by selecting a
small, fixed-size core-template that balances quality and diversity. This
selection uses a differentiable farthest-point sampling with a quality-aware
distance metric, realized by sampling from the Gumbel-Softmax distribution of
quality-aware distances.

The selected core template features are further enhanced using MHA self-
attention. These enhanced features are then used as queries for MHA cross-
attention between the core template and the entire template, enabling the
extraction of additional information that may be missing from the core
template.

Our proposed model achieves state-of-the-art accuracy on IJB-B and IJB-C
datasets, all while maintaining reduced linear time complexity and simple
design. FaceCoresetNet is showcased within the context of face set
recognition, yet its applicability may extend to other domains requiring
efficient feature fusion based on a differentiable selection with a learned
metric. We hope to see an extension of this work to other domains in future
research.

## References

  * Agarwal et al. (2005) Agarwal, P. K.; Har-Peled, S.; Varadarajan, K. R.; et al. 2005.  Geometric approximation via coresets.  _Combinatorial and computational geometry_ , 52(1): 1–30.
  * Campbell and Broderick (2018) Campbell, T.; and Broderick, T. 2018.  Bayesian coreset construction via greedy iterative geodesic ascent.  In _Proceedings of the International Conference on Machine Learning (ICML)_ , 698–706. PMLR.
  * Cao et al. (2018) Cao, Q.; Shen, L.; Xie, W.; Parkhi, O. M.; and Zisserman, A. 2018.  Vggface2: A dataset for recognising faces across pose and age.  In _International Conference on Automatic Face and Gesture Recognition (FGR)_ , 67–74. IEEE.
  * Carion et al. (2020) Carion, N.; Massa, F.; Synnaeve, G.; Usunier, N.; Kirillov, A.; and Zagoruyko, S. 2020.  End-to-end object detection with transformers.  In _Proceedings of the European Conference on Computer Vision (ECCV)_ , 213–229. Springer.
  * Chen et al. (2015) Chen, J.-C.; Ranjan, R.; Kumar, A.; Chen, C.-H.; Patel, V. M.; and Chellappa, R. 2015.  An end-to-end system for unconstrained face verification with deep convolutional neural networks.  In _Proceedings of the IEEE International Conference on Computer Vision Workshops (ICCVW)_ , 118–126.
  * Deng et al. (2019) Deng, J.; Guo, J.; Xue, N.; and Zafeiriou, S. 2019.  ArcFace: Additive angular margin loss for deep face recognition.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 4690–4699.
  * Dosovitskiy et al. (2021) Dosovitskiy, A.; Beyer, L.; Kolesnikov, A.; Weissenborn, D.; Zhai, X.; Unterthiner, T.; Dehghani, M.; Minderer, M.; Heigold, G.; Gelly, S.; Uszkoreit, J.; and Houlsby, N. 2021.  An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale.  _Proceedings of the International Conference on Learning Representations (ICLR)_.
  * Eldar et al. (1997) Eldar, Y.; Lindenbaum, M.; Porat, M.; and Zeevi, Y. Y. 1997.  The farthest point strategy for progressive image sampling.  _IEEE Transactions on Image Processing_ , 6(9): 1305–1315.
  * Feldman, Faulkner, and Krause (2011) Feldman, D.; Faulkner, M.; and Krause, A. 2011.  Scalable Training of Mixture Models via Coresets.  In Shawe-Taylor, J.; Zemel, R.; Bartlett, P.; Pereira, F.; and Weinberger, K., eds., _Advances in Neural Information Processing Systems (NIPS)_ , volume 24. Curran Associates, Inc.
  * fvcore (2023) fvcore. 2023.  fvcore.  https://github.com/facebookresearch/fvcore.
  * Gong, Shi, and Jain (2019) Gong, S.; Shi, Y.; and Jain, A. 2019.  Low quality video face recognition: Multi-mode aggregation recurrent network (MARN).  In _IEEE Conference on Computer Vision and Pattern Recognition Workshops (CVPRW)_ , 0–0.
  * Gong et al. (2019) Gong, S.; Shi, Y.; Kalka, N. D.; and Jain, A. K. 2019.  Video face recognition: Component-wise feature aggregation network (c-fan).  In _International Conference on Biometrics (ICB)_ , 1–8. IEEE.
  * Graves and Graves (2012) Graves, A.; and Graves, A. 2012.  Long short-term memory.  _Supervised sequence labelling with recurrent neural networks_ , 37–45.
  * Guo et al. (2016) Guo, Y.; Zhang, L.; Hu, Y.; He, X.; and Gao, J. 2016.  Ms-celeb-1m: A dataset and benchmark for large-scale face recognition.  In _Proceedings of the European Conference on Computer Vision (ECCV)_ , 87–102. Springer.
  * Har-Peled and Kushal (2005) Har-Peled, S.; and Kushal, A. 2005.  Smaller coresets for k-median and k-means clustering.  In _Proceedings of the twenty-first annual symposium on Computational geometry_ , 126–134.
  * Jang, Gu, and Poole (2016) Jang, E.; Gu, S.; and Poole, B. 2016.  Categorical reparameterization with gumbel-softmax.  _arXiv preprint arXiv:1611.01144_.
  * Kalka et al. (2018) Kalka, N. D.; Maze, B.; Duncan, J. A.; O’Connor, K.; Elliott, S.; Hebert, K.; Bryan, J.; and Jain, A. K. 2018.  Ijb–s: Iarpa janus surveillance video benchmark.  In _2018 IEEE 9th international conference on biometrics theory, applications and systems (BTAS)_ , 1–9. IEEE.
  * Kim, Jain, and Liu (2022) Kim, M.; Jain, A. K.; and Liu, X. 2022.  AdaFace: Quality adaptive margin for face recognition.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 18750–18759.
  * Kim et al. (2022) Kim, M.; Liu, F.; Jain, A.; and Liu, X. 2022.  Cluster and Aggregate: Face Recognition with Large Probe Set.  In _Advances in Neural Information Processing Systems (NIPS)_.
  * Kim and Shin (2022) Kim, Y.; and Shin, B. 2022.  In Defense of Core-set: A Density-aware Core-set Selection for Active Learning.  In _Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining_ , 804–812.
  * Liu et al. (2019) Liu, X.; Guo, Z.; Li, S.; Kong, L.; Jia, P.; You, J.; and Kumar, B. 2019.  Permutation-invariant feature restructuring for correlation-aware image set-based recognition.  In _Proceedings of the IEEE International Conference on Computer Vision (ICCV)_ , 4986–4996.
  * Liu et al. (2018) Liu, X.; Kumar, B.; Yang, C.; Tang, Q.; and You, J. 2018.  Dependency-aware attention control for unconstrained face recognition with image sets.  In _Proceedings of the European Conference on Computer Vision (ECCV)_ , 548–565.
  * Liu, Yan, and Ouyang (2017) Liu, Y.; Yan, J.; and Ouyang, W. 2017.  Quality aware network for set to set recognition.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 5790–5799.
  * Maze et al. (2018) Maze, B.; Adams, J.; Duncan, J. A.; Kalka, N.; Miller, T.; Otto, C.; Jain, A. K.; Niggel, W. T.; Anderson, J.; Cheney, J.; and Grother, P. 2018.  IARPA Janus Benchmark - C: Face Dataset and Protocol.  In _International Conference on Biometrics (ICB)_ , 158–165.
  * Meng et al. (2021) Meng, Q.; Zhao, S.; Huang, Z.; and Zhou, F. 2021.  Magface: A universal representation for face recognition and quality assessment.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 14225–14234.
  * Mirzasoleiman, Bilmes, and Leskovec (2020) Mirzasoleiman, B.; Bilmes, J.; and Leskovec, J. 2020.  Coresets for data-efficient training of machine learning models.  In _International Conference on Machine Learning_ , 6950–6960. PMLR.
  * Mussay et al. (2021) Mussay, B.; Feldman, D.; Zhou, S.; Braverman, V.; and Osadchy, M. 2021.  Data-independent structured pruning of neural networks via coresets.  _IEEE Transactions on Neural Networks and Learning Systems_ , 33(12): 7829–7841.
  * Parkhi, Vedaldi, and Zisserman (2015) Parkhi, O. M.; Vedaldi, A.; and Zisserman, A. 2015.  Deep face recognition.
  * Pooladzandi, Davini, and Mirzasoleiman (2022) Pooladzandi, O.; Davini, D.; and Mirzasoleiman, B. 2022.  Adaptive second order coresets for data-efficient machine learning.  In _Proceedings of the International Conference on Machine Learning (ICML)_ , 17848–17869. PMLR.
  * Roth et al. (2020) Roth, K.; Milbich, T.; Sinha, S.; Gupta, P.; Ommer, B.; and Cohen, J. P. 2020.  Revisiting training strategies and generalization performance in deep metric learning.  In _Proceedings of the International Conference on Machine Learning (ICML)_ , 8242–8252. PMLR.
  * Sinha et al. (2020) Sinha, S.; Zhang, H.; Goyal, A.; Bengio, Y.; Larochelle, H.; and Odena, A. 2020\.  Small-gan: Speeding up gan training using core-sets.  In _Proceedings of the International Conference on Machine Learning (ICML)_ , 9005–9015. PMLR.
  * Vaswani et al. (2017) Vaswani, A.; Shazeer, N.; Parmar, N.; Uszkoreit, J.; Jones, L.; Gomez, A. N.; Kaiser, Ł.; and Polosukhin, I. 2017.  Attention is all you need.  _Advances in Neural Information Processing Systems (NIPS)_ , 30.
  * Wang, Zhao, and Wu (2022) Wang, J.; Zhao, Z.; and Wu, F. 2022.  Set-Based Face Recognition Beyond Disentanglement: Burstiness Suppression With Variance Vocabulary.  In _Proceedings of the 30th ACM International Conference on Multimedia_ , 6125–6135.
  * Wang et al. (2018) Wang, X.; Girshick, R.; Gupta, A.; and He, K. 2018.  Non-local neural networks.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 7794–7803.
  * Whitelam et al. (2017) Whitelam, C.; Taborsky, E.; Blanton, A.; Maze, B.; Adams, J.; Miller, T.; Kalka, N.; Jain, A. K.; Duncan, J. A.; Allen, K.; Cheney, J.; and Grother, P. 2017\.  IARPA Janus Benchmark-B Face Dataset.  In _IEEE Conference on Computer Vision and Pattern Recognition Workshops (CVPRW)_ , 592–600.
  * Xie and Zisserman (2018) Xie, W.; and Zisserman, A. 2018.  Multicolumn networks for face recognition.  _arXiv preprint arXiv:1807.09192_.
  * Yang et al. (2019) Yang, J.; Zhang, Q.; Ni, B.; Li, L.; Liu, J.; Zhou, M.; and Tian, Q. 2019.  Modeling point clouds with self-attention and gumbel subset sampling.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 3323–3332.
  * Yao et al. (2021) Yao, Z.; Ai, J.; Li, B.; and Zhang, C. 2021.  Efficient detr: improving end-to-end object detector with dense prior.  _arXiv preprint arXiv:2104.01318_.
  * Zhang et al. (2022) Zhang, H.; Li, F.; Liu, S.; Zhang, L.; Su, H.; Zhu, J.; Ni, L.; and Shum, H. 2022\.  Dino: Detr with improved denoising anchor boxes for end-to-end object detection.  In _Proceedings of the International Conference on Learning Representations (ICLR)_.
  * Zheng et al. (2021) Zheng, S.; Lu, J.; Zhao, H.; Zhu, X.; Luo, Z.; Wang, Y.; Fu, Y.; Feng, J.; Xiang, T.; Torr, P. H.; et al. 2021.  Rethinking semantic segmentation from a sequence-to-sequence perspective with transformers.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 6881–6890.
  * Zhu et al. (2021) Zhu, Z.; Huang, G.; Deng, J.; Ye, Y.; Huang, J.; Chen, X.; Zhu, J.; Yang, T.; Lu, J.; Du, D.; et al. 2021.  Webface260m: A benchmark unveiling the power of million-scale deep face recognition.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 10492–10502.


# Few-shot Sequence Learning with Transformers

Lajanugen Logeswaran1, Ann Lee2, Myle Ott2, Honglak Lee1,
Marc’Aurelio Ranzato2, Arthur Szlam2
1University of Michigan, 2Facebook AI Research

###### Abstract

Few-shot algorithms aim at learning new tasks provided only a handful of
training examples. In this work we investigate few-shot learning in the
setting where the data points are sequences of tokens and propose an efficient
learning algorithm based on Transformers. In the simplest setting, we append a
token to an input sequence which represents the particular task to be
undertaken, and show that the embedding of this token can be optimized on the
fly given few labeled examples. Our approach does not require complicated
changes to the model architecture such as adapter layers nor computing second
order derivatives as is currently popular in the meta-learning and few-shot
learning literature. We demonstrate our approach on a variety of tasks, and
analyze the generalization properties of several model variants and baseline
approaches. In particular, we show that compositional task descriptors can
improve performance. Experiments show that our approach works at least as well
as other methods, while being more computationally efficient.

##  1 Introduction

The problem of learning a classifier from a handful examples has received
considerable attention in the vision domain under the name of few-shot
learning (Fink, 2005; Fei-Fei et al., 2006). However, less work exists in the
space of few-shot problems involving discrete sequences, such as sequences of
discrete actions in reinforcement learning or sequences of words in natural
language processing. In this work, we study the problem of sequence
classification and modeling in the few-shot regime. Specifically, we assume
there are several training tasks available for learning and, at test time, we
are interested in performing few-shot adaptation to a given new task.

Transformers (Vaswani et al., 2017) have been very successful at modeling
discrete sequences (Barrault et al., 2019; Devlin et al., 2018; Parisotto et
al., 2019). Further, they have been shown to use context tokens appended to an
input to adapt their generations or to switch between different tasks (Lample
et al., 2019; Shen et al., 2019; Zellers et al., 2019; Keskar et al., 2019).
Thus, one might hope that such context tokens could be effectively used in the
meta-learning setting for discrete sequences.

In this work, we show that this is indeed the case. Our approach to few-shot
learning introduces a set of task specific parameters (a task embedding), in
addition to the parameters of the model that are shared among all tasks.
Unlike other approaches that require architectural changes (Houlsby et al.,
2019), task embeddings are simply fed to the input of the transformer.
Learning a new task consists of inferring an appropriate task embedding for
the task, leaving the shared model parameters intact. Towards this end, we
propose a simple training algorithm where the task embedding is found via
gradient based optimization, which is simpler and computationally less
expensive than second order optimization methods (Finn et al., 2017; Zintgraf
et al., 2019).

To summarize, our contributions in this work are as follows. First, we show
that a simple alternating-minimization approach for few-shot learning works
well in combination with the transformer architecture. Second, we show that a
simple yet effective way to condition the transformer with task information is
via input conditioning (i.e., feeding task information as input to the
transformer); this naturally extends to compositional task information. Third,
we introduce a battery of synthetic sequence classification and modeling tasks
to benchmark in a controlled setting various baseline approaches and model
variants for few shot learning of discrete sequences. And finally, we
demonstrate that the proposed approach offers a better trade-off between few-
shot performance and run time cost compared to other baselines, including
meta-learning approaches.

##  2 Problem Definition

We assume a distribution $p_{\text{data}}(\mathcal{T})$ over tasks from which
disjoint sets of training, validation and test sets of tasks are drawn. The
set of training tasks is denoted by
$\\{\mathcal{T}_{i}^{\text{train}}\\}_{i=1}^{N}$, where each task
$\mathcal{T}_{i}^{\text{train}}$ has an associated set of training examples
$\\{(x^{i}_{j},y^{i}_{j})_{j=1}^{N_{i}}\\}$. Validation and test tasks are
defined similarly, except each test task only has $k$ training examples. We
use the training tasks to learn the model parameters and the validation tasks
to determine hyperparameters. The optimal parameters and hyperparameters
identified are used for evaluating average model performance on test tasks.
This involves first (optionally) training on the small set of training
examples accompanied by a test task, followed by testing on the corresponding
test set.

We focus on two types of tasks that involve discrete sequences as inputs and
outputs: sequence classification and transduction. In sequence classification,
the inputs $x$ are sequences and the output $y$ is a discrete categorical
label. In sequence transduction, each task consists of modeling the joint
distribution of a sequence $y$, conditioned on some input context sequence
$x$. The performance metrics for these settings are respectively accuracy and
perplexity, averaged across the test tasks.

##  3 Approach

###  3.1 Architecture

In this work, we explore an adaptation of transformers to the few-shot regime.
Previous works have shown that the behavior of a transformer model can be
conditioned by appending special tokens describing the task to be performed on
the input sequence (Lample et al., 2019; Zellers et al., 2019). We append a
task embedding vector which represents information about the task of interest
to the input sequence of token embeddings. We intend to control the overall
behavior of the model for a particular task by altering the task embeddings
while keeping the rest of the model parameters intact.

For classification tasks, we use a transformer encoder similar to the BERT
model (Devlin et al., 2018). A classification head sits on the final layer
representation of a special token at the beginning of the sequence. We replace
this special token with the task embedding vector $z$ in our model. We use a
transformer decoder architecture for the transduction tasks and append the
task embedding vector to the input sequence similar to the classification
setting.

In both settings we compute a log-likelihood of the form $\log
p(y|x,z;\theta)$, where $x$ is the input sequence, $z$ is the task embedding
and $\theta$ the model parameters. In the classification setting $y$ is a
single categorical value. In the sequence transduction setting, $y$ is a
sequence and the log likelihood decomposes as the sum of the conditional log-
likelihood terms via the chain rule of probability theory: $\log
p(y|x,z;\theta)=\sum_{i}\log p(y_{i}|y_{i-1},\cdots,y_{1},x,z;\theta)$.

Note that in practical applications $\theta$ can be high-dimensional, in the
order of hundreds of millions. Our goal is to alleviate overfitting in the
few-shot regime by adapting only $z$ to learn a new task, where $z$ is a small
vector with at most a few hundred components. Next, we describe how we learn
the model parameters $\theta$ and how we estimate the task embedding $z$ for a
given task.

1

Input : Training tasks
$\mathcal{T}_{1}^{\text{train}},...,\mathcal{T}_{N}^{\text{train}}$

Output : Model parameters $\theta$

2

3repeat

4 Sample a training task: $\mathcal{T}_{i}^{\text{train}}$

5 Sample $N_{i}\geq k$ training examples from the task
$\\{(x^{j},y^{j})_{j=1,\cdots,N_{i}}\\}\sim\mathcal{T}_{i}^{\text{train}}$

6 Initialize: $z_{\mathcal{T}_{i}^{\text{train}}}=0,\Delta\theta=0$

7 while  _loss improves and max number of updates not reached_ do

8
$z_{\mathcal{T}_{i}^{\text{train}}}{\scriptsize\leftarrow}z_{\mathcal{T}_{i}^{\text{train}}}-\nabla_{z_{\mathcal{T}_{i}^{\text{train}}}}\sum_{j}\,-\log
p(y^{j}|x^{j},z_{\mathcal{T}_{i}^{\text{train}}};\theta)$

9 $\Delta\theta\leftarrow\Delta\theta-\nabla_{\theta}\sum_{j=1}^{N_{i}}\,-\log
p(y^{j}|x^{j},z_{\mathcal{T}_{i}^{\text{train}}};\theta)$

10 $\theta\leftarrow\theta+\Delta\theta$

11until  _max training iterations_ ;

Algorithm 1 TAM for k-shot Learning

###  3.2 Training and Inference Algorithm

We train our models with an alternating-minimization scheme similar to Maurer
et al. (2013) and Kumar & Daume III (2012), that can be considered a
simplification of the CAVIA approach in Zintgraf et al. (2019) (or as a
refinement of the “first-order” method in that work). See Algorithm 1 for
pseudo-code. We separate the weights of the network defining the model into
shared weights $\theta$, and per-task weights, as in CAVIA. In our case, the
per-task weights form the embedding $z$, one for each task; while all other
parameters $\theta$ are shared.

Given a few examples from the training task $\mathcal{T}_{i}^{\text{train}}$
(line 3), we alternate training $z_{\mathcal{T}_{i}^{\text{train}}}$ (the task
embedding of the $\mathcal{T}_{i}^{\text{train}}$ task) for a few gradient
descent steps keeping $\theta$ fixed (see line 5 and 6), and then update
$\theta$ based on the optimal task embedding. In practice, however, we found
it helpful to update $\theta$ based on gradients accumulated for the
intermediate values of the task embedding encountered in the inner loop
optimization (line 7). We surmise that this optimization choice helps the
model finding better task embeddings as the whole parameter vector $\theta$ is
updated to account for this search. Task embedding gradient updates (line 6)
are performed until the loss no longer improves or the maximum number of
update steps has been reached. Note that unlike prior methods such as MAML or
CAVIA we do not backpropagate gradients through an optimization process, which
simplifies and speeds up our optimization. We call our method, transformer
trained with Alternating Minimization (TAM) – although the alternating
minimization algorithm could be applied to other architectures as well. At
test time, given a new task $\mathcal{T}^{\text{test}}$,
$z_{\mathcal{T}^{\text{test}}}$ is trained with a few steps of gradient
descent (similar to line 6), with all other parameters held fixed. Since TAM
is trained to optimize task embeddings on the fly, we expect it to find good
embeddings of the new task at test time as well.

##  4 Related work

##### Few-shot learning and meta-learning

There is now a vast literature on learning methods designed for quickly
adapting to new settings. At a coarse level, one can consider classes of
methods that adapt the learning algorithm based on the task (and so are “meta-
learners”) (Schmidhuber, 1987; Hochreiter et al., 2001; Andrychowicz et al.,
2016; Finn et al., 2017; Nichol & Schulman, 2018), or describe model
architectures that can adapt to learn sample-efficiently over a task
distribution (Vinyals et al., 2016; Snell et al., 2017). Many methods have
elements of both of these, e.g. Mishra et al. (2018); Rusu et al. (2019);
Zintgraf et al. (2019).

The method we describe in this work can be considered squarely in the class of
model architectures for sample efficient learning. It is a descendant of
Hinton & Plaut (1987); Schmidhuber (1992); Ba et al. (2016) and is closely
related to Rusu et al. (2019); Zintgraf et al. (2019) in that we pick a subset
of the weights of the model that are task specific (the “fast” weights), and
update them using the training examples for a specific task; but update the
other (“slow”) weights on all training examples for all tasks. Our approach is
closest to Zintgraf et al. (2019), but differs in the way the fast weights are
used by the model. We do not use higher-order gradients for the slow weights,
instead we use an alternating minimization type update.

##### Task transfer for transformers

Our approach is also related to other recent work in natural language
processing. We leverage the particular structure of the transformer
architecture (Vaswani et al., 2017), which has been successful in many NLP
tasks. Several works have shown that adding a token to an input can be used to
switch between different tasks (Lample et al., 2019; Shen et al., 2019;
Zellers et al., 2019; Keskar et al., 2019). Transformer language models
trained on large corpora have also been recently shown to have impressive few-
shot learning capabilities (Brown et al., 2020). More generally, with the
success of methods based on pretraining transformer models (Devlin et al.,
2018), and finetuning on target tasks, there have been several works
discussing how to adapt a pre-trained model without full finetuning (Houlsby
et al., 2019; Stickland & Murray, 2019) but their focus has been on reducing
the number of parameters subject to optimization at finetuning time as opposed
to reducing the number of examples as in this study.

##  5 Experiments

###  5.1 Model and Training Details

In the classification setting, TAM is a bidirectional transformer that takes
the input sequence $x$ and the task embedding $z$ as input, and outputs a
distribution over classes. In the sequence transduction setting, TAM is a
transformer decoder with a causal attention mechanism and takes as additional
input the output sequence $y$ up to the token before the last. In this case
the model is trained to predicted the sequence $y$ at the last $|y|$ (length
of sequence $y$) time steps. Since model parameters are shared across tasks,
TAM needs to leverage the task embedding to perform the tasks well. Both
classification and transduction models are trained with cross-entropy loss.

Unless otherwise specified our transformer has 4 layers with an embedding size
of 128. We use the Adam optimizer (Kingma & Ba, 2015) for both outer and inner
loop optimization. The maximum number of task embedding optimization steps is
set to 25 during training. We train a single model using $N$ samples at
training time, treating $N$ as a hyperparameter and apply it to k-shot
problems with different values of $k$ at test time. The size of the task
embedding was set to match the embedding dimension of the transformer (128).
We discuss more details about hyperparameter choices and how they influence
model performance in section 5.6.

###  5.2 Baselines

##### Task-Agnostic transformer

This baseline uses the same architecture as TAM but is not informed about the
existence of different tasks at training time, i.e., no task embedding is fed
at the input. At test time, the model is fine-tuned on the $k$ training
examples from the test task.

##### Multitask transformer

This is a transformer that is conditioned on the current task both in the
classification and transduction settings. It is identical to TAM except all
parameters including task embeddings are trained by standard back-propagation,
without any alternating minimization.

##### Matching Networks (Vinyals et al., 2016)

We consider Matching Networks only in our classification setting, as it is not
straightforward to use it for transduction. We use a transformer to model the
similarity between a query instance and support set instance which takes the
concatenation of the two sequences as input and outputs a similarity score.
The prediction is a convex sum of the training example labels, the weights
being similarity scores.

##### SNAIL (Mishra et al., 2018)

This model is similar to the task-agnostic transformer except the input is
augmented with the concatenation of all input-output training pairs. For both
Matching Networks and SNAIL, we construct training episodes by sampling $k$
training examples to define a task, to match the test scenario. We train
different models for each $k$-shot problem. Both Matching networks and SNAIL
are trained using the multi-task training loss and applied to test tasks
without any finetuning.

##### MAML (Finn et al., 2017)

All model parameters are trained using MAML, with the same model architecture
as TAM. The entire model is fine-tuned on test tasks.

##### CAVIA (Zintgraf et al., 2019)

Similar to TAM, CAVIA has a set of task-specific parameters and shared
parameters. The training algorithm is similar to MAML, but inner loop updates
are performed on the task-specific parameters as opposed to the entire model.

###  5.3 Sequence Classification and Transduction

Most prior work on few-shot learning have focused on computer vision
benchmarks such as Omniglot (Lake, 2019) and Mini-ImageNet (Vinyals et al.,
2016). In the sequential data setting, Bao et al. (2019) constructed synthetic
benchmarks from existing text datasets but the number of tasks is rather
limited. In this work we construct a new set of benchmarks involving synthetic
sequential data, allowing us to evaluate models in a more controlled setting
after training on a larger number of tasks.

####  5.3.1 Synthetic Benchmarks

We construct a synthetic few-shot classification benchmark as follows. The
benchmark consists of tasks that involve a non-negative integer sequence as
input and a discrete label as output. A task is constructed by applying a
sequence of mathematical transformations to input sequences as follows:
Element-wise transform ($T_{1}$) $\rightarrow$ Subsequence extraction
($T_{2}$) $\rightarrow$ Labeling function $(T_{3})$. The arrows indicate
function composition and the sequence of transformations maps an input
sequence to a single integer. The transformations are defined as $T_{1}\in
S_{1},T_{2}\in S_{2},T_{3}\in S_{3}$ where,  $S_{1}=$ {mul $v$, add $v$, div
$v$, mod $v$}; $S_{2}=$ {(not) multiple of $v$, (not) greater than $v$, (do
not) have exactly $v$ divisors}; $S_{3}=$ {count, min, max, mean, median,
mode, first, last, max-min, middle}, where $v\in\\{1\cdots n\\}$ for some
integer $n$. We randomly generate a large number of sequences $X$ of integers
from $\\{0\cdots N\\}$. We apply the transformation sequence
$T_{1},T_{2},T_{3}$ to these sequences $x\in X$ and get the corresponding
outputs $T_{3}(T_{2}(T_{1}(x)))$. The $C$ most frequent outputs are then
defined to be the $C$ classes of interest. We obtain a uniform amount of data
from each class and discard input sequences for which the output does not
belong to one of the chosen $C$ classes. Cases where at least $C$ distinct
outputs cannot be obtained are discarded. These $C$ classes then constitute a
$C$-way classification task. An example task is mul 2 $\rightarrow$ less than
5 $\rightarrow$ count, where the goal is to count the number of input elements
which, when multiplied by 2, are less than 5 (i.e., count number of input
integers less than 3). The semantics of each of the transforms are defined in
the appendix. We set $C=4$ in our experiments. Vocabulary size and input
sequence length are set to 12 and 5, respectively. Combinations of transforms
that have identical input-output relationship are identified and removed
during task construction. All tasks are thus unique in terms of input-output
mapping.

We also construct two sequence transduction benchmarks. The first benchmark is
constructed in a way similar to the classification tasks where we consider a
sequence of transformations mapping an input sequence to an output sequence
$T_{1}\rightarrow T_{2}\rightarrow T_{3}$, where $T_{1}\in S_{1},T_{2}\in
S_{2},T_{3}\in S_{3}$;  $S_{1}=$ {mul $v$, add $v$, div $v$, mod $v$};
$S_{2}=$ {replace $v$ with $v^{\prime}$, replace $x_{i}$ with $f(x_{i},x_{j})$
}; $S_{3}=$ {sort ascending, sort descending, reverse, swap$(x_{i},x_{j})$,
shift right $v$}, and $v,v^{\prime},i,j$ are integers chosen at random,
$x_{p}$ represents the element at position $p$ in the input sequence, $f$ is a
mathematical function (Eg: $f(a,b)\in\\{a+b,\text{abs}(a-b),b,\cdots\\}$). An
example task is add 2 $\rightarrow$ replace 2 with 1 $\rightarrow$ reverse,
and an (input, output) sample drawn from this task is:
$([0,5,0,3,6],[8,5,1,7,1])$.

Our second transduction benchmark is a path finding task in a grid world (see
appendix B for an illustration). A task is defined by start and end positions
in a square grid of size $N\times N$. Given the locations of obstacles in this
grid, the objective of the task is to find the shortest path connecting start
and end positions that avoids the obstacles. The source and target sequences
correspond to the locations of obstacles and optimal path from start to end
position avoiding the obstacles, respectively.

We use 500, 16, 64 tasks respectively for training, validation and testing for
all three setups. Tasks are unique and randomly assigned to these sets, in
other words we test generalization under the condition of distributional match
between the training and the test set. Each training task has 500 examples.

####  5.3.2 Results

Model | Sequence Classification | Sequence Transduction | Path Finding
---|---|---|---
1 | 5 | 10 | 20 | 1 | 5 | 10 | 20 | 1 | 5 | 10 | 20
Task | 40.50 | 64.75 | 74.25 | 82.50 | 6.27 | 5.26 | 4.71 | 4.01 | 3.17 | 1.75 | 1.55 | 1.39
Agnostic |  $\pm$1.73 |  $\pm$1.29 |  $\pm$1.29 |  $\pm$1.58 |  $\pm$0.17 |  $\pm$0.04 |  $\pm$0.05 |  $\pm$0.07 |  $\pm$0.27 |  $\pm$0.03 |  $\pm$0.02 |  $\pm$0.01
Multitask | 38.75 | 66.00 | 77.50 | 87.50 | 13.80 | 6.80 | 5.18 | 2.91 | 6.39 | 1.98 | 1.64 | 1.44
|  $\pm$0.96 |  $\pm$0.82 |  $\pm$1.29 |  $\pm$0.58 |  $\pm$3.36 |  $\pm$0.26 |  $\pm$1.01 |  $\pm$0.49 |  $\pm$1.96 |  $\pm$0.12 |  $\pm$0.05 |  $\pm$0.02
Matching | 43.00 | 58.75 | 64.50 | 67.00 | – | –
Network |  $\pm$1.15 |  $\pm$2.45 |  $\pm$1.83 |  $\pm$1.41
SNAIL | 43.00 | 44.00 | 68.25 | 67.50 | 2.48 | 3.80 | 4.98 | 4.11 | 2.63 | 1.95 | 3.47 | 3.04
|  $\pm$1.41 |  $\pm$2.00 |  $\pm$1.26 |  $\pm$4.43 |  $\pm$0.38 |  $\pm$0.38 |  $\pm$0.03 |  $\pm$2.94 |  $\pm$0.27 |  $\pm$0.25 |  $\pm$1.56 |  $\pm$1.01
MAML | 39.60 | 63.40 | 71.80 | 78.80 | 6.73 | 5.84 | 5.19 | 4.20 | 5.49 | 2.05 | 1.65 | 1.44
|  $\pm$0.55 |  $\pm$0.89 |  $\pm$0.84 |  $\pm$0.84 |  $\pm$0.16 |  $\pm$0.2 |  $\pm$0.08 |  $\pm$0.10 |  $\pm$0.85 |  $\pm$0.03 |  $\pm$0.01 |  $\pm$0.01
CAVIA | 43.00 | 78.00 | 87.00 | 91.00 | 11.05 | 2.75 | 1.78 | 1.53 | 2.21 | 1.31 | 1.25 | 1.21
|  $\pm$0.58 |  $\pm$1.26 |  $\pm$0.50 |  $\pm$0.58 |  $\pm$2.94 |  $\pm$0.51 |  $\pm$0.14 |  $\pm$0.06 |  $\pm$0.21 |  $\pm$0.03 |  $\pm$0.03 |  $\pm$0.02
TAM | 40.50 | 75.50 | 89.50 | 94.50 | 8.47 | 2.92 | 1.47 | 1.15 | 1.82 | 1.27 | 1.22 | 1.17
|  $\pm$0.82 |  $\pm$0.50 |  $\pm$0.58 |  $\pm$0.82 |  $\pm$1.42 |  $\pm$0.67 |  $\pm$0.18 |  $\pm$0.03 |  $\pm$0.08 |  $\pm$0.01 |  $\pm$0.01 |  $\pm$0.01
Table 1: $k$-shot sequence classification and sequence transduction
experiments on our three benchmarks for $k\in\\{1,5,10,20\\}$. The metric for
sequence classification is average accuracy on test tasks (higher is better).
On the transduction tasks, the performance metric is average perplexity on
test tasks (lower is better). Random performance is at 25% accuracy
(classification) and 12 perplexity points (other two tasks). Entries in
smaller font are error bars, and they are estimated on 4 trials varying the
model initialization.

Table 1 reports the results on this benchmark. In the extreme few-shot setting
($k=1$), all methods perform poorly, although memory based methods such as
matching networks and SNAIL fare the best. However, they start performing
relatively worse when more labelled data is available, where fine-tuning part
or all of the model parameters could be beneficial. Both SNAIL and matching
networks sometimes perform absolutely worse when more labeled examples are
present, suggesting they are failing to effectively use their memory when
confronted with longer sequences. Fine-tuning the whole model, particularly in
the multitask setting, works remarkably well for larger values of $k$,
although the best performance is achieved by TAM, suggesting the need for
sample efficient task adaptation methods. For $k>1$, TAM performs comparably
or better than all baselines, including MAML and CAVIA. Furthermore, TAM is
more efficient to train than CAVIA (see section 5.6).

###  5.4 Compositional Task Representations

Compositional reasoning is arguably an important skill for few-shot learning
(Lake, 2019; Purushwalkam et al., 2019). The underlying assumption is that
there exist primitive skills which can be learned and combined together to
solve entirely new tasks. If a learner can leverage the compositional
structure of the learning task, it may learn with even less labeled data.

In this section we assess how much better TAM works when we expose the
compositional structure of the tasks described in §5.3.1. Specifically, we
assess the ability to learn new tasks which are composed of primitives, some
of which were unseen during training. To present an example from the
classification setting, assume the models know that tasks are composed of
three transforms $T_{1}\in S_{1},T_{2}\in S_{2},T_{3}\in S_{3}$. We henceforth
refer to the elements of $S_{1}\cup S_{2}\cup S_{3}$ as primitives. Further
assume the model never saw the add $k$ primitive during training. Given a new
test task for which $T_{1}=$ add 3 (and $T_{2},T_{3}$ are known primitives
seen during training), we expect the model to infer the concept of add from
the few training examples of the test task.

####  5.4.1 Task Construction

In the compositional setting, we provide models with information about the
primitives used to construct the task. For the classification and transduction
tasks, the training and test tasks are constructed as follows. Assume the set
of primitives available for the three transforms to be $S_{1},S_{2},S_{3}$. We
hold out a subset of primitives $S_{1}^{\prime},S_{2}^{\prime},S_{3}^{\prime}$
respectively from each of these three sets, which shall constitute the unseen
primitives. The training tasks are made up of primitives from
$S_{1}-S_{1}^{\prime},S_{2}-S_{2}^{\prime},S_{3}-S_{3}^{\prime}$, which we
will refer to as seen primitives. The test tasks are made up of seen and
unseen primitives where exactly one primitive is unseen (For instance,
$T_{1}\in S_{1}-S_{1}^{\prime},T_{2}\in S_{2}^{\prime},T_{3}\in
S_{3}-S_{3}^{\prime}$). Model performance is averaged over multiple (8)
different choices of $S_{1}^{\prime},S_{2}^{\prime},S_{3}^{\prime}$.

Model | Sequence Classification | Sequence Transduction | Path Finding
---|---|---|---
1 | 5 | 10 | 20 | 1 | 5 | 10 | 20 | 1 | 5 | 10 | 20
Multitask | 57.50 | 74.00 | 81.00 | 88.5 | 43.32 | 7.50 | 3.48 | 2.16 | 3.08 | 1.62 | 1.32 | 1.23
|  $\pm$3.51 |  $\pm$5.48 |  $\pm$4.24 |  $\pm$2.08 |  $\pm$10.87 |  $\pm$0.43 |  $\pm$0.12 |  $\pm$0.05 |  $\pm$0.61 |  $\pm$0.33 |  $\pm$0.05 |  $\pm$0.02
Matching | 61.25 | 69.50 | 72.25 | 67.5 | – | –
Network |  $\pm$4.35 |  $\pm$5.45 |  $\pm$6.08 |  $\pm$5.92
SNAIL | 63.5 | 71.75 | 76.25 | 80.25 | 7.00 | 6.24 | 6.93 | 17.10 | 1.53 | 1.71 | 3.36 | 4.22
|  $\pm$4.80 |  $\pm$4.86 |  $\pm$2.75 |  $\pm$2.87 |  $\pm$2.03 |  $\pm$0.27 |  $\pm$3.25 |  $\pm$9.66 |  $\pm$0.29 |  $\pm$0.09 |  $\pm$0.57 |  $\pm$0.79
CAVIA | 57.25 | 66.25 | 67.50 | 68.50 | 36.72 | 6.01 | 3.99 | 3.27 | 1.99 | 1.27 | 1.21 | 1.17
|  $\pm$12.09 |  $\pm$14.73 |  $\pm$15.67 |  $\pm$16.42 |  $\pm$8.83 |  $\pm$0.88 |  $\pm$0.50 |  $\pm$0.24 |  $\pm$0.11 |  $\pm$0.01 |  $\pm$0.00 |  $\pm$0.00
TAM | 63.00 | 76.50 | 82.75 | 88.5 | 6.15 | 3.43 | 2.69 | 2.13 | 2.33 | 1.30 | 1.23 | 1.19
(Comp) |  $\pm$5.35 |  $\pm$4.65 |  $\pm$3.86 |  $\pm$2.65 |  $\pm$0.91 |  $\pm$0.05 |  $\pm$0.04 |  $\pm$0.02 |  $\pm$0.18 |  $\pm$0.02 |  $\pm$0.01 |  $\pm$0.01
TAM | 45.25 | 72.5 | 81.5 | 89.75 | 7.80 | 5.08 | 3.60 | 2.42 | 4.37 | 1.27 | 1.17 | 1.11
(Non-comp) |  $\pm$3.59 |  $\pm$3.70 |  $\pm$2.89 |  $\pm$0.96 |  $\pm$0.09 |  $\pm$0.24 |  $\pm$0.15 |  $\pm$0.08 |  $\pm$3.59 |  $\pm$0.01 |  $\pm$0.00 |  $\pm$0.00
Table 2: Compositional models for few-shot sequence classification and
sequence transduction. All models (except non-compositional TAM) get
information on the primitives present in the tasks via extra tokens appended
to the input sequence, except that one such primitive is unseen at test time.
Non-compositional TAM is not given information about primitives, and estimates
a single task embedding instead.

We also define a compositional path-finding task as follows. In addition to
finding the optimal path from start, end positions while avoiding obstacles,
we now require the path to lie on a specified way-point. The locations of the
start, end and way points thus define the primitives that make up a task.
Similar to the previous settings, we hold out sets of values for each of these
points and construct the train/test tasks in an analogous manner.

####  5.4.2 Training

For all the models, a sequence of primitive ids representing the primitives
that make up the task is appended to the input sequence. These primitive
embeddings $\theta_{e}$ are learned along with the other model parameters. To
simulate the testing conditions, at training time we pretend some primitives
are unknown. For the multitask, matching network and SNAIL baselines we learn
an unknown primitive embedding, which is used to initialize embeddings of
unknown primitives encountered at test time. Although the tasks themselves are
harder (because entire primitives are unseen), modeling them is easier because
the primitive information is given to the model. For CAVIA and TAM, we infer
embeddings for unknown primitives on the fly using gradient descent during
train and test. See Algorithm 2 in the appendix for the complete algorithm. We
use 5000 training tasks, and 100 validation and test tasks each. Each training
task has 500 examples.

####  5.4.3 Results

Table 2 summarizes the results in the compositional setting. We observe
similar trends as before for the non-compositional case. Multitask learning
becomes competitive only for larger values of $k$. Vice versa, matching
networks and SNAIL suffer with long sequences (larger values of $k$). TAM
performs at least comparably if not better than methods relying on second
order derivatives like CAVIA. In fact, CAVIA sometimes fail to converge as
shown by the rather large error bars. Finally, the compositional version of
TAM often yields higher accuracy than the corresponding non-compositional
version, showing that the model is able to cleverly leverage the additional
knowledge about a subset of primitives (two out of three) that compose the new
task. Compositionality is particularly helpful with fewer shots (e.g., 1-shot)
– with sufficient training examples (e.g., 20-shot) models benefit less from
compositionality.

###  5.5 Ablation experiments

Model | Sequence Classification (Accuracy)
---|---
1-shot | 5-shot | 10-shot | 20-shot
Input token | 0.41 | 0.76 | 0.89 | 0.94
Adapters | 0.43 | 0.75 | 0.88 | 0.94
Layer Norm | 0.42 | 0.64 | 0.78 | 0.88

Table 3: $k$-shot classification accuracy when plugging the task embedding in
various ways for different values of $k$.

Arch | Training | Sequence Classification (Acc.)
---|---|---
1-shot | 5-shot | 10-shot | 20-shot
LSTM | Multi | 0.38 | 0.57 | 0.72 | 0.87
Alt | 0.35 | 0.78 | 0.83 | 0.85
Transf. | Multi | 0.39 | 0.68 | 0.80 | 0.88
Alt | 0.41 | 0.76 | 0.89 | 0.94

Table 4: k-shot accuracy for different architectures with multitask and the
proposed training algorithms.

##### Where to plug task embeddings

The experiments in the paper so far consider a simple conditioning scheme
where the task embedding appears as an additional embedding in the input
sequence of token embeddings. We compare this against other ways of
incorporating task-specific parameters into the model. Houlsby et al. (2019)
introduce adapter layers, parameter modules that are inserted at every layer
of a pre-trained transformer. An adapter layer down-projects its input,
applies a non-linearity, and up-projects the representation back to the
original size. Our first baseline considers parameters in the adapter layers
as the task embedding. Another popular method for adapting pre-trained
networks to new tasks is adapting parameters in normalization layers (Perez et
al., 2018; Ghiasi et al., 2017). In our second baseline, we consider the scale
and bias parameters in the Layer Normalization layers of the transformer as
the task embedding.

The results in Table 3 on non-compositional classification tasks show that
using adapter layer parameters as task embedding yields similar results to the
simplest conditioning scheme where the task embedding is fed as an additional
input. Using normalization parameters as the task embedding instead performs
slightly worse. This shows that our input conditioning scheme is simple yet
effective.

##### Importance of transformer architecture

The experiments presented in this paper so far have used a transformer
architecture. Although transformers are a natural choice for problems
involving sequences owing to their recent success, the proposed training
algorithm applies equally well to other architectures. We study the impact of
swapping out the transformer with a recurrent model in Table 4 on non-
compositional tasks. We use a bidirectional LSTM with a comparable number of
parameters to our transformer model. The classifier head acts on the final
representation of the final layer of the LSTM. We examine the performance of
the two architectures when trained using both multitasking and the proposed
alternating minimization training algorithm. First, we observe that the
transformer generally performs better than the recurrent model. Second, the
proposed training algorithm yields consistent improvements over the multitask
baseline for the transformer. This shows that the proposed algorithm is
general, but particularly effective when used in conjunction with the
transformer architecture.

##### Visualizing learned task embeddings

In Figure 1 we visualize task embeddings learned by the non-compositional TAM
model in our gridworld task. We visualize the first two principal components
of task embeddings corresponding to tasks which have the same start position.
The projections are color coded by the horizontal and vertical coordinates of
the end position for each task. This shows that the task embeddings have
learned the structure of the tasks.

Figure 1:  2D PCA projections of task embeddings learned by our algorithm for
the gridworld domain. Tasks visualized here have the same start position
(4,4). Points are color coded based on horizontal (left plot) and vertical
(right plot) coordinates of the end position corresponding to each task.

Model | Classification | Transduction | Path-finding
---|---|---|---
Acc., Time | Ppl., Time | Ppl., Time
Multitask |  67.4, 30min |  7.2, 23min |  2.9, 20min
CAVIA | 74.8, 3h | 4.5, 5.3h | 1.5, 3.7h
TAM |  75.0, 2h |  1.5, 2.3h |  1.3, 2.7h

Table 5: Training efficiency: Time taken by each training algorithm to reach
the best model (identified using validation tasks) and corresponding model
performance (non-compositional setting). Performance and time are averaged
across $k\in\\{1,5,10,20\\}$ shots.

###  5.6 Discussion

##### Optimizing Task Embeddings

We observed that both CAVIA and TAM generally attain better performance when
trained with a larger number of inner loop updates. In this work, we use a
maximum of 25 inner loop updates for TAM because it strikes a good balance
between finding an optimal task embedding and containing training time. CAVIA
performed best with 10 inner loop updates, beyond which we hit the
computational limitations of our hardware. We also found that TAM works better
when trained with a number of examples per task much greater than $k$, in our
case 300. All these empirical findings suggest that optimizing for the task
embedding and replacing the second order optimization with TAM’s first order
is an intrinsically difficult problem that requires more iterations and a
larger number of examples.

##### Training Efficiency

We discuss the training efficiency of different models in Table 5. The
multitask baseline is not expensive to train, but it doesn’t perform well on
few-shot scenarios. CAVIA does well especially in the extreme few-shot
scenarios, but has stability issues. TAM is simple, easy to implement,
performs comparably or better than the baselines and trains more efficiently
than CAVIA.

##### First vs. Second Order Gradients

Double backprop to optimize the test optimization has become a standard method
of meta-learning. In the appendix of Finn et al. (2017), and in Zintgraf et
al. (2019) (the “first order” variant), similar approaches to TAM were shown
to perform relatively worse than the methods with second order gradients. In
contrast, in our settings, we have found that first order gradients (via
alternating minimization) are sufficient if done correctly, despite being
simpler and more efficient. Although CAVIA (Zintgraf et al., 2019) sometimes
outperforms TAM, especially for very small numbers of test examples, TAM is
always competitive; with more test examples, TAM is usually superior. TAM
always outperforms MAML (Finn et al., 2017).

##  6 Conclusion

In this work we demonstrate a simple and effective approach to adapt
transformer models to new tasks with limited data. TAM is trained to adapt to
new tasks by inferring a small set of parameters called the task embedding
using gradient descent. On synthetic sequence classification and transduction
benchmarks we constructed TAM yields comparable or superior performance to
approaches relying on second order derivatives, while being computationally
more efficient.

## References

  * Andrychowicz et al. (2016) Andrychowicz, M., Denil, M., Gomez, S., Hoffman, M. W., Pfau, D., Schaul, T., Shillingford, B., and De Freitas, N.  Learning to learn by gradient descent by gradient descent.  In _Advances in neural information processing systems_ , pp. 3981–3989, 2016.
  * Ba et al. (2016) Ba, J., Hinton, G. E., Mnih, V., Leibo, J. Z., and Ionescu, C.  Using fast weights to attend to the recent past.  In _Advances in Neural Information Processing Systems_ , pp. 4331–4339, 2016.
  * Bao et al. (2019) Bao, Y., Wu, M., Chang, S., and Barzilay, R.  Few-shot text classification with distributional signatures.  _arXiv preprint arXiv:1908.06039_ , 2019.
  * Barrault et al. (2019) Barrault, L., Bojar, O., Costa-jussà, M. R., Federmann, C., Fishel, M., Graham, Y., Haddow, B., Huck, M., Koehn, P., Malmasi, S., Monz, C., Müller, M., Pal, S., Post, M., and Zampieri, M.  Findings of the 2019 conference on machine translation (WMT19).  In _Proceedings of the Fourth Conference on Machine Translation (Volume 2: Shared Task Papers, Day 1)_ , pp. 1–61, Florence, Italy, August 2019\. Association for Computational Linguistics.
  * Brown et al. (2020) Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al.  Language models are few-shot learners.  _arXiv preprint arXiv:2005.14165_ , 2020.
  * Devlin et al. (2018) Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Fei-Fei et al. (2006) Fei-Fei, L., Fergus, R., and Perona, P.  One-shot learning of object categories.  _IEEE transactions on pattern analysis and machine intelligence_ , 28(4), 2006.
  * Fink (2005) Fink, M.  Object classification from a single example utilizing class relevance metrics.  In _In Advances in Neural Information Processing Systems_ , pp. 449––456, 2005.
  * Finn et al. (2017) Finn, C., Abbeel, P., and Levine, S.  Model-agnostic meta-learning for fast adaptation of deep networks.  In _Proceedings of the 34th International Conference on Machine Learning-Volume 70_ , pp. 1126–1135. JMLR. org, 2017.
  * Ghiasi et al. (2017) Ghiasi, G., Lee, H., Kudlur, M., Dumoulin, V., and Shlens, J.  Exploring the structure of a real-time, arbitrary neural artistic stylization network.  _arXiv preprint arXiv:1705.06830_ , 2017.
  * Hinton & Plaut (1987) Hinton, G. E. and Plaut, D. C.  Using fast weights to deblur old memories.  In _Proceedings of the ninth annual conference of the Cognitive Science Society_ , pp. 177–186, 1987.
  * Hochreiter et al. (2001) Hochreiter, S., Younger, A. S., and Conwell, P. R.  Learning to learn using gradient descent.  In _International Conference on Artificial Neural Networks_ , pp. 87–94. Springer, 2001.
  * Houlsby et al. (2019) Houlsby, N., Giurgiu, A., Jastrzebski, S., Morrone, B., De Laroussilhe, Q., Gesmundo, A., Attariyan, M., and Gelly, S.  Parameter-efficient transfer learning for nlp.  In _International Conference on Machine Learning_ , 2019.
  * Keskar et al. (2019) Keskar, N. S., McCann, B., Varshney, L. R., Xiong, C., and Socher, R.  Ctrl: A conditional transformer language model for controllable generation.  _arXiv preprint arXiv:1909.05858_ , 2019.
  * Kingma & Ba (2015) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  In _International Conference on Learning Representations_ , 2015.
  * Kumar & Daume III (2012) Kumar, A. and Daume III, H.  Learning task grouping and overlap in multi-task learning.  In _International Conference on Machine Learning_ , 2012.
  * Lake (2019) Lake, B. M.  Compositional generalization through meta sequence-to-sequence learning.  _arXiv preprint arXiv:1906.05381_ , 2019.
  * Lample et al. (2019) Lample, G., Subramanian, S., Smith, E., Denoyer, L., Ranzato, M., and Boureau, Y.-L.  Multiple-attribute text rewriting.  In _International Conference on Learning Representations_ , 2019.
  * Maurer et al. (2013) Maurer, A., Pontil, M., and Romera-Paredes, B.  Sparse coding for multitask and transfer learning.  In _International conference on machine learning_ , pp. 343–351, 2013.
  * Mishra et al. (2018) Mishra, N., Rohaninejad, M., Chen, X., and Abbeel, P.  A simple neural attentive meta-learner.  In _International Conference on Learning Representations_ , 2018.
  * Nichol & Schulman (2018) Nichol, A. and Schulman, J.  Reptile: a scalable metalearning algorithm.  _arXiv preprint arXiv:1803.02999_ , 2, 2018.
  * Parisotto et al. (2019) Parisotto, E., Song, F., Rae, J., Pascanu, R., Gulcehre, C., Jayakumar, S., Jaderberg, M., Kaufman, R. L., Clark, A., Noury, S., Botvinick, M., Heess, N., and Hadsell, R.  Stabilizing transformers for reinforcement learning.  _arXiv:1910.06764_ , 2019.
  * Perez et al. (2018) Perez, E., Strub, F., De Vries, H., Dumoulin, V., and Courville, A.  Film: Visual reasoning with a general conditioning layer.  In _Thirty-Second AAAI Conference on Artificial Intelligence_ , 2018\.
  * Purushwalkam et al. (2019) Purushwalkam, S., Nickel, M., Gupta, A., and Ranzato, M.  Task-driven modular networks for zero-shot compositional learning.  _arXiv preprint arXiv:1905.05908_ , 2019.
  * Rusu et al. (2019) Rusu, A. A., Rao, D., Sygnowski, J., Vinyals, O., Pascanu, R., Osindero, S., and Hadsell, R.  Meta-learning with latent embedding optimization.  In _International Conference on Learning Representations_ , 2019.
  * Schmidhuber (1987) Schmidhuber, J.  _Evolutionary principles in self-referential learning, or on learning how to learn: the meta-meta-… hook_.  PhD thesis, Technische Universität München, 1987.
  * Schmidhuber (1992) Schmidhuber, J.  Learning to control fast-weight memories: An alternative to dynamic recurrent networks.  _Neural Computation_ , 4(1):131–139, 1992.
  * Shen et al. (2019) Shen, T., Ott, M., Auli, M., and Ranzato, M.  Mixture models for diverse machine translation: Tricks of the trade.  In _International Conference on Machine Learning_ , 2019.
  * Snell et al. (2017) Snell, J., Swersky, K., and Zemel, R.  Prototypical networks for few-shot learning.  In _Advances in Neural Information Processing Systems_ , pp. 4077–4087, 2017.
  * Stickland & Murray (2019) Stickland, A. C. and Murray, I.  Bert and pals: Projected attention layers for efficient adaptation in multi-task learning.  In _International Conference on Machine Learning_ , 2019.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I.  Attention is all you need.  In _Advances in neural information processing systems_ , pp. 5998–6008, 2017.
  * Vinyals et al. (2016) Vinyals, O., Blundell, C., Lillicrap, T., Wierstra, D., et al.  Matching networks for one shot learning.  In _Advances in neural information processing systems_ , pp. 3630–3638, 2016.
  * Zellers et al. (2019) Zellers, R., Holtzman, A., Rashkin, H., Bisk, Y., Farhadi, A., Roesner, F., and Choi, Y.  Defending against neural fake news.  In _Neural Information Processing Systems_ , 2019.
  * Zintgraf et al. (2019) Zintgraf, L. M., Shiarlis, K., Kurin, V., Hofmann, K., and Whiteson, S.  Fast context adaptation via meta-learning.  In _International Conference on Machine Learning_ , 2019.

##  Appendix A Sequence transformations used to construct classification and
transduction tasks

In tables 6, 7 we describe the transformations used to construct
classification and transduction tasks, respectively.

|  Transformation |  Description
---|---|---
$S_{1}$ |  mul $v$ |  Elementwise multiply by $v$
add $v$ |  Elementwise add $v$
div $v$ |  Elementwise integer division by $v$
mod $v$ |  Elementwise modulo $v$ operation
$S_{2}$ |  (not) multiple of $v$ |  Extract subset of integers that are (not) multiples of $v$
(not) greater of $v$ |  Extract subset of integers that are (not) greater than $v$
(do not) have exactly $v$ divisors |  Extract subset of integers that (do not) have exactly $v$ divisors
$S_{3}$ |  count |  Sequence length
min |  Smallest integer in sequence
max |  Largest integer in sequence
mean |  Mean of sequence elements
median |  Median of sequence elements
mode |  Mode of sequence elements
first |  First element in sequence
last |  Last element in sequence
max-min |  Difference between largest and smallest elements in sequence
middle |  Element in the middle position of sequence
Table 6: Sequence transformations used to construct classification tasks and their descriptions. Each transformation takes a sequence as input and outputs a sequence (transformations in $S_{1}$ and $S_{2}$), or a single integer (transformations in $S_{1}$). |  Transformation |  Description
---|---|---
$S_{1}$ |  mul $v$ |  Elementwise multiply by $v$
add $v$ |  Elementwise add $v$
div $v$ |  Elementwise integer division by $v$
|  mod $v$ |  Elementwise modulo $v$ operation
$S_{2}$ |  reverse $v$ with $v^{\prime}$ |  Replace all occurrences of $v$ in the sequence with $v^{\prime}$
replace $x_{i}$ with $f(x_{i},x_{j})$ |  Replace element $x_{i}$ with one of the following:
|  {$ax_{i}+b$, $x_{j}$, abs$(x_{i}-x_{j})$, $x_{i}+x_{j}$} where $a,b$ are
integer constants and $x_{i},x_{j}$ are elements of the sequence at position
$i,j$ respectively
$S_{3}$ |  sort ascending |  Sort the sequence in ascending order
sort descending |  Sort the sequence in descending order
reverse |  Reverse the sequence
|  swap$(x_{i},x_{j})$ |  Swap elements at positions $i,j$ of the sequence
|  shift right $v$ |  Cyclic shift the sequence right by $v$ positions
Table 7: Sequence transformations used to construct transduction tasks and
their descriptions. Each transformation takes a sequence as input and outputs
a sequence.

##  Appendix B Path-finding task

###  B.1 Non-compositional path-finding task

We present an example task from the path-finding task below. The grids are
$10\times 10$. The following task is defined by the start position (7, 0) and
end position (1, 4), indicated by the green and red squares, respectively.
Each example in the task corresponds to a particular configuration of
obstacles in the grid. The source sequence represents the locations of
obstacles. The obstacles are represented by the top left position of a
$2\times 2$ blob. The target sequence represents the optimal path from source
to target. Source and target sequences consist of rasterized grid coordinates
(Eg. rasterized start and end positions are 70 and 14, respectively). In
addition, elements of the target sequence have an offset of 100 (Eg.
rasterized position 14 is represented as 114).

  * •

Source: [39, 78, 51, 9, 31, 63, 44, 69], Target: [170, 160, 150, 140, 130,
121, 112, 103, 114]

  * •

Source: [12, 35, 99, 22, 62, 44, 25, 21], Target: [170, 161, 152, 143, 134,
124, 114]

  * •

Source: [90, 99, 1, 96, 34, 50, 94, 31], Target: [170, 171, 162, 152, 143,
133, 123, 114]

###  B.2 Compositional path-finding task

In the compositional setting, we require the optimal path to pass through a
way-point, indicated in yellow in the following grids. A task is thus defined
by a start position, end position and way-point position. The possible values
for each of these three parameters represent the primitives in this
compositional setting.

  * •

Source: [63, 38, 90, 93, 73, 68, 18, 67], Target: [126, 115, 124, 133, 142,
131, 122]

  * •

Source: [95, 60, 95, 70, 23, 34, 83, 85], Target: [126, 115, 104, 113, 122,
131, 142, 131, 122]

  * •

Source: [91, 29, 57, 96, 8, 53, 77, 13], Target: [126, 125, 134, 133, 142,
131, 122]

##  Appendix C Compositional TAM

Algorithm 2 presents the training algorithm for compositional TAM. We draw a
training task $\mathcal{T}^{\text{train}}$ with primitive ids
$T_{1}=i_{1},T_{2}=i_{2},T_{3}=i_{3}$ respectively in line 3. These primitive
ids index into the primitive embedding table $\theta_{e}$. We pretend that one
of the primitives is unknown, and to illustrate the algorithm, we assume
without loss of generality that $T_{2}=i_{2}$ is unknown (line 5). In the
inner loop optimization, we infer an embedding $z$ for this unknown primitive
using gradient descent, while using the primitive embedding table to load the
known primitive embeddings ($\theta_{e}[i_{1}],\theta_{e}[i_{3}]$ in this case
(lines 8, 9)).

1

Input : Training tasks
$\mathcal{T}_{1}^{\text{train}},...,\mathcal{T}_{N}^{\text{train}}$

Output : Model parameters $\theta$, primitive embeddings $\theta_{e}$

2

3$\theta^{\prime}=\theta\cup\theta_{e}$

4repeat

5 Sample training task $\mathcal{T}^{\text{train}}$ with primitive ids
$T_{1}=i_{1},T_{2}=i_{2},T_{3}=i_{3}$

6 Sample $k$ training examples from the task
$\\{(x^{j},y^{j})_{j=1,\cdots,k}\\}\sim\mathcal{T}^{\text{train}}$

7 Pretend one of the primitives (chosen at random) is unknown, say $T_{2}$

8 Initialize $z=0,\Delta\theta^{\prime}=0$

9 while  _loss improves and max iterations not reached_ do

10 $z\leftarrow z-\nabla_{z}\sum_{j=1}^{k}\,-\log
p(y^{j}|x^{j},z_{1}=\theta_{e}[i_{1}],z_{2}=z,z_{3}=\theta_{e}[i_{3}];\theta^{\prime})$

11
$\Delta\theta^{\prime}\leftarrow\Delta\theta^{\prime}-\nabla_{\theta^{\prime}}\sum_{j=1}^{k}-\log
p(y^{j}|x^{j},z_{1}=\theta_{e}[i_{1}],z_{2}=z,z_{3}=\theta_{e}[i_{3}];\theta^{\prime})$

12 $\theta^{\prime}\leftarrow\theta^{\prime}+\Delta\theta^{\prime}$

13until  _max training iterations_ ;

Algorithm 2 Compositional TAM for k-shot Learning

##  Appendix D Model Architecture

Figure 2 shows an illustration of how we use transformers for sequence
classification (left) and sequence transduction (right) problems. In the
classification setting the input is a sequence $(x_{1}\cdots x_{n})$ and the
output is a discrete label $y$. In the transduction setting, the input
$(x_{1}\cdots x_{n})$ and output $(y_{1}\cdots y_{m})$ are sequences. $z$ is
an embedding vector we refer to as the task embedding and appears in the input
to the transformer, in addition to the input sequence. The task embedding $z$
is task specific, and is inferred on the fly for each task during training.
Learning a new task $\mathcal{T}$ at test time involves inferring the
corresponding task embedding $z_{\mathcal{T}}$, leaving the rest of the model
parameters untouched.

Transformer$z$$x_{1}$$x_{2}$$x_{3}$$y$

Transformer$x_{1}$$x_{2}$$z$$y_{1}$$y_{1}$$y_{2}$

Figure 2: Illustration of how we use transformers for sequence classification
(left) and sequence transduction (right) problems.