# One-Shot Imitation Learning with Invariance Matching for Robotic
Manipulation

Xinyu Zhang and Abdeslam Boularias
Rutgers University
Email: {xz653<EMAIL_ADDRESS>
###### Abstract

Learning a single universal policy that can perform a diverse set of
manipulation tasks is a promising new direction in robotics. However, existing
techniques are limited to learning policies that can only perform tasks that
are encountered during training, and require a large number of demonstrations
to learn new tasks. Humans, on the other hand, often can learn a new task from
a single unannotated demonstration. In this work, we propose the Invariance-
Matching One-shot Policy Learning (IMOP) algorithm. In contrast to the
standard practice of learning the end-effector’s pose directly, IMOP first
learns invariant regions of the state space for a given task, and then
computes the end-effector’s pose through matching the invariant regions
between demonstrations and test scenes. Trained on the 18 RLBench tasks, IMOP
achieves a success rate that outperforms the state-of-the-art consistently, by
$4.5\%$ on average over the 18 tasks. More importantly, IMOP can learn a novel
task from a single unannotated demonstration, and without any fine-tuning, and
achieves an average success rate improvement of $11.5\%$ over the state-of-
the-art on 22 novel tasks selected across nine categories. IMOP can also
generalize to new shapes and learn to manipulate objects that are different
from those in the demonstration. Further, IMOP can perform one-shot sim-to-
real transfer using a single real-robot demonstration.

††footnotetext: Project Website: https://mlzxy.github.io/imop/

##  I Introduction

Multi-Task Learning (MTL) is a recent type of learning technique where a
single policy is trained to perform various tasks. With the recent advances in
computer vision architectures [18], a popular MTL strategy in robotics is to
acquire a multi-task control policy through imitation learning from visual
demonstrations [14, 15, 33, 36]. These methods achieved impressive results in
performing challenging manipulation tasks in unstructured 3D environments,
such as RLBench [22]. However, these methods rely on the assumption that
training and testing share the same set of tasks. Moreover, training such
policies on new tasks requires hundreds of demonstrations [15], which results
in the catastrophic forgetting of old tasks [24].

To address these issues, one-shot imitation learning aims to learn on a set of
base tasks and generalize to novel tasks given only a single demonstration for
each novel task and without re-training. However, existing methods rely on
strong assumptions, such as requiring the novel tasks to be limited variations
of the same base tasks [37], requiring the base and novel tasks to have the
same object setup [9], only generalizing specific actions on certain
categories of objects with known 3D models [2], or operating in simple 2D
planar environments [11]. Moreover, most recent works focus on applying
general popular techniques, e.g., transformers and diffusion models, to one-
shot imitation settings, without taking advantage of the particular structure
of robotic manipulation tasks [9, 28, 37, 38].

Therefore, a key question here is: can robots learn a manipulation policy that
not only performs well on base tasks but also generalizes to novel unseen
tasks using a single demonstration and without any fine-tuning? To this end,
we propose IMOP (Invariance Matching One-shot Policy Learning), a new
algorithm that not only outperforms the state-of-the-art on the standard 18
tasks of RLBench ($69.6\%$ average overall success rate compared to $65.1\%$
for the state-of-the-art [14]), but also generalizes to 22 novel tasks with a
single demonstration and without any fine-tuning ($41.3\%$ average success
rate compared to $29.8\%$ for the state-of-the-art). The 22 novel tasks are
selected across nine categories and are substantially different from the base
tasks in terms of objects and task goals [17, 22]. Moreover, we find that IMOP
also generalizes to substantial shape variations, and can manipulate new
objects that are different from those in the demonstration.

Instead of learning the desired end-effector’s pose directly, IMOP learns key
invariant regions of each task, and finds pairwise correspondences between the
invariant regions in a one-shot demonstration and in a given test scene. The
pairwise correspondences are used to analytically compute the desired end-
effector’s pose in the test scene from the least-squares solution of a point-
set registration problem. The invariant region is defined as a set of 3D
points whose coordinates remain invariant when viewed in the end-effector’s
frame, across states that share the same semantic action. We devise a graph-
based invariant region matching network. The invariant regions are located
through neighbor attention [42] from the KNN graphs that connect point clouds
of demonstration and test scenes. The ground-truth invariant regions are
discovered in offline unannotated demonstrations of base tasks.

To summarize, our contributions are threefold. (1) We propose IMOP, a one-shot
imitation learning algorithm for robotic manipulation that learns a universal
policy that is not only successful on base tasks, but also generalizes to
novel tasks using a single unannotated demonstration. (2) We propose a
correspondence-based pose regression method for manipulation tasks, which
predicts robot actions by matching key visual elements, and a graph-based
invariant region matching network on KNN graphs that connect demonstrations
and test scenes. (3) We present a thorough empirical study of the performance
and generalization ability of IMOP on a diverse set of tasks.

##  II Related Work

Learning robotic manipulation from demonstrations. Learning manipulation
policies from offline visual demonstrations has gained increasing attention
following the rapid improvement of vision models [10]. Prior works such as
Transporter networks [40] and CLIPort [32] learned simple pick-and-place tasks
in a 2D planar setting. C2F-ARM [23] and PerAct [33] extend the control
capacity to the 3D space with 6-DoF but learn a separate policy for each
single task. More recently, RVT [15], Hiveformer [17], Act3D [14], and
ChainedDiffuser [36] learned a multi-task policy from demonstrations of
multiple tasks. Yet, these prior approaches can only perform tasks that were
seen during training and cannot generalize to novel tasks upon inference.
Training these models also requires hundreds of demonstrations per task.
Instead, this work aims to develop a multi-task policy that not only performs
well on seen tasks but also generalizes to novel tasks given a single
demonstration.

One-shot imitation learning.  Traditional imitation learning considers
learning a policy for a single task with many expert trajectories provided for
this particular task. One-shot imitation learning aims to learn from
demonstrations of a set of base tasks and generalize to novel tasks using a
single trajectory per new task without further training. Duan et al. [11]
trained a one-shot imitator that generalizes to simple task variations in a 2D
stacking environment. Finn et al. [12] used the same setting in the context of
meta-learning. Dasari and Gupta [9], Xu et al. [37] and Mandi et al. [28]
trained transformer-based one-shot imitators, such as encoding expert
trajectories as prompts. Biza et al. [2] achieved one-shot generalization for
grasping tasks where 3D models of the objects are available. Xu et al. [38]
trained adaptation layers through hyper-networks. Yet, none of these prior
works considered 6D manipulation tasks in the context of one-shot learning.
Moreover, recent works focus on applying existing techniques, such as
transformers and adaptation layers, to one-shot imitation settings, but did
not investigate the inherent task structure of robotic manipulation. In
contrast, this work proposes a one-shot imitation learner for 6D manipulation
by discovering invariant regions of each task.

Invariance and affordance in manipulation.  Incorporating invariance into deep
learning models has been shown to drastically increase data efficiency and
generalization [3, 5, 6, 27, 34]. Graf et al. [16] design image augmentations
to learn viewpoint-invariant features for manipulation. Visual affordance is a
region representation of action possibility [1, 30]. For example, Where2Act
learns a probability map for pushing and pulling at each pixel [31]. In this
work, we propose the concept of invariant region. Instead of camera viewpoint
invariance, we train neural networks to predict regions whose positions remain
invariant to the robot end-effector for a given task. Unlike affordance, the
proposed invariant region is not used to represent action probability but to
transfer actions from demonstrations to test scenes.

##  III One-shot Imitation with Invariance Matching

Figure 1: Example of a pick-and-pour task executed by a real Kuka robot after
observing a single demonstration and without any re-training. The
correspondences between estimated invariant regions of the 3D point cloud in
the demonstration and in the test scene are visualized in red lines. The
invariant regions are predicted and matched by a neural network, trained
offline.

###  III-A A Motivating Example

Figure 1 illustrates an example from our experiments where a Kuka robot is
tasked with picking up a cup and pouring it into a bowl, using a single
demonstration of picking up and pouring a different cup in a different
position. At the core of IMOP lies the capacity to estimate and match
invariant regions of the given task. By finding correspondences of invariant
regions in the demonstration and in the test scene, the demonstrated actions
can be transferred to the test scene. In this example, the invariant regions
are: (1) a set of 3D points on the surfaces of the cups where contact with the
fingertips occurs, and (2), a set of 3D points from the bowl’s point cloud
that capture its spherical concave shape. The invariant regions of the 3D
point cloud are learned offline from unannotated demonstrations, on various
objects and tasks.

###  III-B Formulation of One-shot Manipulation Learning Problem

Figure 2: Examples of tasks, and object-level generalization in various
trials. After training on base tasks, IMOP is evaluated on novel tasks that
are substantially different from the base tasks. Every learned task is
evaluated in multiple trials, each with different object layouts and
orientations. For each novel task, only one recorded trajectory is given as a
demonstration. We summarize the performance of IMOP on both base and novel
tasks in Section IV-A and IV-B. Furthermore, we study the ability of IMOP to
generalize to new objects beyond those in the demonstrations in Section IV-C.

We use $\mathbf{T}$ to denote a set of manipulation tasks,
$\mathbf{T}=\\{\textrm{Task}_{1},\textrm{Task}_{2},\dots,\textrm{Task}_{N_{\textrm{tasks}}}\\}$.
Each manipulation task $\textrm{Task}_{i}\in\mathbf{T}$ is performed in a
Markov Decision Process (MDP) represented by 5-tuple
$\mathcal{M}_{i}=(\mathcal{S},\mathcal{A},\mathbf{P},R_{i},\mu_{i})$, wherein
$\mathcal{S}$ and $\mathcal{A}$ are state and action spaces,
$\mathbf{P}:\mathcal{S}\times\mathcal{A}\mapsto\mathcal{S}$ is a state
transition probability, $R_{i}:\mathcal{S}\mapsto\mathbb{R}$ denotes a reward
function associated with task $\textrm{Task}_{i}$, and $\mu_{i}$ is the
initial state distribution of $\textrm{Task}_{i}$. $\mathbf{T}$ is composed of
a set of base tasks, denoted by $\mathbf{T}^{base}$, and a set of novel tasks,
denoted by $\mathbf{T}^{novel}$. Thus,
$\mathbf{T}=\mathbf{T}^{base}\cup\mathbf{T}^{novel}$. During training, a large
collection of offline demonstrations is provided for the base tasks. The
underlying reward function of each task is unknown and is not used in this
work. During testing, only one successful trajectory
$\tau_{i}=\\{(s_{i},a_{i},s_{i}^{\prime})\\}^{|\tau_{i}|}$ is provided to the
robot for a novel $\textrm{Task}_{i}\in\mathbf{T}^{novel}$, where $|\tau_{i}|$
denotes the number of transitions, $a_{i}$ is the demonstrated action in state
$s_{i}$, and $s{{}_{i}}^{\prime}$ is the resulting next state. The goal is to
learn a one-shot universal policy $\pi(s,\tau_{i})$, which is a multi-task
policy that returns an action $a$ in new state $s$ given a single demonstrated
trajectory $\tau_{i}$ of the desired novel skill. $\tau_{i}$ is provided in an
environment that is different from the one in which $\pi$ is tested.

A state $s\in\mathcal{S}$ is a raw 3D point cloud of the entire scene,
including the robot. Each point in state $s$ is represented in homogeneous
coordinates. An action $a\in\mathcal{A}$ is the desired (i.e., target) 6D pose
of the robot’s end-effector, along with the desired state of the gripper. To
execute an action, low-level robot movements are generated using off-the-shelf
motion planners. More precisely, we define actions as 18-dimensional vectors
$a=(T,\lambda)$, wherein $T\in\mathbb{R}^{4\times 4}$ denotes the target end-
effector pose as a homogeneous matrix and is referred to as action pose, and
$\lambda$ denotes two binary values that indicate whether gripper is open and
collisions are allowed or not.

Figure 3: Overview of the proposed invariant region matching network and
correspondence-based pose regression. Given the observed transition of the
support state $s_{i}$, we build KNN graphs from the scene point cloud and
apply graph self-attention (within $s_{i}$ and $s_{i}^{\prime}$) and cross-
attention (between $s_{i}$ and $s_{i}^{\prime}$) layers. The invariant region
$\mathcal{I}(s_{i})$ is predicted as the set of activated points through a
point-wise sigmoid over the support state $s_{i}$. Next, we apply graph cross-
attention layers between the KNN graphs of $\mathcal{I}(s_{i})$ and query
state $s_{j}$. The correspondence matrix $C$ is predicted by matching the
point-wise features $h_{\mathcal{I}(s_{i})}$ and $h_{s_{j}}$. The action pose
$T_{j}$ of query state $s_{j}$ is analytically solved from the action pose
$T_{i}$ of support state $s_{i}$, correspondence matrix $C$,
$P_{\mathcal{I}(s_{i})}$ and $P_{s_{j}}$ (points in $\mathcal{I}(s_{i})$ and
$s_{j}$), as detailed in Section III-C2.

###  III-C Proposed Method

####  III-C1 Invariant Region Matching Network

We first consider the following problem. Given a task
$\text{Task}_{k}\in\mathbf{T}$, an observed transition
$(s_{i},a_{i},s_{i}^{\prime})$ in a demonstrated trajectory, and a new state
$s_{j}$ such that $s_{i}\equiv s_{j}$, how can action pose $T_{i}$ in state
$s_{i}$ be translated to a new action pose $T_{j}$ to execute in the new state
$s_{j}$? We use $s_{i}\equiv s_{j}$ to indicate that $s_{i}$ and $s_{j}$ share
the same optimal manipulation action. For example, the first manipulation
action of the task ‘open door’ is ‘reach the door handle’. Therefore, states
that describe different scenes with closed doors all share the optimal action
‘reach the door handle’. We refer to $s_{i}$ as the support state and to the
new state $s_{j}$ as the query state.

To answer this question, we define the notion of invariant regions in the
following.

######  Definition III.1.

The invariant region of state $s_{i}$ in $\textrm{Task}_{k}$ is defined as
$\mathcal{I}(s_{i}|\textrm{Task}_{k})=\\{p\in s_{i}|\forall
s_{j}\in\mathcal{S}\textrm{ s.t. }s_{i}\equiv s_{j},\exists q\in
s_{j}:\parallel T^{-1}_{i}p-T^{-1}_{j}q\parallel<\epsilon\\}$, where
$(T_{i},\lambda_{i})=\pi^{*}(s_{i}|\textrm{Task}_{k})$ and
$(T_{j},\lambda_{j})=\pi^{*}(s_{j}|\textrm{Task}_{k})$.

In this definition, $s_{i}$ is a scene point-cloud, and
$\mathcal{I}(s_{i}|\text{Task}_{k})$ is the set of all the 3D points of
$s_{i}$ whose coordinates in frame $T_{i}$ remain invariant across all the
scenes $s_{j}$ that share the same optimal action as the optimal action of
$s_{i}$. Frame $T_{i}$ is the frame of the end-effector’s pre-contact pose
when performing the optimal manipulation action in scene $s_{i}$. For example,
for the task of picking up a cup, the optimal action is to grasp the cup’s
handle, regardless of the location, orientation, or style of the cup. The
learned invariant region will likely be the cup’s handle. To simplify the
notation, we drop $\textrm{Task}_{k}$ and simply refer to the invariant region
as $\mathcal{I}(s_{i})$, while it is implied that $\mathcal{I}(s_{i})$ is the
invariant region of $s_{i}$ for a given $\textrm{Task}_{k}$.

The optimal action pose $T_{j}$ of the query scene $s_{j}$ can be computed by
first predicting the invariant regions $\mathcal{I}(s_{i})$ and
$\mathcal{I}(s_{j})$ of $s_{i}$ and of support scene $s_{j}$, and then
transforming $T_{i}$ according to a matching between $\mathcal{I}(s_{i})$ and
$\mathcal{I}(s_{j})$. The overall framework of our proposed invariant region
prediction and matching network is shown in Figure 3. We first build a KNN
graph for each scene point-cloud by connecting each point to its $k$ nearest
points within the same scene. Next, we apply graph self-attention within each
support scene $s_{i}$, and cross-attention between the KNN graphs of pairs of
consecutive frames $s_{i}$ and $s_{i}^{\prime}$ within the same support
demonstration. In contrast to traditional attention that is computed globally
over all points, graph attention operates within the local neighborhood of
each given point. We use the point transformer layer [35] as the graph
attention operator. The invariant region $\mathcal{I}(s_{i})$ is predicted as
the set of activated points through a point-wise sigmoid over $s_{i}$. Then,
we apply graph cross-attention layers between the KNN graph of
$\mathcal{I}(s_{i})$ and the query state $s_{j}$ to extract the point-wise
features $h_{\mathcal{I}(s_{i})}\in\mathbb{R}^{|\mathcal{I}(s_{i})|\times D}$
and $h_{s_{j}}\in\mathbb{R}^{|s_{j}|\times D}$, where $D$ denotes the size of
feature dimension. Finally, we perform a dual softmax matching [25] between
$h_{\mathcal{I}(s_{i})}$ and $h_{s_{j}}$ to obtain the correspondence matrix
$C\in{[0,1]}^{|\mathcal{I}(s_{i})|\times|s_{j}|}$:

$C=\operatorname{softmax}(h_{\mathcal{I}(s_{i})}\cdot
h_{s_{j}}^{\top})\cdot\operatorname{softmax}(h_{s_{j}}\cdot
h_{\mathcal{I}(s_{i})}^{\top})^{\top}$

where $\operatorname{softmax}$ is applied on each row. The correspondence
matrix $C$ maps each point in the invariant region $\mathcal{I}(s_{i})$ to the
query state $s_{j}$. The matched points in $s_{j}$ constitute
$\mathcal{I}(s_{j})$ but we only need $C$ to predict the action pose $T_{j}$
of the query state $s_{j}$, as explained in Section III-C2. We detail the
training of invariant region matching network in Section III-C4.

####  III-C2 Correspondence-based Pose Regression

The standard practice of 6-DoF pose regression is to obtain the action pose
$T$ from a neural network. However, this approach does not generalize well to
new tasks, as shown in Section IV-B. Instead, we propose to analytically
compute action pose $T_{j}$ of the query state $s_{j}$ by solving the
optimization problem in Equation 1 with a standard least-squares algorithm
[19], as follows,

$T_{j}=\arg\min_{T\in\textrm{SE(3)}}\parallel
TT^{-1}_{i}P_{\mathcal{I}(s_{i})}C-P_{s_{j}}\parallel,$ (1)

where $T_{i}$ is the demonstrated action pose of the support state $s_{i}$,
$P_{\mathcal{I}(s_{i})}$ and $P_{s_{j}}$ are the points in
$\mathcal{I}(s_{i})$ and $s_{j}$, $C$ is the predicted correspondence matrix.
$C$ can be interpreted as an assignment matrix that maps each point in
$\mathcal{I}(s_{i})$ to a point in $s_{j}$. Based on Definition III.1, the
optimal action pose $T_{j}$ is the solution for Equation 1 when the point
mapping in correspondence matrix $C$ produces the minimal overall pairwise
distance after applying the transformation $T_{j}T_{i}^{-1}$. Therefore, the
pose regression problem can be solved by learning to match visual elements
with Equation 1. We use the differentiable Procrustes operator from Lepard
[25] to solve the least-squares problem. As shown in Section IV-E, our
correspondence-based pose regression significantly improves the generalization
performance.

It is worth noting that our correspondence-based pose regression shows a
resemblance to the point cloud registration (PCR) problem [20]. However, there
are three major differences. First, PCR assumes the point clouds are captured
from the same scene with different camera viewpoints, while our method aligns
point clouds from different scenes and objects. Second, PCR aligns the entire
or the majority of two point clouds, while our method finds a sparser matching
between only the invariant regions of some given task. Third, PCR finds the
best matching, while our method uses the matching as an intermediate step to
find the optimal pose of a robot’s end-effector in a new scene.

####  III-C3 State Routing Network

We design the state routing network (shown in Figure 4) to select the support
frame $s_{i}$ in the one-shot demonstration $\tau$, given a query scene
$s_{j}$. We first extract the scene-level features for the query state $s_{j}$
and for each state in $\tau$ using a PTv2 backbone [8]. Next, we follow the
convention of existing work [15] to concatenate the scene-level features with
the low-dimensional internal robot state, including joint positions and
timesteps. Then, we apply cross-attention over the features of multiple
states. Unlike the invariant region matching network, classic attention layers
are used instead of graph attention. Finally, we select the state with the
strongest attention to the query state $s_{j}$ as the support state $s_{i}$.

Figure 4: Overview of the proposed state routing network. Given the query
state $s_{j}$ and one-shot demonstration trajectory $\tau$, we extract the
scene-level features for $s_{j}$ and each state in $\tau$. The scene features
are concatenated with the corresponding low-dimensional robot states, and
applied with cross-attention layers. The state with the strongest attention to
the query state $s_{j}$ is selected as the support state $s_{i}$.

The three techniques presented above form together the Invariance Matching
One-shot Policy Learning (IMOP) algorithm. For each new query state $s_{j}$,
the support state $s_{i}$ is determined from the one-shot demonstration
trajectory $\tau$ by the state routing network. Next, the invariant region
matching network predicts the correspondence matrix $C$ between
$\mathcal{I}(s_{i})$ and $s_{j}$. Finally, the action pose $T_{j}$ is derived
from Equation 1. We directly use the gripper values $\lambda_{i}$ of support
state $s_{i}$ as the value of $\lambda_{j}$ for $s_{j}$. By adapting the
demonstration trajectory $\tau$, IMOP achieves one-shot generalization on
novel tasks.

####  III-C4 Training

To train IMOP, we assume the instance segmentation results are available. We
use the segmentation masks provided by RLBench [22]. Note that IMOP does not
require any segmentation during inference. Following the convention of RVT
[15] and PerAct [33], we represent each state as a macro-step using the
keyframe extraction procedure from C2F-ARM [23]. At each training iteration,
two trajectories $\tau_{i}$ and $\tau_{j}$ are sampled if
$|\tau_{i}|=|\tau_{j}|$. The state routing network is trained with focal loss
[26] to predict that $s_{i}\equiv s_{j}$ (i.e., $s_{i}$ and $s_{j}$ share the
same manipulation action) if $s_{i}$ and $s_{j}$ have the same timestep and
vice versa. Let $\\{c_{i,n}\\}^{N}_{n=1}$ denote $N$ instance segments, where
each segment $c_{i,n}\in s_{i}$ is a set of 3D points of an object instance in
$s_{i}$. To train the invariant region matching network, we estimate the
ground-truth $\mathcal{I}(s_{i})$ in Equation 2 as follows,

$\mathcal{I}(s_{i})=\arg\min_{c_{i,n}}\mathbf{E}_{\begin{subarray}{c}p\in
c_{i,n}\\\ s_{i}\equiv s_{j}\end{subarray}}[\min_{\begin{subarray}{c}q\in
c_{j,m}\\\ \operatorname{cls}(c_{i,n})=\\\
\operatorname{cls}(c_{j,m})\end{subarray}}\parallel
T_{i}^{-1}p-T_{j}^{-1}q\parallel],$ (2)

where $\operatorname{cls}(c_{i,n})$ denotes the object class of the instance
segment $c_{i,n}$. The ground-truth $\mathcal{I}(s_{i})$ is the instance
segment that has the minimal displacement in the frame of the action pose $T$.
The ground-truth correspondence between $s_{i}$ and $s_{j}$ is the set of
point mappings
$\\{(p,q)\in\mathcal{I}(s_{i})\times\mathcal{I}(s_{j})|q=\operatorname*{argmin}_{q\in\mathcal{I}(s_{j})}\parallel
T^{-1}_{i}p-T^{-1}_{j}q\parallel\\}$ such that the point-wise displacement is
minimal in the action pose frame.

TABLE I: Base Task Performance on RLBench. We report the success rate for each task, and measure the average success rate and rank across all tasks. IMOP has the best overall success rate and rank. The success rate of IMOP is measured with an average of 5 runs. The success rates of baselines are reported from respective papers. Method | Avg. | Avg. | Close | Drag | Insert | Meat off | Open | Place | Place | Push
---|---|---|---|---|---|---|---|---|---|---
Success | Rank | Jar | Stick | Peg | Grill | Drawer | Cups | Wine | Buttons
C2F-ARM [23] | 20.1 | 4.9 | 24 | 24 | 4 | 20 | 20 | 0 | 8 | 72
PerAct [33] | 49.4 | 3.6 | 55.2 | 89.6 | 5.6 | 70.4 | 88 | 2.4 | 44.8 | 92.8
RVT [15] | 62.9 | 2.3 | 52 | 99.2 | 11.2 | 88 | 71.2 | 4 | 91 | 100
Act3D [14] | 65.1 | 2.2 | 92 | 92 | 27 | 94 | 93 | 3 | 80 | 99
IMOP (Ours) | 69.6 | 1.9 | 39.2 | 98.4 | 12.0 | 92.0 | 100.0 | 52.8 | 96.0 | 96.0
| Put in | Put in | Put in | Screw | Slide | Sort | Stack | Stack | Sweep to | Turn
| Cupboard | Drawer | Safe | Bulb | Block | Shape | Blocks | Cups | Dustpan | Tap
C2F-ARM [23] | 0 | 4 | 12 | 8 | 16 | 8 | 0 | 0 | 0 | 68
PerAct [33] | 28 | 51.2 | 84 | 17.6 | 74 | 16.8 | 26.4 | 2.4 | 52 | 88
RVT [15] | 49.6 | 88.0 | 91.2 | 48 | 81.6 | 36 | 28.8 | 26.4 | 72 | 93.6
Act3D [14] | 51 | 90 | 95 | 47 | 93 | 8 | 12 | 9 | 92 | 94
IMOP (Ours) | 46.4 | 100.0 | 96.0 | 82.0 | 58.4 | 37.6 | 40.0 | 56.8 | 100.0 | 51.2

##  IV Experiments

We evaluate IMOP on the visual robotic manipulation tasks from RLBench [22], a
standard benchmark used in multi-task learning [21]. We first train and
evaluate our algorithm on the standard 18 RLBench tasks, and then measure its
one-shot generalization ability on 22 novel tasks from 9 categories [17].
Further, we evaluate our method in real-robot experiments through one-shot
sim-to-real transfer. Specifically, we aim to answer the following questions:

  * •

Can IMOP achieve high success rates on base tasks?

  * •

Can IMOP generalize to novel tasks given a single demonstration without
tuning?

  * •

Can IMOP generalize to objects with large shape variations?

  * •

Can IMOP estimate meaningful invariant regions from demonstrations of novel
tasks?

  * •

Can IMOP be trained in simulation and transferred to solve real-robot
manipulation tasks?

Environment. Following the convention [15], we record $128\times 128$ RGB-D
images from the front, left/right-shoulders, and wrist cameras. During
training, we use 100 recorded trajectories per base task. RLBench contains
various task categories, such as pick-and-place, tool use, high-precision
operations, screwing, tasks with visual occlusion, and long-term manipulation.
We choose 22 novel tasks that have different object setups and task goals from
the base ones, according to the task categorization of Hiveformer [17]. For
each novel task, only a single successful trajectory is provided, as a one-
shot demonstration. Each task is evaluated on 25 independent trials, and we
report the average success rate. RLBench generates different object layouts
and orientations in different trials of the same task. For the base tasks, we
use the same test scenes as the ones in [21]. For novel tasks, we generate 25
test scenes for each task. Figure 2 illustrates the task, and object-level
generalizations of IMOP within the RLBench environment. Further details are
provided in the supplementary material.

###  IV-A Performance Comparison on Base Tasks

Baselines. To evaluate the performance of IMOP on base tasks, we test policies
learned from offline demonstrations without the one-shot demonstrations.
Specifically, we consider C2F-ARM [23], PerAct [33], RVT [15], and Act3D [14].
C2F-ARM and PerAct learn separate policies for each task. Similar to our IMOP,
RVT and Act3D learn a single multi-task policy.

Results. Table I compares the success rates of IMOP and other baselines on the
base tasks. C2F-ARM and PerAct use 3D voxel representation. RVT renders point
clouds to virtual images. Similar to IMOP, Act3D uses raw point clouds as
states. Overall, IMOP outperforms all baselines when the best rank and success
rate are averaged across all tasks ($+4.5\%$ average success rate). IMOP has
the highest rank on 10 out of 18 tasks. Remarkably, IMOP achieves a success
rate of $52.8\%$ on the “place cups” task while the existing works only reach
a rate of $2\%$ to $4\%$. Moreover, not only it outperforms existing
techniques, IMOP also generalizes to novel tasks, while the baselines can only
be used on the tasks they were trained on.

Figure 5: Visualization of invariant regions $\mathcal{I}(s_{i})$ as estimated
by our invariant region matching network on selected base and novel tasks. The
invariant regions on $s_{i}$ are highlighted with green masks and yellow
bounding boxes. The end-effector position is marked with red dot.
$\mathcal{I}(s_{i})$ generally covers the target object of the action that is
applied by the robot to transition from state $s_{i}$ to next state
$s^{\prime}_{i}$. For example, the invariant region of ‘turn the lamp on’
covers the ‘lamp switch’ area. The invariant region of ‘pull the rope
straight’ covers one side of the rope and the rope target. It is worth noting
that IMOP only takes point cloud as input and has no access to segmentation
masks.

###  IV-B Generalization to Novel Tasks

TABLE II: Novel Task Performance on RLBench. We report the success rate for each task, and measure the average success rate and rank across all tasks. IMOP has the best overall success rate and rank. Both RVT+FT and RVT+HDT require fine-tuning but IMOP does not. Directly using RVT without fine-tuning results in an average success rate of 0. We do not include this in the table. All success rates are measured with an average of 5 runs.  Method | Avg. | Avg. | Ball in | Rubbish | Scoop | Place | Hit | Block | Slide | Phone | Open | Close
---|---|---|---|---|---|---|---|---|---|---|---|---
Success | Rank | Hoop | in Bin | Spatula | Hanger | Ball | Pyramid | Place | on Base | Box | Lid
RVT+FT | 29.8 | 2 | 42.4 | 50 | 7.2 | 4 | 1.6 | 0 | 0 | 20.4 | 28 | 100
RVT+HDT | 26.9 | 2.4 | 87.6 | 40.4 | 0 | 0 | 0 | 0 | 0 | 0 | 24.8 | 91.2
IMOP (Ours) | 41.3 | 1.5 | 80.0 | 94.0 | 52.0 | 24.4 | 32.0 | 5.2 | 6.8 | 22.8 | 8.0 | 71.2
| Lid off | Lamp | Beat | Remove | Play | Knife on | Straighten | Change | Open | Open | Money | Close
| Pan | On | Buzz | Cups | Jenga | Board | Rope | Clock | Bottle | Door | out Safe | Microwave
RVT+FT | 78.8 | 65.2 | 19.2 | 4 | 40.8 | 14 | 0 | 10.4 | 34.4 | 68 | 22.4 | 44
RVT+HDT | 80.4 | 72 | 16.4 | 0 | 27.6 | 8 | 0 | 36 | 23.2 | 12 | 12 | 59.6
IMOP (Ours) | 100.0 | 12.0 | 20.0 | 36.0 | 100.0 | 20.0 | 16.8 | 24.0 | 81.6 | 41.2 | 42.8 | 21.6

Baselines. To assess the one-shot ability of IMOP, and since there is no
existing one-shot learning method in the literature that was evaluated in
RLBench, we extend the state-of-the-art technique RVT into two baselines,
called RVT+FT and RVT+HDT, that can be used for one-shot learning.
Specifically, since fine-tuning is the most common strategy for transferring
models to new domains [39], we fine-tune RVT models pretrained on the base
tasks with demonstrations of novel tasks and denote this baseline as RVT+FT.
HDT [38] is a recent work that achieves one-shot policy generalization by
training additional adaptation layers for new tasks while freezing original
parameters. We implement the HDT adaptation layers in RVT and denote this
baseline model as RVT+HDT. RVT and Act3D are both the most recent methods for
learning multi-task robot control policies, and have similar performances, as
shown in Table I. We choose RVT over Act3D because the pretrained model of the
latter is not yet available. For both RVT+FT and RVT+HDT, we train the
respective network parameters with the one-shot demonstrations of novel tasks
until convergence and adopt the sampling and 3D perturbation augmentation
strategy from RVT and PerAct.

Results. Table II compares the success rates of IMOP, RVT+FT, and RVT+HDT on
novel task. IMOP outperforms all baselines with the best rank and success rate
averaged across all tasks ($+11.5\%$ over RVT+FT). IMOP has the best rank on
15 out of 22 tasks. Despite using the simple fine-tuning strategy, RVT+FT
performs better than RVT+HDT on novel tasks. However, we observe the
catastrophic forgetting phenomenon [24] in RVT+FT, whose average success rate
on base tasks plummets to $4.6\%$ from $62.9\%$ after fine-tuning on novel
tasks. RVT-HDT maintains the base task performance by freezing the original
parameters but at the cost of having more parameters in the adaptation layers.
Parameter sizes are compared in Table III.

Figure 6: Real-robot manipulation examples of IMOP via one-shot sim-to-real transfer. The correspondences between the demonstration and test scenes are shown in red lines. The invariant regions closely overlap the next object to interact with. We observe that the one-shot generalization capacity of IMOP, trained with simulation data, enables the sim-to-real transfer using a single real-world demonstration without any re-training. TABLE III: Network parameters size comparison between IMOP and one-shot baselines. | Success Rate (%) | Params Size
---|---|---
Method | Base | Novel
RVT+FT | 4.7 | 29.8 | 104M
RVT+HDT | 62.9 | 26.9 | 131M
Ours | 69.6 | 41.3 | 80.8M

Compared to the baselines, IMOP generalizes to novel tasks by matching
invariant regions between the one-shot demonstration and the states
encountered during testing, without further training. Therefore, IMOP does not
suffer from catastrophic forgetting or increased size of the parameters.
Remarkably, both RVT+FT and RVT+HDT completely fail on certain tasks, e.g.,
“block pyramid”, “slide place”, and “straighten rope”, while IMOP still
delivers a non-trivial performance on these tasks. Both “block pyramid” and
“slide place” are categorized as long-horizon tasks and require over 30 and 10
macro-steps, respectively, to finish [17]. The “straighten rope” task requires
operations on ropes, which are drastically different from the objects in the
base tasks. This further indicates the limitation of the existing one-shot
policy generalization strategies and the advantage of our proposed invariance
matching-based policy. Figure 5 visualizes invariant regions on selected base
and novel tasks. It can be seen that the invariant regions generally cover the
target objects of the robot’s actions. Additional visualizations are provided
in the supplementary material.

Analysis of underperformed tasks. Despite the overall superiority of IMOP, the
performance on specific tasks is still worth discussing. As in Table II, IMOP
exhibits a lower success rate on ‘open box’, ‘close microwave’, and ‘close
lid’ compared to RVT+FT and RVT+HDT. These tasks involve manipulating flat
surfaces, such as lids and box covers, as shown by the invariant region of
‘open box’ visualized in Figure 5. The challenge arises from the potential
ambiguity in identifying invariant regions or correspondences on flat surfaces
lacking rich geometric features. IMOP also performs worse than the baselines
on the ‘lamp on’ task, as you can see in the first two columns of Figure 5.
Although the invariant region estimation for this task is good, the area of
the switch is small and hence, the correspondence is sparse compared to the
whole scene, which leads to unstable solutions of the least-square problem in
Equation 1. The success rate of IMOP on long-term tasks, i.e., ‘block pyramid’
and ‘slide place’, outperforms the baseline methods but still has a large room
for improvement. The IMOP mechanism of routing new states in test scenes to
demonstrations is simple yet lacks the failure recovery ability because the
demonstrations only contain successful trajectories. A more detailed analysis
of failure cases is provided in the supplementary material.

###  IV-C Generalization to Large Shape Variations

(a) ‘Spam’ $\rightarrow$ ‘Crackers’

(b)

(c) ‘Coffee jar’ $\rightarrow$ ‘Soup can’

(d)

Figure 7: Visualization of correspondence between different objects for the
picking action. ‘Spam’ $\rightarrow$ ‘Crackers’ denotes using ‘Spam’ as the
demonstration to correspond to ‘Crackers’ during testing. Despite large shape
variations, IMOP is able to estimate correspondences that effectively transfer
actions, as shown by the success rate in Figure 8.

Setup. Previous experiments evaluate the one-shot generalization ability of
IMOP on novel tasks with one demonstration trajectory for each task. Even
though there are large variations in object shapes and categories between base
and novel tasks, the variations within each novel task are limited. For
example, the novel task ‘put the knife on the chopping board’ includes objects
such as a knife, knife holder, and chopping board. None of these objects are
seen during training and they are placed randomly at each trial and sometimes
with color variations, but the object categories and shapes remain unchanged
among various trials, as illustrated in Figure 2. We are interested in whether
IMOP can generalize to objects with large shape variations. To study this, we
evaluate IMOP on objects different from those in demonstrations, which
naturally have different shapes. Specifically, we consider a pick-and-place
setting where the objective is to pick groceries and place them into a
cupboard. The demonstration is given on a certain object, but picking and
placing a different object is required during testing. To simplify our study,
we assume the access of the segmentation mask of the required object during
testing. Therefore, the correspondence between different objects can be
obtained by masking out other regions. To avoid memorization, we remove the
corresponding state pairs from the training data.

Figure 8: Success rate of pick-and-place task using different objects with
large shape variations as demonstrations. For example, the entry at the top-
right corner, i.e., ‘berry-jello, tuna’, denotes the success rate is 30% at
picking and placing ‘tuna can’ while using ‘berry-jello box’ as the
demonstration.

Result. Figure 8 shows the success rate of pick-and-place with demonstrations
on different objects. For example, the entry ‘mustard, coffee’ represents the
success rate of picking and placing a coffee jar with a demonstration on a
mustard bottle. There are 9 objects in total, and all success rates are
measured on an average of 10 trials. It can be seen that IMOP performs
competitively when a different object is used as a demonstration, except in
some cases of the mustard bottle whose shape is drastically different than
others. Using berry-jello and choco-jello as demonstrations delivers the same
or even better results on many objects during testing, e.g., soup, spam,
sugar, and tuna. These two jello boxes are the smallest objects in this setup.

Figure 7 visualizes the correspondence between the invariant regions of
different objects. The corresponding points, visualized as green dots, show a
resemblance to the keypoint-based manipulation methods [13, 29]. kPAM [29] and
kPAM-SC [13] use object keypoints to compose objectives and constraints in the
optimization-based planning for manipulation and demonstrate generalization
over large shape variations in a pick-and-place setting. However, these
methods require the manual design of keypoints for each object category. Our
method IMOP does not require any specification of keypoints or any other form
of labeling. The visualized connections are obtained by simply thresholding
the correspondence matrix.

###  IV-D Real-Robot Manipulation via One-shot Sim-to-Real Transfer

Setup. We evaluate the one-shot sim-to-real performance of IMOP using
simulation demonstrations for the base tasks, and trajectories recorded in the
real-world as the one-shot demonstrations of novel tasks. We adopt five novel
tasks: put toys in bowl, stack blocks, put item in shelf, open bottle, and
pour cup. Contrary to existing work that requires over 10 demonstrations per
task for re-training [15], the one-shot generalization capability of IMOP
enables the sim-to-real transfer using only one demonstration and without re-
training. We deploy the model on a Kuka LBR iiwa robot. We use RBG-D
observations from two RealSense D415 cameras. We use MoveIt [7] for planning
paths that move the arm to the gripper poses predicted by IMOP.

TABLE IV: One-shot sim-to-real transfer performance of IMOP for solving real-world manipulation tasks. Task | Success
---|---
One-shot RVT | IMOP (Ours)
Put toys in bowl | 3/10 | 8/10
Stack blocks | 1/10 | 4/10
Put item in shelf | 2/10 | 5/10
Open bottle | 2/10 | 4/10
Pour cup | 2/10 | 5/10

Result. Table IV shows the success rate of IMOP on the three real-world tasks.
The success rate of each task is the average of independent 10 runs with
different object layouts. Figure 6 visualizes the correspondence between
demonstration and test scenes. For the one-shot RVT, we finetune the
pretrained RVT model with a single demonstration for each task. The sim-to-
real gap is commonly acknowledged as a challenging problem in deploying
control policies trained with simulation data [43]. In our study, we find that
the region matching is robust to this domain shift. However, the sigmoid
output of the invariant region prediction has much smaller values than those
from simulation data. We circumvent this problem by thresholding the sigmoid
output with a much smaller value, $0.1$. This indicates that IMOP has the one-
shot sim-to-real capacity, but still suffers from the sim-to-real gap for the
invariant region estimation. This issue can potentially be addressed by fine-
tuning with real-robot demonstrations or using more cameras to improve the
point cloud quality.

###  IV-E Ablation Studies

TABLE V: Ablation studies on the components of IMOP. Configurations | Success Rate (%)
---|---
Base | Novel
State Routing Network
Using Low-dim robot states | Y | 69.6 | 41.3
N | 63.2 | 33.0
Invariant Region Matching Network
Attention Type | Graph Attention | 69.6 | 41.3
Traditional Attention | 52.8 | 16.4
Matching Strategy |  | Estimate $\mathcal{I}(s_{i})$, then
---
match $\mathcal{I}(s_{i})$ and $s_{j}$
69.6 | 41.3
Directly match $s_{i}$ and $s_{j}$ | 6.7 | -
KNN-Graph (K) | 4 | 58.7 | 32.4
8 | 64.5 | 38.6
16 | 69.6 | 41.3
Regression Method
Correspondence-based Regression (Equation 1) | 69.6 | 41.3
Regression from matched points | 58.2 | 29.4
Conventional pose regression | 65.1 | -

We study the component effects of IMOP in Table V. For the state routing
network, we find the low-dimensional robot state (i.e., joint positions and
timesteps) contributes to the routing accuracy and success rate. This practice
is also adopted by many existing work [33, 15, 14].

For the invariant region matching network, we study the attention type,
matching strategy, and the degree (K) of the point cloud KNN graphs. Table V
shows that KNN graphs of a larger degree provide stronger performance.
However, using traditional attention instead of graph attention, which
connects all points globally, significantly decreases performance. Note that
we only apply traditional attention for region matching and still keep graph
attention for invariant region estimation. This suggests local attention and
propagating information through KNN graphs are effective inductive biases for
3D region matching. In contrast to the framework in Figure 3, which first
estimates $\mathcal{I}(s_{i})$ and then match $\mathcal{I}(s_{i})$ and
$s_{j}$, we investigate the strategy of directly matching $s_{i}$ and $s_{j}$
and predict the correspondence matrix. However, the model fails to predict
actions on all tasks. This indicates the importance of estimating invariant
regions.

We compare our correspondence-based pose regression with two alternative
solutions. The ‘regression from matched points’ trains extra layers to predict
pose from the matched points in correspondence matrix $C$ instead of deriving
the action pose $T_{j}$ with Equation 1. The ‘conventional pose regression’
adopts a standard 3D regression head from Act3D [14] while dropping the region
matching mechanism completely. Table V shows that conventional regression
delivers comparable performance on base tasks but fails to generalize over
novel tasks, and deriving action pose with Equation 1 performs better than
predicting pose from matched points with an extra regression head.

##  V Discussion and Conclusion

We have shown that one-shot novel task generalization can be achieved by
learning to estimate and match key invariant regions in the demonstration and
test scene. The target end-effector pose can be transferred by finding
correspondences between invariant regions. Some open questions are, however,
still worth discussing. IMOP is trained within the RLBench simulation
environment, which offers segmentation masks and a validated key frame
extraction procedure. For real-world data, a more sophisticated key frame
extraction method and the use of open-set object detectors [41, 44] may be
necessary. The proposed idea of transferring actions through matching key
visual elements is general but the current definition of the invariant region
is still closely linked to rigid body transformation. This suggests the
potential for extending the formulation of invariant regions based on more
general motion descriptors, such as warping or flows [25], to improve the
manipulation accuracy on non-rigid objects and tasks with complex dynamics
that cannot be regulated by poses. With the emergence of neural architectures
that are scalable to Internet-scale data [4], the possibility of learning
invariant regions from unlabeled 2D videos that are widely available has yet
to be fully explored. Beyond only leveraging a single demonstration, it is
possible to maintain a pool of demonstrations and transfer actions from the
most relevant state to improve the manipulation performance and reduce error
accumulation under scenes with large variations or requiring failure recovery.

## References

  * Ardón et al. [2020] Paola Ardón, Èric Pairet, Katrin S Lohan, Subramanian Ramamoorthy, and Ronald Petrick.  Affordances in robotic tasks–a survey.  _arXiv preprint arXiv:2004.07400_ , 2020.
  * Biza et al. [2023] Ondrej Biza, Skye Thompson, Kishore Reddy Pagidi, Abhinav Kumar, Elise van der Pol, Robin Walters, Thomas Kipf, Jan-Willem van de Meent, Lawson LS Wong, and Robert Platt.  One-shot imitation learning via interaction warping.  _arXiv preprint arXiv:2306.12392_ , 2023.
  * Brandstetter et al. [2021] Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J Bekkers, and Max Welling.  Geometric and physical quantities improve e (3) equivariant message passing.  _arXiv preprint arXiv:2110.02905_ , 2021.
  * Brown et al. [2020] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Cohen and Welling [2016] Taco Cohen and Max Welling.  Group equivariant convolutional networks.  In _International conference on machine learning_ , pages 2990–2999. PMLR, 2016.
  * Cohen et al. [2018] Taco S Cohen, Mario Geiger, Jonas Köhler, and Max Welling.  Spherical cnns.  _arXiv preprint arXiv:1801.10130_ , 2018.
  * Coleman et al. [2014] David Coleman, Ioan Sucan, Sachin Chitta, and Nikolaus Correll.  Reducing the barrier to entry of complex robotic software: a moveit! case study.  _arXiv preprint arXiv:1404.3785_ , 2014.
  * Contributors [2023] Pointcept Contributors.  Pointcept: A codebase for point cloud perception research.  https://github.com/Pointcept/Pointcept, 2023.
  * Dasari and Gupta [2021] Sudeep Dasari and Abhinav Gupta.  Transformers for one-shot visual imitation.  In _Conference on Robot Learning_ , pages 2071–2084. PMLR, 2021.
  * Dosovitskiy et al. [2020] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Duan et al. [2017] Yan Duan, Marcin Andrychowicz, Bradly Stadie, OpenAI Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojciech Zaremba.  One-shot imitation learning.  _Advances in neural information processing systems_ , 30, 2017.
  * Finn et al. [2017] Chelsea Finn, Tianhe Yu, Tianhao Zhang, Pieter Abbeel, and Sergey Levine.  One-shot visual imitation learning via meta-learning.  In _Conference on robot learning_ , pages 357–368. PMLR, 2017.
  * Gao and Tedrake [2021] Wei Gao and Russ Tedrake.  kpam-sc: Generalizable manipulation planning using keypoint affordance and shape completion.  In _2021 IEEE International Conference on Robotics and Automation (ICRA)_ , pages 6527–6533. IEEE, 2021.
  * Gervet et al. [2023] Theophile Gervet, Zhou Xian, Nikolaos Gkanatsios, and Katerina Fragkiadaki.  Act3d: 3d feature field transformers for multi-task robotic manipulation.  In _Conference on Robot Learning_ , pages 3949–3965. PMLR, 2023.
  * Goyal et al. [2023] Ankit Goyal, Jie Xu, Yijie Guo, Valts Blukis, Yu-Wei Chao, and Dieter Fox.  Rvt: Robotic view transformer for 3d object manipulation.  _arXiv preprint arXiv:2306.14896_ , 2023.
  * Graf et al. [2023] Christian Graf, David B Adrian, Joshua Weil, Miroslav Gabriel, Philipp Schillinger, Markus Spies, Heiko Neumann, and Andras Gabor Kupcsik.  Learning dense visual descriptors using image augmentations for robot manipulation tasks.  In _Conference on Robot Learning_ , pages 871–880. PMLR, 2023.
  * Guhur et al. [2023] Pierre-Louis Guhur, Shizhe Chen, Ricardo Garcia Pinel, Makarand Tapaswi, Ivan Laptev, and Cordelia Schmid.  Instruction-driven history-aware policies for robotic manipulations.  In _Conference on Robot Learning_ , pages 175–187. PMLR, 2023.
  * Han et al. [2022] Kai Han, Yunhe Wang, Hanting Chen, Xinghao Chen, Jianyuan Guo, Zhenhua Liu, Yehui Tang, An Xiao, Chunjing Xu, Yixing Xu, et al.  A survey on vision transformer.  _IEEE transactions on pattern analysis and machine intelligence_ , 45(1):87–110, 2022.
  * Huang et al. [1986] TS Huang, SD Blostein, and EA Margerum.  Least-squares estimation of motion parameters from 3-d point correspondences.  In _Proc. IEEE Conf. Computer Vision and Pattern Recognition_ , volume 10, pages 112–115. IEEE Computer Soc. Press Washington DC, 1986.
  * Huang et al. [2021] Xiaoshui Huang, Guofeng Mei, Jian Zhang, and Rana Abbas.  A comprehensive survey on point cloud registration.  _arXiv preprint arXiv:2103.02690_ , 2021.
  * James and Davison [2022] Stephen James and Andrew J Davison.  Q-attention: Enabling efficient learning for vision-based robotic manipulation.  _IEEE Robotics and Automation Letters_ , 7(2):1612–1619, 2022.
  * James et al. [2020] Stephen James, Zicong Ma, David Rovick Arrojo, and Andrew J Davison.  Rlbench: The robot learning benchmark & learning environment.  _IEEE Robotics and Automation Letters_ , 5(2):3019–3026, 2020.
  * James et al. [2022] Stephen James, Kentaro Wada, Tristan Laidlow, and Andrew J Davison.  Coarse-to-fine q-attention: Efficient learning for visual robotic manipulation via discretisation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 13739–13748, 2022.
  * Kirkpatrick et al. [2017] James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, et al.  Overcoming catastrophic forgetting in neural networks.  _Proceedings of the national academy of sciences_ , 114(13):3521–3526, 2017.
  * Li and Harada [2022] Yang Li and Tatsuya Harada.  Lepard: Learning partial point cloud matching in rigid and deformable scenes.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 5554–5564, 2022.
  * Lin et al. [2017] Tsung-Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár.  Focal loss for dense object detection.  In _Proceedings of the IEEE international conference on computer vision_ , pages 2980–2988, 2017.
  * Lyle et al. [2020] Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy.  On the benefits of invariance in neural networks.  _arXiv preprint arXiv:2005.00178_ , 2020.
  * Mandi et al. [2022] Zhao Mandi, Fangchen Liu, Kimin Lee, and Pieter Abbeel.  Towards more generalizable one-shot visual imitation learning.  In _2022 International Conference on Robotics and Automation (ICRA)_ , pages 2434–2444. IEEE, 2022.
  * Manuelli et al. [2019] Lucas Manuelli, Wei Gao, Peter Florence, and Russ Tedrake.  kpam: Keypoint affordances for category-level robotic manipulation.  In _The International Symposium of Robotics Research_ , pages 132–157. Springer, 2019.
  * Mendonca et al. [2023] Russell Mendonca, Shikhar Bahl, and Deepak Pathak.  Structured world models from human videos.  _arXiv preprint arXiv:2308.10901_ , 2023.
  * Mo et al. [2021] Kaichun Mo, Leonidas J Guibas, Mustafa Mukadam, Abhinav Gupta, and Shubham Tulsiani.  Where2act: From pixels to actions for articulated 3d objects.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 6813–6823, 2021.
  * Shridhar et al. [2022] Mohit Shridhar, Lucas Manuelli, and Dieter Fox.  Cliport: What and where pathways for robotic manipulation.  In _Conference on Robot Learning_ , pages 894–906. PMLR, 2022.
  * Shridhar et al. [2023] Mohit Shridhar, Lucas Manuelli, and Dieter Fox.  Perceiver-actor: A multi-task transformer for robotic manipulation.  In _Conference on Robot Learning_ , pages 785–799. PMLR, 2023.
  * Van der Pol et al. [2020] Elise Van der Pol, Daniel Worrall, Herke van Hoof, Frans Oliehoek, and Max Welling.  Mdp homomorphic networks: Group symmetries in reinforcement learning.  _Advances in Neural Information Processing Systems_ , 33:4199–4210, 2020.
  * Wu et al. [2022] Xiaoyang Wu, Yixing Lao, Li Jiang, Xihui Liu, and Hengshuang Zhao.  Point transformer v2: Grouped vector attention and partition-based pooling.  _Advances in Neural Information Processing Systems_ , 35:33330–33342, 2022.
  * Xian et al. [2023] Zhou Xian, Nikolaos Gkanatsios, Theophile Gervet, Tsung-Wei Ke, and Katerina Fragkiadaki.  Chaineddiffuser: Unifying trajectory diffusion and keypose prediction for robotic manipulation.  In _Conference on Robot Learning_ , pages 2323–2339. PMLR, 2023.
  * Xu et al. [2022] Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan.  Prompting decision transformer for few-shot policy generalization.  In _international conference on machine learning_ , pages 24631–24645. PMLR, 2022.
  * Xu et al. [2023] Mengdi Xu, Yuchen Lu, Yikang Shen, Shun Zhang, Ding Zhao, and Chuang Gan.  Hyper-decision transformer for efficient online policy adaptation.  _arXiv preprint arXiv:2304.08487_ , 2023.
  * Yu et al. [2022] Fuchao Yu, Xianchao Xiu, and Yunhui Li.  A survey on deep transfer learning and beyond.  _Mathematics_ , 10(19):3619, 2022.
  * Zeng et al. [2021] Andy Zeng, Pete Florence, Jonathan Tompson, Stefan Welker, Jonathan Chien, Maria Attarian, Travis Armstrong, Ivan Krasin, Dan Duong, Vikas Sindhwani, et al.  Transporter networks: Rearranging the visual world for robotic manipulation.  In _Conference on Robot Learning_ , pages 726–747. PMLR, 2021.
  * Zhang et al. [2023] Xinyu Zhang, Yuting Wang, and Abdeslam Boularias.  Detect every thing with few examples, 2023.
  * Zhao et al. [2021] Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip HS Torr, and Vladlen Koltun.  Point transformer.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , pages 16259–16268, 2021.
  * Zhao et al. [2020] Wenshuai Zhao, Jorge Peña Queralta, and Tomi Westerlund.  Sim-to-real transfer in deep reinforcement learning for robotics: a survey.  In _2020 IEEE symposium series on computational intelligence (SSCI)_ , pages 737–744. IEEE, 2020.
  * Zhong et al. [2022] Yiwu Zhong, Jianwei Yang, Pengchuan Zhang, Chunyuan Li, Noel Codella, Liunian Harold Li, Luowei Zhou, Xiyang Dai, Lu Yuan, Yin Li, et al.  Regionclip: Region-based language-image pretraining.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 16793–16803, 2022.


2023 Full Paper – MIDL 2023 submission Jun Li 1,Contributed equally
<EMAIL_ADDRESS>
Che Liu<EMAIL_ADDRESS>
Sibo Cheng 2<EMAIL_ADDRESS>
Rossella Arcucci 2,3<EMAIL_ADDRESS>
Shenda Hong 4,5,Corresponding author<EMAIL_ADDRESS>

1 College of Electronic Science and Engineering, Jilin University, Changchun,
China
2 Department of Earth Science and Engineering, Imperial College London, SW7
2AZ, UK
3 Data Science Institute, Department of computing, Imperial College London,
SW7 2AZ, UK
4 National Institute of Health Data Science, Peking University, Beijing, China
5 Institute of Medical Technology, Health Science Center of Peking University,
Beijing, China

# Frozen Language Model Helps ECG Zero-Shot Learning

###### Abstract

The electrocardiogram (ECG) is one of the most commonly used non-invasive,
convenient medical monitoring tools that assist in the clinical diagnosis of
heart diseases. Recently, deep learning (DL) techniques, particularly self-
supervised learning (SSL), have demonstrated great potential in the
classification of ECG. SSL pre-training has achieved competitive performance
with only a small amount of annotated data after fine-tuning. However, current
SSL methods rely on the availability of annotated data and are unable to
predict labels not existing in fine-tuning datasets. To address this
challenge, we propose Multimodal ECG-Text Self-supervised pre-training (METS),
the first work to utilize the auto-generated clinical reports to guide ECG SSL
pre-training. We use a trainable ECG encoder and a frozen language model to
embed paired ECG and automatically machine-generated clinical reports
separately. The SSL aims to maximize the similarity between paired ECG and
auto-generated report while minimize the similarity between ECG and other
reports. In downstream classification tasks, METS achieves around 10%
improvement in performance without using any annotated data via zero-shot
classification, compared to other supervised and SSL baselines that rely on
annotated data. Furthermore, METS achieves the highest recall and F1 scores on
the MIT-BIH dataset, despite MIT-BIH containing different classes of ECG
compared to the pre-trained dataset. The extensive experiments have
demonstrated the advantages of using ECG-Text multimodal self-supervised
learning in terms of generalizability, effectiveness, and efficiency.

###### keywords:

Multimodal self-supervised learning, Zero-shot learning, Language model, ECG,
Signal processing

##  1 Introduction

Figure 1: (a) denotes a supervised learning method for ECG. (b) denotes a
common self-supervised learning method for ECG, i.e. pre-training followed by
fine-tuning. (c) denotes a self-supervised learning method for multimodal ECG-
Text. Zero-shot classification is performed after pre-training is completed.

The electrocardiogram (ECG) is a diagnostic tool that is widely used in
clinical practice [Addison(2005)]. In practice, the ECG is used to detect a
wide range of cardiac conditions, including arrhythmias, heart attacks, and
heart failure [Berkaya et al.(2018)Berkaya, Uysal, Gunal, Ergin, Gunal, and
Gulmezoglu]. Recently, deep learning (DL) methods have shown promising results
in classifying ECG data [Ebrahimi et al.(2020)Ebrahimi, Loni, Daneshtalab, and
Gharehbaghi, Tripathy et al.(2019)Tripathy, Bhattacharyya, and Pachori]. DL
models, such as convolutional neural networks (CNNs) and recurrent neural
networks (RNNs), have been shown to be highly accurate in classifying ECG for
a variety of cardiac conditions [Baloglu et al.(2019)Baloglu, Talo, Yildirim,
San Tan, and Acharya, Singh et al.(2018)Singh, Pandey, Pawar, and Janghel, Xu
et al.(2020)Xu, Jeong, and Li]. However, training DL models in a supervised
manner (see Figure 1 (a)) often requires a large number of high-quality labels
to obtain strong generalization performance [Ebrahimi et al.(2020)Ebrahimi,
Loni, Daneshtalab, and Gharehbaghi]. In addition, some forms of ECG, such as
ST-elevation myocardial infarction, are difficult to detect and often require
manual interpretation of the ECG by trained cardiologists [Ayer and
Terkelsen(2014)]. This work requires a huge effort which is costly and
laborious.

Currently, self-supervised learning (SSL) has achieved impressive performance
on datasets containing a small number of annotations, which provides a
promising solution for unannotated ECG data [Jaiswal et al.(2020)Jaiswal,
Babu, Zadeh, Banerjee, and Makedon, Chou et al.(2020)Chou, Hong, Zhou, Shang,
Song, and Li]. It allows models to mine useful representations of ECG and can
be widely used for various downstream tasks such as abnormality detection and
arrhythmia classification [Lan et al.(2022)Lan, Ng, Hong, and Feng, Mehari and
Strodthoff(2022)]. Nevertheless, existing ECG SSL methods still require a
large amount of annotated data in order to fine-tune them for downstream tasks
(see Figure 1 (b)). This requirement hinders the real-world application of ECG
methods as some heart diseases are rare, which leads to problems with zero-
shot learning. Zero-shot learning means that the model does not need any
annotated samples for unseen categories [Socher et al.(2013)Socher, Ganjoo,
Manning, and Ng]. This is achieved by explicitly learning shared features from
seen samples, and then generalizing them on unseen samples based on
”descriptions” of the unseen categories’ features [Xian et al.(2018)Xian,
Lampert, Schiele, and Akata, Pourpanah et al.(2020)Pourpanah, Abdar, Luo,
Zhou, Wang, Lim, and Wang]. Specifically, such ”descriptions”, are usually
borrowed from external medical domain knowledge, textual ECG reports for
example (see Figure 1 (c)). Zero-shot learning for ECG faces a number of
challenges. The first challenge is the semantic gap, where ECG and text
(automatically machine-generated ECG reports) are two heterogeneous
modalities. ECG is long-term continuous numbers and text is relatively short-
term discrete clinical terminologies [Krishnan and Sowmya Kamath(2018)]. They
are difficult to align and characterize each other [Liang et al.(2022)Liang,
Zadeh, and Morency]. The second challenge is domain adaptation. Zero-shot
learning model may be sensitive to unknown domains, making it difficult to
adapt to new domains or unseen categories and not performing well for
downstream tasks in zero-shot learning. The third challenge is scalability.
Zero-shot learning models need to learn a large number of representations and
apply them to downstream tasks, which leads to a large computational cost for
the model [Wang et al.(2019)Wang, Zheng, Yu, and Miao]. Recently, [Yamaç et
al.(2022)Yamaç, Duman, Adalıoğlu, Kiranyaz, and Gabbouj] and [Bhaskarpandit et
al.(2022)Bhaskarpandit, Gupta, and Gupta] have reached considerable results on
ECG zero-shot classification tasks. However, they pre-trained the model with
supervised learning, which indicates that their methods still require large-
scale annotated ECG for the pre-training stage.

To fully utilize the unannotated data, CLIP [Radford et al.(2021)Radford, Kim,
Hallacy, Ramesh, Goh, Agarwal, Sastry, Askell, Mishkin, Clark, et al.] and
ALIGN [Jia et al.(2021a)Jia, Yang, Xia, Chen, Parekh, Pham, Le, Sung, Li, and
Duerig] first implement multimodal SSL with two individual encoders and use
zero-shot classification as the downstream task to evaluate SSL pre-trained
model performance [Radford et al.(2021)Radford, Kim, Hallacy, Ramesh, Goh,
Agarwal, Sastry, Askell, Mishkin, Clark, et al., Jia et al.(2021b)Jia, Yang,
Xia, Chen, Parekh, Pham, Le, Sung, Li, and Duerig]. Florence [Yuan et
al.(2021)Yuan, Chen, Chen, Codella, Dai, Gao, Hu, Huang, Li, Li, et al.], LiT
[Zhai et al.(2022)Zhai, Wang, Mustafa, Steiner, Keysers, Kolesnikov, and
Beyer], and ALBEF [Li et al.(2021)Li, Selvaraju, Gotmare, Joty, Xiong, and
Hoi] explore the potential of multimodal SSL on large-scale pre-training tasks
[Yuan et al.(2021)Yuan, Chen, Chen, Codella, Dai, Gao, Hu, Huang, Li, Li, et
al., Li et al.(2021)Li, Selvaraju, Gotmare, Joty, Xiong, and Hoi, Zhai et
al.(2022)Zhai, Wang, Mustafa, Steiner, Keysers, Kolesnikov, and Beyer].
Although recent works have achieved substantial progress on an image-text
task, the medical signals-text, such as ECG, has not yet leveraged the
benefits of the multimodal SSL.

To take advantage of multimodal SSL, we propose a novel method to do
Multimodal ECG-TEXT SSL pre-training (METS). The METS model takes the ECG and
the corresponding reported text as input and feeds them into a multimodal
comparative learning framework. The multimodal framework contains a language
component and an ECG encoder to obtain embedding representations of the text
and the ECG respectively. To make full use of the a priori clinical knowledge
in the report text, the report text is fed into a large frozen language model
with Resnet1d-18 as the backbone of the ECG encoder, both of which have a
linear projection head to embed the text and ECG into the same dimension.
Then, the similarity between ECG embedding and text embedding is computed to
minimize the contrast learning loss and obtain a pre-trained model with rich
medical knowledge. probabilities, which can be used to classify the various
categories of ECG.

The main contributions of this paper are summarised as follows:

  * •

Our proposed METS is the first work to apply a large language model for ECG
SSL. The apriori clinical knowledge from the language model can be fully
exploited to help generate ECG medical reports.

  * •

METS is independent of the annotated categories. Even if the external dataset
categories are unseen, the classification can be done directly with zero-shot,
unlike other SSL which require fine-tuning.

  * •

Experiments demonstrate that METS can be adapted to any of the downstream
tasks, e.g. form, rhythm, and superclass, without the need to fine-tune on
different tasks. Besides, METS does not require any annotated data but can
exceed the supervised and SSL methods of fine-tuning with small-scale
annotated data.

##  2 Methods

In this section, we will demonstrate the details of METS. The framework is
shown in Figure 2. METS consists of two components: Multimodal self-supervised
pre-training (Section 2.1), and the zero-shot classification downstream task
(Section 2.2).

Figure 2: A framework for the METS approach. (a) shows a self-supervised pre-
training approach. ECG-text pairs are fed into the model, and after
comparative learning, the ECG encoder learns the parameters. (b) shows the
zero-shot classification task. The corresponding labels are found by computing
ECG and text similarity. (c) shows the visualization of the results of zero-
shot classification.

###  2.1 Multimodal self-supervised pre-training

####  2.1.1 Frozen Pre-trained Language Models

Our approach starts with a large pre-trained language model based on the
transformer architecture. In order for the large language model to fully
understand the report text, we extend the report text into the language model
as a complete sentence input [Radford et al.(2021)Radford, Kim, Hallacy,
Ramesh, Goh, Agarwal, Sastry, Askell, Mishkin, Clark, et al.]. Specifically,
we construct a prompt template for the report “The report of the ECG is
$\textbf{that\ \\{text\\}''}.$ We use a large clinical language model as the
backbone of the text component. ClinicalBert has been pre-trained on all text
from MIMIC III dataset [Alsentzer et al.(2019)Alsentzer, Murphy, Boag, Weng,
Jin, Naumann, and McDermott].

####  2.1.2 ECG Encoder

Our ECG encoder $E_{ecg}$ is based on ResNet1d-18, which modifies the kernel
of ResNet-18 from a 2D patch to a 1D stride, in order to obtain the deep ECG
embeddings $\mathsf{e}$ [He et al.(2016)He, Zhang, Ren, and Sun, Hong et
al.(2020)Hong, Xu, Khare, Priambada, Maher, Aljiffry, Sun, and Tumanov]. This
process can be represented as follows:
$\mathbf{e}=E_{ecg}\left(\mathbf{y}\right)$, where $\mathbf{y}$ is the input
of ECG. Then, a linear projection head $f_{e}$ maps raw embeddings to
$\mathbf{e}_{d}\in\mathbb{R}^{D}$. The embedding dimension of the ECG encoder
is set to be the same as the language model embedding dimension $d$ for
comparison learning.

Inspired by [Tsimpoukelli et al.(2021)Tsimpoukelli, Menick, Cabi, Eslami,
Vinyals, and Hill], we freeze the parameters of the language model (LM) and
use only paired ECG-text data from the PTB-XL dataset to update the parameters
of the ECG encoder during SSL pre-training. This has the advantage of allowing
the ECG encoder to learn rich prior clinical knowledge from the medical
corpus, thus improving the generalization ability of the model. In addition,
the parameters of the language model are frozen to reduce the significant
computational cost of LM parameter updates.

####  2.1.3 Multimodal Contrastive Learning

Following the multimodal contrast learning framework, we treat a pair of
report text and ECG belonging to the same patient as a positive sample pair,
while treating pairs of other patients’ report text and that ECG as negative
sample pairs. We maximize the contrast loss of different pairs
($\mathbf{t_{i}}$, $\mathbf{e_{j}}$) and minimize the contrast loss of the
same pair ($\mathbf{t_{i}}$, $\mathbf{e_{i}}$) to improve the similarity of
the same pair of samples. We first define the similarity between the
representations $\mathbf{t}$ and $\mathbf{e}$ of two modalities in terms of
cosine similarity, as shown in Equation 1.

$\mathbf{}\operatorname{sim}\left(\mathbf{t},\mathbf{e}\right)=\frac{t^{\top}\cdot
e}{\left\|t\right\|\left\|e\right\|}$ (1)

Then, we need to train two contrast loss functions. The first loss function is
the ECG-to-text contrast loss for the $i^{th}$ pair, as shown in Equation 2.

$\ell^{(e\rightarrow
t)}_{\mathbf{i}}=-\log\frac{\exp\left(\operatorname{sim}\left(\mathbf{t_{i}},\mathbf{e_{i}}\right)/\tau\right)}{\sum_{k=1}^{N}\exp\left(\operatorname{sim}\left(\mathbf{t_{i}},\mathbf{e_{j}}\right)/\tau\right)}$
(2)

The initialization of $\tau$ is set to 0.07.Similarly, the text-to-ECG
contrast loss Equation 3 is represented as follows.

$\ell^{(t\rightarrow
e)}_{\mathbf{i}}=-\log\frac{\exp\left(\operatorname{sim}\left(\mathbf{e_{i}},\mathbf{t_{i}}\right)/\tau\right)}{\sum_{k=1}^{N}\exp\left(\operatorname{sim}\left(\mathbf{e_{i}},\mathbf{t_{j}}\right)/\tau\right)}$
(3)

Finally, our training losses are calculated as the average combination of the
two losses for all positive ECG-text pairs in each minibatch, as shown in
Equation 4.

$\mathcal{L}=\frac{1}{N}\sum_{i=1}^{N}\frac{\ell_{\mathbf{i}}^{(e\rightarrow
t)}+\ell_{\mathbf{i}}^{(t\rightarrow e)}}{2}$ (4)

###  2.2 Zero-Shot ECG Classification

In zero-shot classification, a segment of the ECG is used as input. To
evaluate the zero-shot performance of the model on a multi-label
classification task, we extend the discrete labels into full medical
diagnostic statements and feed them into the language model to obtain
embedding representations. Finally, the similarity between ECG embedding and
text embedding is computed to obtain probabilities, which can be used to
classify the various categories of ECG.

##  3 Experiments

###  3.1 Datasets

##### PTB-XL

We use the PTB-XL dataset to train the METS model [Wagner et al.(2020)Wagner,
Strodthoff, Bousseljot, Kreiseler, Lunze, Samek, and Schaeffter]. The PTB-XL
dataset contains 21,837 clinical 12-lead ECG of 10 seconds duration from
18,885 patients, where each ECG segment is paired with the corresponding ECG
reports. The reports are generated by the machine and only describe the ECG
without final diagnosis. The original ECG reports were written in 70.89%
German, 27.9% English, and 1.21% Swedish, and were converted into structured
SCP-ECG statements. The statement sets were assigned to three non-reciprocal
categories: diagnosis, form, and rhythm. Specifically, the dataset consisted
of 71 different statements, broken down into 44 diagnostic statements, 12
rhythmic statements, and 15 formal statements. For the diagnostic labels were
divided into 5 superclasses and 24 subclasses. In the current experiments, we
focused on investigating ECG-text pairs without using any other labels.
Following the experimental setup in [Huang et al.(2021)Huang, Shen, Lungren,
and Yeung, Wang et al.(2022)Wang, Wu, Agarwal, and Sun], we extracted a
multiclass classification dataset, PTB-XL test set, from the test set split.

##### PTB-XL Test Set

The original ECG in the PTB-XL dataset is multi-labeled with diagnostic, form,
and rhythm. In zero-shot downstream task classification, we need to calculate
the similarity of ECG and text to find the most similar target, and multiple
labels for a target can confuse the categories. Therefore, we produce
diagnostic superclass, form, and rhythm test sets to complete the
corresponding zero-shot downstream tasks. There are 1,000 samples on each test
set. Details of the specific split test set are shown in figure 4.

##### MIT-BIH Test Set

We use the MIT-BIH dataset for testing to evaluate the performance of our pre-
trained representation framework for classification on external datasets
[Moody and Mark(2001)]. Please note that we do not pre-train on the MIT-BIH
dataset. Similarly, we produced an MIT-BIH test set following the segmentation
method above. Details of the specific split test set are also shown in figure
4.

Table 1: PTB-XL result on superclass. $\%$ refers to fractions of labels used in the training data. Methods | Accuracy | Precision | Recall | F1
---|---|---|---|---
Self-supervised
random - $5\%$ | 0.581 | 0.438 | 0.421 | 0.429
SimCLR - $5\%$ | 0.648 | 0.545 | 0.443 | 0.485
METS - 0% | 0.842 | 0.694 | 0.626 | 0.657
Supervised
Resnet18 - $100\%$ | 0.894 | 0.811 | 0.745 | 0.776
Table 2: PTB-XL result on form. $\%$ refers to fractions of labels used in the training data. Methods | Accuracy | Precision | Recall | F1
---|---|---|---|---
Self-supervised
random - $5\%$ | 0.603 | 0.364 | 0.342 | 0.351
SimCLR - $5\%$ | 0.660 | 0.446 | 0.471 | 0.456
METS - 0% | 0.734 | 0.537 | 0.503 | 0.518
Supervised
Resnet18 - $100\%$ | 0.724 | 0.520 | 0.508 | 0.509
Table 3: PTB-XL result on rhythm. $\%$ refers to fractions of labels used in the training data. Methods | Accuracy | Precision | Recall | F1
---|---|---|---|---
Self-supervised
random - $5\%$ | 0.627 | 0.435 | 0.442 | 0.438
SimCLR - $5\%$ | 0.697 | 0.516 | 0.565 | 0.549
METS - 0% | 0.746 | 0.576 | 0.612 | 0.593
Supervised
Resnet18 - $100\%$ | 0.790 | 0.664 | 0.607 | 0.633
Table 4: MIT-BIH result. $\%$ refers to fractions of labels used in the training data. Methods | Accuracy | Precision | Recall | F1
---|---|---|---|---
Self-supervised
random - $5\%$ | 0.565 | 0.468 | 0.499 | 0.483
SimCLR - $5\%$ | 0.749 | 0.642 | 0.610 | 0.624
METS - 0% | 0.794 | 0.680 | 0.735 | 0.706
Supervised
Resnet18 - $100\%$ | 0.836 | 0.697 | 0.712 | 0.704

###  3.2 Implementation Details

The models for the transformer were taken from the transformer library [Wolf
et al.(2020)Wolf, Debut, Sanh, Chaumond, Delangue, Moi, Cistac, Rault, Louf,
Funtowicz, et al.]. We took a linear projection head with an output dimension
of 128 and a temperature $\tau$ initialized to 0.07. The ECG encoder is
optimized using Adam optimizer with a learning rate of 1e-3 and weight decay
of 1e-3. We use 50 epochs and a batch size of 32 for pre-training and
downstream tasks. The experiments were conducted using PyTorch 1.7 on NVIDIA
GeForce RTX-3090 GPU, which took about 8 hours.

###  3.3 Baselines

To demonstrate the performance of the METS method, our approach is compared
with the following baselines. (1) ResNet-18 [He et al.(2016)He, Zhang, Ren,
and Sun]: We choose ResNet-18 for showing the performance of small data
fraction fine-tune. (2) SimCLR [Chen et al.(2020)Chen, Kornblith, Norouzi, and
Hinton]: A self-supervised contrast learning model that achieves good
performance in SSL. We compare it with ECG SSL. The temperature parameter of
SimCLR is set to 0.1. For all SSL methods above, we use 5% data for fine-
tuning. (3) Supervised [He et al.(2016)He, Zhang, Ren, and Sun]: We train
ResNet1d-18 in a supervised manner, in order to compare the learning
performance of our method with that of a fully supervised.

###  3.4 Results and Discussion

In this experiment, we assessed the ECG classification using the commonly used
metrics: Accuracy, Precision, Recall, F1. We first performed a zero-shot
classification of the PTB-XL Test set for the diagnostic superclass. We
illustrate the classification results for the diagnostic superclasses in Table
1. It can be found that our method outperforms all other SSL methods with
comparable performance to supervised training. For example, compared to
SimCLR, accuracy and F1 are improved by 11% and 4%, respectively. In contrast,
METS far outperforms other SSL methods in the form classification results
presented in Table 2 and achieves better performance than supervised learning
in accuracy, precision, and F1. As shown in Table 3, METS also achieves good
performance for the rhythm classification task. Overall, our results in PTB-XL
show that the representations learned by METS are more informative than those
of other state-of-the-art SSL methods. This also demonstrates that reports
containing a priori knowledge can improve performance on the metrics.

We evaluated the performance of METS in migration learning. Table 4 shows the
performance comparison under cross-dataset testing. In general, METS
outperforms other state-of-the-art methods and even outperforms supervised
learning. Compared to Table 1, there is a significant improvement in F1 for
METS. This suggests that the features learned by METS are robust and have the
potential to be generalized to other data sources.

##  4 Conclusion

In this paper, we present METS that uses automatically generated clinical
reports to guide ECG pre-training. We pre-train the ECG encoder by applying
the rich medical knowledge from the frozen large language model to the report
text. As a result, our approach is independent of the class of annotated data
and can be directly migrated to any unseen database. We can also do
classification directly with zero-shot, unlike other SSL methods that require
fine-tuning. Our experiments demonstrate that METS can be adapted to various
downstream tasks, e.g. form, rhythm, disease, and abnormality classification.
This means that the METS approach is more effective and efficient.

## Acknowledgement

This work was supported by the National Natural Science Foundation of China
(No.62102008).

## References

  * [Addison(2005)] Paul S Addison.  Wavelet transforms and the ecg: a review.  _Physiological measurement_ , 26(5):R155, 2005\.
  * [Alsentzer et al.(2019)Alsentzer, Murphy, Boag, Weng, Jin, Naumann, and McDermott] Emily Alsentzer, John R Murphy, Willie Boag, Wei-Hung Weng, Di Jin, Tristan Naumann, and Matthew McDermott.  Publicly available clinical bert embeddings.  _arXiv preprint arXiv:1904.03323_ , 2019.
  * [Ayer and Terkelsen(2014)] Antoine Ayer and Christian Juhl Terkelsen.  Difficult ecgs in stemi:: Lessons learned from serial sampling of pre-and in-hospital ecgs.  _Journal of electrocardiology_ , 47(4):448–458, 2014.
  * [Baloglu et al.(2019)Baloglu, Talo, Yildirim, San Tan, and Acharya] Ulas Baran Baloglu, Muhammed Talo, Ozal Yildirim, Ru San Tan, and U Rajendra Acharya.  Classification of myocardial infarction with multi-lead ecg signals and deep cnn.  _Pattern Recognition Letters_ , 122:23–30, 2019.
  * [Berkaya et al.(2018)Berkaya, Uysal, Gunal, Ergin, Gunal, and Gulmezoglu] Selcan Kaplan Berkaya, Alper Kursat Uysal, Efnan Sora Gunal, Semih Ergin, Serkan Gunal, and M Bilginer Gulmezoglu.  A survey on ecg analysis.  _Biomedical Signal Processing and Control_ , 43:216–235, 2018.
  * [Bhaskarpandit et al.(2022)Bhaskarpandit, Gupta, and Gupta] Sathvik Bhaskarpandit, Priyanka Gupta, and Manik Gupta.  Lets-gzsl: A latent embedding model for time series generalized zero shot learning.  _arXiv preprint arXiv:2207.12007_ , 2022.
  * [Chen et al.(2020)Chen, Kornblith, Norouzi, and Hinton] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.  A simple framework for contrastive learning of visual representations.  In _International conference on machine learning_ , pages 1597–1607. PMLR, 2020.
  * [Chou et al.(2020)Chou, Hong, Zhou, Shang, Song, and Li] Yen-hsiu Chou, Shenda Hong, Yuxi Zhou, Junyuan Shang, Moxian Song, and Hongyan Li.  Knowledge-shot learning: An interpretable deep model for classifying imbalanced electrocardiography data.  _Neurocomputing_ , 417:64–73, 2020.
  * [Ebrahimi et al.(2020)Ebrahimi, Loni, Daneshtalab, and Gharehbaghi] Zahra Ebrahimi, Mohammad Loni, Masoud Daneshtalab, and Arash Gharehbaghi.  A review on deep learning methods for ecg arrhythmia classification.  _Expert Systems with Applications: X_ , 7:100033, 2020.
  * [He et al.(2016)He, Zhang, Ren, and Sun] 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.
  * [Hong et al.(2020)Hong, Xu, Khare, Priambada, Maher, Aljiffry, Sun, and Tumanov] Shenda Hong, Yanbo Xu, Alind Khare, Satria Priambada, Kevin Maher, Alaa Aljiffry, Jimeng Sun, and Alexey Tumanov.  Holmes: health online model ensemble serving for deep learning models in intensive care units.  In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, pages 1614–1624, 2020.
  * [Huang et al.(2021)Huang, Shen, Lungren, and Yeung] Shih-Cheng Huang, Liyue Shen, Matthew P Lungren, and Serena Yeung.  Gloria: A multimodal global-local representation learning framework for label-efficient medical image recognition.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 3942–3951, 2021.
  * [Jaiswal et al.(2020)Jaiswal, Babu, Zadeh, Banerjee, and Makedon] Ashish Jaiswal, Ashwin Ramesh Babu, Mohammad Zaki Zadeh, Debapriya Banerjee, and Fillia Makedon.  A survey on contrastive self-supervised learning.  _Technologies_ , 9(1):2, 2020.
  * [Jia et al.(2021a)Jia, Yang, Xia, Chen, Parekh, Pham, Le, Sung, Li, and Duerig] Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoc Le, Yun-Hsuan Sung, Zhen Li, and Tom Duerig.  Scaling up visual and vision-language representation learning with noisy text supervision.  In _International Conference on Machine Learning_ , pages 4904–4916. PMLR, 2021a.
  * [Jia et al.(2021b)Jia, Yang, Xia, Chen, Parekh, Pham, Le, Sung, Li, and Duerig] Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoc Le, Yun-Hsuan Sung, Zhen Li, and Tom Duerig.  Scaling up visual and vision-language representation learning with noisy text supervision.  In _International Conference on Machine Learning_ , pages 4904–4916. PMLR, 2021b.
  * [Krishnan and Sowmya Kamath(2018)] Gokul S Krishnan and S Sowmya Kamath.  A supervised learning approach for icu mortality prediction based on unstructured electrocardiogram text reports.  In _International Conference on Applications of Natural Language to Information Systems_ , pages 126–134. Springer, 2018.
  * [Lan et al.(2022)Lan, Ng, Hong, and Feng] Xiang Lan, Dianwen Ng, Shenda Hong, and Mengling Feng.  Intra-inter subject self-supervised learning for multivariate cardiac signals.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 36, pages 4532–4540, 2022.
  * [Li et al.(2021)Li, Selvaraju, Gotmare, Joty, Xiong, and Hoi] Junnan Li, Ramprasaath Selvaraju, Akhilesh Gotmare, Shafiq Joty, Caiming Xiong, and Steven Chu Hong Hoi.  Align before fuse: Vision and language representation learning with momentum distillation.  _Advances in neural information processing systems_ , 34:9694–9705, 2021.
  * [Liang et al.(2022)Liang, Zadeh, and Morency] Paul Pu Liang, Amir Zadeh, and Louis-Philippe Morency.  Foundations and recent trends in multimodal machine learning: Principles, challenges, and open questions.  _arXiv preprint arXiv:2209.03430_ , 2022.
  * [Mehari and Strodthoff(2022)] Temesgen Mehari and Nils Strodthoff.  Self-supervised representation learning from 12-lead ecg data.  _Computers in Biology and Medicine_ , 141:105114, 2022.
  * [Moody and Mark(2001)] George B Moody and Roger G Mark.  The impact of the mit-bih arrhythmia database.  _IEEE Engineering in Medicine and Biology Magazine_ , 20(3):45–50, 2001.
  * [Pourpanah et al.(2020)Pourpanah, Abdar, Luo, Zhou, Wang, Lim, and Wang] Farhad Pourpanah, Moloud Abdar, Yuxuan Luo, Xinlei Zhou, Ran Wang, Chee Peng Lim, and Xi-Zhao Wang.  A review of generalized zero-shot learning methods.  _arXiv preprint arXiv:2011.08641_ , 2020.
  * [Radford et al.(2021)Radford, Kim, Hallacy, Ramesh, Goh, Agarwal, Sastry, Askell, Mishkin, Clark, et al.] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In _International Conference on Machine Learning_ , pages 8748–8763. PMLR, 2021.
  * [Singh et al.(2018)Singh, Pandey, Pawar, and Janghel] Shraddha Singh, Saroj Kumar Pandey, Urja Pawar, and Rekh Ram Janghel.  Classification of ecg arrhythmia using recurrent neural networks.  _Procedia computer science_ , 132:1290–1297, 2018.
  * [Socher et al.(2013)Socher, Ganjoo, Manning, and Ng] Richard Socher, Milind Ganjoo, Christopher D Manning, and Andrew Ng.  Zero-shot learning through cross-modal transfer.  _Advances in neural information processing systems_ , 26, 2013.
  * [Tripathy et al.(2019)Tripathy, Bhattacharyya, and Pachori] Rajesh Kumar Tripathy, Abhijit Bhattacharyya, and Ram Bilas Pachori.  A novel approach for detection of myocardial infarction from ecg signals of multiple electrodes.  _IEEE Sensors Journal_ , 19(12):4509–4517, 2019\.
  * [Tsimpoukelli et al.(2021)Tsimpoukelli, Menick, Cabi, Eslami, Vinyals, and Hill] Maria Tsimpoukelli, Jacob L Menick, Serkan Cabi, SM Eslami, Oriol Vinyals, and Felix Hill.  Multimodal few-shot learning with frozen language models.  _Advances in Neural Information Processing Systems_ , 34:200–212, 2021.
  * [Wagner et al.(2020)Wagner, Strodthoff, Bousseljot, Kreiseler, Lunze, Samek, and Schaeffter] Patrick Wagner, Nils Strodthoff, Ralf-Dieter Bousseljot, Dieter Kreiseler, Fatima I Lunze, Wojciech Samek, and Tobias Schaeffter.  Ptb-xl, a large publicly available electrocardiography dataset.  _Scientific data_ , 7(1):1–15, 2020.
  * [Wang et al.(2019)Wang, Zheng, Yu, and Miao] Wei Wang, Vincent W Zheng, Han Yu, and Chunyan Miao.  A survey of zero-shot learning: Settings, methods, and applications.  _ACM Transactions on Intelligent Systems and Technology (TIST)_ , 10(2):1–37, 2019.
  * [Wang et al.(2022)Wang, Wu, Agarwal, and Sun] Zifeng Wang, Zhenbang Wu, Dinesh Agarwal, and Jimeng Sun.  Medclip: Contrastive learning from unpaired medical images and text.  _arXiv preprint arXiv:2210.10163_ , 2022.
  * [Wolf et al.(2020)Wolf, Debut, Sanh, Chaumond, Delangue, Moi, Cistac, Rault, Louf, Funtowicz, et al.] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al.  Transformers: State-of-the-art natural language processing.  In _Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations_ , pages 38–45, 2020.
  * [Xian et al.(2018)Xian, Lampert, Schiele, and Akata] Yongqin Xian, Christoph H Lampert, Bernt Schiele, and Zeynep Akata.  Zero-shot learning—a comprehensive evaluation of the good, the bad and the ugly.  _IEEE transactions on pattern analysis and machine intelligence_ , 41(9):2251–2265, 2018.
  * [Xu et al.(2020)Xu, Jeong, and Li] Xue Xu, Sohyun Jeong, and Jianqiang Li.  Interpretation of electrocardiogram (ecg) rhythm by combined cnn and bilstm.  _Ieee Access_ , 8:125380–125388, 2020.
  * [Yamaç et al.(2022)Yamaç, Duman, Adalıoğlu, Kiranyaz, and Gabbouj] Mehmet Yamaç, Mert Duman, İlke Adalıoğlu, Serkan Kiranyaz, and Moncef Gabbouj.  A personalized zero-shot ecg arrhythmia monitoring system: From sparse representation based domain adaption to energy efficient abnormal beat detection for practical ecg surveillance.  _arXiv preprint arXiv:2207.07089_ , 2022.
  * [Yuan et al.(2021)Yuan, Chen, Chen, Codella, Dai, Gao, Hu, Huang, Li, Li, et al.] Lu Yuan, Dongdong Chen, Yi-Ling Chen, Noel Codella, Xiyang Dai, Jianfeng Gao, Houdong Hu, Xuedong Huang, Boxin Li, Chunyuan Li, et al.  Florence: A new foundation model for computer vision.  _arXiv preprint arXiv:2111.11432_ , 2021.
  * [Zhai et al.(2022)Zhai, Wang, Mustafa, Steiner, Keysers, Kolesnikov, and Beyer] Xiaohua Zhai, Xiao Wang, Basil Mustafa, Andreas Steiner, Daniel Keysers, Alexander Kolesnikov, and Lucas Beyer.  Lit: Zero-shot transfer with locked-image text tuning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 18123–18133, 2022.

##  Appendix A Frozen Language Model Details

The ClinicalBert language model trained on all MIMIC notes and was initialized
from BioBERT. The model was trained using code from Google’s BERT repository
on a GeForce GTX TITAN X 12 GB GPU. Model parameters were initialized with
BioBERT. The model was trained using a batch size of 32, a maximum sequence
length of 128, and a learning rate of 5e-5 for pre-training language models.
The models trained on all MIMIC notes were trained for 150,000 steps. The dup
factor for duplicating input data with different masks was set to 5.
Specifically, masked language model probability = 0.15 and max predictions per
sequence = 20 [Alsentzer et al.(2019)Alsentzer, Murphy, Boag, Weng, Jin,
Naumann, and McDermott].

##  Appendix B Example texts of inputted Frozen Language Model

Figure 3: Example texts of inputted Frozen Language Model.

Example texts of inputted Frozen Language Model are shown in figure 4.

First, during the pre-training process of self-supervised learning, in order
for the large language model to understand the report text, we extend the
report text into the language model as a complete sentence input. A sample
inputted text report is shown below.

$\textbf{``The\ report\ of\ the\ ECG\ is\ that\ \\{sinus\ arrhythmia\
otherwise\ normal\ ECG\\}''}.$

Second, in zero-shot classification, we expand the labels into complete
sentences for input into the frozen language model. This may be useful for
specifying categories. Here are a few examples.

In the diagnostic classification task: “The ECG of {X}, a type of diagnostic.”
where X represents the different superclass diagnostic labels, e.g. Normal
ECG, Conduction Disturbance, Myocardial Infarction, etc.

In the form classification task: “The ECG of {Y}, a type of form.” where Y
represents a different form label, such as Abnormal QRS, Non-diagnostic
abnormalities, etc.

##  Appendix C Description of test sets.

Figure 4: Description of test sets.