# Adaptive Cross Batch Normalization for Metric Learning

Thalaiyasingam Ajanthan  Amazon Australia  Australian National University
Matt Ma  Amazon Australia  Anton van den Hengel  Amazon Australia  University
of Adelaide  Stephen Gould  Amazon Australia  Australian National University

###### Abstract

Metric learning is a fundamental problem in computer vision whereby a model is
trained to learn a semantically useful embedding space via ranking losses.
Traditionally, the effectiveness of a ranking loss depends on the minibatch
size, and is, therefore, inherently limited by the memory constraints of the
underlying hardware. While simply accumulating the embeddings across
minibatches has proved useful [40], we show that it is equally important to
ensure that the accumulated embeddings are up to date. In particular, it is
necessary to circumvent the representational drift between the accumulated
embeddings and the feature embeddings at the current training iteration as the
learnable parameters are being updated. In this paper, we model
representational drift as distribution misalignment and tackle it using moment
matching. The result is a simple method for updating the stored embeddings to
match the first and second moments of the current embeddings at each training
iteration. Experiments on three popular image retrieval datasets, namely, sop,
In-shop and DeepFashion2, demonstrate that our approach significantly improves
the performance in all scenarios.

##  1 Introduction

Metric learning has proved useful in solving a variety of problems in computer
vision including image retrieval [43, 14], zero shot learning [3], face
recognition [41], and visual tracking [37]. The main objective in metric
learning is to learn a metric space where the feature embeddings of instances
from the same class lie closer together than those from different classes.

Since metric learning requires comparing the feature embeddings of different
instances, the loss functions used for this problem [7, 34, 2, 19] are
commonly referred to as ranking losses. Such a ranking loss is computed by
comparing the embedding of each instance in the minibatch against the
embeddings of a reference set (usually the minibatch itself). To improve
effectiveness some methods select the most informative samples from the
reference set via sophisticated mining strategies [11, 36].

Increasing the size of the reference set improves the accuracy of the ranking
loss calculations. However, since processing each minibatch results in a model
update, previously computed embeddings are invalidated, and the reference set
embeddings need to be recomputed after every model update. Therefore, the
reference set (and the minibatch) size is limited by the memory and
computational constraints of the underlying hardware, which limits the
accuracy.

Given the above, we seek a scalable approach that allows to use a larger
reference set (preferably as large as the training set) while ensuring that
the corresponding embeddings are up to date. Cross Batch Memory (xbm) [40] is
a recent technique to expand the reference set by accumulating embeddings
across minibatches. While this approach is efficient, it does not ensure that
the embeddings are up to date. Embeddings quickly become outdated as the model
evolves during training, particularly in the early stages. Using outdated
embeddings limits the value in comparing against a larger reference set, and
provides a contradictory supervision signal which may lead to sub-optimal
learning.

We propose here an adaptation of metric learning that explicitly models, and
compensates for representational drift. Specifically, we adopt a Bayesian
framework and model the representational drift as a transformation of the
distribution of the embeddings. To this end, the ideal transformation to
mitigate representational drift is to ensure that the embeddings of the
reference set (_i.e_., cross batch memory) follow the same distribution as the
embeddings of the full dataset at any given iteration. For practical purposes
we represent the empirical distributions of the embeddings using their first
and second moments (_i.e_., mean and standard deviation).

We first introduce Cross Batch Normalization (xbn) which simply adapts the
embeddings in the reference set to match the first and second moments of the
current minibatch embeddings at each training iteration. Then, to better
estimate the dataset statistics at any given iteration, we adopt a Kalman
filter [17] which is well suited to modelling evolving statistics over a
sequence of noisy observations. Thus, we use a Kalman filter to estimate the
mean and standard deviation of the dataset embeddings on the basis of the
minibatch observations and adapt the reference set embeddings. We refer to
this approach as Adaptive Cross Batch Normalization (axbn).

We provide extensive experiments on three popular image retrieval datasets,
namely, sop [29], In-shop [23] and DeepFashion2 [10]. Our results demonstrate
significant improvements over the xbm method in all scenarios, which confirms
that simple moment matching can alleviate representational drift. Furthermore,
while the simpler xbn approach outperforms traditional xbm, the adaptive
version is better in the small batch regime, albeit by a small margin,
providing evidence that classical noise modelling approaches can be useful in
deep learning.

##  2 Preliminaries

In this section, we discuss metric learning in the context of image retrieval,
although it has a wide variety of applications. Let
${\mathcal{D}}=\\{({\mathbf{x}}_{i},y_{i})\\}_{i=1}^{n}$ be the dataset where
${\mathbf{x}}_{i}\in{\mathcal{X}}$ is the input image corresponding to sample
$i$ and $y_{i}\in{\mathcal{Y}}$ is the label. The objective of metric learning
is to learn an embedding function
$f_{{\mathbf{w}}}:{\mathcal{X}}\to\mathbb{R}^{d}$ by optimizing the parameters
${\mathbf{w}}\in\mathbb{R}^{m}$, such that images corresponding to the same
label are clustered together in the $d$-dimensional embedding space. For
notational convenience we write ${\mathbf{z}}\coloneqq
f_{{\mathbf{w}}}({\mathbf{x}})\in\mathbb{R}^{d}$. Ideally, for any triplet
$(a,p,n)$ with $y_{a}=y_{p}\neq y_{n}$ we want
$\|{\mathbf{z}}_{a}-{\mathbf{z}}_{p}\|<\|{\mathbf{z}}_{a}-{\mathbf{z}}_{n}\|$
for some appropriate distance metric $\|\cdot\|$. Once trained, given a set of
gallery images, their corresponding embeddings are computed and stored111The
set of gallery embeddings are sometimes referred to as the index.. During
testing, a new query image is passed through the embedding function and its
k-nearest neighbours in the embedding space are retrieved from the stored
gallery embeddings.

The problem of metric learning can be written as:

$\text{minimize}_{{\mathbf{w}}}\,L({\mathbf{w}};{\mathcal{D}})\coloneqq\frac{1}{n}\sum_{i=1}^{n}\ell({\mathbf{w}},({\mathbf{x}}_{i},y_{i});{\mathcal{D}})\
,$ (1)

where $\ell$ denotes the embedding function composed with a metric learning
loss. Note that in addition to ${\mathbf{w}}$ and the sample
$({\mathbf{x}}_{i},y_{i})$, $\ell$ also depends on the full dataset
${\mathcal{D}}$. Specifically, metric learning requires comparing the
embedding of a sample with all other embeddings to ensure that the samples
belong to the same class are clustered together while other class samples are
pushed away. For example, the commonly used triplet loss [34] computes

$\displaystyle\ell^{\text{triplet}}$
$\displaystyle({\mathbf{w}},({\mathbf{x}}_{i},y_{i});{\mathcal{D}})\coloneqq$
(2)
$\displaystyle\frac{1}{|{\mathcal{T}}_{i}|}\sum_{(i,p,n)\in{\mathcal{T}}_{i}}\big{[}\|{\mathbf{z}}_{i}-{\mathbf{z}}_{p}\|_{2}-\|{\mathbf{z}}_{i}-{\mathbf{z}}_{n}\|_{2}+\alpha\big{]}_{+}\
,$

where ${\mathcal{T}}_{i}$ denotes the set of all triplets for sample $i$ such
that $y_{i}=y_{p}\neq y_{n}$, $\alpha\geq 0$ is a margin, $\|\cdot\|_{2}$ is
the $L_{2}$ norm, and $\left[\cdot\right]_{+}$ projects onto the nonnegative
real numbers.

Since the dataset is large, it is infeasible to compare all the samples in a
dataset at each iteration. To this end, the standard approach is to use a
reference set ${\mathcal{R}}^{k}\subset{\mathcal{D}}$ in place of the whole
dataset for each optimization iteration $k$. Additionally, minibatch based
Stochastic Gradient Descent (sgd) is employed to optimize the loss function.
Putting these two together, the sgd update equation at iteration $k$ can be
written as:

$\displaystyle{\mathbf{w}}^{k+1}$
$\displaystyle={\mathbf{w}}^{k}-\eta^{k}\,\nabla
L\left({\mathbf{w}}^{k};{\mathcal{B}}^{k},{\mathcal{R}}^{k}\right)\ ,$ (3)
$\displaystyle
L\left({\mathbf{w}}^{k};{\mathcal{B}}^{k},{\mathcal{R}}^{k}\right)$
$\displaystyle\coloneqq\frac{1}{|{\mathcal{B}}^{k}|}\sum_{({\mathbf{x}},y)\sim{\mathcal{B}}^{k}}\ell\left({\mathbf{w}}^{k},({\mathbf{x}},y);{\mathcal{R}}^{k}\right)\
.$ (4)

Here, $\eta^{k}>0$ is the learning rate, ${\mathcal{B}}^{k}$ is the minibatch,
and ${\mathcal{R}}^{k}$ is the reference set. In almost all metric learning
scenarios the minibatch itself is used as the reference set, _i.e_.,
${\mathcal{R}}^{k}={\mathcal{B}}^{k}$.

###  2.1 Cross Batch Memory (XBM)

The idea of xbm [40] is to expand the reference set beyond the minibatch. To
alleviate the computational complexity of performing the forward pass on a
larger reference set, the insight is to store the embeddings computed during
the previous iterations and use them as the reference set.

Let $\bar{{\mathcal{B}}}^{k}=\\{({\mathbf{z}}^{k}\coloneqq
f_{{\mathbf{w}}^{k}}({\mathbf{x}}),y)\mid({\mathbf{x}},y)\in{\mathcal{B}}^{k}\\}\subset\mathbb{R}^{d}\times{\mathcal{Y}}$
denote the embeddings and their labels corresponding to the minibatch
${\mathcal{B}}^{k}$ at iteration $k$. Then the analogous reference set of
embeddings $\bar{{\mathcal{R}}}^{k}$ is the union of the most recent $M$
embeddings222For brevity we refer to $\bar{{\mathcal{R}}}^{k}$ as the
reference set embeddings, however it contains the embeddings and the
corresponding labels., where $|\bar{{\mathcal{R}}}^{k}|\leq M$ is the limit on
reference set size. In certain cases, $M$ can be as large as the dataset
itself. These accumulated embeddings $\bar{{\mathcal{R}}}^{k}$ are used to
optimize the loss function, specifically the xbm loss function at each
iteration $k$ can be written as
$L\left({\mathbf{w}}^{k};{\mathcal{B}}^{k},\bar{{\mathcal{R}}}^{k}\right)$.

The benefit of xbm relies heavily on the assumption that the embeddings evolve
slowly. In the xbm paper, it is empirically shown that the features evolve
slowly after a certain number of training iterations (refer to Fig. 3 in
[40]), where feature drift at iteration $k$ is defined as:

$D({\mathbf{x}},k,\Delta
k)\coloneqq\left\|{\mathbf{z}}^{k}-{\mathbf{z}}^{k-\Delta k}\right\|_{2}\ ,$
(5)

where ${\mathbf{z}}^{k}$ is the embedding vector for the image ${\mathbf{x}}$,
$\Delta k$ is the iteration step, and $\|\cdot\|_{2}$ is the $L_{2}$ norm. We
note that this quantity is merely an empirical diagnostic and is not practical
to compute during training to determine whether the features have become
outdated or not. Furthermore, even though the features may change slowly in
absolute terms, the drift accumulates over time (large $\Delta k$), which is
problematic when the minibatch size is small, relative to the size of the
cross batch memory. Altogether this could lead to an inaccurate supervision
signal from the ranking loss.

As noted in [40], the slow-drift assumption is violated in the early phase of
training. While we believe the initial training phase is most important, as
discussed above, the embeddings in the reference set $\bar{{\mathcal{R}}}^{k}$
can become outdated even in the slow-drift regime. Such outdated embeddings
not only limits the full potential of using a larger reference set, but also
provides contradicting supervision signal degrading performance.

Therefore, we believe, it is important to ensure that the embeddings in the
reference set are up to date throughout the whole training process.

##  3 Adapting the Cross Batch Memory

Our idea is to adapt the embeddings in the cross batch memory at each
iteration to circumvent the representational drift between the embeddings in
the reference set $\bar{{\mathcal{R}}}^{k}$ and the current minibatch
$\bar{{\mathcal{B}}}^{k}$. Here the term _representational drift_ refers to
the notion that statistics computed on batches of embeddings vary as training
progresses since $f_{{\mathbf{w}}^{k}}({\mathbf{x}})\neq
f_{{\mathbf{w}}^{k^{\prime}}}({\mathbf{x}})$ where
${\mathbf{w}}^{k^{\prime}}\neq{\mathbf{w}}^{k}$ are model parameters from some
previous iteration $k^{\prime}<k$. We model the representational drift as a
(linear) transformation of the distribution of the embeddings. Then, the ideal
transformation to mitigate representational drift is to ensure that the
embeddings of the reference set (_i.e_., cross batch memory) follow the same
distribution as the embeddings of the full dataset at any given iteration.

For practical purposes, we represent the empirical distributions of embeddings
with their respective first and second moments (_i.e_., mean and standard
deviation). To this end, at any iteration, we intend to ensure that the mean
and standard deviation of the embeddings of the reference set match the mean
and standard deviation of the embeddings of the dataset at that iteration.

###  3.1 Cross Batch Normalization (XBN)

In this section, we outline our algorithm with a simplifying assumption that
the minibatch statistics match the statistics of the dataset. In the
subsequent section, we discuss an approach to circumvent this assumption.

Since we only have access to a minibatch of data at each iteration, we simply
adapt the embeddings of the reference set to have the mean and standard
deviation of the embeddings of the current minibatch. Suppose
$\mathbb{E}[\cdot]$ and $\sigma[\cdot]$ denote the mean and standard deviation
of the embeddings, then Cross Batch Normalization (xbn)can be written as:

$\hat{{\mathbf{z}}}^{k}=\frac{{\mathbf{z}}^{k}-\mathbb{E}\\!\left[\bar{{\mathcal{R}}}^{k}\right]}{\sigma\\!\left[\bar{{\mathcal{R}}}^{k}\right]}\sigma\\!\left[\bar{{\mathcal{B}}}^{k}\right]+\mathbb{E}\\!\left[\bar{{\mathcal{B}}}^{k}\right]\
,\quad\mbox{for each $({\mathbf{z}}^{k},y)\in\bar{{\mathcal{R}}}^{k}$}\ .$ (6)

Here, the division and multiplication are performed elementwise. After the
adaptation, the current batch embeddings are added to the cross batch memory
and the combined set is used as the reference set to compute the metric
learning loss. This updated reference set is stored in memory and used in the
subsequent iteration, ensuring that at every iteration, the cross batch memory
is only an iteration behind the current batch.

Note that, this is just one additional line to the xbm code, but as can be
seen in the experiments, it improves the results significantly.

####  3.1.1 Justification for XBN

Let ${\mathbf{z}}=f_{{\mathbf{w}}}({\mathbf{x}})$ and
${\mathbf{z}}^{\prime}=f_{{\mathbf{w}}^{\prime}}({\mathbf{x}})$ be two
embeddings for the same image ${\mathbf{x}}$ computed using different
parameters of the model, ${\mathbf{w}}$ and ${\mathbf{w}}^{\prime}$,
respectively. Assume that ${\mathbf{z}}^{\prime}=g({\mathbf{z}})$ for some
unknown function $g\coloneqq f_{{\mathbf{w}}^{\prime}}\circ
f_{{\mathbf{w}}}^{-1}$.333The function $g$ may not actually exist if $f$ is
not injective, which is typically the case. Nevertheless, we can still
estimate the approximation to a notional $g$. We can approximate
${\mathbf{z}}^{\prime}$ by the first-order Taylor expansion of $g$ around an
arbitrary point ${\mathbf{z}}_{0}$ as follows:

$\displaystyle{\mathbf{z}}^{\prime}$ $\displaystyle\approx
g({\mathbf{z}}_{0})+\nabla
g({\mathbf{z}}_{0})^{T}({\mathbf{z}}-{\mathbf{z}}_{0})\ ,$ (7)
$\displaystyle={\mathbf{A}}{\mathbf{z}}+{\mathbf{b}}\ .$

To estimate coefficients ${\mathbf{A}}$ and ${\mathbf{b}}$ we need quantities
that can be estimated independently from either the minibatch or the cross
batch memory. By the method of moments [5],

$\displaystyle{\bm{\mu}}^{\prime}$
$\displaystyle=\mathbb{E}\\!\left[{\mathbf{z}}^{\prime}\right]=\mathbb{E}\\!\left[{\mathbf{A}}{\mathbf{z}}+{\mathbf{b}}\right]={\mathbf{A}}{\bm{\mu}}+{\mathbf{b}}\
,$ (8) $\displaystyle{\bm{\Sigma}}^{\prime}$
$\displaystyle=\mathbb{E}\\!\left[({\mathbf{z}}^{\prime}-{\bm{\mu}}^{\prime})({\mathbf{z}}^{\prime}-{\bm{\mu}}^{\prime})^{T}\right]$
(9)
$\displaystyle=\mathbb{E}\\!\left[{\mathbf{A}}({\mathbf{z}}-{\bm{\mu}})({\mathbf{z}}-{\bm{\mu}})^{T}{\mathbf{A}}^{T}\right]={\mathbf{A}}{\bm{\Sigma}}{\mathbf{A}}^{T}\
.$

Here, ${\bm{\mu}}$ and ${\bm{\Sigma}}$ denote the mean and the covariance and
the expectations are taken over samples drawn from the same distribution.
However, in practice, we estimate ${\bm{\mu}}$ and ${\bm{\Sigma}}$ from the
reference set $\bar{{\mathcal{R}}}^{k}$, and ${\bm{\mu}}^{\prime}$ and
${\bm{\Sigma}}^{\prime}$ from the minibatch $\bar{{\mathcal{B}}}^{k}$, at each
iteration $k$. Solving for ${\mathbf{A}}$ and ${\mathbf{b}}$ we have,

$\displaystyle{\mathbf{A}}$
$\displaystyle=\left({\bm{\Sigma}}^{\prime}{\bm{\Sigma}}^{-1}\right)^{\frac{1}{2}}\
,$ (10) $\displaystyle{\mathbf{b}}$
$\displaystyle={\bm{\mu}}^{\prime}-\left({\bm{\Sigma}}^{\prime}{\bm{\Sigma}}^{-1}\right)^{\frac{1}{2}}{\bm{\mu}}\
.$ (11)

Assuming independence in components of ${\mathbf{z}}$ we can consider scalar
equations:

$\displaystyle z^{\prime}_{j}$
$\displaystyle=A_{jj}\,z_{j}+b_{j}=\frac{\sigma^{\prime}_{j}}{\sigma_{j}}z_{j}+\mu^{\prime}_{j}-\frac{\sigma^{\prime}_{j}}{\sigma_{j}}\mu_{j}=\sigma^{\prime}_{j}\frac{z_{j}-\mu_{j}}{\sigma_{j}}+\mu^{\prime}_{j}\
,$ (12)

where $\sigma_{j}\coloneqq\Sigma_{jj}^{\frac{1}{2}}$ is the standard
deviation. This is same as our adaptation formula (Eq. (6)). Here, we are
taking a linear approximation of the transformation function $g$.
Nevertheless, higher-order Taylor series expansions, requiring higher-order
moments to be calculated, may give better results with increased complexity,
but we do not consider that direction further in this paper.

As we shown above, xbn is a moment matching approach that considers the first
and second moments of the features. This can be shown to minimize the kl-
divergence between the reference set and minibatch embeddings when the
distributions are assumed to be Gaussian [20]. This provides further
theoretical justification for our approach.

###  3.2 Adaptive Cross Batch Normalization (AXBN)

Note that, we intend to circumvent the representational drift by adapting the
statistics of the reference set to match the statistics of the dataset at each
iteration. Since we only have access to a minibatch of data at each iteration,
we only have a noisy estimation of the dataset statistics. Therefore, we adopt
a Bayesian framework to model the process of estimating dataset statistics
from minibatches. With the dataset statistics in hand, we can simply transform
the reference set embeddings to match the estimated dataset statistics.

####  3.2.1 Kalman Filter based Estimation of Dataset Statistics

The Kalman filter is a special case of recursive Bayesian filter where the
probability distributions are assumed to be Gaussians [17, 32]. To this end,
we first briefly review Bayesian filtering and then turn to the Kalman filter.

##### Recursive Bayesian Filter.

Let ${\mathbf{u}}$ be the random variable correspond to a dataset statistic
(_e.g_., mean or standard deviation of the embeddings) that we want to
estimate and ${\mathbf{v}}$ be that statistic computed using the minibatch of
data. Bayesian filtering assumes the Markovian property where the probability
of the current state given the previous state is conditionally independent of
the other earlier states [6]. With this assumption, let
$p({\mathbf{u}}_{k}\mid{\mathbf{u}}_{0:k-1})=p({\mathbf{u}}_{k}\mid{\mathbf{u}}_{k-1})$
be the process noise distribution and
$p({\mathbf{v}}_{k}\mid{\mathbf{u}}_{k})$ be the measurement noise
distribution at iteration $k$. Since ${\mathbf{u}}$ is hidden and we only get
to observe ${\mathbf{v}}$, one can estimate the dataset statistic at iteration
$k$ given the past observations as:

$p({\mathbf{u}}_{k}\mid{\mathbf{v}}_{1:k-1})=\int
p({\mathbf{u}}_{k}\mid{\mathbf{u}}_{k-1})p({\mathbf{u}}_{k-1}\mid{\mathbf{v}}_{1:k-1})d{\mathbf{u}}_{k-1}\
.$ (13)

Now, with the current observation ${\mathbf{v}}_{k}$ the updated estimation
becomes:

$p({\mathbf{u}}_{k}\mid{\mathbf{v}}_{1:k})=\frac{p({\mathbf{v}}_{k}\mid{\mathbf{u}}_{k})p({\mathbf{u}}_{k}\mid{\mathbf{v}}_{1:k-1})}{p({\mathbf{v}}_{k}\mid{\mathbf{v}}_{1:k-1})}\
.$ (14)

Once $p({\mathbf{u}}_{k}\mid{\mathbf{v}}_{1:k})$ is estimated, the mean of the
distribution can be taken as the estimated value. This is a general form of
Bayesian filtering and depending on the application, certain assumptions have
to be made (_e.g_., the type of probability distributions, their initial
values, _etc_.)

##### Kalman Filter.

As noted earlier, the Kalman filter [17] assumes the distributions are
Gaussians and for simplicity we assume a static state model. Let us consider
the case where we estimate the mean of the embeddings of the dataset at
iteration $k$. Then ${\mathbf{u}}_{k},{\mathbf{v}}_{k}\in\mathbb{R}^{d}$. Let
$\hat{{\mathbf{u}}}_{k}$ be the estimate of ${\mathbf{u}}_{k}$ and
${\mathbf{Q}},{\mathbf{R}},{\mathbf{P}}\in\mathbb{R}^{d\times d}$ be the
process covariance, the measurement covariance and the estimation covariance,
respectively. Now, the distributions used in Eq. (14) take the following
forms:

$\displaystyle p({\mathbf{u}}_{k}\mid{\mathbf{u}}_{k-1})$
$\displaystyle={\mathcal{N}}({\mathbf{u}}_{k-1},{\mathbf{Q}}_{k})\ ,$ (15)
$\displaystyle p({\mathbf{v}}_{k}\mid{\mathbf{u}}_{k})$
$\displaystyle={\mathcal{N}}({\mathbf{u}}_{k},{\mathbf{R}}_{k})\ ,$ (16)
$\displaystyle p({\mathbf{u}}_{k-1}\mid{\mathbf{v}}_{1:k-1})$
$\displaystyle={\mathcal{N}}(\hat{{\mathbf{u}}}_{k-1},{\mathbf{P}}_{k-1})\ ,$
(17)

where ${\mathcal{N}}(\mu,\Sigma)$ denotes the multi-variate Gaussian
distribution. The Kalman filtering steps are then:

$\displaystyle\hat{{\mathbf{u}}}_{k,k-1}=\hat{{\mathbf{u}}}_{k-1}\
,\quad\mbox{predicted state}$ (18)
$\displaystyle{\mathbf{P}}_{k,k-1}={\mathbf{P}}_{k-1}+{\mathbf{Q}}_{k}\
,\quad\mbox{predicted variance}$ (19)
$\displaystyle{\mathbf{K}}_{k}={\mathbf{P}}_{k,k-1}({\mathbf{P}}_{k,k-1}+{\mathbf{R}}_{k})^{-1}\
,\quad\mbox{Kalman gain}$ (20)
$\displaystyle\hat{{\mathbf{u}}}_{k}=\hat{{\mathbf{u}}}_{k,k-1}+{\mathbf{K}}_{k}({\mathbf{v}}_{k}-\hat{{\mathbf{u}}}_{k,k-1})\
,\quad\mbox{updated state}$ (21)
$\displaystyle{\mathbf{P}}_{k}=({\mathbf{I}}-{\mathbf{K}}_{k}){\mathbf{P}}_{k,k-1}\
.\quad\mbox{updated variance}$ (22)

Here, $\hat{{\mathbf{u}}}_{k,k-1}$ and ${\mathbf{P}}_{k,k-1}$ denote the
intermediate state and noise variance estimates. In our approach, we assume
that the dimensions are independent so that ${\mathbf{Q}}$, ${\mathbf{R}}$,
and ${\mathbf{P}}$ are diagonal matrices (and each dimension can be processed
independently), which has computational advantages for high dimensional
spaces.

For our estimation model, we need to initialize the estimation variance
${\mathbf{P}}$ and most importantly the process noise ${\mathbf{Q}}$ and
measurement noise ${\mathbf{R}}$ needs to be identified. Even though there are
some heuristic approaches to estimate these noise covariances [28], for
simplicity, we treat them as constant hyperparameters in our implementation
and use the same hyperparameter value for all dimensions to limit the number
of hyperparameters. In summary, there are three hyperparameters $p_{0},q,$ and
$r$, where
${\mathbf{P}}_{0}=p_{o}\,{\mathbf{I}},{\mathbf{Q}}_{k}=q\,{\mathbf{I}}$ and
${\mathbf{R}}_{k}=r\,{\mathbf{I}}$.

It is known that the initial estimation variance $p_{0}$ is not important as
the Kalman filter converges quickly. Moreover, for scalar systems, it can be
shown that the static system depends on ${\lambda=r/q}$ [9]. Therefore,
essentially there is only one hyperparameter to be tuned. In our case,
measurement noise is inversely proportional to the minibatch size. Therefore,
we use $r/|{\mathcal{B}}^{k}|$ as the measurement noise, tune $r$ only, and
the same value can be used with different minibatch sizes. Note that, when the
measurement noise is assumed to be zero ($r=0$), this process specializes to
directly using the minibatch statistics (refer to Section 3.1). Our final
update in its simplest form is provided in the appendix.

####  3.2.2 Adaptation

Suppose $\hat{{\bm{\mu}}}_{k}$ and $\hat{{\bm{\sigma}}}_{k}$ be the estimated
mean and variance of the dataset at iteration $k$ from the Kalman filter. Then
the adaptation can be written as:

$\hat{{\mathbf{z}}}^{k}=\frac{{\mathbf{z}}^{k}-\mathbb{E}\\!\left[\bar{{\mathcal{R}}}^{k}\right]}{\sigma\\!\left[\bar{{\mathcal{R}}}^{k}\right]}\hat{{\bm{\sigma}}}_{k}+\hat{{\bm{\mu}}}_{k}\
,\qquad\mbox{for each $({\mathbf{z}}^{k},y)\in\bar{{\mathcal{R}}}^{k}$}\ ,$
(23)

where $\mathbb{E}\\!\left[\bar{{\mathcal{R}}}^{k}\right]$ and
$\sigma\\!\left[\bar{{\mathcal{R}}}^{k}\right]$ denote the mean and standard
deviation of embeddings in the reference set. After the adaptation, the
current batch embeddings are added to the cross batch memory and the combined
set is used as the reference set to compute the metric learning loss.

We would like to emphasize that our approach does not add any overhead to the
memory and computational requirements of the xbm method. Furthermore, our
adaptation is a simple change to the xbm code which results in significant
performance improvements.

##  4 Related Work

##### Metric Learning Methods.

Metric learning is a popular problem in machine learning with applications
varying from image retrieval [43] to visual tracking [37]. Recent advances in
metric learning are mostly on improving the effectiveness of learning by
modifying the loss function and/or the example mining strategy. Based on how
embeddings of different instances are compared against each other, the loss
functions can be broadly categorized into 1) pair/triplet based approaches [7,
34, 35, 39, 18], 2) methods that directly optimize average precision [4, 2],
and 3) proxy-based losses [25, 31, 19]. It is worth noting that the
differences in these loss functions impact the learning effectiveness,
however, they are theoretically equivalent.

Apart from the loss, example mining strategy is important to ensure that the
deep learning model focuses more on the informative examples. In addition to
the popular (semi)-hard mining, there are many sophisticated strategies have
been developed recently [15, 11, 36, 39]. We only mentioned a few works here
and we refer the interested reader to [26, 27] for a comprehensive list of
metric learning losses and example mining strategies.

In contrast to these works, we focus on expanding the reference set which
allows us to compare more examples across minibatches while benefiting from
these recent advances.

##### Using External Memory.

Using external memory in metric learning is not new [38, 21, 46, 40]. However,
an important distinction is that we use an external memory to expand the
reference set and ensure that those embeddings are kept up to date so that
they can be used for effective learning. Similarly, in self-supervised
learning, Momentum Contrast (MoCo) [12] uses an external memory along with a
separate encoding network to compute the features and uses momentum based
updates to ensure slow-drift of encoder parameters. Another conceptually
similar approach is cross iteration batch normalization [44, 16], which
estimates minibatch statistics across training iterations using first order
information. In contrast to both of these methods, we simply store the feature
embeddings at each iteration and directly tackle the representational drift by
adapting the distribution parameters for effective metric learning.

##  5 Experiments

We first describe the experimental setup and the datasets and then discuss the
results.

###  5.1 Experimental Setup

We follow the standard approach and use an ImageNet [33] pretrained ResNet-50
[13] backbone and set the embedding dimension $d$ to $512$. ResNet-50
architecture is used as is and therefore batch normalization [16] layers are
enabled. We implemented our code in the PyTorch framework [30] and made use of
the Pytorch Image Models (timm) [42] library for training pipeline including
data augmentations and pretrained weights. For metric learning specific
components including the implementation of xbm, we used the Pytorch Metric
Learning (pml) [27] library.

For all experiments we use the supervised contrastive loss [18] with a pair
margin miner, where the default values are used, _i.e_., pos_margin=$0.2$ and
neg_margin=$0.8$ [27]. Default data augmentations in the timm library are used
along with RandAugment [8] profile rand-m9-mstd0.5, reprob is set to $0.2$ and
mixed precision training. The embeddings are L2-normalized and cosine
similarity is used as the distance metric for training and evaluation. We used
AdamW optimizer [24] with initial learning rate $0.0001$ and the learning rate
is multiplied by $0.33$ at every $15$ epochs. We train for a total of $50$
epochs and the best model with respect to Recall@1 on the validation set is
selected. We used a custom sampler to ensure that there are four images per
class in each minibatch.

To ensure that the feature embeddings are stable, for all methods, we employ a
pre-warmup stage to finetune the randomly initialized last layer444Note that
the rest of the network is pretrained on ImageNet. which projects the backbone
features to a $512$ dimensional embedding. Specifically, we finetune this
layer for $2$ epochs with the standard supervised constrastive loss. For this
stage, sgd with the learning rate of $0.001$ is used.

For our axbn method, the noise hyperparameters are set as follows: $q=1$,
$p_{0}=1$, and $r=0.01$. We did a small grid search on the sop dataset for
batch size $64$ to obtain $r$ and other parameters are not tuned. The same
value of $r$ is used for all datasets and all batch sizes. In addition to
this, we also tune the interval at which to update the Kalman gain (Eq. (20)).
We found the value $100$ to work well for sop and In-shop, and the value $10$
is used for DeepFashion2. While it may possible to tune these hyperparameters
for each setting individually to squeeze out a bit more performance, we did
not do that in our experiments. All our experiments are performed on a V100
gpu in the aws cloud.

As noted earlier, our technique is a simple modification to the xbm approach
and we implemented it as a custom loss in the pml library. Our code will be
released upon publication.

###  5.2 Datasets

We evaluated on three large-scale datasets for few-shot image retrieval
following the standard data splits provided by the respective papers.

#####  Stanford Online Products.

Stanford Online Products (sop) [29] contains 120,053 online product images in
22,634 categories. There are only 2 to 10 images for each category. Following
[29], we use 59,551 images (11,318 classes) for training, and 60,502 images
(11,316 classes) for testing.

#####  In-shop Clothes Retrieval.

In-shop Clothes Retrieval (In-shop) [23] contains 72,712 clothing images of
7,986 classes. Following [23], we use 3,997 classes with 25,882 images as the
training set. The test set is partitioned to a query set with 14,218 images of
3,985 classes, and a gallery set having 3,985 classes with 12,612 images.

#####  Deep Fashion 2.

Deep Fashion 2 (DeepFashion2) [10] consumer to shop retrieval dataset contains
217,778 cloth bounding boxes that have valid consumer to shop pairs. We use
ground truth bounding boxes in training and testing. We follow the evaluation
protocol described in [10] and use the validation set for evaluation. It has
36,961 items in the gallery set and 12,377 items in the query set. We consider
a pair is unique if the pair_id, category_id, and style match, following the
original paper. To this end, there are about 31,517 unique labels in the
training set, 7,059 labels in the gallery set, and 3,128 labels in the query
set. Note that, in contrast to the other two datasets, DeepFashion2 is much
larger and there is a domain gap between the images in query and gallery sets.

Figure 1: Top: Recall@1 _vs_. batch size where cross batch memory size is fixed to $50\%$ (sop and In-shop) or $100\%$ (DeepFashion2) of the training set. Bottom: Recall@1 _vs_. cross batch memory size with batch size is set to $64$. In all cases, our algorithms significantly outperform xbm and the adaptive version is better than the simpler xbn method especially for smaller batch sizes.  | Algorithm | sop | In-shop | DeepFashion2
---|---|---|---|---
| Recall@1 | Recall@10 | Recall@1 | Recall@10 | Recall@1 | Recall@10
|  No-xbm | $75.94\pm 0.03$ | $89.74\pm 0.05$ | $88.76\pm 0.15$ | $97.76\pm 0.13$ | $36.45\pm 0.06$ | $61.44\pm 0.31$
| xbm | $76.80\pm 2.65$ | $89.34\pm 1.90$ | $86.17\pm 0.18$ | $96.58\pm 0.28$ | $41.22\pm 2.96$ | $62.24\pm 3.54$
Ours | xbn | $\underline{80.62}\pm 0.17$ | $\underline{91.85}\pm 0.11$ | $\underline{91.49}\pm 0.07$ | $\underline{98.31}\pm 0.10$ | $\underline{45.12}\pm 0.08$ | $\textbf{66.32}\pm 0.17$
axbn | $\textbf{80.73}\pm 0.30$ | $\textbf{91.98}\pm 0.15$ | $\textbf{91.51}\pm 0.20$ | $\textbf{98.35}\pm 0.09$ | $\textbf{45.33}\pm 0.28$ | $\underline{66.26}\pm 0.63$
Table 1: Summary of results for a particular setting where batch size is $64$ and the cross batch memory size is $50\%$ (sop and In-shop) or $100\%$ (DeepFashion2) of the training set. This corresponds to a point in each plot provided in the top row of Fig. 1. The experiments are repeated three times and the mean and standard deviation are reported. Best numbers are in bold and the second best numbers are underlined. In all cases, our methods significantly outperform xbm. The larger standard deviations for xbm indicate its training instability.  | Algorithm | sop | In-shop
---|---|---|---
| Recall@1 | Recall@10 | Recall@1 | Recall@10
|  No-xbm | $75.94\pm 0.03$ | $89.74\pm 0.05$ | $88.76\pm 0.15$ | $97.76\pm 0.13$
|  No-xbm 512 | $79.63\pm 0.13$ | $91.76\pm 0.02$ | $90.39\pm 0.01$ | $98.18\pm 0.03$
| xbm | $76.80\pm 2.65$ | $89.34\pm 1.90$ | $86.17\pm 0.18$ | $96.58\pm 0.28$
|  xbm * | $79.53\pm 0.09$ | $91.64\pm 0.01$ | $90.91\pm 0.10$ | $98.20\pm 0.03$
Ours | xbn | $\underline{80.62}\pm 0.17$ | $\underline{91.85}\pm 0.11$ | $\underline{91.49}\pm 0.07$ | $\underline{98.31}\pm 0.10$
axbn | $\textbf{80.73}\pm 0.30$ | $\textbf{91.98}\pm 0.15$ | $\textbf{91.51}\pm 0.20$ | $\textbf{98.35}\pm 0.09$
Table 2: Additional comparisons on smaller datasets where batch size is $64$
and memory size is $50\%$ of the training set. Here, No-xbm512 denotes No-xbm
with batch size $512$ and xbm* denotes adding the loss on the minibatch to the
xbm loss to stabilize training. The experiments are repeated three times and
the mean and standard deviation are reported. Best numbers are in bold and the
second best numbers are underlined. In all cases, our methods clearly
outperform both versions of xbm.

###  5.3 Results

Following the paradigm of [26], we compare our approach against the original
xbm method and the one without any cross batch memory while keeping everything
else the same for fair comparison. In this way, we can clearly demonstrate the
benefits of ensuring the reference set (_i.e_., cross batch memory) is up to
date.

In Fig. 1, we plot the performance of different methods by varying the batch
size and varying the cross batch memory size. We can clearly see that our
methods significantly outperform xbm in all cases, validating our hypothesis
that it is important to ensure that the reference set embeddings are up to
date. Furthermore, while our simpler xbn method is powerful, our adaptive
method yields slightly better performance for the smaller minibatches where
the sampling noise is higher.

In Table 1, we summarise results for a particular setting with batch size $64$
and reference set size is $50\%$ of the training set for sop and In-shop and
$100\%$ for DeepFashion2. We repeat the experiments three times and report the
mean and standard deviation of the results. In all cases, our methods
significantly outperform xbm, confirming the merits of tackling
representational drift. Furthermore, the xbm results have large standard
deviations for repeated experiments, which we believe indicates its training
instability due to outdated embeddings.

####  5.3.1 Additional Comparisons

On the smaller sop and In-shop datasets we perform further experiments to
understand the benefits of a larger and up to date reference set. To this end,
in Table 2 we provide the results of No-xbm version with a larger minibatch
size. Even with the $8\times$ larger minibatch, the performance is worse than
our approach. Note that, larger minibatch size significantly increases the
required gpu memory, whereas storing the embeddings adds little overhead [40].

Furthermore, we include the performance of a modified xbm method where the xbm
loss is summed with the loss on the minibatch to stabilize training. This is a
trick used in the original xbm code555https://github.com/msight-tech/research-
xbm but was not mentioned in the paper. While this helps the xbm performance,
it is still worse than both of our methods. Clearly this trick diminishes the
value of using a larger reference set and confirms our hypothesis that it is
necessary to ensure that the reference set embeddings are up to date.

In the plots and table above, xbm is sometimes worse than not using the cross
batch memory. We believe this is due to outdated embeddings providing
inaccurate supervision signal. Note that, when the loss on minibatch is added
to the xbm loss, performance improves and the standard deviation of repeated
experiments decreases. For our methods, we did not observe any improvements
when adding the minibatch loss as our adaptation handles the representational
drift effectively.

More experiments comparing other normalization techniques are provided in the
appendix.

Figure 2: Maximum and mean feature drift on a minibatch of samples on the sop
dataset for xbm and our xbn method, and the corresponding performance on the
validation set. Our method has much smaller feature drift for most of the
training period compared to xbm and yields improved performance on the
validation set.

####  5.3.2 Feature Drift Diagnostics

We measure feature drift (refer to Eq. (5)) for xbm and our xbn method for the
sop dataset in Fig. 2. We fix $\Delta k=1$ and measure the drift on a
minibatch of training samples with batch size $64$. The maximum or mean drift
is computed on the minibatch and the respective quantities are averaged over
the epoch. Our method has much smaller feature drift for the most part of the
training compared to xbm and yields improved performance on the validation
set.

Note that, xbm* also reduces feature drift similar to our method, however, its
performance is inferior to our approach as reported in Table 2 and Fig. 2. We
hypothesize that, even though xbm* shows slow-drift, adding minibatch loss
diminishes the value of using a larger reference set, leading to inferior
performance.

As opposed to our approach which directly tackles feature drift, the slow-
drift phenomenon is emergent for xbm* (also for No-xbm as shown in [40]), and
the reason is not well understood. We believe relying on a principled approach
to handle feature drift while optimizing the correct loss is valuable and as
shown in the experiments our approach yields superior performance.

We would like to clarify that low feature drift does not always mean high
accuracy. A trivial example for this is, when one does not optimize the
parameters, the drift will be zero, however, the accuracy will also be low.
Therefore, when comparing feature drift, all other factors (_e.g_., loss,
optimiser, learning rate, architecture, _etc_.) should be kept the same across
the compared methods. Since the loss is different between xbm* and xbn, merely
comparing feature drift is meaningless and the validation accuracy should also
be considered.

##  6 Discussion

We have introduced an efficient approach that adapts the embeddings in the
cross batch memory to tackle representational drift. This enables us to
effectively expand the reference set to be as large as the full training set
without significant memory overhead. Our simple xbn approach significantly
outperforms the standard xbm method, demonstrating the importance of ensuring
that the reference set embeddings are up to date. Furthermore, the adaptive
version (labelled axbn), which uses Kalman filter to model the noisy
estimation process, performs slightly better than the simpler xbn method for
smaller batch sizes. We believe the methods offer useful additions to the
standard toolset for metric learning due to their simplicity and
effectiveness. Indeed, other applications where accumulation of examples over
many minibatches is needed may also benefit from our approach.

###  6.1 Limitations and Societal Impact

xbn applies a linear transformation function to tackle representational drift.
Nevertheless, as noted earlier, it is possible to incorporate higher-order
terms and/or cross-correlation among the elements of the embeddings.
Additionally, in axbn, it is not clear how to choose the Kalman filter noise
variances as they are problem dependent. Currently we treat them as
hyperparameters. However, it would be interesting to come up with an automatic
mechanism to obtain these variances depending on the dataset and architecture.
We also made some simplifying assumptions for this method such as independence
among the dimensions. While our resulting method works well in our
experiments, the full potential of our approach may be limited due to these
assumptions.

Our method focuses on improving the learning effectiveness in metric learning
by ensuring the cross batch memory embeddings are up to date. As with any
method that optimises for performance on the training data, our method may
amplify dataset bias, however, methods that combat dataset bias may also
benefit from the better training signal. Overall, our method may contribute to
the societal impact of deep metric learning, both positive and negative.

## References

  * [1] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.  Layer normalization.  arXiv preprint arXiv:1607.06450, 2016.
  * [2] Andrew Brown, Weidi Xie, Vicky Kalogeiton, and Andrew Zisserman.  Smooth-ap: Smoothing the path towards large-scale image retrieval.  In European Conference on Computer Vision, pages 677–694. Springer, 2020.
  * [3] Maxime Bucher, Stéphane Herbin, and Frédéric Jurie.  Improving semantic embedding consistency by metric learning for zero-shot classiffication.  In European Conference on Computer Vision, pages 730–746. Springer, 2016.
  * [4] Fatih Cakir, Kun He, Xide Xia, Brian Kulis, and Stan Sclaroff.  Deep metric learning to rank.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 1861–1870, 2019.
  * [5] George Casella and Roger L Berger.  Statistical inference.  Cengage Learning, 2021.
  * [6] Zhe Chen et al.  Bayesian filtering: From kalman filters to particle filters, and beyond.  Statistics, 182(1):1–69, 2003.
  * [7] Sumit Chopra, Raia Hadsell, and Yann LeCun.  Learning a similarity metric discriminatively, with application to face verification.  In 2005 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR’05), volume 1, pages 539–546. IEEE, 2005.
  * [8] Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le.  Randaugment: Practical automated data augmentation with a reduced search space.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pages 702–703, 2020.
  * [9] Simone Formentin and Sergio Bittanti.  An insight into noise covariance estimation for kalman filter design.  IFAC Proceedings Volumes, 47(3):2358–2363, 2014.
  * [10] Yuying Ge, Ruimao Zhang, Xiaogang Wang, Xiaoou Tang, and Ping Luo.  Deepfashion2: A versatile benchmark for detection, pose estimation, segmentation and re-identification of clothing images.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5337–5345, 2019.
  * [11] Ben Harwood, Vijay Kumar BG, Gustavo Carneiro, Ian Reid, and Tom Drummond.  Smart mining for deep metric learning.  In Proceedings of the IEEE International Conference on Computer Vision, pages 2821–2829, 2017.
  * [12] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick.  Momentum contrast for unsupervised visual representation learning.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9729–9738, 2020.
  * [13] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016.
  * [14] Xinwei He, Yang Zhou, Zhichao Zhou, Song Bai, and Xiang Bai.  Triplet-center loss for multi-view 3d object retrieval.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1945–1954, 2018.
  * [15] Alexander Hermans, Lucas Beyer, and Bastian Leibe.  In defense of the triplet loss for person re-identification.  arXiv preprint arXiv:1703.07737, 2017.
  * [16] Sergey Ioffe and Christian Szegedy.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In International conference on machine learning, pages 448–456. PMLR, 2015.
  * [17] Rudolph Emil Kalman.  A new approach to linear filtering and prediction problems.  1960\.
  * [18] Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan.  Supervised contrastive learning.  Advances in Neural Information Processing Systems, 33:18661–18673, 2020.
  * [19] Sungyeon Kim, Dongwon Kim, Minsu Cho, and Suha Kwak.  Proxy anchor loss for deep metric learning.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3238–3247, 2020.
  * [20] Gerhard Kurz, Florian Pfaff, and Uwe D Hanebeck.  Kullback-leibler divergence and moment matching for hyperspherical probability distributions.  In 2016 19th International Conference on Information Fusion (FUSION), pages 2087–2094. IEEE, 2016.
  * [21] Suichan Li, Dapeng Chen, Bin Liu, Nenghai Yu, and Rui Zhao.  Memory-based neighbourhood embedding for visual recognition.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 6102–6111, 2019.
  * [22] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 10012–10022, 2021.
  * [23] Ziwei Liu, Ping Luo, Shi Qiu, Xiaogang Wang, and Xiaoou Tang.  Deepfashion: Powering robust clothes recognition and retrieval with rich annotations.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1096–1104, 2016.
  * [24] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  arXiv preprint arXiv:1711.05101, 2017.
  * [25] Yair Movshovitz-Attias, Alexander Toshev, Thomas K Leung, Sergey Ioffe, and Saurabh Singh.  No fuss distance metric learning using proxies.  In Proceedings of the IEEE International Conference on Computer Vision, pages 360–368, 2017.
  * [26] Kevin Musgrave, Serge Belongie, and Ser-Nam Lim.  A metric learning reality check.  In European Conference on Computer Vision, pages 681–699. Springer, 2020.
  * [27] Kevin Musgrave, Serge Belongie, and Ser-Nam Lim.  Pytorch metric learning, 2020.
  * [28] Brian J Odelson, Murali R Rajamani, and James B Rawlings.  A new autocovariance least-squares method for estimating noise covariances.  Automatica, 42(2):303–308, 2006.
  * [29] Hyun Oh Song, Yu Xiang, Stefanie Jegelka, and Silvio Savarese.  Deep metric learning via lifted structured feature embedding.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4004–4012, 2016.
  * [30] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  Advances in neural information processing systems, 32, 2019.
  * [31] Qi Qian, Lei Shang, Baigui Sun, Juhua Hu, Hao Li, and Rong Jin.  Softtriple loss: Deep metric learning without triplet sampling.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 6450–6458, 2019.
  * [32] Sam Roweis and Zoubin Ghahramani.  A unifying review of linear gaussian models.  Neural computation, 11(2):305–345, 1999.
  * [33] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al.  Imagenet large scale visual recognition challenge.  International journal of computer vision, 115(3):211–252, 2015\.
  * [34] Florian Schroff, Dmitry Kalenichenko, and James Philbin.  Facenet: A unified embedding for face recognition and clustering.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 815–823, 2015.
  * [35] Kihyuk Sohn.  Improved deep metric learning with multi-class n-pair loss objective.  Advances in neural information processing systems, 29, 2016.
  * [36] Yumin Suh, Bohyung Han, Wonsik Kim, and Kyoung Mu Lee.  Stochastic class-based hard example mining for deep metric learning.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 7251–7259, 2019.
  * [37] Ran Tao, Efstratios Gavves, and Arnold WM Smeulders.  Siamese instance search for tracking.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1420–1429, 2016.
  * [38] Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Daan Wierstra, et al.  Matching networks for one shot learning.  Advances in neural information processing systems, 29, 2016.
  * [39] Xun Wang, Xintong Han, Weilin Huang, Dengke Dong, and Matthew R Scott.  Multi-similarity loss with general pair weighting for deep metric learning.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5022–5030, 2019.
  * [40] Xun Wang, Haozhi Zhang, Weilin Huang, and Matthew R Scott.  Cross-batch memory for embedding learning.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 6388–6397, 2020.
  * [41] Yandong Wen, Kaipeng Zhang, Zhifeng Li, and Yu Qiao.  A discriminative feature learning approach for deep face recognition.  In European conference on computer vision, pages 499–515. Springer, 2016.
  * [42] Ross Wightman.  Pytorch image models.  https://github.com/rwightman/pytorch-image-models, 2019.
  * [43] Fan Yang, Ajinkya Kale, Yury Bubnov, Leon Stein, Qiaosong Wang, Hadi Kiapour, and Robinson Piramuthu.  Visual search at ebay.  In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pages 2101–2110, 2017.
  * [44] Zhuliang Yao, Yue Cao, Shuxin Zheng, Gao Huang, and Stephen Lin.  Cross-iteration batch normalization.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12331–12340, 2021.
  * [45] Hongwei Yong, Jianqiang Huang, Deyu Meng, Xiansheng Hua, and Lei Zhang.  Momentum batch normalization for deep learning with small batch size.  In European Conference on Computer Vision, pages 224–240. Springer, 2020.
  * [46] Zhun Zhong, Liang Zheng, Zhiming Luo, Shaozi Li, and Yi Yang.  Invariance matters: Exemplar memory for domain adaptive person re-identification.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 598–607, 2019.

Algorithm |  L2 \+ BN |  L2 \+ MBN |  L2 \+ LN | L2
---|---|---|---|---
Recall@1 | Recall@10 | Recall@1 | Recall@10 | Recall@1 | Recall@10 | Recall@1 | Recall@10
No-xbm | $48.96$ | $65.06$ | $48.96$ | $65.06$ | $75.77$ | $89.54$ | $75.94$ | $89.74$
xbm | $48.96$ | $65.06$ | $48.96$ | $65.06$ | $79.51$ | $91.17$ | $76.80$ | $89.34$
xbn | $48.96$ | $65.06$ | $48.96$ | $65.06$ | 80.58 | 91.77 | 80.62 | 91.85
Table 3: Experiments with BN, MBN, and LN for emebedding normalization and
compared against our method on the sop dataset. While networks with BN and MBN
yield poor results, LN improves over the standard xbm method. Nevertheless,
our approach (xbn) improves further and the best performance of our method is
attained when only L2 is used for embedding normalization. Note, the best
results for BN and MBN are achieved during the warmup stage and hence
identical.

##  Appendix A Comparison to other Normalizations

We experiment with Batch Normalization (BN) [16], Momentum Batch Normalization
(MBN) [45], and Layer Normalization (LN) [1] to normalize the feature
embeddings and the results on the sop dataset with batch size $64$ are
reported in Table 3. Here, the nomalization approach is preceded by L2
normalization to mimic the protocol of our approach. For BN and MBN the best
validation performance is obtained during the warmup stage and we found the
networks with BN/MBN for embedding normalization untrainable. While networks
with BN and MBN yield poor results, LN improves over the standard xbm method.
Nevertheless, our approach (xbn) improves further and the best performance of
our method is attained when neither BN nor LN is used for embedding
normalization.

Even though we are unaware of any work that used BN to normalize feature
embeddings, we used BN for embedding normalization in various settings
(with/without L2 normalization, freezing/not freezing BN parameters, and
with/without tracking running statistics) and the resulting networks were
untrainable for a range of learning rates and attained poor results. We
hypothesize that BN is not designed to be used for embedding normalization and
it would require deeper analysis and modifications to it if one wants to use
it. Similar behaviour is observed even for MBN where we used the recommended
hyperparameters and tested with two batch sizes: $64$ (with $m_{0}=128$) and
$16$ (with $m_{0}=32$). In both the settings the behaviour is similar to BN.

Algorithm | CIFAR-10 | CIFAR-100
---|---|---
MoCo | $82.96$ | 56.02
MoCo +xbn | 83.22 | $55.90$

Figure 3: MoCo with and without our xbn approach. The left two figures show
the the training curves for the first $50$ epochs and the right table shows
the final $k$-NN accuracy after $200$ epochs. xbn clearly improves the
training of MoCo early in the training demonstrating the benefit of adapting
stored embeddings even when a momentum encoder is used.

##  Appendix B MoCo with XBN

MoCo [12] is a self-supervised learning approach that uses an external memory
similar to xbm and uses a momentum based encoder to handle representational
drift. To understand if our approach is beneficial despite the presence of the
momentum encoder, we experimented MoCo on CIFAR-10/100 datasets with our xbn
approach to adapt the stored embeddings. Specifically, we followed the MoCo
example provided in pml 666https://github.com/KevinMusgrave/pytorch-metric-
learning/blob/master/examples/notebooks/MoCoCIFAR10.ipynb including all
hyperparameter values except the batch size, which is set $128$.

As shown in Fig. 3, our approach indeed improves the training of MoCo
especially early in the training (_e.g_., $\sim 4\%$ improvement at epoch $15$
for both datasets) where features are adapted quickly. Note this improvement
is on top of the momentum encoder where the momentum parameter is set to
$0.99$. Even though the gap is reduced towards the end of training, this
clearly demonstrates the benefits of our approach outside of metric learning.

##  Appendix C More on the AXBN Method

Since we have made several simplifying assumptions to the Kalman filter based
noise estimation, we provide the update in its simplest form in Algorithm 1.

Note that, Exponential Moving Average (ema) can be thought of as a special
case of our final axbn approach where the Kalman gain is replaced with a
constant throughout training. To this end, we perform an experiment with ema
on the sop dataset with batch size $64$ using ResNet-50 architecture. We
varied the momentum parameter within the range $\\{0.1,0.2,...,0.9\\}$ (where
$0$ corresponds to xbn) and the best performance is obtained when the momentum
parameter is $0.1$. The results are reported in Table 4.

Algorithm 1 axbn update at iteration $k$

1:$\bar{{\mathcal{R}}}^{k-1},\bar{{\mathcal{B}}}^{k},\hat{{\bm{\mu}}}_{k-1},\hat{{\bm{\sigma}}}_{k-1},K_{k-1},p_{k-1},q,r$

2:$\bar{{\mathcal{R}}}^{k},\hat{{\bm{\mu}}}_{k},\hat{{\bm{\sigma}}}_{k},K_{k},p_{k}$

3:$p_{k,k-1}\leftarrow p_{k-1}+q$ $\triangleright$ Predicted noise estimate

4:$K_{k}\leftarrow p_{k,k-1}/\left(p_{k,k-1}+r\right)$ $\triangleright$ Kalman
gain

5:$p_{k}\leftarrow(1-K_{k})\,p_{k,k-1}$ $\triangleright$ Updated noise
estimate

6:$\hat{{\bm{\mu}}}_{k}\leftarrow\hat{{\bm{\mu}}}_{k-1}+K_{k}\left(\mathbb{E}\\!\left[\bar{{\mathcal{B}}}^{k}\right]-\hat{{\bm{\mu}}}_{k-1}\right)$
$\triangleright$ Mean est.

7:$\hat{{\bm{\sigma}}}_{k}\leftarrow\hat{{\bm{\sigma}}}_{k-1}+K_{k}\left(\sigma\\!\left[\bar{{\mathcal{B}}}^{k}\right]-\hat{{\bm{\sigma}}}_{k-1}\right)$
$\triangleright$ Variance est.

8:for $({\mathbf{z}}^{k-1},y)\in\bar{{\mathcal{R}}}^{k-1}$ do$\triangleright$
Normalization

9:
${\mathbf{z}}^{k}\leftarrow\left({\mathbf{z}}^{k-1}-\mathbb{E}\left[\bar{{\mathcal{R}}}^{k-1}\right]\right)\hat{{\bm{\sigma}}}_{k}/\sigma\left[\bar{{\mathcal{R}}}^{k-1}\right]+\hat{{\bm{\mu}}}_{k}$

10:end for

11:$\bar{{\mathcal{R}}}^{k}\leftarrow\\{({\mathbf{z}}^{k},y)\\}$
$\triangleright$ Store the updated reference set

Algorithm | sop
---|---
Recall@1 | Recall@10
xbn | $\underline{80.62}$ | $91.85$
xbn +ema | $80.48$ | $\underline{91.95}$
axbn | 80.73 | 91.98
Table 4: Results with ema version of our approach. The results are competitive to both of our methods but estimating the Kalman gain performs slightly better.  | Algorithm |  Batch size = $16$ |  Batch size = $32$
---|---|---|---
| Recall@1 | Recall@10 | Recall@1 | Recall@10
|  No-xbm | $77.80$ | $90.70$ | $81.81$ | $92.91$
| xbm | $61.18$ | $76.95$ | $86.99$ | $95.39$
Ours | xbn | $\underline{86.25}$ | $\underline{94.89}$ | $\underline{87.59}$ | $\underline{95.57}$
axbn | 86.40 | 94.96 | 87.74 | 95.66
Table 5: SwinTransformer results on the sop dataset with batch size $16$ and
$32$. As expected, in both cases our approaches clearly outperform xbm and our
adaptive version axbn is slightly better than the simple xbn method.

As discussed in Sec. 5.3, the AXBN approach is useful when the sampling noise
due to minibatches is high. There are two factors that affect the sampling
noise: 1) minibatch size, and 2) the stability of the embedding network. The
adaptive method may not be necessary when one has a stable embedding network
(a network that shows slow feature-drift) and/or the minibatch size is large.

| Algorithm | sop
---|---|---
| Recall@1 | Recall@10
|  No-xbm | $78.69$ | $\underline{91.21}$
| xbm | $78.38$ | $90.39$
Ours | xbn | $\underline{79.96}$ | $91.02$
axbn | 80.31 | 91.31
Table 6: Results with batch size $256$. The gap between different methods
decreased compared to smaller batch sizes.

##  Appendix D Large Model Experiments

All the experiments in the main paper are provided with ResNet-50
architecture. Here, we experiment with much larger SwinTransformer [22].
Similar to other experiments in the paper ImageNet pretrained weights are used
to initialize the network. The results are provided in Table 5. The behaviour
is similar to ResNet-50 and our approaches clearly outperform the original
xbm. However, even with small batch sizes the improvement due to the adaptive
version is marginal. We hypothesize that these large models are stable
pretrained models and therefore the noise in minibatch based feature
statistics is small, and Kalman filter based noise estimation does not improve
significantly.

##  Appendix E Larger Batch Sizes

We performed an experiment with batch size $256$ on the sop dataset (maximum
batch size allowed by our A10G gpu with 25GB memory) and the results are
reported in Table 6. From these results and Fig. 1 top-row in the main paper,
one may extrapolate that larger batch sizes tend to lead to smaller gap
between the methods. This is expected as when the batch size is large enough,
cross batch memory may not be required.

Nevertheless, the point of our paper is not about when to use cross batch
memory (which has already been established in [40]), rather in cases where
cross batch memory is relevant, our approach is the most effective way to use
it.