# iART: Learning from Demonstration for Assisted Robotic Therapy Using LSTM

Shrey Pareek1 and Thenkurussi Kesavadas2 Research supported by the National
Science Foundation under Grant No. 1502339.1 Shrey Pareek is with the
Department of Industrial and Enterprise Systems Engineering, University of
Illinois, Urbana-Champaign, IL-USA<EMAIL_ADDRESS>Thenkurussi
Kesavadas is with the Department of Industrial and Enterprise Systems
Engineering, University of Illinois, Urbana-Champaign, IL-USA,
<EMAIL_ADDRESS>
###### Abstract

In this paper, we present an intelligent Assistant for Robotic Therapy (iART),
that provides robotic assistance during 3D trajectory tracking tasks. We
propose a novel LSTM-based robot learning from demonstration (LfD) paradigm to
mimic a therapist’s assistance behavior. iART presents a trajectory agnostic
LfD routine that can generalize learned behavior from a single trajectory to
any 3D shape. Once the therapist’s behavior has been learned, iART enables the
patient to modify this behavior as per their preference. The system requires
only a single demonstration of 2 minutes and exhibits a mean accuracy of
91.41% in predicting, and hence mimicking a therapist’s assistance behavior.
The system delivers stable assistance in realtime and successfully reproduces
different types of assistance behaviors.

##  I INTRODUCTION

Traditional stroke therapy is hospital-centric and involves a team of doctors
and therapists designing repetitive therapy exercises for a patient. The
success of these programs relies heavily on the ability of the patient to
practice the prescribed exercises at home. In these situations, robot-based
tele-rehabilitation systems serve as reliable tools for home-based therapy.

Robot-based tele-rehabilitation systems [1] comprise of a robotic system that
can assist/resist a patient’s motions as they perform therapy exercises using
a simulation system. These systems also allow a remote therapist to observe
and assist the patient over a network and modify the therapy tasks and/or
robotic assistance/resistance [2].

The robotic device may be programmed to automatically assist/resist the
patient in the absence of a therapist. The selection of the type of robotic-
assistance is a crucial yet challenging aspect of robotic rehabilitation. A
task rendered too difficult due to the presence of inadequate robotic
assistance may induce anxiety [3] in the patient; forcing them to quit the
rehabilitation task early. Conversely, excessive assistance in a task may not
challenge the patient enough and can lead to boredom [3]. Further, this may
lead to over-reliance on the robotic assistance which may be inhibitory to the
rehabilitation outcome [4, 5]. Thus, the choice of the assistance mechanism is
crucial to the success of the rehabilitative routine.

One such commonly used assistance mechanism is the error reduction (ER)
strategy. ER refers to an assistance strategy that minimizes the tracking
error in a trajectory tracking task. These strategies are usually dependent on
the instantaneous tracking error and follow strict threshold-based rules. In
other words, this strategy constrains the subject within a fixed no-error zone
around the reference trajectory. If the subject deviates outside this zone,
the robotic device provides a corrective force and guides them back inside
this zone. Mathematically, if the tracking error is larger than a threshold
$r$ (called the radius of the no-error zone), robotic assistance is switched
on. Conversely, if the tracking error is lower than this threshold, assistance
is switched off (see Appendix for details). Such rule-based methods only use
the instantaneous position error as a metric for deciding whether or not to
assist the patient; and do not take into account the patient’s performance
history. This makes the selection of the size of the no-error zone, very
crucial and challenging.

In our experiments with robot-based trajectory tracking, we observed that
subjects tend to stay at the boundary of the virtual no-error zone, as at the
boundary, the haptic device provides minimal assistance enabling the subjects
to correctly follow the trajectory with minimal effort and low tracking error.
The plots in Fig. 1 show the absolute tracking error v/s time (blue line)
exhibited by a subject under two assistance conditions (ER and human
therapist) for the same reference trajectory. The red dots denote instances
when the robotic assistance was switched on; and the dotted black line
represents the radius ($r$) of the no-error zone (see Appendix).

It is evident that under the rule-based assistance, the subject tends to stay
close to the boundary of the no-error zone. The robotic assistance is switched
on for around $50\%$ ($\sim 7.5s$) of the duration of the experiment. In the
second case, the therapist assists the subject for about $10\%$ ($\sim 6.5s$)
of the experiment and allows them to reduce the tracking error on their own
accord before intervening.

Further, under ER, the subject completes the tracking task in 15 seconds
whereas they require a little more than a minute in the second case for the
same task. Conventional wisdom dictates that lower task completion time
signifies superior performance. However, it is evident that the lower
completion time observed in the first case is due to the presence of excessive
robotic assistance. This increased dependency on the robotic assistance may be
inhibitory to the therapy outcome [4].

In other words, although rule-based approaches might improve patient
performance during assisted training, they may hinder their ability to learn
the desired skills for successful task completion in the absence of this
assistance. Hence, the choice of when to assist the patient needs be made
carefully to ensure that skills learned during robot assisted rehabilitation
are transferred to real-world scenarios.

Figure 1: Tracking behavior demonstrated under rule-based assistance (left)
and therapist assistance (right).

##  II ROBOT LEARNING FROM DEMONSTRATION

We propose the use of a robot learning from demonstration (LfD) strategy to
learn a therapist’s assistance behavior (see Fig. 1). A key challenge in
robotics is the problem of learning a mapping between the state space and
actions. This mapping, commonly referred to as policy, describes a set of
rules that enables a robot to select an optimal action based on its current
state. Although these rules or policies can be generated by hand (rule-based
methods); this process can be exhaustive and usually only works reliably for
certain scenarios [6].

LfD refers to a methodology wherein the policy is learned from demonstrations
or examples provided by a teacher (in case of robotic rehabilitation the
therapist is the teacher). These examples refer to a set of state-action pairs
recorded during a demonstration phase \- wherein the teacher demonstrates the
desired robot behavior. LfD algorithms then use these examples to derive a
policy that reliably reproduces the demonstrated behavior. LfD does not
require a large library of state-action pairs as the demonstrations can be
focused to areas of the state-space that will actually be encountered during
the task execution. Finally, since LfD relies on the assumption that expert
demonstrations are based on the optimal policy to be reproduced; the learning
process is faster.

In this paper, we approach LfD as a supervised learning problem to learn a
therapist’s assistance behavior through a set of state-action pairs. We
propose an intelligent Assistant for Robotic Therapy (iART), that enables
home-based robotic therapy for stroke rehabilitation. iART refers to a LSTM-
based [7] LfD paradigm that monitors patient-therapist interactions and
generates a patient-centric assistance platform using a single 2-minute
trajectory tracking task. The system does not require any external sensors
such as brain-computer-interface [8] or surface electromyography [9]. The
proposed method is trajectory agnostic, which implies it can learn from a
single trajectory and be generalized for multiple scenarios. Use of LSTM takes
into account the historical performance of the patient while providing
assistance, leading to the realization of an intuitive assistive paradigm.
Finally, iART enables the patient to modify the behavior of the system as per
their requirement using an iterative learning process. A video demo of iART
may be found here111https://youtu.be/Gc-VVE-lJaA.

###  II-A LfD in Robotic Rehabilitation

Although LfD has been widely implemented in mainstream robotics [6], it is a
relatively nascent concept in the field of assisted stroke rehabilitation.
Tavakoli’s research group [10, 11, 12] have pioneered the use of LfD in
robotic rehabilitation. Their algorithms rely on learning the average
impedance and/or position tracking behavior exhibited by an expert while
conducting a therapy task. Here, average position tracking refers to the
tracking behavior obtained over multiple trials of the same trajectory
tracking task. Similarly, the average impedance behavior can be used to model
the kinesthetic behavior of a therapist [12].

This average behavior is then used as a reference to which the patient’s
performance in realtime is compared. Any deviations from this reference are
used to adjust the task difficulty/assistance for the patient [10]. Jung et
al. [13] use Latent Dirichlet Allocation to model demonstrations of a
therapist in the form of a generative process for autonomous therapy. However,
this methodology is task/trajectory specific and the therapist needs to train
the system for multiple different trajectories in order to develop a
practically viable system. This makes the demonstration process highly
exhaustive and limits the scalability of the system.

Additionally, the above methods only learn the average behavior of the
therapist and do not take into account the preferences of the patient. While
adherence to the therapy regime prescribed by the therapist is essential for
successful rehabilitation; an ideal system should focus on learning the
therapist-patient interaction rather than just the therapist’s behavior.

iART makes two critical distinctions from the above methodologies. First,
instead of learning the average therapist behavior and providing assistance
based on deviations from this average behavior, iART learns when a therapist
chooses to assist a patient as they perform a trajectory tracking task. This
enables iART to be trajectory agnostic, making it scalable. However, this
means that iART does not learn the impedance behavior of a therapist [10, 12]
and only learns ’when-to-assist’ and not ’how-to-assist’. This is a limitation
of the proposed methodology. Second, iART focuses on observing therapist-
patient interactions rather than just the therapist’s behavior. This enables
iART to be tuned as per the specific requirements of a patient. However, this
requires individual demonstrations for every patient. We now describe iART’s
algorithm in detail.

Figure 2: Schematic representation of iART.

##  III iART

Figure 3: Experimental setup.

iART uses LSTM to learn a therapist’s assistance behavior. LSTM is a class of
neural networks that takes into account the past history through a sequence of
features. LSTM has the ability to memorize information over a series of
timesteps using the concept of memory cells. $c_{k}\in\mathbb{R}^{n}$
represents this long term memory at timestep $k$. $\hat{c}_{k}$ represents a
candidate state at $k$. During training LSTM learns whether to overwrite a
memory cell, retrieve it, or keep it for the next time step. The formulation
for the output state $h$ at a future timestep $k+1$ is given by [14]:

$\displaystyle\hat{c}_{k+1}$ $\displaystyle=$ $\displaystyle
tanh(w_{c}[h_{k},x_{k+1}],b_{c})$ $\displaystyle c_{k+1}$ $\displaystyle=$
$\displaystyle\sigma(w_{f}[h_{k},x_{k+1}],b_{f})\odot
c_{k}+\sigma(w_{i}[h_{k},x_{k+1}],b_{i})\odot\hat{c}_{k+1}$ $\displaystyle
h_{k+1}$ $\displaystyle=$
$\displaystyle\sigma(w_{o}[h_{k},x_{k+1}],b_{o})\odot tanh(c_{k+1})$ (1)

where, $\tanh(\cdot)$ is the tangent hyperbolic function and
$\sigma(\cdot)=\exp(\cdot)/(1+\exp(\cdot))$ is the sigmoid function. $i,f,$
and $o$ refer to the input, forget and output gates, respectively. $W_{x}$ and
$b_{x}$ are the weights and biases associated with these gates. The symbol
$\odot$ denotes the Hadamard product.

Traditionally, LSTM finds applications in scenarios such as natural language
processing and time-series prediction. We propose a novel paradigm of using
LSTM in LfD for robotic rehabilitation. Consider the LfD methodologies
described in Section II-A. These methods rely on the instantaneous behavior of
the patient in order to decide the degree of assistance to be supplied. If the
system observes deviations from the average behavior, it immediately supplies
robotic assistance to the patient. This greedy approach does not allow the
patient to reduce the tracking errors on their own accord and may increase the
patient’s dependence on the robot.

On the other hand, a therapist assisting the patient observes the patient’s
behavior over a brief period of time before supplying any assistance. They may
allow the patient to return to the desired trajectory autonomously. If
however, the patient continues to struggle and cannot execute the task
properly, the therapist may decide to intervene and assist the patient. Use of
LSTM enables iART to capture this history-based patient-therapist interaction
and successfully reproduce it for the realization of an intuitive assistive
paradigm. Non-history-based classifiers such as support vector machine or
neural networks provide assistance based on the instantaneous performance of
the subjects. In our initial experiments, LSTM outperformed all other types of
non-history-based classifiers. Overall, iART’s architecture comprises of three
key phases (Fig. 2):

  1. 1.

Demonstration Phase: refers to the data collection phase wherein the patient
performs the tracking task using a robotic device and the therapist decides
whether or not to supply robotic assistance to the patient based on their
performance.

  2. 2.

Learning Phase: involves the use of an LSTM-based neural network to learn the
therapist’s assistance behavior observed during the demonstration phase.

  3. 3.

DAgger Phase: enables the patient/therapist to test the realtime assistance
behavior learned by iART and iteratively change the behavior if needed.

We now describe these modules individually.

###  III-A Demonstration Phase

The demonstration phase involves a shared therapist-patient interaction
session during which a 3D reference trajectory is chosen. The patient is then
required to track this reference path using the end-effector of a robotic
device to the best of their ability (Fig. 3). By default, the robot does not
offer any assistance during the execution of these trajectories. If the expert
deems that robotic assistance needs to be applied, they can instruct the robot
(using a keyboard key press) to guide the subject back to the closest point on
the trajectory (denoted by $\mathbf{x_{l}}$ in Fig. 10). Assistance is
provided in accordance with the ER rule described via (2). iART eliminates the
’no-error zone’ parameter ($r$) and instead assists the subject to the closest
point on the trajectory ($x_{l}$). In other words, $r$ can be set to zero and
the error threshold that decides whether assistance needs to be provided is
replaced by the non-linear LSTM function that depends on multiple performance
measures or features (described later).

Once the therapist decides that adequate assistance has been provided to the
patient, they can disable the robotic assistance using another key press. On
average, the demonstration phase lasts for 2 minutes. Switching the assistance
on would lead to a sudden jump in the robotic assistance from 0 to a high
value, and may ’jerk’ the subject’s arm around. The $K_{p}$ gain in (2) was
chosen so as to minimize this effect while providing adequate assistance. In
the future however, the gain would be gradually increased from zero to a
maximum value (and vice versa) to completely eliminate this jump in robotic
force.

During this phase, various local states are logged at a sampling rate of 30
Hz. These states are - tracking error in x ($e_{x}$), y ($e_{y}$), and z
($e_{z}$) direction, error magnitude ($e$), radius of curvature ($r_{c}$),
velocity magnitude ($v$), and a boolean describing whether the subject is
tracking the trajectory or navigating to the start point of the trajectory
after completion of a trial ($is_{track}$). Mathematically, the states are
denoted by a vector $\mathbf{s}=[e_{x},e_{y},e_{z},e,r_{c},v,is_{track}]$.

####  III-A1 State-Action Pairs

The use of local states enables us to train a trajectory agnostic LfD model.
Consider the use of x, y, z error. Using tracking error as a feature as
opposed to the absolute position of the robot end-effector is analogous to
using distance information from LIDAR or RADAR signals to localize neighboring
vehicles during autonomous driving. The assistance behavior depends on the
distance from the reference trajectory and not the location of the end-
effector in the world-space. Velocity magnitude captures the ability of a
patient to reduce any tracking errors. For instance, a lower velocity with
high tracking error may denote the inability of the patient to return the end
effector closer to the reference trajectory. Use of radius of curvature
captures the assistance behavior w.r.t. the complexity of the trajectory.
Straight sections of the trajectory (large radius of curvature) are generally
easier to traverse and would exhibit lower tracking error; signifying need of
assistance even when the error magnitude may be small. Conversely, curved
sections (small radius of curvature) would allow more leeway in terms of the
allowable error before the assistance in enabled.

It should be noted that the effects of the states such as velocity magnitude
or curvature on assistance described above have been only included as an
intuition to explain the reasoning behind the choice of these features. Since
we are using an LSTM (with highly non-linear dynamics), it is impossible to
identify the true role that these states might play during the actual learning
process. Local states at each epoch are associated with a corresponding action
- robotic assistance On or Off. The action is denoted by $a\in\\{0,1\\}$. This
set of state-action pairs ($\mathbf{s},a$) is used for supervised learning
during the learning phase.

###  III-B Learning Phase

Learning phase refers to the training phase of iART wherein an LSTM model is
used to mimic the therapist-patient interaction recorded during the
demonstration phase. The learned model is then used to automatically supply
realtime adaptive assistance to the patient in the absence of the therapist.

The LSTM architecture comprises of one hidden LSTM layer with 100 neurons,
each with a hyperbolic tangent (tanh) activation function. The output layer
uses a single neuron with sigmoid activation function. A batch size of 32 is
used. Finally, a mean-squared-error (MSE) loss function is used with the Adam
optimizer and the model is trained for 400 epochs.

As mentioned earlier, the network uses a set of local state-action pairs as
the labelled dataset for supervised learning. Since LSTM requires historic
data, we use a moving window of 1 sec (or 30 data points, since the data are
sampled at a rate of 30 Hz) to provide this past information. This means that
the network observes the patient behavior over a moving window of 1 sec (with
96.66% or a 29 step overlap between subsequent windows) to determine whether
or not to provide assistance to the subject. The final output is given as the
probability of assistance ($P(A)$). If $P(A)>0.5$, robotic assistance is
supplied to the patient and vice-versa.

###  III-C DAgger Phase

The focus of iART is the development of a patient-centric assistance platform.
To this end, iART enables the patient (or therapist) to modify the assistance
behavior learned by iART as per their need through iART’s DAgger module.

The DAgger module is inspired by the Dataset Aggregation (DAgger) methodology
proposed in [15]. DAgger is an iterative imitation learning algorithm that
aims at mimicking the expert’s policy $\pi_{e}$, iteratively. Briefly, in the
first iteration, a dataset of expert examples $\chi$ is collected and used to
extract the approximate policy $\hat{\pi_{1}}$. Then at iteration $n$, it uses
policy $\hat{\pi_{n-1}}$ (i.e. the previous best policy) to collect more
examples and adds them to the dataset $\chi$. The augmented dataset is then
used to estimate the next policy $\pi_{n}$ that best mimics the expert’s
policy $\pi_{e}$ 222Detailed explanation of DAgger may be found in the
original implementation at [15].

####  III-C1 Modifying DAgger For Retraining

In this paper, we modify the original DAgger implementation to enable a
patient to change the assistance behavior of pre-trained models. For instance,
consider the straight line tracking task between two points $\mathbf{p_{1}}$
and $\mathbf{p_{2}}$ as shown in Fig. 10. While the patient tracks the
straight line between $\mathbf{p_{1}}$ and $\mathbf{p_{2}}$, the therapist
decides whether or not to provide assistance. Once the patient navigates to
the end point of the line ($\mathbf{p_{2}}$), they are required to return to
the starting point ($\mathbf{p_{1}}$) and continue the tracking task. The
therapist may choose to assist the patient as they traverse their path back to
the starting point. In this case, if the robotic assistance is switched on,
the patient will be guided directly to $\mathbf{p_{1}}$ 333 In other words,
the desired point $\mathbf{x_{d}}$ in (2) will be set as $\mathbf{p_{1}}$..

Consider the following scenario - during the demonstration phase (Section
III-A) the therapist always assists the patient as they return to the starting
point $\mathbf{p_{1}}$. The LSTM model described in Section III-B uses a
dataset $\chi$ that exhibits this behavior to learn a representation of this
policy $\hat{\pi_{1}}$. This is analogous to the first iteration of DAgger
described above. Over time the patient’s performance improves and they no
longer require assistance while returning to $\mathbf{p_{1}}$.

The DAgger module enables the patient to modify this behavior. The patient can
override iART’s actions by pressing a key on the keyboard. In this case, since
iART will assist the patient by guiding them back to $\mathbf{p_{1}}$, the
patient overrides this action to not assist during the return phase. These
overridden data points $\chi_{o}$ are then used to augment the original
dataset $\chi$. This is a modification of the original DAgger implementation,
wherein all data points are used to augment the dataset.

The LSTM network in then retrained on the augmented dataset. Additionally, the
overridden data points $\chi_{o}$ are assigned higher weightage during the
learning task to adapt quickly to the new desired behaviour. We use a ratio of
$\beta$:1 ($\beta=20$ in this case) to assign higher weightage to the
overridden data points. This means that, mis-classifications on the overridden
data points $\chi_{o}$, are penalized $\beta$ times more by the loss function
( $\beta^{2}$ times in case of MSE) compared to the data points in the
original dataset $\chi$. The new policy $\pi_{n}$ is then learned on the
augmented dataset (see Algorithm 1).

Algorithm 1 iART Algorithm: Learning Expert Policy $\pi_{e}$

1: Initialize $\chi\leftarrow\phi$

2: Demonstration Phase:

3: Collect expert state-action pairs $[\mathbf{s},a]\rightarrow\chi$.

4: Learning Phase:

5: Train LSTM on $\chi$ to learn approximate policy $\hat{\pi_{1}}$.

6: DAgger Phase:

7: for $n=2,3,\dots$ do

8: Get dataset $\chi_{o}$ of overridden state-action pairs.

9: Aggregate data sets: $\chi\leftarrow\chi\cup\beta\chi_{o}$.

10: Learning Phase:

11: Learn new policy $\hat{\pi_{n}}$ using LSTM on updated $\chi$.

12: end for

13: Return chosen policy $\hat{\pi_{i}}$ for $i\in 1,2,\dots$.

##  IV EXPERIMENTAL EVALUATION

We studied the efficacy and feasibility of iART through four different
experiments. Nine healthy subjects (average age 27.22 years; 7 male; 2 female;
8 right handed; 1 left handed) who served as the ’patient’ were recruited for
the experiments. One of the authors served as the ’therapist’.

The experiments involved a Unity-based 3D trajectory tracking task [16, 17]
wherein the subject uses the end-effector of a robotic/haptic device
(Geomagic® TouchTM) as a writing stylus controller. The subjects were required
to track 3D reference trajectories chosen from Fig. 4. Use of 3D trajectories
on a 2D computer monitor increased the complexity of the task and led to large
(forced) tracking errors despite the use of healthy subjects in the
experiment. This was intentional as it forced subjects to make errors and thus
enabled the therapist to provide assistance.

Figure 4: Reference 3D curves.

###  IV-A Imitation Learning

The first experiment was designed to study the imitation learning aspects of
iART. During the experiment, as the subjects followed the reference
trajectory, the therapist toggled the robotic assistance on or off using a
keyboard key press. Assistance was supplied according to the ER rule described
in Section III-A. Two curves were randomly chosen from Fig. 4 to serve as the
reference trajectory. These shapes were chosen as they provided a balance
between complexity and covering basic motor primitives. The subjects were
required to track each trajectory twice. State-action pairs collected during
the execution of the first trajectory served as the training dataset $\chi$
for the LSTM network. A different model was trained for each subject. Data
from the second trajectory were used for testing the trajectory-agnostic
nature of iART. The results from this experiment are summarized in Section
V-A. The $K_{p}$ and $K_{d}$ gains in (2) were set to 4 and 0.001,
respectively.

###  IV-B Realtime Testing

The second experiment tested the realtime prediction capabilities of iART. 7
out of the original 9 subjects were recruited again for this experiment. The
trained model obtained from the first experiment was then used for realtime
assistance prediction using Python as subjects tracked a reference trajectory
chosen from Fig. 4. Local states were streamed to from Unity3D to Python using
Transmission Control Protocol. Once the assistance was predicted using Python,
the assistance behavior was communicated to the robotic device using a
simulated key press. During this trial, the therapist used a key press (that
had no influence on the robotic assistance) to indicate whenever they would
have provided assistance to the user. In other words, while the actual
assistance was provided through iART, corresponding assistance behavior of the
therapist was recorded simultaneously as ground truth to verify the realtime
performance of iART in mimicking the therapist.

###  IV-C DAgger Phase

Next, we tested the retraining aspect of iART through the DAgger Phase. The
experiment was repeated with one of the subjects. During the demonstration
phase, the therapist adopted an ’assist-too-often’ strategy wherein, they
assisted the subject for very small deviations from the reference trajectory.
Additionally, the subject was assisted by the therapist while returning to the
starting point (Section III-C). An LSTM model was trained using this behavior
and a new reference trajectory was chosen for realtime testing and retraining.

During the DAgger phase, the subject was instructed to modify the assistance
behavior of the model using iART’s override methodology. They were instructed
to retrain the network to prevent it from assisting while the subject
navigated to the starting point and to assist for larger deviations from the
trajectory as opposed to the learned behavior of ’assist-too-often’.

###  IV-D Special Cases

We designed two special cases to test the imitation learning capabilities of
iART using one of the subjects. These case were:

#####  Assist-too-often

The demonstration phase involved the therapist assisting the subject for small
deviations from the reference trajectory (same as above).

##### Assist-on-stop

In this case, the therapist assisted the subject whenever they stopped moving
the end effector, irrespective of the tracking error. The models trained using
these behaviors were then tested in realtime on different reference
trajectories.

##  V RESULTS AND DISCUSSIONS

In this section, we provide the results for the aforementioned experimental
evaluation. The efficacy of iART is represented through classification
accuracy in the offline and realtime setting and graphical representations to
enable visual comparisons between the assistance behavior demonstrated by the
therapist and that learnt by iART. We also conducted a paired t-test on the
percentage time during which the assistance was switched on by the therapist
and iART as a metric to validate the realtime performance of iART. Another
statistical analysis was conducted to compare the velocity and position error
at which iART and therapist assistance are switched on to understand special
behaviors learnt by iART and also to validate the DAgger module.

###  V-A Imitation Learning

Across the nine subjects, the system demonstrated a mean accuracy of $91.41\%$
(range 85.77%-94.71%) in predicting (and mimicking) the assistance behavior of
the therapist. The system demonstrated a mean accuracy of $90.21\%$ in
identifying the assistance-off action (true negatives) and $96.24\%$ in
identifying the assistance-on action (true positives).

Figure 5: Classification results on one of the test trajectories.

We present the prediction results for one of the test trajectories in Fig 5.
Remember that the test trajectory was different from the one used during the
demonstration and learning phase. The therapist’s actions are denoted by blue
squares and the dots represent iART’s predicted actions. Cases where iART’s
actions matched the therapist’s are labelled in green and mis-classifications
are marked as red. It can be observed that the system performs reliably
throughout the experiment. Most mis-classifications occur at the point where
the assistance is switched from off to on, and vice-versa. These mis-
classifications are usually off by a few time-steps (a few milliseconds) and
can be ignored for practical purposes. It is clear from the above observations
that iART provides a reliable methodology to mimic a therapist’s assistance
behavior and also highlights its trajectory agnostic nature.

###  V-B Realtime Testing

Figure 6: Realtime assistance comparison between therapist and iART for two
different subjects.

The system demonstrated a realtime tracking accuracy of 81.85% (range 76.48% -
89.31%) across 7 subjects. The ground truth for this case was collected from
the therapist while iART provided the actual assistance. The mean accuracy for
these 7 subjects in the offline setting was 88.41% (Section V-A). Despite the
lower accuracy demonstrated by iART in realtime, the assistance pattern was
similar to the one exhibited by the ground truth collected from the therapist.

Fig. 6 demonstrates the performance of iART for two different subjects. In
both cases, iART’s assistance behavior (solid red line) was similar to the
therapist’s (dotted blue line). The assistance switching behavior was stable
and did not switch from one state to another too often. This enabled a smooth
assistance mechanism. One difference was the therapist’s tendency to switch
the assistance on a little earlier when compared with iART. Additionally, iART
tends to assist over shorter intervals compared to the assistance provided by
the therapist. This may be due to the increased familiarity of the subject
with the system and hence the need for shorter bursts of assistance. Since at
this point the subject had been using the system for a while, some learning
effects may be present, which would have reduced their assistance
requirements. The subjects also reported that they were unable to distinguish
the assistance behavior of iART from the original experiment involving an
expert.

We also conducted a paired t-test on the percentage time for which assistance
was switched on by the therapist and iART across the 7 subjects. The two
demonstrated no significant differences ($t(6)=0.92$, $p=0.39$). While the
therapist kept the assistance on for an average of 24.81% (std 5.91%) of the
time across the 7 subjects, iART kept it on for 26.71% (std 5.44%). This
reaffirms the ability of iART to learn and recreate the therapist’s assistance
behavior in realtime.

We also report iART’s performance when tested on the same reference trajectory
as used for training. In this case the ground truth refers to therapist
assistance collected during the initial demonstration phase. Fig. 7
demonstrates the tracking behavior exhibited by the subject under therapist
(green) and iART (blue) assistance for the same reference trajectory (black).
The green and blue circles denote instances where the robotic assistance was
switched on by the therapist and iART, respectively.

iART was able to successfully mimic the therapist’s assistance behavior.
Additionally, in this particular case, it was observed that the subject
required assistance at similar phases of the trajectory. For instance, in both
cases the therapist and iART assisted the subject at the mid-way point of the
trajectory. Similar behavior was exhibited towards the last quarter of the
trajectory.

Figure 7: Tracking behavior executed by a subject under therapist and iART
assistance for the same reference trajectory. Figure 8: Box plots
demonstrating the position error for which the robotic assistance is switched
on under: assist-too-often therapist behavior (left); iART’s learned assist-
too-often behavior (middle); and iART modified using the DAgger phase to
assist for larger deviations (right). The red asterisks denote significant
differences.

###  V-C DAgger Phase

The goal of the experiment was to modify the assist-too-often behavior of the
therapist and iART. Fig. 8 shows the box plots for the tracking error
magnitude during which the robotic assistance is switched on. The first two
plots depict the therapist behavior and the learned iART assistance. It is
evident from the plots that both the therapist and iART assist for very small
tracking errors as described by the assist-too-often strategy. The final plot
describes the modified iART behavior wherein the subject modified the learned
iART behavior to assist for larger deviation.

It was observed that the subject was able to successfully retrain the model to
learn the new desired behavior. A Welch’s t-test was conducted between the
three pairs viz. Therapist-iART, Therapist-Modified iART and iART-Modified
iART. Both Therapist-Modified iART and iART-Modified iART demonstrated
significant differences at a p-value $<0.02$. reaffirming iART’s ability to
retrain models.

Additionally, the subject was instructed to modify the assistance mechanism to
not assist as they navigated to the starting point. The subject was able to
successfully modify this behavior. The original model assisted for 62.18% of
the time-steps as the subject navigated to the start point. In the modified
behavior, this was reduced to 28.25% of the time-steps.

###  V-D Special Cases

Finally, we report iART’s ability to mimic special behaviors exhibited by the
therapist.

##### Assist-too-often

As explained in the previous section and Fig. 8, iART can successfully mimic
this behavior.

##### Assist-on-stop

For this case, we present the box plots describing the cursor velocity at
which the assistance is toggled from off to on. Fig. 9 shows this behavior
under three different settings. The first plot describes the assistance
behavior of the therapist. Remember, the therapist switches on the assistance
whenever the end-effector comes to a near stop (low velocity). Compared with
the plot for iART, it was observed that both cases exhibited a similar spread
of cursor velocity, indicating that iART can successfully learn this behavior.
Additionally, we present the assistance start velocity for a generic case
wherein assistance was supplied based on the subject’s position error and
overall performance. Compared to the previous cases, the spread of velocity
covers a wider range for this generic case.

Figure 9: Box plots demonstrating the end-effector velocity at for which the
robotic assistance is switched from off to on under different settings.

##  VI CONCLUSION AND FUTURE DIRECTIONS

In this paper, we have demonstrated the use of LSTM as a novel paradigm for
robot LfD in robotic rehabilitation. We have presented a system capable of
mimicking a therapist’s assistance behavior with relatively short
demonstrations of 2 minutes. As opposed to conventional LfD techniques, our
method is trajectory-agnostic and can generalize learned behavior from a
single trajectory to any 3D shape. The system can be re-trained based on the
preferences of the patient using a modified DAgger-based methodology. iART
demonstrates a mean accuracy of $91.41\%$ in predicting, and hence mimicking a
therapist’s assistance behavior. The system delivers stable performance in
realtime and can be retrained easily to modify pre-learnt assistance
behaviors.

While the use of local features enables iART to generalize its assistance
behavior across multiple trajectories, this generalization fails to capture
the critical dynamics that may be associated with a specific shape. Although,
the use of radius of curvature aims at capturing this information, it is not
the most effective measure and may fail in the presence of obstacles on the
path.

The current system provides assistance in cases wherein subjects deviate from
the reference trajectory and is unable to guide subjects requiring assistance
along the trajectory. We are in the process of developing a similar LSTM-based
methodology that can learn this behavior using therapist demonstrations. The
new system would not only predict ’when-to-assist’ but will also decide ’how-
to-assist’ by predicting whether to guide the subject towards the trajectory
or along it.

## References

  * [1] B. R. Brewer, S. K. McDowell, and L. C. Worthen-Chaudhari, “Poststroke upper extremity rehabilitation: a review of robotic systems and clinical results,” _Topics in stroke rehabilitation_ , vol. 14, no. 6, pp. 22–44, 2007.
  * [2] S. Pareek, P. Chembrammel, and T. Kesavadas, “Development and evaluation of haptics-based rehabilitation system,” in _Medical Robotics (ISMR), 2018 International Symposium on_. IEEE, 2018, pp. 1–6.
  * [3] J. Nakamura and M. Csikszentmihalyi, “The concept of flow,” in _Flow and the foundations of positive psychology_. Springer, 2014, pp. 239–263.
  * [4] E. T. Wolbrecht, V. Chan, D. J. Reinkensmeyer, and J. E. Bobrow, “Optimizing compliant, model-based robotic assistance to promote neurorehabilitation.” _IEEE transactions on neural systems and rehabilitation engineering : a publication of the IEEE Engineering in Medicine and Biology Society_ , vol. 16, no. 3, pp. 286–297, jun 2008.
  * [5] E. T. Esfahani, S. Pareek, P. Chembrammel, M. Ghobadi, and T. Kesavadas, “Adaptation of rehabilitation system based on user’s mental engagement,” in _ASME 2015 International Design Engineering Technical Conferences and Computers and Information in Engineering Conference_. American Society of Mechanical Engineers, 2015, pp. V003T14A017–V003T14A017.
  * [6] B. D. Argall, S. Chernova, M. Veloso, and B. Browning, “A survey of robot learning from demonstration,” _Robotics and autonomous systems_ , vol. 57, no. 5, pp. 469–483, 2009.
  * [7] F. A. Gers, J. Schmidhuber, and F. Cummins, “Learning to forget: Continual prediction with lstm,” 1999.
  * [8] S. Pareek, “Cognitive adaptation for human computer interaction,” Ph.D. dissertation, State University of New York at Buffalo, 2015.
  * [9] S. Pareek, H. Manjunath, E. T. Esfahani, and T. Kesavadas, “Myotrack: Realtime estimation of subject participation in robotic rehabilitation using semg and imu,” _IEEE Access_ , vol. 7, pp. 76 030–76 041, 2019.
  * [10] M. Maaref, A. Rezazadeh, K. Shamaei, R. Ocampo, and T. Mahdi, “A bicycle cranking model for assist-as-needed robotic rehabilitation therapy using learning from demonstration,” _IEEE Robotics and Automation Letters_ , vol. 1, no. 2, pp. 653–660, 2016.
  * [11] M. Najafi, M. Sharifi, K. Adams, and M. Tavakoli, “Robotic assistance for children with cerebral palsy based on learning from tele-cooperative demonstration,” _International Journal of Intelligent Robotics and Applications_ , vol. 1, no. 1, pp. 43–54, 2017.
  * [12] J. Fong, H. Rouhani, and M. Tavakoli, “A therapist-taught robotic system for assistance during gait therapy targeting foot drop,” _IEEE Robotics and Automation Letters_ , vol. 4, no. 2, pp. 407–413, 2019.
  * [13] H.-T. Jung, R. G. Freedman, T. Foster, Y.-K. Choe, S. Zilberstein, and R. A. Grupen, “Learning therapy strategies from demonstration using latent dirichlet allocation,” in _Proceedings of the 20th International Conference on Intelligent User Interfaces_. ACM, 2015, pp. 432–436.
  * [14] M. Amrouche, D. S. Anand, A. Lekić, V. R. Royo, E. T. Chai, D. M. Stipanović, B. Murmann, and C. J. Tomlin, “Long short-term memory neural network equilibria computation and analysis,” 2018.
  * [15] S. Ross, G. Gordon, and D. Bagnell, “A reduction of imitation learning and structured prediction to no-regret online learning,” in _Proceedings of the fourteenth international conference on artificial intelligence and statistics_ , 2011, pp. 627–635.
  * [16] S. Pareek, P. Chembrammel, J. K. Nguyen, and T. Kesavadas, “Position error-based identification of subject participation in robotic-rehabilitation,” in _2018 7th IEEE International Conference on Biomedical Robotics and Biomechatronics (Biorob)_. IEEE, 2018, pp. 432–438.
  * [17] S. Pareek, H. Manjunath, E. T. Esfahani, and T. Kesavadas, “Myotrack: Tracking subject participation in robotic rehabilitation using semg and imu,” in _2019 International Symposium on Medical Robotics (ISMR)_. IEEE, 2019, pp. 1–7.

Figure 10: A representation of the virtual no-error zone.

The error threshold $r$ may be imagined as a virtual ’no-error zone’ around
the reference trajectory. For instance, the virtual no-error zone for a
straight line connecting two points $\mathbf{p_{1}}$ and $\mathbf{p_{2}}$
would be a cylinder of radius $r$ (see Fig. 10). Here, robotic assistance is
applied to guide the subject back to the closest point on the cylinder’s
surface ($\mathbf{x_{d}}$) using the following rule:

$\displaystyle\textbf{u}(t)=\begin{cases}K_{p}[\textbf{x}_{d}(t)-\textbf{x}(t)]-K_{d}\dot{\textbf{x}}(t),&\text{if
}d>r\\\ 0,&\text{otherwise}\end{cases}$ (2)

Where, $d$ is the Euclidean distance between the current point $\mathbf{x}(t)$
and closest point on the line $\mathbf{x_{l}}(t)$. $K_{p}$ and $K_{d}$ refer
to proportional and derivative gains for the P-d controller.


# Learning Symbolic Representations Through Joint GEnerative and
DIscriminative Training

Emanuele Sansone
Department of Computer Science
KU Leuven
Leuven, 3001, Belgium
<EMAIL_ADDRESS>
&Robin Manhaeve
Department of Computer Science
KU Leuven
Leuven, 3001, Belgium
<EMAIL_ADDRESS>

###### Abstract

We introduce GEDI, a Bayesian framework that combines existing self-supervised
learning objectives with likelihood-based generative models. This framework
leverages the benefits of both GEnerative and DIscriminative approaches,
resulting in improved symbolic representations over standalone solutions.
Additionally, GEDI can be easily integrated and trained jointly with existing
neuro-symbolic frameworks without the need for additional supervision or
costly pre-training steps. We demonstrate through experiments on real-world
data, including SVHN, CIFAR10, and CIFAR100, that GEDI outperforms existing
self-supervised learning strategies in terms of clustering performance by a
significant margin. The symbolic component further allows it to leverage
knowledge in the form of logical constraints to improve performance in the
small data regime.

##  1 Introduction

Recently, neuro-symbolic learning has received attention as a new approach for
integrating symbolic-based and sub-symbolic methods based on neural networks.
This integration provides new capabilities in terms of perception and
reasoning. Currently, neuro-symbolic solutions rely either on costly pre-
training methods or on additional supervision at the symbolic representation
level provided by the neural network, in order to effectively utilize
subsequent learning feedback from the logical component (Manhaeve et al.,
2018). This traditional top-down learning paradigm is subject to the problem
of representational collapse. To gain a clearer understanding of the problem,
imagine we have a tuple of three images, each of which contains a single digit
(e.g., $<3,5,8>$). Along with this, we have information about the logical
relationships between these digits (e.g., the third digit is the sum of the
first two). Note that this task introduces less supervision compared to the
digit addition experiment typically used in neuro-symbolic systems (i.e. the
information about the sum is not provided). Current neuro-symbolic solutions
can easily solve the task by mapping all input data onto the same symbol 0 and
clearly solve the constrained task.

In this study, we present a bottom-up representation learning that can
naturally integrate with, and leverage the information in logical constraints.
We demonstrate that several existing self-supervised learning techniques and
likelihood-based generative models can be unified within a coherent Bayesian
framework called GEDI (Sansone & Manhaeve, 2022). The model leverages the
complementary properties of discriminative approaches, which are suitable for
representation learning, and of generative approaches, which capture
information about the underlying density function generating the data, to
learn better symbolic representations and support logical reasoning.
Importantly, GEDI has two main advantages: it can be easily extended to the
neuro-symbolic setting to address the collapse problem and it can also allow
for learning symbolic representations in the small data regime, which is
currently out of reach for existing self-supervised learning techniques.

##  2 GEDI Model

Model. Let us introduce the random quantities used in the model shown in
Figure 1: (i) $x\in\Omega$, where $\Omega$ is a compact subset of
$\mathbb{R}^{d}$, represents a data vector drawn independently from an unknown
distribution $p(x)$ (for instance an image), (ii) $x^{\prime}\in\Omega$
represents a transformed version of $x$ using a stochastic data augmentation
strategy $\mathcal{T}(x^{\prime}|x)$ (obtained by adding for instance noisy or
cropping the original image) (iii) $\xi\in\mathbb{R}^{h}$ is the latent
representation of an input data point obtained from an encoder network (the
latent representation of the original image) (iv) $w\in\mathcal{S}^{h-1}$,
where $\mathcal{S}^{h-1}$ is a $h-1$ dimensional unit hypersphere, is the
embedding vector of an input data point (obtained from the latent
representation using a network called projection head), while (v)
$y\in\\{1,\dots,c\\}$ is the symbolic representation of an input data point
defined over $c$ categories (namely the cluster label obtained by an output
layer defined over the embedding representation).

(a) PGM

(b) Block model

Figure 1: GEDI model. (a) shows the corresponding probabilistic graphical
model (PGM). (b) shows the different modules of GEDI, namely the encoder, the
projector head and an output module computing the cosine similarity between
the embedding representation and the cluster centers.

The corresponding probabilistic graphical model is given in Figure 1(a).
Importantly, the generative process (solid arrows) is defined using the
following conditional densities, namely: $p(w|x)=\mathcal{N}(w|0,I)$, viz. a
multivariate Gaussian with zero mean and identity covariance,
$p(\xi)=\mathcal{N}(\xi|0,I)$, $p(x^{\prime}|x,\xi)=\mathcal{T}(x^{\prime}|x)$
and $p(y|x)=\text{Softmax}(out(proj(enc(x))))$, where
$enc:\Omega\rightarrow\mathbb{R}^{h}$ is an encoder used to compute the latent
representation, $proj:\mathbb{R}^{h}\rightarrow\mathcal{S}^{h-1}$ is a
projector head used to compute the embedding representation, and $out$
computes the cosine similarity between the embedding representation $w$ and
the column vectors of a matrix of parameters $U\in\mathbb{R}^{h\times c}$
known as the cluster centers/prototypes (Caron et al., 2020). The inference
process (dashed arrows) is given by the following conditional densities:
$q(w|x)=\mathcal{N}(w|0,\Sigma)$, where
$\Sigma=\sum_{i=1}^{n}(w_{i}-\bar{w})(w_{i}-\bar{w})^{T}+\beta I$ is an
unnormalized sample covariance matrix computed over the embedding
representations of $n$ data points, $\beta$ is a positive scalar used to
ensure that $\Sigma$ is positive-definite, $\bar{w}=1/n\sum_{i=1}^{n}w_{i}$ is
the mean of the embedding representations; $q(\xi|enc(x)-enc(x^{\prime}),I)$
assesses the level of invariance between the latent representation of the
input data and its augmented version; finally
$q(y|x)=\text{SK}(out(proj(enc(x^{\prime}))))$ defines a distribution over
cluster/prototype assignments leveraging the Sinkhorn-Knopp algorithm (SK).
Please refer to the work of Caron et al. (2020) for further details.

Objective. Our training objective is based on an evidence lower bound on the
negative entropy, derived from the probabilistic graphical model of Figure
1(a), namely:

$E_{p(x)}\\{\log p(x)\\}\geq\underbrace{-CE(p,p_{\Psi})}_{\text{Generative
term}}+\underbrace{\mathcal{L}_{NF}(\Theta)+\mathcal{L}_{DI}(\Theta)}_{\text{Self-
supervised learning terms}}$ (1)

where $CE(p,p_{\Psi})$ is the cross-entropy between the unknown distribution
$p$ and a generative model $p_{\Psi}$, equivalently seen as the negative data
log-likelihood of the generative model $p_{\Psi}$. We define
$p_{\Psi}(x)=e^{-u^{T}enc(x)}/\Gamma(\Psi)$ as an energy-based model, where
$\Psi$ includes both $u\in\mathbb{R}^{h}$ and the encoder parameters.
Additionally,

$\displaystyle\mathcal{L}_{NF}(\Theta)$
$\displaystyle=-\underbrace{\mathbb{E}_{p(x)}\\{KL(q(w|x)\|p(w))\\}}_{\text{Decorrelation
term}}-\underbrace{\mathbb{E}_{p(x)\mathcal{T}(x^{\prime}|x)}\\{KL(q(\xi|x,x^{\prime})\|p(\xi))\\}}_{\text{Invariance
term}}$ (2)

where the first and the second addends promote decorrelated features in the
embedding representation and latent representations that are invariant to data
augmentations, respectively. Finally,

$\displaystyle\mathcal{L}_{DI}(\Theta)$
$\displaystyle\geq\mathbb{E}_{p(x)\mathcal{T}(x^{\prime}|x)}\\{\mathbb{E}_{q(y|x^{\prime})}\\{\log
p(y|x;\Theta)\\}+H_{q}(y|x^{\prime})\\}$ (3)

where $H_{q}(y|x^{\prime})$ is the entropy computed over $q(y|x^{\prime})$ and
$\Theta$ includes all parameters of the encoder, projector head and the output
layer of our model. Intuitively, the first addend in Eq. 3 forces the symbolic
representations of the input data and its augmented version to be similar,
whereas the second addend enforces uniformity on the cluster assignments, so
as to avoid that all representations collapse to a single cluster. It is
important to mention that the two objectives in Eqs. 2 and 3 are general
enough to cover several proposed criteria in the literature of negative-free
and cluster-based self-supervised learning (cf. Appendix A) (Sansone &
Manhaeve, 2022). Interestingly, the objective in Eq. 1 provides a natural
unification between generative and discriminative models based on self-
supervised learning. Learning the GEDI model proceeds using standard gradient
descent by maximizing Eq. 1 (more details about the training procedure are
provided in Appendix C).

##  3 Experiments

We perform experiments to evaluate the discriminative performance of GEDI and
its competitors, namely an energy-based model JEM (Grathwohl et al., 2020),
which is trained with persistent contrastive divergence (similarly to our
approach) and 2 self-supervised baselines, viz. a negative-free approach based
on Barlow Twins (Zbontar et al., 2021) and a discriminative one based on SwAV
(Caron et al., 2020). The whole analysis is divided into two main experimental
settings, the first one based on real-world data, including SVHN, CIFAR-10 and
CIFAR-100, and the second one based on a neural-symbolic learning task in the
small data regime constructed from MNIST. We use existing code both as a basis
to build our solution and also to run the experiments for the different
baselines. In particular, we use the code from Duvenaud et al. (2021) for
training energy-based models and the repository from da Costa et al. (2022)
for all self-supervised baselines. Implementation details as well as
additional experiments are reported in the Appendices.

###  3.1 SVHN, CIFAR-10, CIFAR-100

Table 1: Clustering performance in terms of normalized mutual information on test set (SVHN, CIFAR-10, CIFAR-100). Higher values indicate better clustering performance. We observe unstable training for SwAV on CIFAR-100. We report the best performance achieved out of 10 experiments. Dataset | JEM | Barlow | SwAV | GEDI | Gain
---|---|---|---|---|---
SVHN | 0.04 | 0.20 | 0.24 | 0.39 | +0.15
CIFAR-10 | 0.04 | 0.22 | 0.39 | 0.41 | +0.02
CIFAR-100 | 0.05 | 0.46 | 0.69∗ | 0.72 | +0.03

We consider three well-known computer vision benchmarks, namely SVHN, CIFAR-10
and CIFAR-100. We use a simple 8-layer Resnet network for the backbone encoder
for both SVHN and CIFAR-10 (around 1M parameters) and increase the hidden
layer size for CIFAR-100 (around 4.1M parameters) following Duvenaud et al.
(2021). We use a MLP with a single hidden layer for $proj$ (the number of
hidden neurons is twice the size of the input vector), we choose $h=256$ for
CIFAR-100 and $h=128$ for all other cases. Additionally, we use data
augmentation strategies commonly used in the self-supervised learning
literature, including color jitter, and gray scale conversion to name a few.
We train JEM, Barlow, SwAV and GEDI for $100$ epochs using Adam optimizer with
learning rate $1e-4$ and batch size $64$. Further details about the
hyperparameters are available in Appendix F. We evaluate the clustering
performance against the ground truth labels by using the Normalized Mutual
Information (NMI) score.

We report all quantitative performance in Table 1. Specifically, we observe
that JEM fails to solve the clustering task for all datasets. This is quite
natural, as JEM is a purely generative approach, mainly designed to perform
implicit density estimation. Barlow Twins achieves inferior performance to
SwAV, due to the fact that is not a cluster-based self-supervised learning
approach. On the contrary, we observe that GEDI is able to outperform all
other competitors, thanks to the exploitation of the complementary properties
of both generative and self-supervised models. Indeed, the discriminative
component in GEDI leverages the information about the underlying data manifold
structure learnt by the generative part, thus improving the learning of the
symbolic representation. In Appendix G we provide an ablation study to assess
the importance of the different loss terms involved in Eq. 1. Additionally, we
conduct experiments on linear probe evaluation, generation and OOD detection
tasks commonly used in the literature of self-supervised learning and energy-
based models. Results are reported in Appendix H.

###  3.2 Neural-symbolic setting

For the final task, we consider applying the proposed method to a neural-
symbolic setting. For this, we borrow an experiment from DeepProbLog Manhaeve
et al. (2018). In this task, each example consists of a three MNIST images
such that the value of the last one is the sum of the first two, e.g.
$\vbox{\hbox{\includegraphics[height=9.0pt]{img/MNIST/3.png}}}+\vbox{\hbox{\includegraphics[height=9.0pt]{img/MNIST/5.png}}}=\vbox{\hbox{\includegraphics[height=9.0pt]{img/MNIST/8.png}}}$.
This can thus be considered a minimal neural-symbolic tasks, as it requires a
minimal reasoning task (a single addition) on top of the image classification
task. This task only contains positive examples and requires a minimal
modification to the probabilistic graphical model, as shown in Figure 2. We
use the inference mechanism from DeepProbLog to calculate the probability that
this sum holds, and optimize this probability using the cross-entropy loss
function, which is optimized along with the other loss functions. For this
setting, this coincides with the Semantic Loss function (Xu et al., 2018). To
be able to calculate the probability of this addition constraint, we need the
classification probabilities for each digit.

(a) PGM

(b) Block model

Figure 2: Neuro-symbolic task requiring to learn the correct symbolic
representation of the digits given only a tuple of images and the
corresponding logical constraint. In this setting, $n=3$ data points and $q$
is a Boolean random variable used to detect if the logical constraint is
satisfied.

It is a specifically interesting use case for neural-symbolic learning, since
when only the probability is optimized, the neural network tends to collapse
onto the trivial solution of classifying each digit as a $0$ (i.e.
$y_{1}=y_{2}=y_{3}=0$ in Figure 2). This is a logically correct but
undesirable solution. Optimizing the discriminative objective should prevent
this collapse. We hypothesize that a neural network can be trained to
correctly classify MNIST digits by using both GEDI and the logical constraint.
Since the MNIST is an easy dataset, we focus on the small data regime, and see
whether the logical constraint is able to provide additional information. The
hyperparameters are identical to those used in Section 3.1. Further details
about the hyperparameters are dependent on the data regime, and are available
in Appendix I.

We evaluate the model by measuring the accuracy and NMI of the ResNet model on
the MNIST test dataset for different numbers of training examples. The results
are shown in Table 2. Here, N indicates the number of addition examples, which
each have $3$ MNIST digits. As expected, the DeepProbLog baseline from
Manhaeve et al. (2018) completely fails to classify MNIST images. It has
learned to map all images to the class $0$, as this results in a very low loss
when considering only the logic, resultsing in an accuracy of $0.10$ and an
NMI of $0.0$. The results also show that, without the NeSy constraint, the
accuracy is low for all settings. The NMI is higher, however, and increases as
there is more data available. This is expected, since the model is still able
to learn how to cluster from the data. However, it is unable to correctly
classify, as there is no signal in the data that is able to assign the correct
label to each cluster. By including the constraint loss, the accuracy
improves, as the model now has information on which cluster belongs to which
class. Furthermore, it also has a positive effect on the NMI, as we have
additional information on the clustering which is used by the model. These
results show us that the proposed method is beneficial to learn to correctly
recognize MNIST images using only a weakly-supervised constraint, whereas
other NeSy methods fail without additional regularization. Furthermore, we
show that the proposed method can leverage the information offered by the
constraint to further improve the NMI and classification accuracy.

Table 2: The accuracy and NMI of GEDI on the MNIST test set after training on the addition dataset, both with and without the NeSy constraint. Additionally, we use DeepProbLog (Manhaeve et al., 2018) as a baseline without using our GEDI model. We trained each model 5 times and report the mean and standard deviation. | Without GEDI | Without constraint | With constraint
---|---|---|---
N | Acc. | NMI | Acc. | NMI | Acc. | NMI
100 | $0.10\pm 0.00$ | $0.00\pm 0.00$ | $0.08\pm 0.03$ | $0.28\pm 0.03$ | $\mathbf{0.25\pm 0.03}$ | $\mathbf{0.41\pm 0.03}$
1000 | $0.10\pm 0.00$ | $0.00\pm 0.00$ | $0.09\pm 0.02$ | $0.47\pm 0.10$ | $\mathbf{0.52\pm 0.26}$ | $\mathbf{0.86\pm 0.06}$
10000 | $0.10\pm 0.00$ | $0.00\pm 0.00$ | $0.17\pm 0.12$ | $0.68\pm 0.09$ | $\mathbf{0.98\pm 0.00}$ | $\mathbf{0.97\pm 0.01}$

#### Acknowledgments

This research is funded by TAILOR, a project from the EU Horizon 2020 research
and innovation programme under GA No 952215. This research received also
funding from the Flemish Government under the “Onderzoeksprogramma Artificiële
Intelligentie (AI) Vlaanderen” programme.

## References

  * Bardes et al. (2022a) A. Bardes, J. Ponce, and Y. Lecun.  VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning.  In _ICLR_ , 2022a.
  * Bardes et al. (2022b) A. Bardes, J. Ponce, and Y. LeCun.  VICREGL: Self-Supervised Learning of Local Visual Features.  In _NeurIPS_ , 2022b.
  * Caron et al. (2018) M. Caron, P. Bojanowski, A. Joulin, and M. Douze.  Deep Clustering for Unsupervised Learning of Visual Features.  In _ECCV_ , 2018.
  * Caron et al. (2020) M. Caron, I. Misra, J. Mairal, P. Goyal, P. Bojanowski, and A. Joulin.  Unsupervised Learning of Visual Features by Contrasting Cluster Assignments.  In _NeurIPS_ , 2020.
  * Chen et al. (2020) T. Chen, S. Kornblith, M. Norouzi, and G. Hinton.  A Simple Framework for Contrastive Learning of Visual Representations.  In _ICML_ , 2020.
  * Cuturi (2013) M. Cuturi.  Sinkhorn Distances: Lightspeed Computation of Optimal Transport.  In _NeurIPS_ , 2013.
  * da Costa et al. (2022) V. G. T. da Costa, E. Fini, M. Nabi, N Sebe, and E. Ricci.  Solo-learn: A Library of Self-supervised Methods for Visual Representation Learning.  _JMLR_ , 23(56):1–6, 2022.
  * den Oord et al. (2018) A. Van den Oord, Y. Li, and O. Vinyals.  Representation Learning with Contrastive Predictive Coding.  In _arXiv_ , 2018.
  * Du & Mordatch (2019) Y. Du and I. Mordatch.  Implicit Generation and Generalization in Energy-Based Models.  _arXiv_ , 2019.
  * Duvenaud et al. (2021) D. Duvenaud, J. Kelly, K. Swersky, M. Hashemi, M. Norouzi, and W. Grathwohl.  No MCMC for Me: Amortized Samplers for Fast and Stable Training of Energy-Based Models.  In _ICLR_ , 2021.
  * Ermolov et al. (2021) A. Ermolov, A. Siarohin, E. Sangineto, and N. Sebe.  Whitening for Self-Supervised Representation Learning.  In _ICML_ , 2021.
  * Grathwohl et al. (2020) W. Grathwohl, K.-C. Wang, J.-H. Jacobsen, D. Duvenaud, M. Norouzi, and K. Swersky.  Your Classifier is Secretly an Energy Based Model and You Should Treat It Like One.  In _ICLR_ , 2020.
  * Kim & Ye (2022) B. Kim and J. C. Ye.  Energy-Based Contrastive Learning of Visual Representations.  In _NeurIPS_ , 2022.
  * Lee (2022) K. Lee.  Prototypical Contrastive Predictive Coding.  In _ICLR_ , 2022.
  * Liu et al. (2022) X. Liu, Z. Wang, Y. Li, and S. Wang.  Self-Supervised Learning via Maximum Entropy Coding.  In _NeurIPS_ , 2022.
  * Manhaeve et al. (2018) R. Manhaeve, S. Dumancic, A. Kimmig, T. Demeester, and L. De Raedt.  DeepProbLog: Neural Probabilistic Logic Programming.  In _NeurIPS_ , 2018.
  * Nalisnick et al. (2019) E. Nalisnick, A. Matsukawa, Y. W. Teh, D. Gorur, and B. Lakshminarayanan.  Do deep generative models know what they don’t know?  In _ICLR_ , 2019.
  * Nijkamp et al. (2019) E. Nijkamp, M. Hill, S. C. Zhu, and Y. N. Wu.  Learning non-convergent non-persistent short-run mcmc toward energy-based model.  In _NeurIPS_ , 2019.
  * Nijkamp et al. (2020) E. Nijkamp, M. Hill, T. Han, S.-C. Zhu, and Y. N. Wu.  On the Anatomy of MCMC-Based Maximum Likelihood Learning of Energy-Based Models.  In _AAAI_ , 2020.
  * O. Henaff (2020) Olivier O. Henaff.  Data-Efficient Image Recognition with Contrastive Predictive Coding.  In _ICML_ , 2020.
  * O.Chapelle et al. (2010) O.Chapelle, B. Schölkopf, and A. Zien.  Semi-Supervised Learning. Adaptive Computation and Machine Learning.  _Methods_ , 2010.
  * Ozsoy et al. (2022) S. Ozsoy, S. Hamdan, S. Arik, D. Yuret, and A. T. Erdogan.  Self-Supervised Learning with an Information Maximization Criterion.  In _NeurIPS_ , 2022.
  * Sansone & Manhaeve (2022) E. Sansone and R. Manhaeve.  GEDI: GEnerative and DIscriminative Training for Self-Supervised Learning.  In _arXiv_ , 2022.
  * Xie et al. (2016) J. Xie, Y. Lu, S. C. Zhu, and Y. N. Wu.  A Theory of Generative Convnet.  In _ICML_ , 2016.
  * Xu et al. (2022) H. Xu, H. Xiong, and G. J. Qi.  K-Shot Contrastive Learning of Visual Features with Multiple Instance Augmentations.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 2022.
  * Xu et al. (2018) J. Xu, Z. Zhang, T. Friedman, Y. Liang, and G. Broeck.  A Semantic Loss Function for Deep Learning with Symbolic Knowledge.  In _ICML_ , 2018.
  * Zbontar et al. (2021) J. Zbontar, L. Jing, I. Misra, Y. LeCun, and S. Deny.  Barlow Twins: Self-Supervised Learning via Redundancy Reduction.  In _ICML_ , 2021.

##  Appendix A Bayesian Interpretation of Self-Supervised Learning Objectives

(a) Contrastive (CT)

(b) Discriminative (DI)

(c) Negative-Free (NF)

(d) GEDI

Figure 3: Probabilistic graphical models for the different classes of self-
supervised learning approaches. White and grey nodes represent hidden and
observed vectors/variables, respectively. Solid arrows define the generative
process, whereas dashed arrows identify auxiliary posterior
densities/distributions.

We distinguish self-supervised learning approaches according to three
different classes: 1) contrastive, 2) cluster-based (or discriminative) and,
3) negative-free (or non-contrastive) methods. Fig. 3 shows the corresponding
probabilistic graphical models for each of the self-supervised learning
classes. From Fig. 3, we can also see how the different random quantities, we
have introduced in Section 2, are probabilistically related to each other.
Importantly, here we consider to have a set of input samples drawn
independently from $p$ and use it $i$ to index different samples, viz.
$i=\\{1,\dots,n\\}$

We are now ready to give an interpretation of the different SSL classes.

###  A.1 Contrastive SSL

Contrastive self-supervised learning can be interpreted in probabilistic terms
using the graphical model in Fig. 3(a). In particular, we can define the
following conditional density:

$p(\ell_{i}|x_{i};\Theta)=\frac{e^{sim(g(x_{\ell_{i}}),g(x_{i}))/\tau}}{\sum_{j=1}^{n}e^{sim(g(x_{j}),g(x_{i}))/\tau}}$
(4)

where $sim$ is a similarity function, $\tau>0$ is temperature parameter used
to calibrate the uncertainty for $p(\ell_{i}|x_{i};\Theta)$ and
$\Theta=\\{\theta,\\{x_{i}\\}_{i}^{n}\\}$ is the set of parameters, including
the parameters of the embedding function and the observed data. The learning
criterion can be obtained from the expected log-likelihood computed on the
observed random quantities (and using the factorization provided by the
graphical model), namely:

$\displaystyle\mathbb{E}_{p(x_{1:n},\ell_{1:n})}\big{\\{}\log
p(x_{1:n},\ell_{1:n};\Theta)\big{\\}}$ (5)
$\displaystyle=\mathbb{E}_{\prod_{j=1}^{n}p(x_{j})\delta(\ell_{j}-j)}\big{\\{}\log\prod_{i=1}^{n}p(x_{i})p(\ell_{i}|x_{i};\Theta)\big{\\}}$
$\displaystyle=\sum_{i=1}^{n}\mathbb{E}_{p(x_{i})}\big{\\{}\log
p(x_{i})\big{\\}}+\mathbb{E}_{\prod_{j=1}^{n}p(x_{j})\delta(\ell_{j}-j)}\bigg{\\{}\sum_{i=1}^{n}\log
p(\ell_{i}|x_{i};\Theta)\bigg{\\}}$
$\displaystyle=\underbrace{\sum_{i=1}^{n}\mathbb{E}_{p(x_{i})}\big{\\{}\log
p(x_{i})\big{\\}}}_{\text{Negative entropy term,
$-H_{p}(x_{1:n})$}}+\underbrace{\mathbb{E}_{\prod_{j=1}^{n}p(x_{j})}\bigg{\\{}\sum_{i=1}^{n}\log
p(\ell_{i}=i|x_{i};\Theta)\bigg{\\}}}_{\text{Conditional log-likelihood term
$\mathcal{L}_{CT}(\Theta)$}}$ (6)

where $\delta$ in the second equality is a delta function and we use for
instance $x_{1:n}$ as a compact way to express $x_{1},\dots,x_{n}$. From Eq.
(6), we observe that the expected log-likelihood can be rewritten as the sum
of two quantities, namely a negative entropy and a conditional log-likelihood
terms. However, only the second addend in Eq. (6) is relevant for maximization
purposes over the parameters $\theta$. Importantly, we can now show that the
conditional log-likelikood term in Eq. (6), coupled with the definition
provided in Eq. (4), corresponds to the notorious InfoNCE objective (den Oord
et al., 2018). To see this, let us recall InfoNCE:

$\displaystyle\text{InfoNCE}\propto\mathbb{E}_{\prod_{j=1}^{n}p(x_{j},z_{j})}\bigg{\\{}\sum_{i=1}^{n}\log\frac{e^{f(x_{i},x_{i})}}{\sum_{k=1}^{n}e^{f(x_{k},x_{i})}}\bigg{\\}}$

By choosing $p(x,z)=p(x)\delta(z-g(x))$ and $f(x,z)=sim(g(x),z)/\tau$ we
recover the conditional log-likelihood term in Eq. (6). Importantly, other
contrastive objectives, such as CPC (O. Henaff, 2020), SimCLR (Chen et al.,
2020), ProtoCPC (Lee, 2022), KSCL (Xu et al., 2022) to name a few, can be
obtained once we have a connection to InfoNCE.

###  A.2 Discriminative/Cluster-Based SSL

Cluster-based SSL can be interpreted in probabilistic terms using the
graphical model in Fig. 3(b). In particular, we can define the following
conditional density:

$p(y_{i}|x_{i};\Theta)=\frac{e^{U_{:y_{i}}^{T}G_{:i}/\tau}}{\sum_{y}e^{U_{:y}^{T}G_{:i}/\tau}}$
(7)

where $U\in\mathbb{R}^{h\times c}$ is a matrix111We use subscripts to select
rows and columns. For instance, $U_{:y}$ identify $y-$th column of matrix $U$.
of $c$ cluster centers, $G=[g(x_{1}),\dots,g(x_{n})]$ is a matrix of
embeddings of size $h\times n$ and $\Theta=\\{\theta,U\\}$ is the set of
parameters, including the ones for the embedding function and the cluster
centers. The learning criterion can be obtained in a similar way to what we
have done previously for contrastive methods. In particular, we have that

$\displaystyle\mathbb{E}_{p(x_{1:n})}\\{\log p(x_{1:n};\Theta)\\}$ (8)
$\displaystyle=-H_{p}(x_{1:n})+\mathbb{E}_{p(x_{1:n})\mathcal{T}(x_{1:n}^{\prime}|x_{1:n})}\left\\{\log\sum_{y_{1:n}}p(y_{1:n}|x_{1:n};\Theta)\right\\}$
$\displaystyle=-H_{p}(x_{1:n})+\mathbb{E}_{p(x_{1:n})\mathcal{T}(x_{1:n}^{\prime}|x_{1:n})}\Big{\\{}\log\sum_{y_{1:n}}\frac{q(y_{1:n}|x_{1:n}^{\prime})}{q(y_{1:n}|x_{1:n}^{\prime})}p(y_{1:n}|x_{1:n};\Theta)\\}$
$\displaystyle\geq-
H_{p}(x_{1:n})-\mathbb{E}_{p(x_{1:n})\mathcal{T}(x_{1:n}^{\prime}|x_{1:n})}\\{KL(q(y_{1:n}|x_{1:n}^{\prime})\|p(y_{1:n}|x_{1:n}))\\}$
$\displaystyle=\underbrace{\sum_{i=1}^{n}\mathbb{E}_{p(x_{i})}\big{\\{}\log
p(x_{i})\big{\\}}}_{\text{Negative entropy term, $-H_{p}(x_{1:n})$}}$
$\displaystyle\quad+\underbrace{\sum_{i=1}^{n}\mathbb{E}_{p(x_{i})\mathcal{T}(x_{i}^{\prime}|x_{i})}\\{\mathbb{E}_{q(y_{i}|x_{i}^{\prime})}\log
p(y_{i}|x_{i};\Theta)+H_{q}(y_{i}|x_{i}^{\prime})\\}}_{\text{Discriminative
term $\mathcal{L}_{DI}(\Theta)$}}$ (9)

where $q(y_{i}|x_{i}^{\prime})$ is an auxiliary distribution. Notably,
maximizing the discriminative term is equivalent to minimize the $KL$
divergence between the two predictive distributions $p(y_{i}|x_{i})$ and
$q(y_{i}|x_{i}^{\prime})$, thus learning to predict similar category for both
sample $x_{i}$ and its augmented version $x_{i}^{\prime}$, obtained through
$\mathcal{T}$. Importantly, we can relate the criterion in Eq. (9) to the
objective used in optimal transport Cuturi (2013), by substituting Eq. (7)
into Eq. (9) and adopting a matrix format, namely:

$\displaystyle\mathcal{L}_{DI}(\Theta;\mathcal{T})$
$\displaystyle=\frac{1}{\tau}\bigg{\\{}\mathbb{E}_{p(x_{1:n})\mathcal{T}(x_{i}^{\prime}|x_{i})}\\{Tr(QU^{T}G)\\}+\tau\mathbb{E}_{p(x_{1:n})\mathcal{T}(x_{i}^{\prime}|x_{i})}\\{H_{Q}(y_{1:n}|x_{1:n}^{\prime})\\}\bigg{\\}}$
(10)

where $Q=[q(y_{1}|x_{1}^{\prime}),\dots,q(y_{n}|x_{n}^{\prime})]^{T}$ is a
prediction matrix of size $n\times c$ and $Tr(A)$ is the trace of an arbitrary
matrix $A$. Note that a naive maximization of $\mathcal{L}_{DI}(\Theta)$ can
lead to obtain trivial solutions like the one corresponding to uniformative
predictions, namely
$q(y_{i}|x_{i}^{\prime})=p_{\gamma}(y_{i}|x_{i})=\text{Uniform}(\\{1,\dots,c\\})$
for all $i=1,\dots,n$. Fortunately, the problem can be avoided and solved
exactly using the Sinkhorn-Knopp algorithm, which alternates between
maximizing $\mathcal{L}_{DI}(\Theta)$ in Eq. (10) with respect to $Q$ and with
respect to $\Theta$, respectively. This is indeed the procedure used in
several cluster-based SSL approaches, like DeepCluster Caron et al. (2018) and
SwAV Caron et al. (2020), to name a few.

###  A.3 Non-contrastive/Negative-Free SSL

We can provide a probabilistic interpretation also for negative-free SSL using
the graphical model shown in Fig. 3(c). Note that, for the sake of simplicity
in the graph and in the following derivation, the latent variables ($w$ and
all $\xi_{i}$) are considered independent of each other only for the
generation process. In reality, one should consider an alternative but
equivalent model, using a generating process including also the edges
$x_{i}\rightarrow w$, $\xi_{i}\rightarrow x_{i}^{\prime}$ and defining
$p(w|x_{i})=p(w)$ and
$p(x_{i}^{\prime}|x_{i},\xi_{i})=\mathcal{T}(x_{i}^{\prime}|x_{i})$ for all
$i=1\dots,n$. Based on these considerations, we can define the prior and our
auxiliary and inference densities for the model in the following way:

$\displaystyle p(w)$ $\displaystyle=\mathcal{N}(w|0,I)$ $\displaystyle
p(\xi_{i})$ $\displaystyle=\mathcal{N}(\xi_{i}|0,I)$ $\displaystyle
q(w|x_{1:n};\Theta)$ $\displaystyle=\mathcal{N}(w|0,\Sigma)$ $\displaystyle
q(\xi_{i}|x_{i},x_{i}^{\prime};\Theta)$
$\displaystyle=\mathcal{N}(\xi_{i}|enc(x_{i})-enc(x_{i}^{\prime}),I)$ (11)

where $\mathcal{N}(\cdot|\mu,\Sigma)$ refers to a multivariate Gaussian
density with mean $\mu$ and covariance $\Sigma$, $I$ is an identity matrix,
$\Sigma=\sum_{i=1}^{n}(g(x_{i})-\bar{g})(g(x_{i})-\bar{g})^{T}+\beta I$,
$\beta$ is positive scalar used to ensure the positive-definiteness of
$\Sigma$, $\bar{g}=1/n\sum_{i=1}^{n}g(x_{i})$ and $\Theta=\\{\theta\\}$.
Importantly, while $q(w|x_{1:n})$ in Eq. (11) is used to store the global
statistical information of the data in the form of an unnormalized sample
covariance, $q(\xi_{i}|x_{i},x_{i}^{\prime})$ is used to quantify the
difference between a sample and its augmented version in terms of their latent
representation. Similarly to previous SSL classes and by reusing definitions
in Eq. (11), we can devise the learning criterion in the following way:

$\displaystyle\mathbb{E}_{p(x_{1:n})}\\{\log p(x_{1:n};\Theta)\\}$
$\displaystyle\geq-
H_{p}(x_{1:n})-\mathbb{E}_{p(x_{1:n})}\\{KL(q(w|x_{1:n};\Theta)\|p(w))\\}$
$\displaystyle\quad-\mathbb{E}_{p(x_{1:n})\mathcal{T}(x_{1:n}^{\prime}|x_{1:n})}\\{KL(q(\xi_{1:n}|x_{1:n},x_{1:n}^{\prime};\Theta)\|p(\xi_{1:n}))\\}$
$\displaystyle=-H_{p}(x_{1:n})-\mathbb{E}_{p(x_{1:n})}\\{KL(q(w|x_{1:n};\Theta)\|p(w))\\}$
$\displaystyle\quad-\sum_{i=1}^{n}\mathbb{E}_{p(x_{i})\mathcal{T}(x_{i}^{\prime}|x_{i})}\\{KL(q(\xi_{i}|x_{i},x_{i}^{\prime};\Theta)\|p(\xi_{i}))\\}$
$\displaystyle\propto\underbrace{\sum_{i=1}^{n}\mathbb{E}_{p(x_{i})}\big{\\{}\log
p(x_{i})\big{\\}}}_{\text{Negative entropy term,
$-H_{p}(x_{1:n})$}}\underbrace{-\mathbb{E}_{p(x_{1:n})}\left\\{\frac{Tr(\Sigma)}{2}{-}\frac{\log|\Sigma|}{2}\right\\}}_{\text{Negative-
free term, $\mathcal{L}_{NF}(\Theta)$ (first part)}}$
$\displaystyle\qquad\underbrace{-\sum_{i=1}^{n}\mathbb{E}_{p(x_{i})\mathcal{T}(x_{i}^{\prime}|x_{i})}\left\\{\frac{dist(x_{i},x_{i}^{\prime})}{2}\right\\}}_{\text{Negative-
free term, $\mathcal{L}_{NF}(\Theta)$ (second part)}}$ (12)

where $dist(x,x^{\prime})=\|enc(x)-enc(x^{\prime})\|^{2}$ and $|A|$ computes
the determinant of an arbitrary matrix $A$. Notably, the maximization of
$\mathcal{L}_{NF}(\Theta)$ promotes both decorrelated embedding features, as
the first two addends in $\mathcal{L}_{NF}(\Theta)$ (obtained from the first
KL term in Eq. (12)) force $\Sigma$ to become an identity matrix, as well as
representations that are invariant to data augmentations, thanks to the third
addend in $\mathcal{L}_{NF}(\Theta)$. It is important to mention that
$\mathcal{L}_{NF}(\Theta)$ in Eq. (12) recovers two recent negative-free
criteria, namely CorInfoMax (Ozsoy et al., 2022) and MEC (Liu et al., 2022).
We can also relate $\mathcal{L}_{NF}(\Theta)$ to other existing negative-free
approaches, including Barlow Twins (Zbontar et al., 2021), VicReg (Bardes et
al., 2022a; b) and W-MSE (Ermolov et al., 2021).

##  Appendix B Unifying Generative and SSL Models: A General Recipe (GEDI)

In all three classes of SSL approaches (see Eqs. (6),(9) and (12)), the
expected data log-likelihood can be lower bounded by the sum of two
contribution terms, namely a negative entropy $-H_{p}(x_{1:n})$ and a
conditional log-likelihood term, chosen from
$\mathcal{L}_{CT}(\Theta),\mathcal{L}_{DI}(\Theta)$ and
$\mathcal{L}_{NF}(\Theta)$. A connection to generative models emerges by
additionally lower bounding the negative entropy term, namely:

$\displaystyle-H_{p}(x_{1:n})$ $\displaystyle=\mathbb{E}_{p(x_{1:n})}\\{\log
p(x_{1:n})\\}$ $\displaystyle=\sum_{i=1}^{n}\mathbb{E}_{p(x_{i})}\\{\log
p(x_{i})\\}$ $\displaystyle=\sum_{i=1}^{n}\left[\mathbb{E}_{p(x_{i})}\\{\log
p_{\Psi}(x_{i})\\}+KL(p(x_{i})\|p_{\Psi}(x_{i}))\right]$
$\displaystyle\geq\underbrace{\sum_{i=1}^{n}\mathbb{E}_{p(x_{i})}\\{\log
p_{\Psi}(x_{i})\\}}_{\text{$-CE(p,p_{\Psi})$}}$ (13)

where $p_{\Psi}(x)$ is a generative model parameterized by $\Psi$. Notably,
the relation in (13) can be substituted in any of the objectives previously
derived for the different SSL classes, thus allowing to jointly learn both
generative and SSL models. This leads to a new GEnerative and DIscriminative
family of models, which we call GEDI. Importantly, any kind of likelihood-
based generative model (for instance variational autoencoders, normalizing
flows, autoregressive or energy-based models) can be considered in GEDI. In
this work, we argue that much can be gained by leveraging the GEDI
integration. Notably, there has been a recent work EBCLR (Kim & Ye, 2022)
integrating energy-based models with contrastive SSL approaches. Here, we show
that EBCLR represents one possible instantiation of GEDI. For instance, let us
consider contrastive SSL and observe that the conditional density in Eq. (4)
can be decomposed into a joint and a marginal densities (similarly to what is
done in (Grathwohl et al., 2020)):

$\displaystyle p(\ell,x;\Theta)$
$\displaystyle=\frac{e^{sim(g(x_{\ell}),g(x))/\tau}}{\Gamma(\Theta)}$
$\displaystyle p(x;\Theta)$
$\displaystyle=\frac{\sum_{j=1}^{n}e^{sim(g(x_{\ell_{j}}),g(x))/\tau}}{\Gamma(\Theta)}$
$\displaystyle=\frac{e^{-\underbrace{\left(-\log\sum_{j=1}^{n}e^{sim(g(x_{\ell_{j}}),g(x))/\tau}\right)}_{\text{$E(x;\Theta)$}}}}{\Gamma(\Theta)}$
(14)

where $E(x,\Theta)$ defines the energy score of the marginal density. Now, by
choosing $p_{\Psi}(x)=p(x;\Theta)$ and
$sim(z,z^{\prime})=-\|z-z^{\prime}\|^{2}$ in Eq. (14), one recovers the exact
formulation of EBCLR (Kim & Ye, 2022).

##  Appendix C Whole Model and Training Algorithm

Figure 4 shows the whole GEDI architecture and how to compute the different
losses. Importantly, GEDI training relies on a newly introduced data
augmentation strategy, called DAM.

Figure 4: Diagram of the whole model. Grey boxes are shared among different
rows.

Data Augmentation based on Manifold structure (DAM). The routine uses
generated samples obtained by ”walking” on the approximated data manifold
induced by the energy-based model and it assigns the same label of the
original points to these samples through the discriminative loss
$\mathcal{L}_{DI}(\Theta)$, thus enforcing the manifold assumption, commonly
used in semi-supervised learning (O.Chapelle et al., 2010). DAM consists of an
iterative procedure starting from an original training data point $x$, drawn
from $p$, and generating new samples $x^{\prime}$ along the approximated data
manifold induced by $p_{\Psi}$. At each iteration $t$, the algorithm performs
two main operations: Firstly, it locally perturbs a sample $x^{t}$ by randomly
choosing a vector $\Delta^{t}$ on a ball of arbitrarily small radius
$\epsilon>0$, viz. $B_{\epsilon}$, and secondly, it projects the perturbed
sample $x^{t}+\Delta^{t}$ back onto the tangent plane of the approximated data
manifold $\mathcal{M}_{p_{\Psi}}$ using the following update rule:

$\displaystyle x^{t}\leftarrow
x^{t}+\underbrace{\Delta^{t}-\left(\frac{\nabla_{x}p_{\Psi}(x^{t}+\Delta^{t})^{T}\Delta^{t}}{\|\nabla_{x}p_{\Psi}(x^{t}+\Delta^{t})\|^{2}}\right)\nabla_{x}p_{\Psi}(x^{t}+\Delta^{t})}_{\text{$\Delta_{\parallel}^{t}$}}$
(15)

A visual interpretation as well as a complete description of the strategy are
provided in Figure 5 and Algorithm 1, respectively.

Input: $x_{1:n},p_{\Psi},\epsilon,T$;

Output: $x_{1:n}^{\prime}$;

Sample $\Delta^{0}$ uniformly at random from $B_{\epsilon}$;

For $t=0,\dots,T$;

Evaluate $\nabla_{x}p_{\Psi}(x_{i}^{t}+\Delta^{t})$ for all $i=1,\dots,n$;

Update $x_{i}^{t}$ using Eq. (15) for all $i=1,\dots,n$;

$\Delta^{t}\leftarrow\Delta^{0}$;

$x_{i}^{\prime}\leftarrow x_{i}^{T}$ for all $i=1,\dots,n$;

Return $x_{1:n}^{\prime}$;

Algorithm 1 Data Augmentation based on Manifold structure (DAM). Formulation
based on a batch of samples of size $n$.

Learning a GEDI model. The data augmentation strategy proves effective when
the energy-based model approximates well the unknown data density $p(x)$.
Consequently, we opt to train our GEDI model using a two-step procedure, where
we first train the energy-based model to perform implicit density estimation
and subsequently train the whole model by maximizing the objective in Eq. (1).
Regarding the first stage, we maximize only the generative term in Eq. (1)
whose gradient is given by the following relation:

$\displaystyle-\nabla_{\Psi}CE(p,p_{\Psi})$
$\displaystyle=\mathbb{E}_{p(x)}\left\\{\nabla_{\Psi}\log
e^{u^{T}enc(x)}\right\\}$ $\displaystyle\quad-\nabla_{\Psi}\log\Gamma(\Psi)$
$\displaystyle=\mathbb{E}_{p(x)}\\{\nabla_{\Psi}\log e^{u^{T}enc(x)}\\}$
$\displaystyle\quad\mathbb{E}_{p_{\Psi}(x)}\left\\{\nabla_{\Psi}\log
e^{u^{T}enc(x)}\right\\}$ (16)

Figure 5: Data augmentation strategy exploiting the information about the
manifold structure $\mathcal{M}_{p_{\Psi}}$ and the vector field
$\nabla_{x}p_{\Psi}$ induced by the energy-based model $p_{\Psi}$. A local
perturbation $\Delta^{t}$ of a point $x^{t}$ is projected back onto the
tangent plane of the manifold by using the gradient information. The strategy
is applied iteratively starting from $x\equiv x^{0}$ up to $x^{\prime}\equiv
x^{T}$.

where the first and the second expectations in Eq. (16) are estimated using
the training and the generated data, respectively. To generate data from
$p_{\Psi}$, we use a sampler based on Stochastic Gradient Langevin Dynamics
(SGLD), thus following recent best practices to train energy-based models (Xie
et al., 2016; Nijkamp et al., 2019; Du & Mordatch, 2019; Nijkamp et al.,
2020).

Input: $x_{1:n}$, $\text{Iters}_{1}$, $\text{Iters}_{2}$, $T$, $\epsilon$,
SGLD and Adam optimizer hyperparameters;

Output: Trained model $g$;

# Step 1;

For $\text{iter}=1,\dots,\text{Iters}_{1}$;

Generate samples from $p_{\Psi}$ using SGLD;

$\Psi\leftarrow\text{Adam}$ maximizing $-CE(p,p_{\Psi})$;

# Step 2;

For $\text{iter}=1,\dots,\text{Iters}_{2}$;

$x_{1:n}^{\prime}\leftarrow\text{DAM}(x_{1:n},p_{\Psi},\epsilon,T)$;

Generate samples from $p_{\Psi}$ using SGLD;

$\Psi,\Theta\leftarrow\text{Adam}$ maximizing Eq. (1);

Return $g$;

Algorithm 2 GEDI Training. Formulation based on a batch of samples of size
$n$.

Regarding the second stage, we maximize the whole objective in Eq. (1).
Specifically, at each training iteration, we run the DAM routine to obtain the
augmented sample $x^{\prime}$. We also use the augmented samples from the
stochastic data augmentation strategy $\mathcal{T}$. Both augmentations are
used to compute $\mathcal{L}_{DI}(\Theta)$ through the differentiable
clustering prodecure of SwAV (Caron et al., 2020). The learning process is
summarized in Algorithm 2.

Computational requirements. Compared to traditional SSL training, and more
specifically to SwAV, our learning algorithm requires additional operations
and therefore increased computational requirements (but constant given $T$ and
$\text{Iters}_{1}$). Indeed, (i) we need an additional training step (Step 1)
to pre-train a generative model to approximate the unknown data density and to
ensure the proper working of DAM in the second step, (ii) we need to generate
samples from $p_{\Psi}$ to continue learning the generative model in Step 2
and (ii) we also need to run $T$ additional forward and backward passes
through the energy-based model to run the DAM strategy at each iteration of
the GEDI training.

##  Appendix D Motivating Toy Examples for DAM

Table 3: Clustering performance in terms of normalized mutual information (NMI) on test set (moons and circles). Higher values indicate better clustering performance. Mean and standard deviations are computed from 5 different runs. GEDI uses $T=10$ moves in DAM. Dataset | JEM | SwAV | GEDI | Gain
---|---|---|---|---
Moons | 0.0$\pm$0.0 | 0.8$\pm$0.1 | 0.98$\pm$0.02 | +0.18
Circles | 0.0$\pm$0.0 | 0.0$\pm$0.0 | 1.00$\pm$0.01 | +1.00

We consider two well-known synthetic datasets, namely moons and circles. We
use a multi-layer perceptron (MLP) with two hidden layers (100 neurons each)
for $enc$ and one with a single hidden layer (4 neurons) for $proj$, we choose
$h=2$ and we choose $\mathcal{T}(x^{\prime}|x)=N(0,\sigma^{2}I)$ with
$\sigma=0.03$ as our data augmentation strategy. We train JEM, SwAV and GEDI
for $7k$ iterations using Adam optimizer with learning rate $1e-3$. Further
details about the hyperparameters are available in the Appendix E. We evaluate
the clustering performance both qualitatively, by visualizing the cluster
assignments using different colors, as well as quantitatively, by using the
Normalized Mutual Information (NMI) score. Furthermore, we conduct an ablation
study for the different components of GEDI.

We report all quantitative performance in Table 3. Specifically, we observe
that JEM fails to solve the clustering task for both datasets. This is quite
natural, as JEM is a purely generative approach, mainly designed to perform
implicit density estimation. SwAV can only solve the clustering task for the
moons dataset, highlighting the fact that it is not able to exploit the
information from the underlying density used to generate the data. However,
GEDI can recover the true clusters in both datasets. This is due to the fact
that GEDI uses the information from the generative component through DAM to
inform the cluster-based one. Consequently, GEDI is able to exploit the
manifold structure underlying data. Figure 6 provide some examples of
predictions by SwAV and GEDI on the two datasets. From the figure, we can
clearly see that GEDI can recover the two data manifolds up to a permutation
of the labels.

We conduct an ablation study to understand the impact of the different
components of GEDI. We compare four different versions of GEDI, namely the
full version (called simply GEDI), GEDI trained without
$\mathcal{L}_{NF}(\Theta)$ (called no NF), GEDI trained without the first
stage and also without $\mathcal{L}_{NF}(\Theta)$ (called no NF, no train. 1)
and GEDI trained without $\mathcal{L}_{NF}(\Theta)$ using two different
encoders for computing the discriminative and the generative terms in our
objective (called no NF, 2 enc.). From the results in Figure 8, we can make
the following observations: (i) DAM plays a crucial role to recover the
manifold structure, as the performance increase with the number of steps $T$
for GEDI, no NF and no NF, 2 enc.. However, the strategy is not effective when
we don’t use the first training stage, as demonstrated by the results obtained
by no NF, no train. 1 on circles. This confirms our original hypothesis that
DAM requires a good generative model; (ii) When looking at the results
obtained by no NF and no NF, 2 enc., we observe that there is a clear
advantage, especially on circles, by using a single encoder, highlighting the
fact that the integration between generative and self-supervised models is
effective not only at the objective level but also at the architectural one;
(iii) From the comparison between GEDI and no NF on circles, we observe
improved performance and reduced variance. This suggests that
$\mathcal{L}_{NF}(\Theta)$ complements the other terms in our objective and
therefore contributes to driving the learning towards desired solutions.

(a) Ground truth

(b) SwAV

(c) GEDI

(d) Ground truth

(e) SwAV

(f) GEDI

Figure 6: Qualitative visualization of the clustering performance for the
different strategies (only SwAV and GEDI are shown) on moons (a-c) and on
circles (d-f) datasets. Colors identify different cluster predictions.

(a) Moons

(b) Circles

Figure 7: NMI achieved by JEM, SwAV and GEDI on moons and circles dataset. The
curves are obtained by choosing different values of $T$ for DAM, namely
$T\in\\{1,2,5,8,10\\}$.

(a) Moons

(b) Circles

Figure 8: Ablation study for GEDI on moons and circles dataset. The curves are
obtained by choosing different values of $T$ for DAM, namely
$T\in\\{1,2,5,8,10\\}$.

##  Appendix E Hyperparameters for Synthetic Data

For the backbone $enc$, we use a MLP with two hidden layers and 100 neurons
per layer, an output layer with 2 neurons and LeakyReLU activation functions.
For the projection head $proj$, we use a MLP with one hidden layer and 4 and
an output layer with 2 neurons (batch normalization is used in all layers) and
final $L_{2}$ normalization. For JEM, we use an output layer with only one
neuron. All methods use a batch size of 400. Baseline JEM (following the
original paper):

  * •

Number of iterations $100K$

  * •

Learning rate $1e-3$

  * •

Optimizer Adam $\beta_{1}=0$, $\beta_{2}=0.9$

  * •

SGLD steps $1$

  * •

Buffer size 10000

  * •

Reinitialization frequency $0.05$

  * •

SGLD step-size $1$

  * •

SGLD noise $0.01$

And for self-supervised learning methods, please refer to Table 5.

##  Appendix F Hyperparameters for SVHN, CIFAR-10, CIFAR-100

Table 4: Resnet architecture. Conv2D(A,B,C) applies a 2d convolution to input with B channels and produces an output with C channels using stride (1, 1), padding (1, 1) and kernel size (A, A). Name | Layer | Res. Layer
---|---|---
Block 1 | Conv2D(3,3,F) | AvgPool2D(2)
LeakyRELU(0.2)
Conv2D(3,F,F) | Conv2D(1,3,F) no padding
AvgPool2D(2)
| Sum
Block 2 | LeakyRELU(0.2) |
Conv2D(3,F,F) |
LeakyRELU(0.2) |
Conv2D(3,F,F) |
AvgPool2D(2) |
Block 3 | LeakyRELU(0.2) |
Conv2D(3,F,F) |
LeakyRELU(0.2) |
Conv2D(3,F,F) |
Block 4 | LeakyRELU(0.2) |
Conv2D(3,F,F) |
LeakyRELU(0.2) |
Conv2D(3,F,F) |
AvgPool2D(all) |

For the backbone $enc$, we use a ResNet with 8 layers as in Duvenaud et al.
(2021), where its architecture is shown in Table 4. For the projection head
$proj$, we use a MLP with one hidden layer and $2*F$ neurons and an output
layer with $F$ neurons (batch normalization is used in all layers) and final
$L_{2}$ normalization. $F=128$ for SVHN, CIFAR-10 (1 million parameters) and
$F=256$ for CIFAR-100 (4.1 million parameters). For JEM, we use the same
settings of Duvenaud et al. (2021). All methods use a batch size of 64.
Baseline JEM (following the original paper):

  * •

Number of epochs $100$

  * •

Learning rate $1e-4$

  * •

Optimizer Adam

  * •

SGLD steps $20$

  * •

Buffer size 10000

  * •

Reinitialization frequency $0.05$

  * •

SGLD step-size $1$

  * •

SGLD noise $0.01$

  * •

Data augmentation (Gaussian noise) $0.03$

And for self-supervised learning methods, please refer to Table 5.

Table 5: Hyperparameters (in terms of sampling, optimizer, objective and data augmentation) used in all experiments. Class | Name param. | SVHN | CIFAR-10 | CIFAR-100 | Addition
---|---|---|---|---|---
Data augment. | Color jitter prob. | 0.1 | 0.1 | 0.1 | 0.1
Gray scale prob. | 0.1 | 0.1 | 0.1 | 0.1
Random crop | Yes | Yes | Yes | Yes
Additive Gauss. noise (std) | 0.03 | 0.03 | 0.03 | 0.3
Random horizontal flip | No | Yes | Yes | No
Training 1
SGLD | SGLD iters | 20 | 20 | 20 | 10
Buffer size | 10k | 10k | 10k | 10k
Reinit. frequency | 0.05 | 0.05 | 0.05 | 0.05
SGLD step-size | 1 | 1 | 1 | 1
SGLD noise | 0.01 | 0.01 | 0.01 | 0.01
Optimizer | Batch size | 64 | 64 | 64 | 60
$\text{Iters}_{1}$ | 100k | 70k | 70k | Sec. I
Adam $\beta_{1}$ | 0.9 | 0.9 | 0.9 | 0.9
Adam $\beta_{2}$ | 0.999 | 0.999 | 0.999 | 0.999
Learning rate | $1e-4$ | $1e-4$ | $1e-4$ | $1e-4$
Training 2
SGLD | SGLD iters | idem | idem | idem | idem
Buffer size | idem | idem | idem | idem
Reinit. frequency | idem | idem | idem | idem
SGLD step-size | idem | idem | idem | idem
SGLD noise | idem | idem | idem | idem
Optimizer | Batch size | idem | idem | idem | idem
$\text{Iters}_{2}$ | idem | idem | idem | idem
Adam $\beta_{1}$ | idem | idem | idem | idem
Adam $\beta_{2}$ | idem | idem | idem | idem
Learning rate | idem | idem | idem | idem
DAM | $\epsilon$ | 0.03 | 0.03 | 0.03 | 0.03
$T$ | 10 | 10 | 10 | 10
Weights for losses | $-CE(p,p_{\Psi})$ | 1 | 1 | 1 | 1
$\mathcal{L}_{NF}(\Theta)$ | 1/batch | 1/batch | 1/batch | 1/batch
$\mathcal{L}_{DI}(\Theta,\mathcal{T})$ | 1000 | 1000 | 1000 | 1000
$\mathcal{L}_{DI}(\Theta,p_{\Psi})$ | 500 | 500 | 500 | 500
$\mathcal{L}_{NeSY}(\Theta)$ | - | - | - | 0/3000
Table 6: Ablation study for clustering performance in terms of normalized mutual information on test set (SVHN, CIFAR-10, CIFAR-100). Higher values indicate better clustering performance. Dataset | no NF, 2 enc. (Ours) | no NF (Ours) | GEDI (Ourss) | Gain
---|---|---|---|---
SVHN | 0.21 | 0.31 | 0.39 | +0.08
CIFAR-10 | 0.38 | 0.40 | 0.41 | +0.00
CIFAR-100 | 0.75 | 0.76 | 0.72s | -0.04
Table 7: Supervised linear evaluation in terms of accuracy on test set (SVHN, CIFAR-10, CIFAR-100). The linear classifier is trained for 100 epochs using SGD with momentum, learning rate $1e-3$ and batch size 100. Dataset | JEM | Barlow | SwAV | No NF, 2 enc. | No NF | GEDI
---|---|---|---|---|---|---
SVHN | 0.20 | 0.84 | 0.44 | 0.29 | 0.52 | 0.56
CIFAR-10 | 0.23 | 0.63 | 0.53 | 0.48 | 0.50 | 0.51
CIFAR-100 | 0.03 | 0.35 | 0.14 | 0.12 | 0.15 | 0.15
Table 8: Generative performance in terms of Frechet Inception Distance (FID) (SVHN, CIFAR-10, CIFAR-100). The lower the values the better the performance are. Dataset | JEM | Barlow | SwAV | No NF, 2 enc. | No NF | GEDI
---|---|---|---|---|---|---
SVHN | 166 | 454 | 489 | 158 | 173 | 208
CIFAR-10 | 250 | 413 | 430 | 209 | 265 | 236
CIFAR-100 | 240 | 374 | 399 | 210 | 237 | 244
Table 9: OOD detection in terms of AUROC on test set (CIFAR-10, CIFAR-100). Training is performed on SVHN. Dataset | JEM | Barlow | SwAV | No NF, 2 enc. | No NF | GEDI
---|---|---|---|---|---|---
CIFAR-10 | 0.75 | 0.43 | 0.21 | 0.76 | 0.97 | 0.94
CIFAR-100 | 0.75 | 0.5 | 0.28 | 0.75 | 0.94 | 0.93
Table 10: OOD detection in terms of AUROC on test set (SVHN, CIFAR-100). Training is performed on CIFAR-10. Dataset | JEM | Barlow | SwAV | No NF, 2 enc. | No NF | GEDI
---|---|---|---|---|---|---
SVHN | 0.43 | 0.31 | 0.24 | 0.43 | 0.31 | 0.31
CIFAR-100 | 0.54 | 0.56 | 0.51 | 0.53 | 0.53 | 0.55
Table 11: OOD detection in terms of AUROC on test set (SVHN, CIFAR-10). Training is performed on CIFAR-100. Dataset | JEM | Barlow | SwAV | No NF, 2 enc. | No NF | GEDI
---|---|---|---|---|---|---
SVHN | 0.48 | 0.43 | 0.50 | 0.47 | 0.32 | 0.38
CIFAR-10 | 0.48 | 0.42 | 0.46 | 0.47 | 0.49 | 0.50

##  Appendix G Ablation Study on SVHN, CIFAR-10, CIFAR-100

We conduct an ablation study to understand the impact of the different
components of GEDI. We compare four different versions of GEDI, namely the
full version (called simply GEDI), GEDI trained without
$\mathcal{L}_{NF}(\Theta)$ (called no NF), GEDI trained without the first
stage and also without $\mathcal{L}_{NF}(\Theta)$ (called no NF, no train. 1)
and GEDI trained without $\mathcal{L}_{NF}(\Theta)$ using two different
encoders for computing the discriminative and the generative terms in our
objective (called no NF, 2 enc.). Results are shown in Table 6.

##  Appendix H Additional Experiments on SVHN, CIFAR-10, CIFAR-100

We conduct a linear probe evaluation of the representations learnt by the
different models Table 7. These experiments provide insights on the
capabilities of learning representations producing linearly separable classes.
From Table 7, we observe a large difference in results between Barlow and
SwAV. Our approach provides interpolating results between the two approaches.

We also evaluate the generative performance in terms of Frechet Inception
Distance (FID). From Table 8, we observe that GEDI outperforms all self-
supervised baselines by a large margin, achieving comparable performance to
JEM.

Additionally, we evaluate the performance in terms of OOD detection, by
following the same methodology used in Grathwohl et al. (2020). We use the OOD
score criterion proposed in Grathwohl et al. (2020), namely
$s(x)=-\|\frac{\partial\log p_{\Psi}(x)}{\partial x}\|_{2}$. From Table 9, we
observe that GEDI achieves almost optimal performance. While these results are
exciting, it is important to mention that they are not generally valid.
Indeed, when training on CIFAR-10 and performing OOD evaluation on the other
datasets, we observe that all approaches achieve similar performance both on
CIFAR-100 and SVHN, suggesting that all datasets are considered in-
distribution, see Table 10. A similar observation is obtained when training on
CIFAR-100 and evaluating on CIFAR-10 and SVHN, see Table 11. Importantly, this
is a phenomenon which has been only recently observed by the scientific
community on generative models. Tackling this problem is currently out of the
scope of this work. For further discussion about the issue, we point the
reader to the works in Nalisnick et al. (2019).

##  Appendix I Details on the MNIST addition experiment.

We now discuss the details on the MNIST addition experiment.

###  I.1 Hyperparameters

For the backbone $enc$, we use a ResNet with 8 layers as in Duvenaud et al.
(2021), where its architecture is shown in Table 4. For the projection head
$proj$, we use a MLP with one hidden layer and 256 neurons and an output layer
with 128 neurons (batch normalization is used in all layers) and final $L_{2}$
normalization. The number of epochs for both training phases for the three
settings, i.e. $100$ examples, $1000$ examples and $10000$ examples are $100$,
$30$ and $5$ epochs respectively. These were selected by the point at which
the loss curves flatten out.

###  I.2 Data generation

The data was generated by uniformly sampling pairs $a,b$ such that $0\leq
a\leq 9$, $0\leq b\leq 9$ and $0\leq a+b\leq 9$. For each triplet $(a,b,c)$,
we assigned to $a,b,c,$ random MNIST images with corresponding labels, without
replacement.

###  I.3 Calculating the constraint

To calculate the constraint, we group the three images of each triplet
consecutively in the batches, hence why the batch size is a multiple of 3. To
calculate the probability of the constraint, we used an arithmetic circuit
compiled from the DeepProbLog program that implements this constraint Manhaeve
et al. (2018).