11institutetext: 1School of E.C.E., National Technical University of Athens,
15773, Athens, Greece
2Department of Computer Science TU Darmstadt, 64289, Darmstadt, Germany
Email: {ismarougkas<EMAIL_ADDRESS>{pkoutras<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
# How to track your dragon:
A Multi-Attentional Framework for real-time RGB-D 6-DOF Object Pose Tracking

Isidoros Marougkas 11 Petros Koutras 11 Nikos Kardaris 11 Georgios Retsinas 11
Georgia Chalvatzaki 22
Petros Maragos 11

###### Abstract

We present a novel multi-attentional convolutional architecture to tackle the
problem of real-time RGB-D 6D object pose tracking of single, known objects.
Such a problem poses multiple challenges originating both from the objects’
nature and their interaction with their environment, which previous approaches
have failed to fully address. The proposed framework encapsulates methods for
background clutter and occlusion handling by integrating multiple parallel
soft spatial attention modules into a multitask Convolutional Neural Network
(CNN) architecture. Moreover, we consider the special geometrical properties
of both the object’s 3D model and the pose space, and we use a more
sophisticated approach for data augmentation during training. The provided
experimental results confirm the effectiveness of the proposed multi-
attentional architecture, as it improves the State-of-the-Art (SoA) tracking
performance by an average score of 34.03 $\%$ for translation and 40.01 $\%$
for rotation, when tested on the most complete dataset designed, up to date,
for the problem of RGB-D object tracking. Code will be available in:
https://github.com/ismarou/How_to_track_your_Dragon

###### Keywords:

Pose, Tracking, Attention, Geodesic, Multi-Task

##  1 Introduction

Robust, accurate and fast object pose estimation and tracking, i.e. estimation
of the object’s 3D position and orientation, has been a matter of intense
research for many years. The applications of such an estimation problem can be
found in Robotics, Autonomous Navigation, Augmented Reality, etc. Although the
Computer Vision community has consistently studied the problem of object pose
estimation and tracking for decades, the recent spread of affordable and
reliable RGB-D sensors like Kinect, along with advances in Deep Learning (DL)
and especially the use of CNNs as the new SoA image feature extractors, led to
a new era of research and a re-examination of several problems, with central
aim the generalization over different tasks. CNNs have achieved ground-
breaking results in 2D problems like object classification, object detection
and segmentation. Thus, it has been tempting to the research community to
increasingly use them in the more challenging 3D tasks, renouncing traditional
algorithms.

The innate challenges of object pose estimation from RGB-D streams include
background clutter, occlusions (both static, from other objects present in the
scene, and dynamic, due to possible interactions with a human user),
illumination variation, sensor noise, image blurring (due to fast movement)
and appearance changes as the object viewpoint alters. Moreover, one should
account for the pose ambiguity, which is a direct consequence of the object’s
own geometry, in possible symmetries, the challenges of proper parameter
representation of rotations and the inevitable difficulties that an effort of
forging a model faces, when extracting information about the 3D scene geometry
from 2D-projected images.

Figure 1: Top row:(Left to right) The ”Dragon” model, an estimated pose in the
”Hard Interaction” scenario for each of the SoA [8] (light blue) and our
(pink) approaches and an example frame pair of Foreground extraction (up) and
Occlusion handling (down) attention maps which are learned by minimizing the
two auxiliary binary cross entropy losses. The following tradeoff occurs: as
the occlusion increases, foreground attention, which focuses on the moving
parts of the scene (i.e. the hand and the object), gets blurrier, while
occlusion attention gets sharper and shifts focus from the object center to
its body parts. Bottom row:(Left to right) Translational and Rotational error
plots of the SoA [8] (blue) and (our) approaches, for the “$75\%$ Vert.
Occlusion”, the “Rotation Only” and the “Hard Interaction” scenario,
respectively. Grey regions stand for intervals of high occlusion and green
ones for rapid movement.

In this paper, we build upon previous works [7, 8], in order to face a series
of those challenges that have not been fully resolved, so far. Thus, our main
contributions are:

  * •

An explicit background clutter and occlusion handling mechanism that leverages
spatial attentions and provides an intuitive understanding of the tracker’s
region of interest at each frame, while boosting its performance. To the best
of our knowledge, this is the first such strategy, that explicitly handles
these two challenges, is incorporated into a CNN-based architecture, while
achieving real-time performance. Supervision for this mechanism is extracted
by fully exploiting the synthetic nature of our training data.

  * •

The use of a novel multi-task pose tracking loss function, that respects the
geometry of both the object’s 3D model and the pose space and boosts the
tracking performance by optimizing auxiliary tasks along with the principal
one.

  * •

SoA real-time performance in the hardest scenarios of the benchmark dataset
[8], while achieving lower translation and rotation errors by an average of
$34.03\%$ for translation and $40.01\%$ for rotation.

Accordingly, we provide the necessary methodological design details and
experimental results that justify the importance of the proposed method in the
challenging object pose tracking problem.

##  2 Related Work

Previous works attempt to tackle the problem using DL, focusing on two
different directions: per-frame pose estimation (or, else, “tracking by
detection”) and temporal tracking.

####  2.0.1 Tracking-by-Detection

The first family of proposed approaches in literature processes each video
frame separately, without any feedback from the estimation of the previous
timeframe. In [36], Xiang et al. constructed a CNN architecture that estimates
binary object masks and then predicts the object class and its translation and
rotation separately, while in [19], Kehl et al. extended the Single Shot
Detection (SSD) framework [24] for 2D Object detection by performing discrete
viewpoint classification for known objects. Finally, they refine their initial
estimations via ICP [32] iterations. In [38], Zakharov et al. proposed a CNN
framework that uses RGB images for pixel-wise object semantic segmentation in
a mask-level. Following this, UV texture maps are estimated to extract dense
correspondences between 2D images and 3D object models minimizing cross
entropy losses. Those correspondences are used for pose estimation via P’n’P
[22]. This estimation is, ultimately, inserted as a prior to a refinement CNN
that outputs the final pose prediction. In PVNet [31], Peng et al. perform
per-pixel voting-based regression to match 3D object coordinates with
predefined keypoints inside the object surface, in order to handle occlusions.
In [30], Pavlakos et al. extract semantic keypoints in single RGB images with
a CNN and incorporate them into a deformable shape model. In [35], Wohlhart et
al. employed a supervised contrastive convolutional framework to disentangle
descriptors of different object instances and impose proportional distances to
different poses of the same object. In [33], Sundermeyer et al. built a self-
supervised Augmented Auto-Encoder that predicts 3D rotations only from
synthetic data. In [29], Park et al. trained an adversiarially guided Encoder-
Decoder to predict pixel-wise coordinates in a given image and then fed them
to a P’n’P algorithm. More recently, iPose [18] is one of the attempts the
philosophy of which is the closest to ours. Its authors segment binary masks
with a pretrained MaskRCNN [11] to extract background clutter and occluders
and they map 2D pixels to dense 3D object coordinates, which, in turn, are
used as input to a P’n’P geometric optimization. Our attention modules have
the same effect, but are computationally cheaper than MaskRCNN, as they relax
the requirement for hard segmentation.

####  2.0.2 Temporal Tracking

The second category under study is temporal tracking, where feedback is
utilized, to allow for skipping steps without prior knowledge of the previous
pose. Garon et al. [7, 8], formulated the tracking problem exclusively as a
learning one, by generating two streams of synthetic RGB-D frame pairs from
independent viewpoints and regressing the pose using a CNN. Li et al.[23]
initialized a similar CNN architecture using a FlowNet2 [17] backbone and
fused its two streams by subtraction. In DeepTAM [39], the training was
performed with an Optical flow-based regularization term and the production of
multiple heterogenous pose hypotheses was encouraged. Those hypotheses were
bootstrapped in the final layer. Last but not least, Deng et al.[4] extended
the framework of [33] by combining it with a Rao-Blackwellized Particle Filter
[5]. In brief, they randomly sampled 2D bounding boxes to crop RGB images,
infer 3D translation probabilities from their dimensions and search for the
closest saved rotated sample provided by the Autoencoder of [33]. Then, the
particles were weighted according to the orientation probabilities and were
prepared for the next sampling iteration.

Figure 2: Overview of the proposed CNN architecture for object pose tracking.

##  3 Methodology

###  3.1 Problem Formulation

Our problem consists in estimating the object pose $\mathbb{P}$, which is
usually described as a rigid 3D transformation w.r.t. a fixed coordinate
frame, namely an element of the Special Euclidean Lie group in 3D: $SE(3)$. It
can be disentangled into two components; a rotation matrix R, which is an
element of the Lie Group SO(3) and a translation vector
$\mathbf{t}\in\mathbb{R}^{3}$. However, Brégier et al. [2] proposed a broader
definition for the object pose, which can be considered as a family of rigid
transformations, accounting for the ambiguity caused by possible rotational
symmetry, noted as G $\in$ SO(3). We leverage this augmented mathematical
definition for introducing a relaxation to the pose space $\mathscr{C}$
definition:

$\mathscr{C}=\Biggl{\\{}\mathbb{P}\quad|\quad\mathbb{P}=\left[\begin{array}[]{c|c}R\cdot
G&\mathbf{t}\\\
\hline\cr\mathbf{0}^{T}&1\end{array}\right],\mathbf{t}\in\mathbb{R}^{3},R\in\textnormal{SO(3)},G\in\textnormal{SO(3)}\Biggr{\\}}.$

(1)

For example, as stated in [2], the description of the pose of an object with
spherical symmetry requires just 3 numbers: ($t_{x,y,z}$), as G can be any
instance of SO(3) with the imprinted shape of the object remaining the same.
Obviously, for asymmetrical objects, G=$\mathbb{I}_{3}$.

###  3.2 Architecture Description

The proposed architecture is depicted in Fig.2. Our CNN inputs two RGB-D
frames of size $150\times 150$: I(t),$\hat{{\textnormal{I}}}$(t) (with I(t)
being the “Observed” and $\hat{\textnormal{I}}$(t) the “Predicted” one) and
regresses an output pose representation $\Delta\mathbf{p}\in\mathbb{R}^{9}$,
with 3 parameters for translation ($\hat{t}_{x,y,z}\in[-1,1]$) and 6 for
rotation. The first two layers of the “Observed” stream are initialized with
the weights of a ResNet18[13], pretrained on Imagenet [3], to narrow down the
real-synthetic domain adaptation gap, as proposed in [14]. Since Imagenet
contains only RGB images, we initialize the weights of the Depth input
modality with the average of the weights corresponding to each of the three
RGB channels. Contrary to [14], we find beneficial not to freeze those two
layers during training. The reason is that we aim to track the pose of the
single objects we train on and not to generalize to unseen ones. So,
overfitting to that object’s features helps the tracker to focus only on
distinguishing the pose change. The weights that correspond to the Depth
stream are, of course, not frozen in either case. To the output of the second
“Observed” layer, we apply spatial attention for foreground extraction and
occlusion handling and we add their corresponding output feature maps with the
one of the second layer, along with a Residual connection [12] from the first
layer. As a next step, we fuse the two streams by concatenating their feature
maps and pass this concatenated output through three sequential Fire modules
[16], all connected with residual connections [13].

Background and Occlusion Handling: After our first “Observed” Fire layer, our
model generates an attention weight map by using a Fire layer dedicated to
occlusion handling and foreground extraction, respectively, followed by a
$1\times 1$ convolution that squeezes the feature map channels to a single one
(and normalized by softmax). Our goal is to distil the soft foreground and
occlusion segmentation masks from the hard binary ground-truth ones (that we
keep from augmenting the object-centric image with random backgrounds and
occluders) in order to have their estimations available during the tracker’s
inference. To this end, we add the two corresponding binary cross entropy
losses to our overall loss function. We argue our design choice of using two
attention modules, as after experimentation, we found that assigning a clear
target to each of the two modules is more beneficial, rather than relying on a
single attention layer to resolve both challenges (see Sect.4.3.1), an
observation also reported in [18].

Rotation representation: From a mathematical standpoint, immediate regression
of pose parameters [8] with an Euclidean loss is suboptimal: while the
translation component belongs to the Euclidean space, the rotation component
lies on a non-linear manifold of SO(3). Thus, it is straightforward to model
the rotation loss using a Geodesic metric [15, 10] on SO(3), i.e. the length,
in radians, of the minimal path that connects two of its elements:
$\Delta\hat{R},\Delta R$:

${\small L_{Rot}(\Delta\hat{R},\Delta
R_{GT})=d^{(Geod)}_{Rot}(\Delta\hat{R},\Delta
R_{GT})=\arccos\Big(\frac{\Tr\big(\Delta\hat{R}^{T}\cdot\Delta
R_{GT}\big{missing})-1}{2}\Big{missing}).}$ (2)

In order to minimize the rotation errors due to ambiguities caused by the
parameterization choice, we employ the 6D continuous rotation representation
that was introduced in [40]:
$\Delta\mathbf{r}=(\Delta\mathbf{r_{x}}^{T},\Delta\mathbf{r_{y}}^{T})$, where
$\Delta\mathbf{r_{x/y}}\in\mathbb{R}^{3}$. Given $\Delta\mathbf{r}$, the
matrix $\Delta
R=(\Delta\mathbf{R_{x}},\Delta\mathbf{R_{y}},\Delta\mathbf{R_{z}})^{T}$ is
obtained by:

$\begin{split}&\Delta\mathbf{R_{x}}=N(\Delta\mathbf{r_{x}})\\\
&\Delta\mathbf{R_{y}}=N[\Delta\mathbf{r_{y}}-(\Delta\mathbf{R^{T}_{x}}\cdot\mathbf{r_{y})\cdot\Delta\mathbf{R_{x}}})]\\\
&\Delta\mathbf{R_{z}}=\Delta\mathbf{R_{x}}\times\Delta\mathbf{R_{y}}\end{split}$
(3)

where $\Delta\mathbf{R_{x/y/z}}\in\mathbb{R}^{3}$,
$N(\cdot)=\frac{(\cdot)}{\norm{(\cdot)}}$ is the normalization function.
Furthermore, as it has already been discussed in [2], each 3D rotation angle
has a different visual imprint regarding each rotation axis. So, we multiply
both rotation matrices with an approximately diagonal Inertial Tensor
$\Lambda$, calculated on the object model’s weighted surface and with respect
to its center mass, in order to assign a different weight to each rotational
component. We note here that since we want that matrix product to still lie in
SO(3), we perform a Gramm-Schmidt orthonormalization on the Inertial Tensor
$\Lambda$ before right-multiplying it with each rotation matrix. Finally, we
weigh the translation and rotation losses using a pair of learnable weights
$\mathbf{v}=[v_{1},v_{2}]^{T}$ that are trained along with the rest of the
network’s parameters using Gradient Descent-based optimization, as proposed in
[20].

Symmetric Object Handling In the special case of symmetric objects, we
disentangle the ambiguities inserted due to this property from the core of
rotation estimation. We classify such ambiguities to two distinctive
categories: a continuous set of rotational ambiguities (the unweaving of which
we incorporate into our loss function) and a discrete set of reflective
ambiguities that appear due to our rotation representation choice and are
handled heuristically.

Normally, the presence of the RGB input cue would break any symmetry
ambiguities as their origin is its Depth counterpart that depicts the object’s
3D shape. However, our preliminary experiments have shown that this is only
partially true since the tracker places more emphasis to the Depth cue, in
general, in its effort to estimate the object’s pose. This inclination keeps
the ambiguities present during inference and incurs the need to explicitly
model this disentanglement in the loss function formulation.

Symmetric Object Handling: Discrete Reflective Symmetry

By replacing the 3D Euler rotational parameter regression of Garon et al.[8],
with its 6D continuous counterpart, we face the extra problem of being unable
to constrain the network’s rotational output. This has a severely negative
effect to objects with reflective symmetry, as there are configurations in
which the tracker predicts unacceptable values for one or more rotational
components. For example, in Fig. 3 we observe that the “Cookie Jar” symmetric
object has been turned upside down with respect to both the prediction of the
previous timeframe and the ground truth pose . As a result, we end up with
adding significant extra errors during the mean rotational error calculation
over the total motion length, since that discrepancy is not limited to a
single frame, but is accumulated as we proceed to the following frames, due to
the temporal nature of the tracker, until it is reset. In order to handle this
challenge, we employ the following heuristic algorithm: for each Euler
rotational component we calculate the angular distance
$d^{(o)}_{A}(\hat{r}_{i}(t),\hat{r}_{i}(t-1))$ (in degrees), between the
current and the previous timestep. If a $d^{(o)}_{A}$ exceeds a certain
threshold (here, it is set to $100^{o}$), then the value of this particular
Euler angle is set to $\hat{r_{i}}(t-1)$. Then, we may choose to perform a
second (or more) iterative forward CNN pass(es), if the time constraints of
our broader application allow so.

Figure 3: Comparison of the Rotational Estimation for the 127th frame of the
”Hard Interaction Scenario” without/with our proposed reflective symmetry
handling algorithm. (a) The ground-truth prediction for frame 127. (b) The
rotation estimation of the 126th frame. Rotational error between the
prediction and ground truth is small. (c) An erroneous rotational estimation
of the tracker for frame 127. The prediction is discarded and replaced by the
one of the previous timeframe. The error is kept small and (although larger
than in the 126th frame) reflective ambiguity is not propagated to the next
timestep via the previous pose feedback rendering.

Symmetric Object Handling: Continuously Rotational Symmetry In order to form
G, we train a batch (of size $B_{2}$) of separate Euler angle triplets
$\mathbf{\hat{g}}\in\mathbb{R}^{3}$. At each timestep, one of them is selected
and it is converted to a rotation matrix $\hat{G^{*}}$, which gets right-
multiplied with $\Delta\hat{R}$ before being weighted by the parameters of
$\Lambda_{(G.S.)}$ (see eq.4). Aiming to select the appropriate parameter from
the batch, we train a linear classification layer on top of the first fully
connected layer of the tracker. Moreover, we encourage the symmetry triplets
to be as uniform as possible by incorporating an appropriate penalty to the
overall loss function. To test our approach for the symmetric object case, we
used the cylindrical “Cookie Jar” model of [8], the shape of which has only
one axis of continuous symmetry. Consequently, we estimate a single rotational
symmetry parameter, that of the object-centric z-axis (and keep the rest to
0). On the other hand, in the previous case, we define all three of its axes
as axes of reflective symmetry. Before the conversion, that parameter is
passed through a tanh function and multiplied by $\pi$ to constrain its
values.

####  3.2.1 Overall Loss

As a result, our overall tracking loss function is formulated as:

$\begin{gathered}\footnotesize
L_{Track}(\Delta\hat{\mathbb{P}},\Delta\mathbb{P})=e^{(-v_{1})}\cdot\textit{\footnotesize{MSE}}[(\Delta\hat{\mathbf{t}},\Delta\mathbf{t})]+v_{1}+v_{2}+\\\
+e^{(-v_{2})}\cdot\textit{arcos}\left(\frac{\Tr\left((\Delta\hat{R}\cdot\hat{G}^{*}\cdot\Lambda_{(G.S.)})^{T}\cdot(\Delta
R\cdot\Lambda_{(G.S.)})\right)-1}{2}\right)\normalsize\end{gathered}$ (4)

Using a similar external multi-task learnable weighting scheme
($\mathbf{s}=[s_{1},s_{2},s_{3}]^{T}$) as in eq.(4), we combine our primary
learning task, pose tracking, with the two auxiliary ones: clutter and
occlusion handling. Both $\mathbf{s}$ and $\mathbf{v}$ are initialized to
$\mathbf{0}$.

$\begin{gathered}Loss=e^{(-s_{1})}\cdot L_{Track}+e^{(-s_{2})}\cdot
L_{Unoccl}+e^{(-s_{3})}\cdot L_{Foregr}+s_{1}+s_{2}+s_{3}\end{gathered}$ (5)

For objects with continuously rotational symmetry the loss becomes:

$Loss^{(Symm)}=Loss+e^{(-s_{4})}\Big{(}\frac{1}{B}\sum_{b=1}^{B}\frac{1}{\xi_{b}}\Big{)}+s_{4},\textnormal{
with }$ (6)
$\xi_{b}=\frac{1}{B_{2}(B_{2}-1)}\sum_{j=1}^{B_{2}}\sum_{k\neq
j}d^{(Geod)}_{Rot}(\hat{G}_{k},\hat{G}_{j}),$ (7)

The extra term added to the multi-task loss is a penalty that guides the
classification layer(s) to select the proper rotational symmetry parameter(s)
at each timestep. It encourages the geodesic rotational distances between all
pairs of parameters in the batch to be maximum and, thus, ultimately converge
to as a uniform distribution as possible. Here, we train $B_{2}$=64 such
parameters for each continuous rotational component.

###  3.3 Data Generation and Augmentation

Following [7], for our network (Fig.2), we generate two synthetic RGB-D pairs
I(t),$\hat{\textnormal{I}}(t)$, but we alter their sampling strategy using the
Golden Spiral approach [21], and we modify the augmentation procedure of [7,
8] as follows: Firstly, we blend the object image with a background image,
sampled from a subset of the SUN3D dataset [37]. We also mimic the procedure
of [7, 8] in rendering a 3D hand model-occluder on the object frame with
probability $60\%$. A twist we added, is preparing our network for cases of
$100\%$ occlusion, by completely covering the object by the occluder for
$15\%$ of the occluded subset. Note that both the foreground and unoccluded
object binary masks are kept during both of these augmentation procedures.
Hence, we can use them as ground truth segmentation signals for clutter
extraction and occlusion handling in our auxiliary losses to supervise the
corresponding spatial attention maps. We add to the “Observed” frame pair
I(t): (i) Gaussian RGB noise, (ii) HSV noise, (iii) blurring (to simulate
rapid object movement), (iv) depth downsampling and (v) probabilistic dropout
of one of the modalities, all with same parameters as in [8]. With a
probability of $50\%$, we change the image contrast, using parameters
$\alpha\sim{U(0,3)}$, $\beta\sim{U(-50,50)}$ (where U($\cdot$) is a uniform
distribution) and gamma correction $\gamma\sim{U(0,2)}$ with probability
$50\%$, to help generalize over cases of illumination differences between
rendered and sensor generated images. Instead of modelling the noise added to
the “Observed” Depth modality with an ad-hoc Gaussian distribution as in [8],
we consider the specific properties of Kinect noise [28] and model it with a
3D Gaussian noise (depending on depth and the ground truth object pose), used
for simulating the reality gap between synthetic and real images. Its
distribution consists of a product of an z-Axial:
$n_{A}\sim{\mathscr{N}(0,\sigma_{A})}$ and two z-Lateral:
$n_{L_{x}}\sim{\mathscr{N}(0,\sigma_{L_{x}})}$,
$n_{L_{y}}\sim{\mathscr{N}(0,\sigma_{L_{y}})}$ 1D distributions that vary with
the object depth z and its angle around the y-axis, $\theta_{y}$: with
standard deviations $\sigma_{A}$,$\sigma_{L_{x}}$,$\sigma_{L_{y}}$
respectively. The rest of the preprocessing follows [7].

##  4 Evaluation and Results

###  4.1 Implementation Details

We use ELU activation functions, a minibatch size of 128, Dropout with
probability 0.3, Adam optimizer with decoupled weight decay [26] by a factor
$1e^{-5}$, learning rate $1e^{-3}$ and a scheduler with warm restarts [25]
every 10 epochs. All network weights with ELU activation function (except
those transferred from ResNet18 [13]) are initialized via a uniform K.He [12]
scheme, while for all those with a symmetric one we use a corresponding Xavier
[9] distribution. Since the Geodesic distance suffers from multiple local
minima [6], following [27], we first warm-up the weights, aiming first to
minimize the LogCosh[1] loss function for 25 epochs. Then, we train until
convergence, minimizing the loss (5). The average training time is 12 hours in
a single GeForce 1080 Ti GPU.

Object Attributes Size Symmetry Shape Texture Distinctive parts “Dragon”
Medium No Complex Rich Yes “Cookie Jar” Medium Rotoreflective Simple Poor and
Repetitive No Dog Medium No Complex Almost None Yes “Lego” Small No Complex
Rich and Repetitive No Watering Can Big No Simple Poor Yes

Table 1: Characteristics of the five objects we test our approach on. The fact
that there are no two identical items validates the generalization
capabilities of our tracker.

###  4.2 Dataset and Metrics

We test our approach on all the Interaction scenarios and the highest
percentage Occlusion scenarios of [8], which are considered the most
challenging. As in [8], we initialize our tracker every 15 frames, and use the
same evaluation metrics.

Due to limited computational resources, we produced only 20,000 samples, whose
variability covers the pose space adequately enough to verify the validity of
our experiments, both for the ablation study and the final experimentation. In
the following tables, similarly to [8], we report the mean and standard
deviation of our error metrics, as well, as the overall tracking failures
(only in the final experimentation section).

###  4.3 Ablation Study

In this section, we discuss our main design choices and we demonstrate
quantitative results that led to their selection.

####  4.3.1 Hierarchy choices for the attention modules

Here, we justify the need for both attention modules of our architecture (Fig.
2). We build upon the network proposed by [8] and we firstly introduce a
single convolutional attention map just for occlusion handling . Then, we
explore the possibility for a seperate attentional weighting of the “Observed”
feature map for foreground extraction, prior to the occlusion one, and we,
finally, leverage both in parallel and add their resulting maps altogether.

Translational Error(mm) Rotational Error(degrees) Garon et al. [8] 34.38 $\pm$
24.65 36.38 $\pm$ 36.31 Only occlusion 17.60 $\pm$ 10.74 37.10 $\pm$ 35.08
Hierarchical clutter $\&$ occlusion 14.99 $\pm$ 9.89 39.07 $\pm$ 33.22
Parallel clutter $\mathbf{\&}$ occlusion $\mathbf{14.35}\pm 10.21$
$\mathbf{34.28\pm 29.81}$

Table 2: Comparison of different attentional foreground/occlusion handling
configurations added to the baseline architecture of Garon et al.[8].

Rotational Error(degrees) Garon et al. [8] 36.38 $\pm$ 36.31 Rotational MSE
46.55 $\pm$ 40.88 Geod. 37.69 $\pm$ 35.39 Geod.+[40] 14.90 $\pm$ 21.76
Geod.+[40]+$\Lambda_{(G.S.)}$ $\mathbf{9.99\pm 13.76}$

Table 3: The evolution of the proposed rotation loss, on the baseline
architecture of Garon et al. [8] (without our proposed attention modules).

Translational Error(mm) Rotational Error(degrees) Garon et al. [8] 48.58 $\pm$
38.23 36.38 $\pm$ 36.31 Steady Weights 11.83 $\pm$ 8.94 10.98 $\pm$ 16.74
Recursive Batch Standarization [34] 13.97 $\pm$ 10.23 14.76 $\pm$ 19.24
Learnable Weights [20] 11.63 $\mathbf{\pm}$ 8.79 8.31 $\pm$ 6.76

Table 4: Comparison of different multi-task weighting schemes.

The comparison of Table 4 establishes not only the need for both attentional
modules in our design, but also that the parallel layout is the optimal one.
We can observe the effect of parallel connection in Fig.1 as both attentions
present sharp peaks. We can, also, observe a visual tradeoff between the
parallel attentions: while the object is not occluded (either in steady state
or when moving), the module responsible for foreground extraction is
highlighted more intensely than the occlusion one. As the object gets more and
more covered by the user’s hand, the focus gradually shifts to the module
responsible for occlusion handling. Note that this is not an ability we
explicitly train our network to obtain, but rather a side effect of our
approach, which fits our intuitive understanding of cognitive visual tracking.
Moreover, although our supervising signals are uniform, both attentional
modules learn to highlight specific keypoints of interest during the tracker’s
inference.
We, also, observe another interesting property: the tracker learns to handle
self-occlusion patterns as well. For example, for the “Dragon” object of
Fig.1, it learns to ignore the one wing when it is in front and focuses on the
other one, at the back, if its appearance is more distinctive of the pose. The
same can be said for the legs of the Dog model of Fig.7. This is a clue that
this module has, indeed, learned the concept of occlusions and has not
overfitted to the shape of the user’s hand, the occluder model that was used
for training.

####  4.3.2 Contributions of the rotation Loss components

We demonstrate the value of every component included in our rotation loss
(leaving symmetries temporarily out of study), by: (i) regressing only the
rotational parameters with the baseline architecture of [8], (ii), replacing
the MSE loss with the Geodesic one, (iii), replacing the rotation
parameterization of [8] with the continuous one of (3), and, (iv) including
the Inertial Tensor weighting of each rotational component.

Translational Error(mm) Rotational Error(degrees)
$d^{(o)}_{A}(\hat{r}_{z},r^{GT}_{z})$ Garon et al.[8] 10.75 $\pm$ 6.89 23.53
$\pm$ 18.85 10.60 $\pm$ 38.90 Ours 10.87 $\pm$ 8.14 20.55 $\pm$ 18.06 4.60
$\pm$ 35.42 Ours+Unique,frozen learnable symmetry parameter 11.38 $\pm$ 8.94
10.98 $\pm$ 16.74 2.98 $\pm$ 25.07 Ours+Regression of learnable symmetry
parameter 13.03 $\pm$ 6.88 17.25 $\pm$ 12.40 2.84 $\pm$ 24.95 Ours+Mean of a
batch of frozen, learnable symmetry parameters 11.98 $\pm$ 9.23 13.84 $\pm$
11.87 2.16 $\pm$ 29.25 Ours+Optimal selection out of a batch of frozen,
learnable symmetry parameters 10.43 $\mathbf{\pm}$ 6.63 9.57 $\mathbf{\pm}$
10.01 2.07 $\pm$ 24.52

Table 5: Comparison of different selection methods for the continuously
rotational symmetry parameter.

Table 4 indicates the value that translation estimation brings to rotation
estimation, as when the former’s regression is excluded, the latter’s
performance decreases. Moreover, Table 4 justifies our progressive design
selections in formulating our rotation loss, as with the addition of each
ambiguity modelling, the 3D rotation error decreases, starting from
$46.55^{o}\pm 40.88^{o}$ and reaching $9.99^{o}\pm 13.76^{o}$.

####  4.3.3 Weighting the Multi-Task Loss

Here, we explore various weighting schemes of the multiple loss functions of
our approach, both the main and the auxiliary ones. Our first approach is the
crude addition of the tracking and the two Binary Cross Entropy losses . A
second one is standarizing the three losses by subtracting their batchwise
means and dividing by their batchwise standard deviations, that we calculate
using the Welford algorithm[34]. Lastly, we consider the learnable weighting
strategy that we, ultimately, utilize . In quantitative comparison of Table4
the scheme of [20] emerges as the clear favourite.

####  4.3.4 Comparison of Continuously Rotational Symmetry Regression methods
(“Cookie Jar” model)

Figure 4: Comparison of the SoA[8] (light blue) and our (pink) approaches for
the “Cookie Jar” in 3 scenarios: “Translation Only”, “Rotation Only” and “Full
Interaction”.

Figure 5: Comparison of the SoA[8] (light blue) and our (pink) approaches for
the “Lego” in 3 scenarios: “Translation Only”,’Full’ and “Hard Interaction”.

In our effort to disentangle the rotation estimation and the continuously
rotational symmetries, we try different configurations for optimally choosing
the appropriate parameter(s): (i) learning a unique symmetry parameter over
all possible pose changes in the training set and keeping it frozen during
inference, (ii) regressing a different one per pose pair, (iii) learning a
batch of them and taking their average during inference and, ultimately, (iv)
selecting the optimal from the aforementioned batch using an appropriate
classification layer while encouraging the values of this batch to be as
uniform as possible, at the same time. This freedom of ours resides from the
fact that the minimization of the tracking loss w.r.t. the symmetry matrix
$\hat{G^{*}}$ (see [2]) does not explicitly impose a global-solution
constraint. This time, the comparison is done w.r.t. the full approach of
Sect.3 that does not account for symmetries and, apart from the classic 3D
translational and rotational errors, we, also, report the Euler angle error
for the z-component. As we can see in Table 5, our proposed approach is the
“golden medium” between accuracy and robustness since both its mean and
standard deviation are the lowest across all metrics.

###  4.4 Experimental Results

According to our ablation study, we proceed to merge our parallel attention
modules with the Geodesic rotation loss of eq.(4), along with the remaining
elements of Sect.3. We evaluate our method on five objects of dataset [8]: the
“Dragon”, the “Cookie Jar”, the “Dog”, the “Lego” block and the “Watering
Can”, aiming for maximum variability (see Table 1). In Fig. 1-7, we plot the
3D translational and rotational errors in three randomly selected scenarios
for each object both for the SoA and our tracker.

Figure 6: Comparison of the SoA[8] (light blue) and our (pink) approaches for
the “Dog” in 3 scenarios: “$75\%$ Vertical Occlusion”, “Rotation Only” and
“Hard Interaction”.

Figure 7: Comparison of the SoA[8] (light blue) and our (pink) approaches for
the “Watering Can” in 3 scenarios:“$75\%$ Horizontal Occlusion”, “Rotation
Only” and “Full Interaction”.

Evidently, the object most benefited by our methodological improvements is the
“Dragon”. Since its geometry is the most complex, its texture is rich and it
has several distinctive parts that stand out of the user’s grip, both our
geometric modelling and the parallel attention modules find their best
application in this case. When the user’s hand occludes parts of the “Dragon”,
the attention shifts to its body parts of interest that stand out of the grip,
like its neck, wings or tail (Fig.1). For the symmetric “Cookie Jar”, the
differences between our method and the baseline are lower. The attentions’
effect is less prominent here since this model is of simpler, symmetric shape
and poorer texture. This replaces the distinctive clues of the dragon case
with ambiguities, denying the corresponding modules of the ability to easily
identify the pose. Alongside the “Lego” model, they make the most out of the
reflective symmetry handling algorithm as they avoid large abrupt errors that
propagate to future frames. Although our CNN primarily focuses on the object’s
shape, appearance seems to play a significant role in its predictions, as
well, since the errors of the “Dog” and the “Watering Can” models, the less
textured ones, decrease more mildly. The foreground attention map aids
disentangling the “Dog” model from its background, a table of the same color,
in the “75$\%$ Occlusion” scenarios. On the other hand, for the “Watering
Can”, the most ambiguities are presented when viewpoint-induced symmetries
appear, with the effects of our modelling declining in this case.

The accuracy of our approach exceeds that of the SoA [8], across all objects
and in almost all scenarios, especially for 3D rotations. According to Table6,
both errors are generally lower (in terms of both mean and standard deviation)
and our tracker fails equally or less often. It presents aggravated errors in
fast object motions more rarely than [8] and handles both static and dynamic
high-percentage occlusion patterns better. Also, it not only keeps track of
the object’s 3D position under severe occlusions, but extends this property to
3D rotations, as well. Although more computationally intense than [8], it runs
in 40 fps.

Approach 75$\%$ Horizontal Occlusion 75$\%$ Vertical Occlusion Translational
Error(mm) Rotational Error(degrees) Fails Translational Error(mm)
Rotational(degrees) Fails Garon et al.[8] (’Dragon) 16.02 $\pm$ 8.42 18.35
$\pm$ 11.71 13 18.20 $\pm$ 11.81 14.66 $\pm$ 12.98 13 Ours(“Dragon”)
$\mathbf{12.68}$ $\mathbf{\pm}$ 11.49 $\mathbf{13.00}$ $\mathbf{\pm}$
$\mathbf{9.14}$ $\mathbf{10}$ $\mathbf{12.87\pm 10.49}$ $\mathbf{13.14\pm
8.85}$ 8 Garon et al.[8](“Cookie Jar”) 21.27 $\pm$ 9.74 21.90 $\pm$ 13.97 17
20.77 $\pm$ 6.88 24.86 $\pm$ 13.64 20 Ours(“Cookie Jar”) 9.51 $\pm$ 4.17 15.48
$\pm$ 9.50 15 20.97 $\pm$ 7.32 16.14 $\pm$ 10.06 15 Ours+Symm.(“Cookie Jar”)
$\mathbf{6.37}$ $\mathbf{\pm}$ $\mathbf{2.14}$ 7.22 $\pm$ 3.97 $\mathbf{11}$
$\mathbf{19.01}$ $\mathbf{\pm}$ 7.53 13.00 $\pm$ 7.49 $\mathbf{14}$ Garon et
al.[8](Dog) 37.96 $\pm$ 23.39 47.94 $\pm$ 31.55 21 32.84 $\pm$ 34.07 22.44
$\pm$ 13.60 21 Ours(Dog) 24.43 $\pm$ 18.92 17.24 $\pm$ 12.41 25 36.53 $\pm$
22.39 12.67 $\pm$ 7.95 20 Garon et al.[8](“Lego”) 68.25 $\pm$ 46.97 40.04
$\pm$ 47.37 28 40.04 $\pm$ 47.37 35.30 $\mathbf{\pm}$ 31.32 20 Ours(“Lego”)
72.04 $\pm$ 34.10 18.41 $\pm$ 13.84 28 12.92 $\pm$ 5.73 12.92 $\pm$ 9.02 20
Garon et al.[8](Watering Can) 21.59 $\pm$ 11.32 23.99 $\pm$ 16.95 14 32.76
$\pm$ 24.12 26.74 $\pm$ 19.05 18 Ours(Watering Can) 20.71 $\pm$ 10.24 17.00
$\pm$ 18.99 13 17.66 $\pm$ 17.95 13.46 $\pm$ 10.43 12

Approach Translation Interaction Rotation Interaction Translational Error(mm)
Rotational Error(degrees) Fails Translational Error(mm) Rotational(degrees)
Fails Garon et al.[8] (“Dragon”) 41.60 $\pm$ 39.92 11.55 $\pm$ 15.58 15 23.86
$\pm$ 17.44 27.21 $\pm$ 22.40 15 Ours(“Dragon”) $\mathbf{11.05}$
$\mathbf{\pm}$ $\mathbf{8.20}$ $\mathbf{3.55}$ $\mathbf{\pm}$ $\mathbf{2.27}$
$\mathbf{1}$ $\mathbf{9.37\pm 6.07}$ $\mathbf{7.86\pm 6.69}$ 2 Garon et
al.[8](“Cookie Jar”) 20.43 $\pm$ 25.44 17.19 $\pm$ 12.99 16 10.75 $\pm$ 5.89
23.53 $\pm$ 18.85 19 Ours(“Cookie Jar”) 8.64 $\pm$ 8.23 8.31 $\pm$ 5.97 5
10.87 $\pm$ 8.14 20.55 $\pm$ 18.06 16 Ours+Symm.(“Cookie Jar”) $\mathbf{8.09}$
$\pm$ $\mathbf{7.67}$ $\mathbf{5.83}$ $\mathbf{\pm}$ $\mathbf{5.50}$
$\mathbf{3}$ $\mathbf{9.98}$ $\mathbf{\pm}$ 10.63 $\mathbf{13.84}$
$\mathbf{\pm}$ $\mathbf{11.87}$ $\mathbf{16}$ Garon et al.[8](Dog) 58.87 $\pm$
71.86 16.42 $\pm$ 13.51 20 11.16 $\pm$ 10.28 20.00 $\pm$ 21.31 17 Ours(Dog)
$\mathbf{21.64}$ $\mathbf{\pm}$ $\mathbf{22.78}$ $\mathbf{9.27}$
$\mathbf{\pm}$ $\mathbf{8.03}$ $\mathbf{14}$ $\mathbf{10.68}$ $\mathbf{\pm}$
$\mathbf{7.53}$ $20.07$ $\mathbf{\pm}$ $\mathbf{19.29}$ $\mathbf{17}$ Garon et
al.[8](“Lego”) 27.90 $\pm$ 23.53 11.89 $\pm$ 18.50 29 16.42 $\pm$ 10.90 17.83
$\pm$ 15.90 32 Ours(“Lego”) $\mathbf{22.66}$ $\mathbf{\pm}$ 24.58
$\mathbf{9.08}$ $\mathbf{\pm}$ $\mathbf{7.60}$ $\mathbf{12}$ $\mathbf{10.13}$
$\mathbf{\pm}$ $\mathbf{6.79}$ $\mathbf{7.22}$ $\mathbf{\pm}$ $\mathbf{4.55}$
$\mathbf{4}$ Garon et al.[8](Watering Can) 24.95 $\pm$ 42.91 13.26 $\pm$ 11.34
16 13.14 $\pm$ $\mathbf{8.99}$ 22.19 $\pm$ 25.93 15 Ours(Watering Can)
$\mathbf{24.30}$ $\mathbf{\pm}$ $\mathbf{21.51}$ 8.79 $\pm$ 6.35 $\mathbf{16}$
12.22 $\pm$ 9.46 18.66 $\pm$ 15.51 $\mathbf{15}$

Apporach Full Interaction Hard Interaction Translational Error(mm) Rotational
Error(degrees) Fails Translational Error(mm) Rotational(degrees) Fails Garon
et al.[8] (“Dragon”) 35.23 $\pm$ 31.97 34.98 $\pm$ 29.46 18 34.38 $\pm$ 24.65
36.38 $\pm$ 36.31 17 Ours(“Dragon”) $\mathbf{10.31}$ $\mathbf{\pm}$
$\mathbf{8.66}$ $\mathbf{6.40}$ $\mathbf{\pm}$ $\mathbf{4.52}$ $\mathbf{1}$
$\mathbf{11.63\pm 8.79}$ $\mathbf{8.31\pm 6.76}$ 2 Garon et al.[8](“Cookie
Jar”) 13.06 $\pm$ 9.35 31.78 $\pm$ 23.78 24 15.78 $\pm$ 10.43 24.29 $\pm$
20.84 15 Ours(“Cookie Jar”) 17.03 $\pm$ 11.94 22.24 $\pm$ 20.86 21 15.29 $\pm$
16.06 16.73 $\pm$ 14.79 11 Ours+Symm.(“Cookie Jar”) 14.63 $\mathbf{\pm}$ 11.19
$\mathbf{15.71}$ $\mathbf{\pm}$ $\mathbf{13.80}$ $\mathbf{21}$
$\mathbf{14.96}$ $\mathbf{\pm}$ $\mathbf{9.06}$ $\mathbf{15.00}$
$\mathbf{\pm}$ $\mathbf{13.20}$ $\mathbf{8}$ Garon et al.[8](Dog) 37.73 $\pm$
42.32 20.77 $\pm$ 19.66 23 23.95 $\pm$ 38.86 24.38 $\pm$ 26.39 20 Ours(Dog)
$\mathbf{24.88}$ $\mathbf{\pm}$ $\mathbf{35.85}$ $28.52$ $\pm$ $25.38$
$\mathbf{20}$ $\mathbf{19.32}$ $\mathbf{\pm}$ $\mathbf{15.97}$
$\mathbf{19.72}$ $\mathbf{\pm}$ $\mathbf{20.17}$ $\mathbf{19}$ Garon et
al.[8](“Lego”) 30.96 $\pm$ 31.44 22.10 $\pm$ 20.20 20 30.71 $\pm$ 42.62 36.38
$\pm$ 34.99 20 Ours(“Lego”) $\mathbf{23.58}$ $\mathbf{\pm}$ $\mathbf{27.73}$
$\mathbf{11.80}$ $\pm$ $\mathbf{12.28}$ 13 $\mathbf{16.47}$ $\mathbf{\pm}$
$\mathbf{12.95}$ $\mathbf{14.29}$ $\mathbf{\pm}$ $\mathbf{11.68}$
$\mathbf{11}$ Garon et al.[8](Watering Can) 33.76 $\pm$ 37.62 40.16 $\pm$
35.90 26 28.31 $\pm$ 19.49 23.04 $\pm$ 24.27 28 Ours(Watering Can)
$\mathbf{19.82}$ $\mathbf{\pm}$ $\mathbf{19.98}$ $\mathbf{28.76}$
$\mathbf{\pm}$ $\mathbf{30.27}$ $\mathbf{26}$ $\mathbf{18.03}$ $\mathbf{\pm}$
$\mathbf{14.99}$ $\mathbf{19.57}$ $\mathbf{\pm}$ $\mathbf{17.47}$
$\mathbf{23}$

Table 6: 3D Translational and Rotational errors and overall tracking failures
in six different scenarios for five employed objects.

##  5 Conclusion

In this work, we propose a CNN for fast and accurate single object pose
tracking. We perform explicitly modular design of clutter and occlusion
handling and we account for the geometrical properties of both the pose space
and the object model during training. As a result, we reduce both SoA pose
errors by an average of $34.03\%$ for translation and 40.01$\%$ for rotation
for a variety of objects with different properties. Our tracker exceeds the
SoA performance in challenging scenarios with high percentage occlusion
patterns and rapid movement and we gain an intuitive understanding of our
artificial tracking mechanism.

## Acknowledgements

This research has been co-financed by the European Union and Greek national
funds through the Operational Program Competitiveness, Entrepreneurship and
Innovation, under the call RESEARCH CREATE INNOVATE (project code:
T1EDK-01248, acronym: i-Walk).

## References

  * [1] Belagiannis, V., Rupprecht, C., Carneiro, G., Navab, N.: Robust optimization for deep regression. In: Proceedings of the IEEE international conference on computer vision. pp. 2830–2838 (2015)
  * [2] Brégier, R., Devernay, F., Leyrit, L., Crowley, J.L.: Defining the pose of any 3d rigid object and an associated distance. Int. J. of Comp. Vision (IJCV) 126(6), 571–596 (2018)
  * [3] Deng, J., Dong, W., Socher, R., Li, L.J., Li, K., Fei-Fei, L.: Imagenet: A large-scale hierarchical image database. In: Proc. IEEE Conf. on Computer Vision and Pattern Recognition (CVPR). pp. 248–255. Ieee (2009)
  * [4] Deng, X., Mousavian, A., Xiang, Y., Xia, F., Bretl, T., Fox, D.: Poserbpf: A rao-blackwellized particle filter for 6d object pose tracking. arXiv preprint arXiv:1905.09304 (2019)
  * [5] Doucet, A., De Freitas, N., Murphy, K., Russell, S.: Rao-blackwellised particle filtering for dynamic bayesian networks. arXiv preprint arXiv:1301.3853 (2013)
  * [6] Fletcher, T.: Terse notes on riemannian geometry (2010)
  * [7] Garon, M., Lalonde, J.F.: Deep 6-dof tracking. IEEE transactions on visualization and computer graphics 23(11), 2410–2418 (2017)
  * [8] Garon, M., Laurendeau, D., Lalonde, J.F.: A framework for evaluating 6-dof object trackers. In: Proc. European Conf. on Computer Vision (ECCV). pp. 582–597 (2018)
  * [9] Glorot, X., Bengio, Y.: Understanding the difficulty of training deep feedforward neural networks. In: Proceedings of the thirteenth international conference on artificial intelligence and statistics. pp. 249–256 (2010)
  * [10] Hartley, R., Trumpf, J., Dai, Y., Li, H.: Rotation averaging. Int. J. of Comp. Vision (IJCV) 103(3), 267–305 (2013)
  * [11] He, K., Gkioxari, G., Dollár, P., Girshick, R.: Mask r-cnn. In: Proc. IEEE Int. Conf. on Computer Vision (ICCV). pp. 2961–2969 (2017)
  * [12] He, K., Zhang, X., Ren, S., Sun, J.: Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In: Proc. IEEE Int. Conf. on Computer Vision (ICCV). pp. 1026–1034 (2015)
  * [13] He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: Proc. IEEE Conf. on Computer Vision and Pattern Recognition (CVPR). pp. 770–778 (2016)
  * [14] Hinterstoisser, S., Lepetit, V., Wohlhart, P., Konolige, K.: On pre-trained image features and synthetic images for deep learning. In: Proc. European Conf. on Computer Vision (ECCV). pp. 0–0 (2018)
  * [15] Huynh, D.Q.: Metrics for 3d rotations: Comparison and analysis. J. Math. Imaging and Vision 35(2), 155–164 (2009)
  * [16] Iandola, F.N., Han, S., Moskewicz, M.W., Ashraf, K., Dally, W.J., Keutzer, K.: Squeezenet: Alexnet-level accuracy with 50x fewer parameters and $<$0.5mb model size. arXiv:1602.07360 (2016)
  * [17] Ilg, E., Mayer, N., Saikia, T., Keuper, M., Dosovitskiy, A., Brox, T.: Flownet 2.0: Evolution of optical flow estimation with deep networks. In: Proc. IEEE Conf. on Computer Vision and Pattern Recognition (CVPR). pp. 2462–2470 (2017)
  * [18] Jafari, O.H., Mustikovela, S.K., Pertsch, K., Brachmann, E., Rother, C.: ipose: instance-aware 6d pose estimation of partly occluded objects. In: Asian Conference on Computer Vision. pp. 477–492. Springer (2018)
  * [19] Kehl, W., Manhardt, F., Tombari, F., Ilic, S., Navab, N.: Ssd-6d: Making rgb-based 3d detection and 6d pose estimation great again. In: Proc. IEEE Int. Conf. on Computer Vision (ICCV). pp. 1521–1529 (2017)
  * [20] Kendall, A., Gal, Y., Cipolla, R.: Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In: Proc. IEEE Conf. on Computer Vision and Pattern Recognition (CVPR). pp. 7482–7491 (2018)
  * [21] Leopardi, P.C.: Distributing points on the sphere: partitions, separation, quadrature and energy. Ph.D. thesis, University of New South Wales, Sydney, Australia (2007)
  * [22] Lepetit, V., Moreno-Noguer, F., Fua, P.: Epnp: An accurate o(n) solution to the pnp problem. Int. J. of Comp. Vision (IJCV) 81 (02 2009). https://doi.org/10.1007/s11263-008-0152-6
  * [23] Li, Y., Wang, G., Ji, X., Xiang, Y., Fox, D.: Deepim: Deep iterative matching for 6d pose estimation. In: Proc. European Conf. on Computer Vision (ECCV). pp. 683–698 (2018)
  * [24] Liu, W., Anguelov, D., Erhan, D., Szegedy, C., Reed, S., Fu, C.Y., Berg, A.C.: Ssd: Single shot multibox detector. In: Proc. European Conf. on Computer Vision (ECCV). pp. 21–37. Springer (2016)
  * [25] Loshchilov, I., Hutter, F.: Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983 (2016)
  * [26] Loshchilov, I., Hutter, F.: Fixing weight decay regularization in adam. arXiv preprint arXiv:1711.05101 (2017)
  * [27] Mahendran, S., Ali, H., Vidal, R.: 3d pose regression using convolutional neural networks. In: Proceedings of the IEEE International Conference on Computer Vision Workshops. pp. 2174–2182 (2017)
  * [28] Nguyen, C.V., Izadi, S., Lovell, D.: Modeling kinect sensor noise for improved 3d reconstruction and tracking. In: 2012 Second International Conference on 3D Imaging, Modeling, Processing, Visualization & Transmission. pp. 524–530. IEEE (2012)
  * [29] Park, K., Patten, T., Vincze, M.: Pix2pose: Pixel-wise coordinate regression of objects for 6d pose estimation. In: Proceedings of the IEEE International Conference on Computer Vision. pp. 7668–7677 (2019)
  * [30] Pavlakos, G., Zhou, X., Chan, A., Derpanis, K.G., Daniilidis, K.: 6-dof object pose from semantic keypoints. In: 2017 IEEE international conference on robotics and automation (ICRA). pp. 2011–2018. IEEE (2017)
  * [31] Peng, S., Liu, Y., Huang, Q., Zhou, X., Bao, H.: Pvnet: Pixel-wise voting network for 6dof pose estimation. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. pp. 4561–4570 (2019)
  * [32] Segal, A., Haehnel, D., Thrun, S.: Generalized-icp. In: Robotics: science and systems. vol. 2, p. 435. Seattle, WA (2009)
  * [33] Sundermeyer, M., Marton, Z.C., Durner, M., Brucker, M., Triebel, R.: Implicit 3d orientation learning for 6d object detection from rgb images. In: Proceedings of the European Conference on Computer Vision (ECCV). pp. 699–715 (2018)
  * [34] Welford, B.: Note on a method for calculating corrected sums of squares and products. Technometrics 4(3), 419–420 (1962)
  * [35] Wohlhart, P., Lepetit, V.: Learning descriptors for object recognition and 3d pose estimation. In: Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 3109–3118 (2015)
  * [36] Xiang, Y., Schmidt, T., Narayanan, V., Fox, D.: Posecnn: A convolutional neural network for 6d object pose estimation in cluttered scenes. arXiv preprint arXiv:1711.00199 (2017)
  * [37] Xiao, J., Owens, A., Torralba, A.: Sun3d: A database of big spaces reconstructed using sfm and object labels. In: Proc. IEEE Int. Conf. on Computer Vision (ICCV). pp. 1625–1632 (2013)
  * [38] Zakharov, S., Shugurov, I., Ilic, S.: Dpod: Dense 6d pose object detector in rgb images. arXiv preprint arXiv:1902.11020 (2019)
  * [39] Zhou, H., Ummenhofer, B., Brox, T.: Deeptam: Deep tracking and mapping. In: Proc. European Conf. on Computer Vision (ECCV). pp. 822–838 (2018)
  * [40] Zhou, Y., Barnes, C., Lu, J., Yang, J., Li, H.: On the continuity of rotation representations in neural networks. In: Proc. IEEE Conf. on Computer Vision and Pattern Recognition (CVPR). pp. 5745–5753 (2019)


# Generalization Measures
for Zero-Shot Cross-Lingual Transfer

Saksham Bassi, Duygu Ataman, Kyunghyun Cho
Department of Computer Science
New York University
New York, NY 10011, USA
<EMAIL_ADDRESS>

###### Abstract

A model’s capacity to generalize its knowledge to interpret unseen inputs with
different characteristics is crucial to build robust and reliable machine
learning systems. Language model evaluation tasks lack information metrics
about model generalization and their applicability in a new setting is
measured using task and language-specific downstream performance, which is
often lacking in many languages and tasks. In this paper, we explore a set of
efficient and reliable measures that could aid in computing more information
related to the generalization capability of language models in cross-lingual
zero-shot settings. In addition to traditional measures such as variance in
parameters after training and distance from initialization, we also measure
the effectiveness of sharpness in loss landscape in capturing the success in
cross-lingual transfer and propose a novel and stable algorithm to reliably
compute the sharpness of a model optimum that correlates to generalization.
111Code: https://anonymous.4open.science/r/strikegen-7288

##  1 Introduction

Generalization enables models to use prior knowledge to reasonably respond to
previously unseen stimuli. Although traditional machine learning evaluation is
performed based on a preselected set of prediction or generation tasks,
accuracy on many public benchmarks may often not be sufficient to extensively
assess the ability to perform well in new settings. Therefore, a majority of
researchers have found it worthwhile to investigate measures that could
evaluate the generalization capability of models with properties, such as VC
dimension Vapnik & Chervonenkis (1971), cross-entropy Shannon (1948),
complexity Mohri et al. (2012) or variation in parameters during training
Nagarajan & Kolter (2019). Among these, recent findings support the smoothness
in the loss curvature to correlate best with generalization capability
Chaudhari et al. (2019); Petzka et al. (2021); Kaddour et al. (2022),
motivating the development of learning methods that induce smoothness in the
learning trajectory such that the model becomes more robust; either through
data perturbation Jiang et al. (2020a); Aghajanyan et al. (2021); Liang et al.
(2021); Hua et al. (2021); Park et al. (2022); Zheng et al. (2021); Wang et
al. (2021); Huang et al. (2021) or by integrating the measure directly to the
optimization objective Izmailov et al. (2018); Jastrzebski et al. (2021); Cha
et al. (2021); Foret et al. (2021); Hu et al. (2022); Zaken et al. (2022);
Stickland & Murray (2021). However it might often not be straightforward to
compute such measures in high-dimensional feature space in a stable fashion
Nachum et al. (2024).

As models get larger and cover more languages, the possibility of improving
the applicability of NLP systems in many under-resourced languages gets more
promising. An essential requirement in studying the dynamics of cross-lingual
knowledge transfer is to have an evaluation methodology that can reliably
measure the model’s capability in generalization of knowledge under different
scenarios. There is a common hypothesis that states that a model demonstrating
an extended flat optimum area of low loss value surrounding the minimized loss
is indicative of better generalization capability. In this work, we study the
above hypothesis and present the first study to provide methods that can be
used for measuring the cross-lingual generalization capability of language
models.

  * •

We pick prominent measures that were previously shown to correlate well with
generalization performance Jiang et al. (2020b), such as the Frobenius
distance of the learned parameters after training Nagarajan & Kolter (2019),
the margin between model predictions and true labels Wei et al. (2018) and
sharpness in loss minima to test applicability to zero-shot cross-lingual
generalization measurement Keskar et al. (2017); Foret et al. (2021).

  * •

We also extend the formulation of state-of-the-art sharpness computation
methods Keskar et al. (2017); Foret et al. (2021) to provide a sharpness
prediction algorithm such that the optimization of the parameters can converge
in a more stable fashion.

##  2 Related work

Loss-landscape Minima One of the most promising indicators of generalization
capability to date seems to be related to the form of the loss landscape, in
particular, the sharpness in the loss curvature. A potential reason for this
fallback is traced to stochastic gradient descent (SGD) Bottou (2012) methods
which often fall into sharp minima of the loss surface Keskar et al. (2017);
Chaudhari et al. (2019); Wang et al. (2021). Although clear conclusions on the
relationship between sharpness and generalization performance, such as whether
sharper Dinh et al. (2017) vs. flatter Li et al. (2018); Keskar et al. (2017)
minima would generally yield better generalization, are still due. The main
idea behind these methods is that their objective is to explicitly find flat
minima, often using stochastic averaging methods Polyak & Juditsky (1992);
Izmailov et al. (2018), mini-max or sharpness-aware minimization methods,
which can be computed by direct formulation based on the Hessian matrix of the
loss function Chaudhari et al. (2019); Petzka et al. (2021) or Monte-Carlo
approximations of the minimizer’s neighborhood Foret et al. (2021); Cha et al.
(2021).

Adversarial optimization Comparison of two approaches finds that for NLP
tasks, mini-max methods are more competitive over averaging-based optimization
Kaddour et al. (2022). Jastrzebski et al. (2021) hypothesize that regularizing
the trace of the Fisher information matrix amplifies the implicit bias of SGD,
which prevents memorization. The Fisher information Fisher (1925) measures
local curvature, so a smaller trace implies a flatter minimum, which gives the
model more freedom to reach an optimum. Instead of explicitly minimizing the
values of parameters, Foret et al. (2021) propose minimizing both loss and
sharpness while optimizing the parameters such that they lie in neighborhoods
with low loss values. Perturbation is an auxiliary objective that encourages
the model predictions to be similar in the vicinity of the observed training
samples Englesson & Azizpour (2021), usually by penalizing the KL-divergence
between the probability distribution of the perturbed and normal model.
Perturbations can be adversarial inputs Jiang et al. (2020a) or inputs with
Gaussian or uniform noise Aghajanyan et al. (2021). To improve cross-lingual
generalization, translations of the input generated by machine translation
systems were used as perturbed input Wang et al. (2021); Zheng et al. (2021).
Other work also has found the benefit of enforcing consistency for
perturbations within the model in addition to the input distribution Liang et
al. (2021); Hua et al. (2021).

##  3 Methodology

In this study, we undertake the development of a methodology that could
benefit an accurate assessment of the generalization capability of models for
the purpose of cross-lingual knowledge transfer into under-resourced
languages. This section first presents approaches to improving generalization
performance and the selected measures that provide stable results for
measuring zero-shot cross-lingual transfer performance.

###  3.1 Sharpness-based Optimization

We chose the following objective functions as fine-tuning methods for a given
pre-trained model as a means of comparison since their main purpose is to
enhance the generalization and robustness of models. Following the work of
Stickland & Murray (2021), as the two most prominent approaches to mini-max
optimization, we include Sharpness-Aware Minimization (SAM) Foret et al.
(2021) and regularization with Fisher Information Matrix (FIM) Jastrzebski et
al. (2021) in our evaluation study on cross-lingual generalization.We also
include Multi-view Subword Regularization (MVR) as a perturbation-based
optimization method Wang et al. (2021) which induces stochasticity into the
shared subword vocabulary across languages for easing cross-lingual transfer.

SAM Foret et al. (2021) works on the principle of a mini-max objective
function: $\min_{w}\max_{\|\epsilon\|_{2}<\rho}L(w+\epsilon)$, which
essentially means the optimizing function tries to minimize the maximum loss
value in a given radius in loss landscape. Therefore, SAM states that it tries
to seek ”parameters lying in uniformly low-loss neighborhoods”.

Fisher Penalty is defined as explicitly penalizing the trace of the Fisher
information matrix (FIM). Jastrzebski et al. (2021), Stickland & Murray (2021)
observed penalizing FIM during training correlates to better generalization
performance. It can be written mathematically as
$\frac{1}{n}\sum_{i}^{n}\nabla L(x_{i},y_{i})$ where $L(x_{i},y_{i})$ is the
loss at the data point $(x_{i},y_{i})$.

MVR Wang et al. (2021) function on the concept of consistency regularization
where the divergence between the model predictions on deterministic and
probabilistic segmentation inputs is minimized. The objective function is
formulated as

$\displaystyle\sum_{i=1}^{N}$ $\displaystyle\Big{(}-\frac{1}{2}\log
p(y_{i}|\hat{x}_{i})-\frac{1}{2}\log p(y_{i}|x_{i}^{\prime})$ (1)
$\displaystyle+\lambda
D(p(y_{i}|\hat{x}_{i})\;\|\;p(y_{i}|x_{i}^{\prime}))\Big{)}$ (2)

where the first term is the model loss on deterministic segmentation of the
$i^{th}$ data sample (most probable segmentation), the second term is the
model loss on probabilistic segmentation of the $i^{th}$ data sample (random
segmentation) and the third term is the KL divergence between these two output
predictions. This technique influences the model to be consistent on the
predictions of different input types which successively motivates the model to
be more adversarially robust.

###  3.2 Generalization Measures

Our study aims to investigate which type and characteristics of methods would
best correlate with better performance in generalization, in this case, zero-
shot cross-lingual transfer. We are especially interested in confirming the
applicability of the flatness hypothesis for cross-lingual generalization. In
order to assess whether a flat optimum loss-scape region corresponds to
generalization, we essentially break down the experiment to measure two
things, flatness, and generalization, such that their correlation can be
measured.

Jiang et al. (2020b) conducted an extensive study on image classification
tasks using generalization measures such as flatness-based measures (sharpness
metrics), margin and norm-based metrics (based on parameter norms and distance
from initial weights) to find correlations between measures and model
performance which supported the usability of measures. These measures can be
useful to explore the capabilities of language models to transfer knowledge
from high-resource languages to low-resource ones.

#### Margin

Higher certainty in predicting the correct label leads to a model that is
robust to perturbations and unseen examples. Margin is the distinction between
model prediction for ground truth label and the next highest prediction
probability. We use an average based margin formula defined by Wei et al.
(2018) to calculate margin values on the entire test set. Jiang et al. (2020b)
observed that the margin was directly proportional to better generalization in
the image classification tasks. Margin is

$\frac{1}{n}\sum_{i}^{n}\Big{(}f_{y_{i}}(x_{i})-\max_{j\neq
y_{i}}f_{j}(x_{i})\Big{)}$

where $x_{i}$ is the $i^{th}$ input to model, $y_{i}$ is the ground truth
label, $f(.)$ is the model function. A larger value of the margin of a model
on a given dataset would mean higher confidence in the model to predict the
correct label - including unseen examples (from languages not included in
fine-tuning).

#### Sharpness of optimum

In simpler terms, we can define sharpness as the change in the model loss
value at two neighboring points in the model weights plane. It can also be
loosely interpreted as the inverse of the maximum radius the loss function can
sustain a low loss value at the optimum. Sharpness-based measures resulted in
the highest correlation with generalization in Jiang et al. (2020b).

Jiang et al. (2020b) formulates the sharpness to be

$\phi=\frac{\|W-W_{0}\|_{2}^{2}\log(2\omega)}{4\alpha^{2}}+\log\frac{m}{\sigma}+10$

such that $\max_{|u_{i}|\leq\alpha}L(f_{W+u})<0.1$, where $\alpha$ is the
maximum radius in the model’s loss landscape possible, $W$ and $W_{0}$ are the
models finetuned weights and model initial weights respectively, $\omega$ is
the number of parameters, $m$ is the total number of observations, $\sigma$ is
the standard deviation of Gaussian noise added. In this work, as we are
comparing models with the same architecture (considering mBERT only), on the
same dataset, we can remove the constants, and simplify the equation further
for comparative analysis.

$\phi=\frac{\|W-W_{0}\|_{2}^{2}}{4\alpha^{2}}$

Intuitively, if the radius of the low-loss region in the loss-landscape
($\alpha$) is small, that means the model has a higher loss value near the
optimum, which would mean the landscape of the optimum is not flat. We can
relate this to resulting in an unstable prediction when having perturbations
in either the data or model weights. Jiang et al.’s formula didn’t result in
stable results for our experimental set-up which might be because the ascent
steps taken to optimize the $\alpha$ value resulted in either having a large
or a very small final $\alpha$. The values of $\alpha$ occurred at extreme
points because the algorithm was using a binary search method and whenever
optimal $\alpha$ was not found, the search algorithm stopped with the final
$\alpha$ value at either of the extreme points. The correlation results of the
above sharpness method are shown in Table 3.

We present an alternative definition (inclined with sharpness measure
mentioned in the works of Keskar et al., and Foret et al.),
$\phi_{\text{difference}}$ that removes the need to optimize $\alpha$ by
calculating the difference between loss values at two points in the optimum
region, formulated as

$\phi_{\text{difference}}=L(f_{W^{\prime}})-L(f_{W})$

where $W^{\prime}$ is $W+\epsilon$ ($\epsilon$ being Gaussian noise) and $W$
is the optimum weight parameters. The details of our definition are in
Algorithm 1 and performs calculation at about roughly 5-10 times faster than
Jiang et al.’s algorithm for a given batch size of 8.

Algorithm 1 Difference-based sharpness algorithm

1:$w_{0}=\text{original$\\_$weight}$

2:$w=w_{0}+\epsilon$ $\triangleright$ Small noise added to avoid zero gradient

3:$\Delta w=\nabla L(w)$

4:$w^{\prime}=w+n\Delta w$

5:$p=\lambda\times\|w^{\prime}\|_{F}$ $\triangleright$ $\lambda$ is small like
0.05

6:if $\|w^{\prime}-w_{0}\|>p$ then

7: $w^{\prime}=w_{0}+\frac{(w^{\prime}-w_{0})}{\|w^{\prime}-w_{0}\|}\cdot p$

8:end if

9:$\phi_{\text{difference}}=L(w^{\prime})-L(w_{0})$

##  4 Experiments

For comparison, we implement each Sharpness-based optimization as a fine-
tuning objective on the multilingual mBERT base variant (`bert-base-
multilingual-cased` from huggingface) Devlin et al. (2019) in addition to mT5
model (`google/mt5-small`) Xue et al. (2021). We use a linear classification
layer of size 768x3 where the output dimension is equal to the number of
labels. We adopt a two-step training approach in our experiments. First, we
fine-tune the model on the English language part of the XNLI dataset to
optimize the model to learn the task specifically in English. Subsequently, we
perform a zero-shot transfer of the fine-tuned model on the rest of the 14
languages to facilitate an evaluation of the generalization of models.

###  4.1 Data, Model details, and Settings

For this work, we used the XNLI dataset Conneau et al. (2018) that includes
data samples from the MultiNLI dataset Williams et al. (2018) and their
translated versions in 14 different languages (Arabic ”ar”, Bulgarian ”bg”,
German ”de”, Greek ”el”, Spanish ”es”, French ”fr”, Hindi ”hi”, Russian ”ru”,
Swahili ”sw”, Thai ”th”, Turkish ”tr”, Urdu ”ur”, Vietnamese ”vi”, Chinese
”zh”). We only train the models on the English (”en”) subset of the dataset.
We use the data of these 14 languages only for inference and evaluation of the
models.

We fine-tune pretrained mBERT models for 15 epochs each with a batch size of
32, with a learning rate of $2\times 10^{-5}$, and select best checkpoint on
validation. The objective function we use for the baseline model with the
classification layer is the AdamW optimizer Loshchilov & Hutter (2019) with
cross-entropy loss, the mBERT+FIM model has an additional loss as Fisher
Penalty, the mBERT+SAM model uses the SAM optimizer and mBERT+MVR uses the MVR
algorithm for fine-tuning. We use the hyperparameters and code presented in
XTREME222https://github.com/google-research/xtreme and MVR
codebase333https://github.com/cindyxinyiwang/multiview-subword-regularization.
We run the models with 8 random seeds and present the average performance of
these models (Figure 6). In Algorithm 1, the amount of Gaussian noise we add
to model weights during calculating sharpness is controlled using a scale that
we empirically find (among [0.001, 0.005, 0.01, 0.02]) for each model, with
$n$ equal to 0.05.

We fine-tuned the MT5 model (`google/mt5-small` using Huggingface’s library
over 15 epochs. The XNLI dataset was processed using a function to tokenize
inputs, and the optimizer utilized was Adafactor with a learning rate
scheduler. Adafactor optimizer’s ability to adapt learning rates is helpful
with larger models like T5 in multi-lingual settings. We trained the model
with a batch size of 8, accumulating gradients over 4 steps.

Additional experiments were run on PAWS-X dataset Yang et al. (2019) which has
7 languages: German ”de”, English ”en”, Spanish ”es”, French ”fr”, Japanese
”ja”, Korean ”ko”, Chinese ”zh”. We use similar experimentation of fine-tuning
on english and doing a zero-shot transfer on 6 other languages as defined
above for this dataset. We used Huggingface’s models: mBERT (`bert-base-
multilingual-cased`), RoBERTa (`roberta-base`), and XLM (`xlm-mlm-en-2048`)
using Adam optimizers.

### Results

To evaluate how each of the selected measures correlates with cross-lingual
generalization, we first compare these measures with held-out test accuracy.
In Table 1, we present the correlation coefficients (using `numpy.corrcoef`)
of margin vs. accuracy and sharpness vs. accuracy. We notice that having a
higher margin is exceptionally correlated to achieving great performance on
unseen language data. Hence, we assume the margin to indicate the
generalization performance of a given model. Similarly, sharpness captures a
noteworthy negative correlation with test performance.

Model | Correlation with Accuracy
---|---
| Margin | Sharpness
Baseline | 0.801 | -0.845
mBERT+MVR | 0.818 | -0.793
mBERT+SAM | 0.874 | -0.584
mBERT+FIM | 0.954 | -0.671
mT5 + Adafactor | 0.912 | -0.410
Table 1: Correlation coefficients between Margin & Test Accuracy, and
Sharpness & Test Accuracy on the XNLI dataset.

We notice similar results by extending our similar experimentation to
Paraphrase Identification, PAWS-X dataset Yang et al. (2019) with 3 different
models: mBERT (`bert-base-multilingual-cased`), RoBERTa (`roberta-base`) Liu
et al. (2019), and XLM (`xlm-mlm-en-2048`) CONNEAU & Lample (2019) and analyze
the validity of the flatness hypothesis, i.e. a flat optimum neighborhood
would lead to a generalized model. In Figure 1, we confirm the strong
relationship between Margin (indicating generalization) and Sharpness
(indicating flatness) even when compared across all models and metrics,
suggesting flat neighborhoods of model optimum can help in achieving higher
margin values which correlate to better generalization. More findings about
visualizations are in Appendix A.1.

Model | Correlation with Accuracy
---|---
| Margin | Sharpness
mBERT | 0.998 | -0.289
RoBERTa | 0.997 | -0.708
XLM-R | 0.995 | -0.622
Table 2: Correlation coefficients between Margin and Sharpness with Test
Accuracy on the PAWS-X dataset. Figure 1: Scatter plot of Margin values and
Sharpness ($\phi_{\text{difference}}$) values for each mBERT model (on XNLI
dataset) with different objectives language-wise to show the relationship
between sharpness and generalization. Figure 2: Scatter plot of difference-
based sharpness measure with test performance for all models combined.

We can interpret sharpness as the inverse of flatness, providing us the
verdict that flatness of the minimum in which the fine-tuned model is, would
help the model perform better on unseen language data. When we evaluate
similar models trained with different objectives across languages, we observe
that the relationships between measures are likely dependent on the
optimization objective functions used during fine-tuning. In coherence with
both Figure 1 and 2, overall, we see that min-max based optimization methods
including FIM and SAM, have the lowest sharpness values, compared to the
baseline and the regularization method MVR.

##  5 Conclusion

Enabling cross-lingual knowledge transfer is an important step towards
extending the applicability of NLP models to more languages. Despite recent
efforts to develop better optimization methods for improving the
generalization of language models in new languages or domains; these
techniques try different types of methods to achieve higher performance such
as sharpness-based minimizations, reducing gradient of loss functions, or
consistency regularization. Evaluating these techniques thoroughly without a
standardized methodology remains a difficult task. This work aims to uncover
insights into how to measure cross-lingual generalization by exploring
suitable measures that work well under different settings. Our experiments
studying model loss landscape and parameter properties find strong
relationships between the margin, sharpness in the loss minima neighborhood,
and zero-shot cross-lingual downstream task performance, both on validation
and test sets, supporting strong applicability to evaluate models before
deploying them in new languages.

##  6 Limitations

The algorithm presented in our paper, the difference-based sharpness measure,
is a great novelty for more robust sharpness computation, however, we would
like to acknowledge that a few variables in the algorithm still require tuning
heuristically, including the noise scale and the multiplication coefficient
required to compute the projected radius. Secondly, the mean-based margin
distance is only applicable to classification tasks. Due to the limited scope
of this project, we leave the development of generalization measures more
suitable for generative tasks to future work.

## References

  * Aghajanyan et al. (2021) Armen Aghajanyan, Akshat Shrivastava, Anchit Gupta, Naman Goyal, Luke Zettlemoyer, and Sonal Gupta.  Better fine-tuning by reducing representational collapse.  In _International Conference on Learning Representations_ , 2021.
  * Bottou (2012) Léon Bottou.  Stochastic gradient descent tricks.  _Neural Networks: Tricks of the Trade: Second Edition_ , pp. 421–436, 2012.
  * Cha et al. (2021) Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park.  Swad: Domain generalization by seeking flat minima.  _Advances in Neural Information Processing Systems_ , 34:22405–22418, 2021.
  * Chaudhari et al. (2019) Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina.  Entropy-sgd: Biasing gradient descent into wide valleys.  _Journal of Statistical Mechanics: Theory and Experiment_ , (12):124018, 2019.
  * CONNEAU & Lample (2019) Alexis CONNEAU and Guillaume Lample.  Cross-lingual language model pretraining.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.  URL https://proceedings.neurips.cc/paper_files/paper/2019/file/c04c19c2c2474dbf5f7ac4372c5b9af1-Paper.pdf.
  * Conneau et al. (2018) Alexis Conneau, Ruty Rinott, Guillaume Lample, Adina Williams, Samuel R. Bowman, Holger Schwenk, and Veselin Stoyanov.  Xnli: Evaluating cross-lingual sentence representations.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_. Association for Computational Linguistics, 2018.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pp. 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.  doi: 10.18653/v1/N19-1423.
  * Dinh et al. (2017) Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio.  Sharp minima can generalize for deep nets.  In _International Conference on Machine Learning_ , pp. 1019–1028. PMLR, 2017.
  * Englesson & Azizpour (2021) Erik Englesson and Hossein Azizpour.  Generalized jensen-shannon divergence loss for learning with noisy labels.  _Advances in Neural Information Processing Systems_ , 34:30284–30297, 2021.
  * Fisher (1925) Ronald Aylmer Fisher.  Theory of statistical estimation.  In _Mathematical proceedings of the Cambridge philosophical society_ , volume 22, pp. 700–725. Cambridge University Press, 1925.
  * Foret et al. (2021) Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur.  Sharpness-aware minimization for efficiently improving generalization.  In _International Conference on Learning Representations_ , 2021.
  * Hu et al. (2022) Edward J Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al.  Lora: Low-rank adaptation of large language models.  In _International Conference on Learning Representations_ , 2022.
  * Hua et al. (2021) Hang Hua, Xingjian Li, Dejing Dou, Chengzhong Xu, and Jiebo Luo.  Noise stability regularization for improving bert fine-tuning.  In _Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pp. 3229–3241, 2021.
  * Huang et al. (2021) Kuan-Hao Huang, Wasi Ahmad, Nanyun Peng, and Kai-Wei Chang.  Improving zero-shot cross-lingual transfer learning via robust training.  In Marie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih (eds.), _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ , pp. 1684–1697, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.  doi: 10.18653/v1/2021.emnlp-main.126.  URL https://aclanthology.org/2021.emnlp-main.126.
  * Izmailov et al. (2018) P Izmailov, AG Wilson, D Podoprikhin, D Vetrov, and T Garipov.  Averaging weights leads to wider optima and better generalization.  In _34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018_ , pp. 876–885, 2018.
  * Jastrzebski et al. (2021) Stanislaw Jastrzebski, Devansh Arpit, Oliver Astrand, Giancarlo B Kerg, Huan Wang, Caiming Xiong, Richard Socher, Kyunghyun Cho, and Krzysztof J Geras.  Catastrophic fisher explosion: Early phase fisher matrix impacts generalization.  In Marina Meila and Tong Zhang (eds.), _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pp. 4772–4784. PMLR, 18–24 Jul 2021.
  * Jiang et al. (2020a) Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Tuo Zhao.  Smart: Robust and efficient fine-tuning for pre-trained natural language models through principled regularized optimization.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pp. 2177–2190, 2020a.
  * Jiang et al. (2020b) Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio.  Fantastic generalization measures and where to find them.  In _International Conference on Learning Representations_ , 2020b.
  * Kaddour et al. (2022) Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt J Kusner.  When do flat minima optimizers work?  _Advances in Neural Information Processing Systems_ , 35:16577–16595, 2022.
  * Keskar et al. (2017) Nitish Shirish Keskar, Jorge Nocedal, Ping Tak Peter Tang, Dheevatsa Mudigere, and Mikhail Smelyanskiy.  On large-batch training for deep learning: Generalization gap and sharp minima.  In _5th International Conference on Learning Representations, ICLR 2017_ , 2017.
  * Li et al. (2018) Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein.  Visualizing the loss landscape of neural nets.  _Advances in neural information processing systems_ , 31, 2018.
  * Liang et al. (2021) Lijun Liang, Juntao Li, Yue Wang, Qi Meng, Tao Qin, Wei Chen, Min Zhang, Tie-Yan Liu, et al.  R-drop: Regularized dropout for neural networks.  _Advances in Neural Information Processing Systems_ , 34:10890–10905, 2021.
  * Liu et al. (2019) Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov.  Roberta: A robustly optimized BERT pretraining approach.  _CoRR_ , abs/1907.11692, 2019.  URL http://arxiv.org/abs/1907.11692.
  * Loshchilov & Hutter (2019) Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_ , 2019.
  * Mohri et al. (2012) Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar.  _Foundations of Machine Learning_.  The MIT Press, 2012.  ISBN 026201825X.
  * Nachum et al. (2024) Ido Nachum, Jonathan Shafer, Thomas Weinberger, and Michael Gastpar.  Fantastic generalization measures are nowhere to be found.  In _The Twelfth International Conference on Learning Representations_ , 2024.  URL https://openreview.net/forum?id=NkmJotfL42.
  * Nagarajan & Kolter (2019) Vaishnavh Nagarajan and J. Zico Kolter.  Generalization in deep networks: The role of distance from initialization.  _arXiv:1901.01672 [cs.LG]_ , 2019.
  * Park et al. (2022) Jungsoo Park, Gyuwan Kim, and Jaewoo Kang.  Consistency training with virtual adversarial discrete perturbation.  In _Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pp. 5646–5656, 2022.
  * Petzka et al. (2021) Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley.  Relative flatness and generalization.  _Advances in neural information processing systems_ , 34:18420–18432, 2021.
  * Polyak & Juditsky (1992) Boris T Polyak and Anatoli B Juditsky.  Acceleration of stochastic approximation by averaging.  _SIAM journal on control and optimization_ , 30(4):838–855, 1992.
  * Shannon (1948) Claude E Shannon.  A mathematical theory of communication.  _The Bell system technical journal_ , 27(3):379–423, 1948.
  * Stickland & Murray (2021) Asa Cooper Stickland and Iain Murray.  Regularising fisher information improves cross-lingual generalisation.  In _Proceedings of the 1st Workshop on Multilingual Representation Learning_ , pp. 238–241, 2021.
  * Vapnik & Chervonenkis (1971) V. N. Vapnik and A. Ya. Chervonenkis.  On the uniform convergence of relative frequencies of events to their probabilities.  _Theory of Probability & Its Applications_, 16(2):264–280, 1971.
  * Wang et al. (2021) Xinyi Wang, Sebastian Ruder, and Graham Neubig.  Multi-view subword regularization.  In _Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pp. 473–482, 2021.
  * Wei et al. (2018) Colin Wei, Jason D. Lee, Qiang Liu, and Tengyu Ma.  On the margin theory of feedforward neural networks.  _arXiv:1810.05369 [stat.ML]_ , 2018.
  * Williams et al. (2018) Adina Williams, Nikita Nangia, and Samuel Bowman.  A broad-coverage challenge corpus for sentence understanding through inference.  In _Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers)_ , pp. 1112–1122. Association for Computational Linguistics, 2018.
  * Xue et al. (2021) Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami Al-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel.  mT5: A massively multilingual pre-trained text-to-text transformer.  In Kristina Toutanova, Anna Rumshisky, Luke Zettlemoyer, Dilek Hakkani-Tur, Iz Beltagy, Steven Bethard, Ryan Cotterell, Tanmoy Chakraborty, and Yichao Zhou (eds.), _Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pp. 483–498, Online, June 2021. Association for Computational Linguistics.  doi: 10.18653/v1/2021.naacl-main.41.  URL https://aclanthology.org/2021.naacl-main.41.
  * Yang et al. (2019) Yinfei Yang, Yuan Zhang, Chris Tar, and Jason Baldridge.  PAWS-X: A cross-lingual adversarial dataset for paraphrase identification.  In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan (eds.), _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_ , pp. 3687–3692, Hong Kong, China, November 2019. Association for Computational Linguistics.  doi: 10.18653/v1/D19-1382.  URL https://aclanthology.org/D19-1382.
  * Zaken et al. (2022) Elad Ben Zaken, Yoav Goldberg, and Shauli Ravfogel.  Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers)_ , pp. 1–9, 2022.
  * Zheng et al. (2021) Bo Zheng, Li Dong, Shaohan Huang, Wenhui Wang, Zewen Chi, Saksham Singhal, Wanxiang Che, Ting Liu, Xia Song, and Furu Wei.  Consistency regularization for cross-lingual fine-tuning.  In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_ , pp. 3403–3417, 2021.

##  Appendix A Appendix

###  A.1 Visualization results

In Figure 3, we create scatter plots for mBERT models where in each scatter
plot, we plot the model’s margin based on the validation set for each
language, and we plot the accuracy of that model on the test set on the XNLI
dataset. We observe that the margin measure exhibits a consistent correlation
with test performance across all the models analyzed.

(a) Correlation visualization for Baseline(mBERT + AdamW)

(b) Correlation visualization for mBERT + MVR model

(c) Correlation visualization for mBERT + SAM model

(d) Correlation visualization for mBERT + FIM model

Figure 3: Scatter plots of margin of individual models and their corresponding
performance on test set language-wise on XNLI dataset.

As can be seen in the scatter plots for sharpness (proposed difference-based
sharpness) and accuracy in Figure 4, findings further indicate a negative
correlation between sharpness and test performance, suggesting that lower
sharpness values are associated with better generalization, represented as
model performance on unseen data.

Previous work Nagarajan & Kolter (2019); Jiang et al. (2020b) suggests that a
lower Frobenius distance from initialization would lead to better
generalization. As Figure 5 shows, we fail to observe a strong direct
relationship between generalization and Frobenius distance from
initialization. However, the model trained with Fisher Penalty as an
additional objective function that has a high distance from initialization
overall performed poorly than others. We also see that models trained with
Fisher Penalty, SAM, and MVR optimizers tend to be more stable than the
baseline model, with Fisher Penalty resulting in the most stable model when
trained multiple times (with different seeds, see Figure 6), and SAM achieving
generally the best average zero-shot task accuracy.

(a) Correlation visualization for Baseline (mBERT + AdamW)

(b) Correlation visualization for mBERT + MVR model

(c) Correlation visualization for mBERT + SAM model

(d) Correlation visualization for mBERT + FIM model

Figure 4: Scatter plots of the proposed difference-based sharpness
($\phi_{\text{difference}}$) of individual models and their corresponding
performance on test set language-wise on XNLI Dataset. Figure 5: Scatter plot
of Frobenius distance from initialization and Test accuracy for each model
type (trained multiple times independenty). Figure 6: Average test performance
(and deviations) of models when trained multiple times with different seeds.

###  A.2 Additional experiments

To compare all models together, using the difference-based sharpness measure,
on language-wise performance, we observe it is dependent on the learning
algorithms used during training in Figure 2.

Model | Correlation coefficient of $\alpha$ sharpness with accuracy
---|---
Baseline | 0.249
mBERT+MVR | -0.471
mBERT+SAM | -0.166
mBERT+FIM | -0.440
Table 3: Correlation coefficients between $\alpha$-Sharpness Jiang et al.
(2020b) & Test Accuracy on the XNLI dataset.

(a) Correlation visualization for Baseline (mBERT + AdamW)

(b) Correlation visualization for mBERT + MVR model

(c) Correlation visualization for mBERT + SAM model

(d) Correlation visualization for mBERT + FIM model

Figure 7: Scatter plots of Jiang et al. (2020b) based $\alpha$-sharpness
measure (we are only considering $\frac{1}{\alpha^{2}}$ here) of individual
models and their corresponding performance on test set language-wise.

We used the Jiang et al.’s $\alpha$-based sharpness algorithm for the
experiment and optimized the threshold loss values for our experimental
setting. The results of the correlation coefficient (using `numpy.corrcoef`)
for $\alpha$-based sharpness and test accuracy are shown in Table 3 and Figure
7. We notice that $\alpha$-based sharpness values occur at extreme points (for
example, for mBERT+FIM model, sharpness values are low whereas for the
Baseline or mBERT+MVR model, sharpness values are much larger). Apart from
being a computationally expensive algorithm, we failed to see a strong
relationship of $\alpha$-based sharpness with performance in Baseline and
mBERT+SAM models.