# Test-time Fourier Style Calibration for Domain Generalization

Xingchen Zhao1 Chang Liu1 Anthony Sicilia2 Seong Jae Hwang2 &Yun Fu 1
1Northeastern University
2University of Pittsburgh {zhao.xingc<EMAIL_ADDRESS>{anthonysicilia<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

The topic of generalizing machine learning models learned on a collection of
source domains to unknown target domains is challenging. While many domain
generalization (DG) methods have achieved promising results, they primarily
rely on the source domains at train-time without manipulating the target
domains at test-time. Thus, it is still possible that those methods can
overfit to source domains and perform poorly on target domains. Driven by the
observation that domains are strongly related to styles, we argue that
reducing the gap between source and target styles can boost models’
generalizability. To solve the dilemma of having no access to the target
domain during training, we introduce Test-time Fourier Style Calibration (TF-
Cal) for calibrating the target domain style on the fly during testing. To
access styles, we utilize Fourier transformation to decompose features into
amplitude (style) features and phase (semantic) features. Furthermore, we
present an effective technique to Augment Amplitude Features (AAF) to
complement TF-Cal. Extensive experiments on several popular DG benchmarks and
a segmentation dataset for medical images demonstrate that our method
outperforms state-of-the-art methods.

##  1 Introduction

Over the past few years, deep convolutional neural networks (CNNs) have been
shown to perform excellently in a variety of computer vision applications.
However, CNNs have demonstrated a drastic performance drop when both the
training (source domain) and testing (target domain) data come from different
distributions. This incapability causes numerous problems in real-world ﬁelds
such as autonomous driving and medical imaging analysis. As a consequence,
domain adaptation (DA) has been extensively researched as a mechanism of
transferring knowledge from source domains to labeled/unlabeled target
domains. Yet, DA approaches rely on known target domains at training time,
which fails to generalize well to unknown testing target domains. Domain
generalization (DG) has lately attracted much attention as a more challenging
setting for enhancing CNNs’ generalizability to unknown testing data.
Specifically, DG is to train models that are resistant to distribution shifts
between domains (i.e., domain shifts) so that models can generalize well when
evaluated on any previously unseen target domain, provided that several source
domains are available for training. To address the problem, various methods
were proposed and achieved promising results, including adversarial learning
Zhou et al. (2020); Sicilia et al. (2021), meta learning Li et al. (2018a);
Balaji et al. (2018), and data augmentation Xu et al. (2021); Zhou et al.
(2021). Among these, approaches based on data augmentation have gained
popularity due to their simplicity and effectiveness, which can expose models
to a wide variety of source domains to learn generalizable representation.
However, most existing data augmentation approaches cannot perform well on
test data that are out of the augmented source data distribution. This begs
the question of how we can reduce domain gaps when the augmented source
domains are insufficient to cover the target domain.

In this paper, we propose a novel yet simple method called Test-time Fourier
Style Calibration (TF-Cal) to calibrate the target style features at
inference-time. Our method is driven by the assumption that the lack of
generalizability is due to the difference in styles of images (e.g., sketches,
cartoon, photo) between the source and target Zhou et al. (2021). Based on
this idea, we employ Fourier transformation to decompose latent features from
early CNN layers into amplitude features and phase features. The amplitude
encodes the intensities (style information), while the phase encodes the
spatial positions (semantic information) Piotrowski and Campbell (1982). To
achieve test-time style calibration, we propose to convexly calibrate the
target amplitude features using the source amplitude prototype estimated at
train-time. Hence we narrow the style gap between the target and source
domains on the fly when evaluating, while retaining the key semantic feature
for classification. We additionally show that increasing the diversity of
features can complement TF-Cal to further boost models’ generalizability based
on the convex hull theory proposed by Albuquerque et al. (2019). Inspired by
Xu et al. (2021) which augments amplitude in image space, we propose
augmenting amplitude in feature space (AAF) that diversifies higher-level
styles, while preserving semantics. Then we have the Test-time Augmentated
Fourier Style Calibration (TAF-Cal) by combining TF-Cal and AAF. Moreover, our
algorithm does not require domain labels, allowing for more data collection
flexibility.

The proposed work is evaluated on multiple DG benchmarks with different
backbones on image classification and medical image segmentation. Through
extensive experiments, we demonstrate that our method significantly improves
the generalizability of CNNs and outperforms multiple state-of-the-art
methods. We also show that TF-Cal is useful in conjunction with augmentation
techniques. Finally, we present qualitative analysis and quantitative ablation
studies to verify the effectiveness of our method.

##  2 Related Work

###  2.1 Domain Generalization

DG works in recent years can be roughly categorized into: domain-invariant
learning, meta-learning and data augmentation approaches.

Domain-invariant learning aims to map input data to features that are
invariant to source domain shift in order to be resilient to unseen target
domain shift. The domain adversarial neural network (DANN) Ganin et al.
(2016), which was originally developed for DA, consists of a domain classifier
trained to discriminate between the source and target distributions, and a
feature extractor trained to learn features that are not only beneficial for
the task (e.g., classifcation) but also capable of confusing the domain
classifier. Albuquerque et al. (2019); Sicilia et al. (2021) leverage the DANN
to align the representation of source domains to learn the domain invariance
in DG. Li et al. (2018b) propose to use Maximum Mean Discrepancy with
adversarial autoencoder to align domain distributions, mapping the aligned
distribution to a prior distribution adversarially. Another line of works
focus on minimizing contrastive loss to learn domain-invariant representation.
MASF Dou et al. (2019) minimizes contrastive losses with positive pair of same
class and negative pair of different class.

Meta-learning based approach MLDG Li et al. (2018a) simulates domain shift by
dividing data from multiple sources into meta-train and meta-test sets. Balaji
et al. (2018) propose a meta-regularizer learned to be robust to domain shift
by meta-learning. Dou et al. (2019) regularize semantic structure of features
while optimizing models by meta-learning. While meta-learning methods can
simulate the domain shift during training by dividing source domains to meta-
train and meta-test sets, it is still likely that the simulated domain shift
cannot cover the target domain shift.

Data augmentation method is most related to our work, attempting to augment
the training images or features to enhance the generalization performance on
unseen data. Zhou et al. (2021) probabilistically mix the instance statistics
of training features. Xu et al. (2021) propose to augment amplitude of images
with a regularization to form consistency between original and augmented
images. To generate diverse novel training images, Zhou et al. (2020) design
an image generation network with domain-adversarial training plus optimal
transport based metrics. A concurrent work Zhang et al. (2021) has a similar
motivation as ours to augment/modify target data during evaluation, developing
a multi-view regualizeried meta-learning that encourages test-time
augmentation with multi-view. Although these methods achieve favorable results
on unseen targets, they might fail when the augmented source cannot cover the
unseen target. Our method mitigates this issue by calibrating the target with
a source prototype, allowing the calibrated target to be covered by the
source.

Figure 1: Our method’s overview. The overall pipeline is illustrated in (a).
The images are fed into the feature extractor, and the early layer’ features
are then decomposed into amplitude and phase features. We then augment the
amplitude (AAF). Following that, the model uses TF-Cal to calibrate the
amplitude features. Next, the modified amplitude with the unmodified phase are
inversed into the spatial domain. Note that the AAF is only used at train-
time. Remark: For visualization purpose, we reconstruct the amplitude and
phase with a constant to the spatial domain. (b) illustrates TF-Cal, where
solid circles represent the original source and dashed circles represent the
augmented source.

##  3 Methodology

###  3.1 Problem Definition

We define $\mathcal{X}$ and $\mathcal{Y}$ as the input and label space. Given
$K$ source domains
$\mathcal{D}_{s}=\left\\{D^{1}_{s},D^{2}_{s},\ldots,D^{K}_{s}\right\\}(1\leq
i\leq K)$ that are available at train-time, we indicate that each $D^{i}_{s}$
contains $N$ input and label pairs $(x_{j},y_{j})_{j=1}^{N_{i}}$ where
${x}\in\mathcal{X},y\in\mathcal{Y}$ and $N_{i}$ is the number of inputs in the
source domain $D^{i}_{s}$. The goal of DG is to train a function learned from
data of $\mathcal{D}_{s}$ to perform well on an unseen target domain
$\mathcal{D}_{t}$ which share the same label space $\mathcal{Y}$ as
$\mathcal{D}_{s}$. The function can be defined as
$R_{\theta,\sigma}=T_{\sigma}\circ F_{\theta}$, where $T_{\sigma}$ is a task-
specific network parameterized by $\sigma$ and
$F_{\theta}=\left\\{F^{1}_{\theta},F^{2}_{\theta},\ldots,F^{L}_{\theta}\right\\}(1\leq
i\leq L)$ which is a feature encoder with $L$ layers parameterized by
$\theta$.

###  3.2 Background

To reduce the domain shift, we need to first understand the difference between
two distributions which are measured by $\mathcal{H}$-divergence proposed by
Ben-David et al. (2009):

$d_{\mathcal{H}}(\mathcal{D}_{s},\mathcal{D}_{t})=2\sup_{h\in\mathcal{H}}\left|\operatorname{Pr}_{x\sim{\mathcal{D}_{s}}}[h(x)=1]-\operatorname{Pr}_{x\sim{\mathcal{D}_{t}}}[h(x)=1]\right|$
(1)

where classifier $h:\mathcal{X}\rightarrow\\{0,1\\}$. Following that,
Albuquerque et al. (2019) defines the convex hull $\Lambda_{s}$ of
$\mathcal{D}_{s}$ that is a set of mixture of source domain distributions:

$\Lambda_{s}=\left\\{\sum_{i=1}^{K}\pi_{i}\mathcal{D}_{s}^{i}\mid\pi\in\Delta_{K-1}\right\\}$
(2)

where the $\pi$ represents non-negative coefficient in the $(K-1)$-dimensional
simplex $\Delta_{K-1}$. Then, an ideal case
$\bar{\mathcal{D}}_{t}\in\Lambda_{s}$ is assumed that the ideal target
$\bar{\mathcal{D}}_{t}$ lies in the source domain convex hull $\Lambda_{s}$.
Under this assumption, the risk $\epsilon_{t}(h)$ on the target domain is
upper-bounded Albuquerque et al. (2019) by:

$\epsilon_{t}(h)\leq\sum_{i=1}^{K}\pi_{i}\epsilon^{i}_{s}(h)+\gamma+\zeta$ (3)

where, on the right hand side, the first term corresponds to the risks over
all source domains, and the second term
$\gamma=d_{\mathcal{H}\Delta\mathcal{H}}(\bar{\mathcal{D}}_{t},D_{t})$ is the
$\mathcal{H}$-divergence between the ideal target $\bar{\mathcal{D}}_{t}$ and
the real target $D_{t}$, and the third term
$\zeta=\sup_{D_{s}^{\prime},D_{s}^{\prime\prime}\in\Lambda_{s}}d_{\mathcal{H}\Delta\mathcal{H}}\left(D_{s}^{\prime},D_{s}^{\prime\prime}\right)$
which is the largest $\mathcal{H}$-divergence between any pair of source
domains. $\mathcal{H}\Delta\mathcal{H}$ corresponds to $\left\\{h(x)\oplus
h^{\prime}(x)\mid h,h^{\prime}\in\mathcal{H}\right\\}$. The first term can be
minimized by empirical risk minimization (ERM), and the second term is hard to
minimized due to having no access to the target domain at training, and the
third term can be minimized by removing the source domain-discrimintive
information which is the style information in the context of this paper.
According to Cha et al. (2021), ERM is a strong baseline to learn source
domain invariance compared to other methods such as DANN Ganin et al. (2016).

###  3.3 Motivation

As the Eq.(3) states, a model can achieve acceptable performance on the target
domain by using ERM to train on the source domains when the ideal target
domain $\bar{\mathcal{D}}_{t}$ is assumed to be covered by source convex hull
$\Lambda_{s}$. However, in reality, this assumption typically does not hold.
This is the reason why domain-adversarial based methods Zhou et al. (2020);
Albuquerque et al. (2019) may fail in DG as they only reduce source
divergence. Therefore, augmenting images or features of source domains can
possibly extend the $\Lambda_{s}$ (i.e., $\gamma\rightarrow 0$), which
motivating a variety of works Zhou et al. (2021); Sicilia et al. (2021) to
diversify training data. Still, as illustraed in Fig.1b, it is likely that the
augmented source distributions (dashed circles) may not accurately reflect the
previously unseen target domain (red rhombi). This difficulty arises mostly
due to the absence of target data during training. Instead of extending
$\Lambda_{s}$ arbitrarily, we aim to pull $D_{t}$ closer to $\bar{D_{t}}$ with
the help from source data.

Therefore, it motivates us to calibrate the target features at test-time based
on what a model learned from source features. By doing so, we pull the target
closer to the source convex hull $\Lambda_{s}$, and reduce the difference
between $D_{t}$ and $\bar{D_{t}}$. To further encourage the reduction in
source-target divergence, we augment the source features to expand the convex
hull.

###  3.4 Method Overview

Our method TF-Cal is to calibrate style features. To this end, we use Fourier
transformation to decompose features into phase features that contain
semantics and amplitude features that contain styles. Then we let a model to
learn with a style calibrating function $\mathcal{C}$ (TF-Cal) for reducing
the gap between the target domain and source domains prototype. We further
introduce an augmenting function $\mathcal{A}$ (AAF) for expanding the source
convex hull by diversifying style features. The model then is trained in a
two-stage manner: in the early stage, we train the model with AAF to learn a
great diversity of style features, and in the late stage, we train the model
with both AAF and TF-Cal (i.e., TAF-Cal) that uses the estimated source
prototype from diverse styles. During testing, we leverage the learned
$\mathcal{C}$ to calibrate target styles. Our method is easy to implement as a
plug-and-play module without adding any additional network. The overall
pipeline is illustrated in Fig.1a and TF-Cal is illustrated in Fig.1b, and the
detail of the method will be explained as follows.

###  3.5 Fourier Transformation of Features

Given a latent feature tensor $\mathbf{z}^{i}=F^{i}(\mathbf{z}^{i-1})$ with
$\mathbf{z}^{i}\in\mathbb{R}^{N\times C\times H\times W}$ computed from $i$-th
layer of the feature extractor, we want to extract style information from the
feature. In our implementation, we choose the early layers since they can
extract style information well Zhou et al. (2021). Fourier transformation
$\mathcal{F}(\mathbf{z}^{i})$, which can be implemented using Fast Fourier
transformation (FFT) Nussbaumer (1982), is able to decompose the feature
$\mathbf{z}^{i}$ into the phase spectrum $\mathcal{PH}(\mathbf{z}^{i})$
encoding the spatial information and amplitude spectrum
$\mathcal{AM}(\mathbf{z}^{i})$ encoding the intensity information. As
illustrated in Fig.1, the amplitude component of the horse features depicts
the intensity of texture, while phase component preserves the shape of object.
Thus, it is natural to associate the amplitude with styles and the phase with
semantics. Since our method only modifies the amplitude, the semantic
information in the phase will not be affected.

The Fourier transformation $\mathcal{F}(\mathbf{z}^{i})$ is done on each
channel $C$ of $\mathbf{z}^{i}$:

$\mathcal{F}(\mathbf{z}^{i})(m,n)=\sum_{h=0}^{H-1}\sum_{w=0}^{W-1}\mathbf{z}^{i}(h,w)e^{-j2\pi}\left(m\frac{h}{H}+n\frac{w}{W}\right)$
(4)

where $\mathcal{F}(\mathbf{z}^{i})(m,n)$ represents the value of the features
in the frequency space corresponding to the coordinates $m$ and $n$. With
$\mathcal{F}$, we can retrieve the amplitude and phase features with real
$\operatorname{Re}$ and imaginary $\operatorname{Im}$ parts:

$\displaystyle\mathcal{PH}(\mathbf{z}^{i})=\arctan\frac{\operatorname{Im}(\mathcal{F}(\mathbf{z}^{i})(m,n))}{\operatorname{Re}(\mathcal{F}(\mathbf{z}^{i})(m,n))}$
(5)
$\displaystyle\mathcal{AM}(\mathbf{z}^{i})=\sqrt{\operatorname{Im}(\mathcal{F}(\mathbf{z}^{i})(m,n))^{2}+\operatorname{Re}(\mathcal{F}(\mathbf{z}^{i})(m,n))^{2}}$

To reconstruct phase and amplitude spectrum features back into the spatial
feature space, we can use the inverse Fourier Transformation
$\mathcal{F}^{-1}(\mathcal{AM}(\mathbf{z}^{i}),\mathcal{PH}(\mathbf{z}^{i}))=\mathbf{z}^{i}$.

###  3.6 Test-time Fourier Style Calibration

#### Estimating Source Prototype

The first step in our algorithm TF-Cal is to estimate the source domain
prototype $\mathbf{Pt}_{s}\in\mathbb{R}^{1\times C\times H\times W}$ that will
be used for the calibrating function $\mathcal{C}$. The $\mathbf{Pt}_{s}$ can
be estimated by taking the average over the amplitude features of all data
points $N_{D}$ from source domains along the batch dimension:

$\mathbf{Pt}_{s}=\frac{1}{N_{D}}\sum_{n=0}^{N_{D}-1}\mathcal{AM}(\mathbf{z}^{i}_{s})$
(6)

Since the $\mathbf{Pt}_{s}$ is estimated via all training samples, it will be
a proper representation of the centroid of source styles. We leverage
$\mathbf{Pt}_{s}$ as the anchor to pull a style closer to the source convex
hull $\Lambda_{s}$. In practice, we store the
$\mathcal{AM}(\mathbf{z}^{i}_{s})$ from each mini-batch in a memory bank, and
compute the mean of all amplitude features from the bank at the end of one
epoch. We will verify calibrating style with source prototype instead of a
random prototype in the ablation studies section.

#### Training with the Style Calibrating Function

The TF-Cal is proposed to reduce the divergence between source and target
styles during testing. We need to first train a model with a calibrating
function $\mathcal{C}$ with the input of the source prototype
$\mathbf{Pt}_{s}$ and source amplitude features
$\mathcal{AM}(\mathbf{z}^{i}_{s})$. Therefore, we let the model’s parameters
be associated with the source prototype that is the common information shared
by the target and source styles. The $\mathcal{C}$ is used to train the model
by linearly interpolating between $\mathcal{AM}(\mathbf{z}^{i}_{s})$ and the
$\mathbf{Pt}_{s}$:

$\mathcal{C}(\mathbf{Pt}_{s},\mathcal{AM}(\mathbf{z}^{i}_{s}))=\eta\mathbf{Pt}_{s}+(1-\eta)\mathcal{AM}(\mathbf{z}^{i}_{s})$
(7)

where $\eta$ is the hyperparameter to control the strength of calibration. The
calibrating function $\mathcal{C}$ is activated with a probability $p_{cal}$
during training. Before the training is finished, we store the source
prototype $\mathbf{Pt}_{s}$ from the final epoch for the usage of calibrating
target style at test-time.

#### Calibrating the Target Style on the Fly

During the evaluation, we leverage the $\mathcal{C}$ function to calibrate the
target style $\mathcal{AM}(\mathbf{z}^{i}_{t})$ with the source prototype
$\mathbf{Pt}_{s}$:

$\mathcal{C}(\mathbf{Pt}_{s},\mathcal{AM}(\mathbf{z}^{i}_{t}))=\tau\mathbf{Pt}_{s}+(1-\tau)\mathcal{AM}(\mathbf{z}^{i}_{t})$
(8)

where the $\tau$ is the hyperparamter to control the strength of calibration
during testing. The reason for controlling $\tau$ is that, due to CNNs’
inductive biases toward styles Geirhos et al. (2019), target amplitude may
contain useful style information for classification. We set the calibration
probability $p_{cal}$ to 1 during testing. Our method requires only one test
sample $\mathbf{z}^{i}_{t}\in\mathbb{R}^{1\times C\times H\times W}$ for
calibration which is realistic in practical applications. Through calibration,
the target style is pulled closer to the source style, helping the model to
achieve a similar performance on target data as it does on source data. The
advantage of our method is that we let the target and source styles share the
$\mathbf{Pt}_{s}$ that is controllable and at hand. Thus, we can offset the
source-target divergence with the guidance from the source prototype at test-
time. Instead, other DG methods only rely on the source data, hoping to cover
the target distribution or learn domain-invariance that overfits to source
domains during training.

###  3.7 Augmenting Amplitude Features

As we can see from Eq.(7) and Eq.(8), the common part of the input is the
$\mathbf{Pt}_{s}$ and the different part is the amplitude features
$\mathcal{AM}(\mathbf{z}^{i}_{t})$ and $\mathcal{AM}(\mathbf{z}^{i}_{s})$.
Therefore, to further boost the generalizability with TF-Cal, we want
$\mathcal{AM}(\mathbf{z}^{i}_{t})$ and $\mathcal{AM}(\mathbf{z}^{i}_{s})$ to
be as similar as possible. An effective way is to augment amplitude in feature
space (AAF) hoping to expand the source convex hull $\Lambda_{s}$ to decrease
the source-target divergence. Therefore, based on the strategy of mixing and
swapping amplitude of images Xu et al. (2021), we propose to mix and swap
amplitude feature pairs
($\mathcal{AM}(\mathbf{z}^{i}_{s}),\mathcal{AM}(\mathbf{z}^{i^{\prime}}_{s})$)
randomly from any source domains:

$\mathcal{A}(\mathcal{AM}(\mathbf{z}^{i}_{s}),\mathcal{AM}(\mathbf{z}^{i^{\prime}}_{s}))=\delta\mathcal{AM}(\mathbf{z}^{i}_{s})+(1-\delta)\mathcal{AM}(\mathbf{z}^{i^{\prime}}_{s})$
(9)

where $\delta\sim\operatorname{Beta}(\alpha,\alpha)$. For swapping amplitude
between two samples, we set $\delta$ to 0. Compared to Xu et al. (2021) that
augment styles in the image space, AAF utilizes higher level style information
from features with much more channels, which provides additional statistics
for diversifying styles. In our implementation, we choose amplitude features
from early layer of feature extractors since they encode style information
about the inputs. Our results in the ablation studies demonstrate that
augmenting amplitude on feature is better than on image, which empirically
verify that AAF helps TF-Cal.

##  4 Experiments

###  4.1 Datasets and Settings

Our method is evaluated on three DG benchmarks: (1) PACS Li et al. (2017) is
an object recognition dataset that contains 4 different domains (i.e., Photo,
Art-Painting, Cartoon, Sketch), including 9,991 images and 7 classes. We
follow the train-val split provided by Li et al. (2017). The domain shift is
primarily contributed by the style difference. (2) Office-Home Venkateswara et
al. (2017) is an object recognition dataset that contains 4 different domains
(i.e., Art, Clipart, Product, Real-World), including 15,550 images and 65
classes. The dataset is split randomly into 90% training set and 10%
validation set. The domain shift is primarily contributed by the style and
viewpoint differences, with less diversity than PACS. (3) MICCAI WMH Challenge
Kuijf et al. (2019) is a medical image segmentation dataset for White Matter
Hyperintensity that contains 3 different domains (i.e., Utrecht, Amsterdam,
Singapore), including 20 subjects in each domain. For each subject, there are
47 slices in Utrecht and Singapore, and 82 slices in Amsterdam. The domain
shift is mainly due to scanner and protocol differences.

For evaluation, we use the leave-one-domain-out protocol. In other words, the
model is trained on source domains, and is evaluated on a held-out target
domain. For PACS and OfficeHome, we report the classification accuracy. For
WMH Challenge, we report the Dice Similarity Coefficient (DSC) (higher is
better). All results are averaged over 3 runs with different random seeds. We
use DeepAll as the ERM baseline with no DG mechanisms.

###  4.2 Implementation Details

(1) For PACS and OfficeHome, we employ the ResNet He et al. (2016) pretrained
on the ImageNet as the backbone. The networks are trained with the learning
rate of 1e-3 for feature extractors, the learning rate of 1e-4 for
classifiers, batch size of 64, and epochs of 50. We use the features from the
second residual block as the input for AAF and TF-Cal. We adopt the same data
augmentation as Matsuura and Harada (2020). (2) For WMH Challenge, we use the
U-Net Ronneberger et al. (2015) trained from scratch as the backbone with the
the learning rate of 2e-4, batch size of 30, and epochs of 300. We use the
features from the second downsamling layer as the input for our method. We
randomly augment each training slices by rotation, scaling, and shearing. All
networks are optimized by SGD with weight decay of 5e-4, and all learning
rates are decayed by 0.1 after 80% of the epochs. For all experiments: For TF-
Cal, we set the calibration strength $\eta$ and $\tau$ to 0.5, and set the
$p_{cal}$ to 0.5. For AAF activated with a probability of 0.5, we set
$\delta\sim\operatorname{Beta}(0.2,0.2)$. The AAF is applied for all epochs,
and TF-Cal is trained after 70% of the epochs.

###  4.3 Evaluation on PACS

From the results in Table 1, our method outperforms the SOTA methods on both
ResNet-18 and ResNet-50, improving the DeepAll by a large margin, which
demonstrates the effectiveness of our method on unseen target domains. In
comparison to the meta-learning methods MLDG Li et al. (2018a) and MASF Dou et
al. (2019), as well as the domain-adversarial method MMLD Matsuura and Harada
(2020), our method clearly outperforms them significantly with simple feature
transformations and no additional network. With the advantage of alleviating
the overfit to source domains, our method also surpasses data augmentation
methods DDAIG Zhou et al. (2020) and MixStyle Zhou et al. (2021), as well as a
more recent method, RSC Huang et al. (2020). FACT Xu et al. (2021) is most
related to our work in terms of Fourier transformation and augmentation.
Still, our method outperforms FACT: (1) TF-Cal can calibrate target style at
test time, which solves the problem of training with no knowledge of target
domains. (2) AAF can encourage more diverse styles of source domains by
providing additional signals in higher-level style information via features
from early CNN layers. MVRML Zhang et al. (2021) is relevant to our work in
terms of test-time augmentation because it relies on augmenting target images
with multiple views. Our method beats MVRML since we can calibrate target
images with guidance from the source prototype.

###  4.4 Evaluation on Office-Home

As Table 2 shows, our method has the best performance. For this dataset, we
have similar observations and conclusions as the PACS. Our method shows
improvement consistently over all target domains. Note that, Clipart has the
most style difference compared to other domains. Our method demonstrates its
ability to calibrate the target style to offset the source-target shifts,
which improves the performance on Clipart by a large margin over the DeepAll.

###  4.5 Evaluation on WMH Challenge

We evaluate our method on the medical image segmentation task, which is
important for real-world applications and is often ignored by other DG works.
Our method outperforms other related methods: (1) BigAug Zhang et al. (2020)
applies heavy data augmentation. (2) UDA Li et al. (2020) is an unsupervised
domain adaptation method. (3) MixDANN Zhao et al. (2021) is a method that
combines mixup and DANN. We notice that the DeepAll fails to generalize to the
Utrecht due to the significant scanner-difference. Our method significantly
improves the overall performance of Utrecht by a considerable margin compared
to previous methods.

Method | Art | Cartoon | Photo | Sketch | Avg
---|---|---|---|---|---
ResNet-18
DeepAll | 77.95 | 77.92 | 95.68 | 69.56 | 80.28
MLDG [2017] | 79.50 | 77.30 | 94.30 | 71.50 | 80.65
MASF [2019] | 80.29 | 77.17 | 94.99 | 71.69 | 81.04
MMLD [2019] | 81.28 | 77.16 | 96.09 | 72.29 | 81.83
DDAIG [2020] | 84.20 | 78.10 | 95.30 | 74.70 | 83.08
MixStyle [2021] | 84.10 | 78.80 | 96.10 | 75.90 | 83.73
FACT [2021] | 85.37 | 78.38 | 95.15 | 79.15 | 84.51
RSC [2020] | 83.43 | 80.31 | 95.99 | 80.85 | 85.15
MVRML [2021] | 84.59 | 79.22 | 95.38 | 83.55 | 85.69
TAF-Cal (Ours) | 85.70 | 82.55 | 96.07 | 82.60 | 86.73
ResNet-50
DeepAll | 84.91 | 78.58 | 97.78 | 73.89 | 83.79
MASF [2019] | 82.89 | 80.49 | 95.01 | 72.29 | 82.67
MVRML [2021] | 87.36 | 83.12 | 96.41 | 84.11 | 87.75
RSC [2021] | 87.89 | 82.16 | 97.92 | 83.35 | 87.83
FACT [2021] | 89.63 | 81.77 | 96.75 | 84.46 | 88.15
TAF-Cal (Ours) | 90.53 | 84.91 | 97.84 | 83.83 | 89.28

Table 1: Evaluation accuracy (%) on PACS.

Method | Art | Clipart | Product | Real | Avg
---|---|---|---|---|---
DeepAll | 59.93 | 48.80 | 74.13 | 75.42 | 64.57
RSC [2020] | 58.42 | 47.90 | 71.63 | 74.54 | 63.12
DDAIG [2020] | 59.20 | 52.30 | 74.60 | 76.00 | 65.53
MixStyle [2021] | 58.70 | 53.40 | 74.20 | 75.90 | 65.55
MVRML [2021] | 61.44 | 51.16 | 74.72 | 77.72 | 66.26
FACT [2021] | 60.34 | 54.85 | 74.48 | 76.55 | 66.56
TAF-Cal (Ours) | 61.52 | 54.98 | 74.87 | 76.97 | 67.09

Table 2: Evaluation accuracy (%) on OfficeHome with ResNet-18.

Method | Utrecht | Amsterdam | Singapore | Avg
---|---|---|---|---
DeepAll | 0.183 | 0.619 | 0.781 | 0.528
BigAug [2020] | 0.534 | 0.691 | 0.711 | 0.645
UDA [2020] | 0.529 | 0.737 | 0.782 | 0.683
MixDANN [2021] | 0.694 | 0.700 | 0.839 | 0.744
TAF-Cal (Ours) | 0.736 | 0.732 | 0.860 | 0.776

Table 3: Evaluation DSC on WMH Challenge with U-Net.

AAF | TF-Cal (Train) | TF-Cal (Test) | Art | Cartoon | Photo | Sketch | Avg
---|---|---|---|---|---|---|---
- | - | - | 77.95 | 77.92 | 95.68 | 69.56 | 80.28
✓ | - | - | 85.35 | 78.63 | 96.23 | 77.58 | 84.45
- | ✓ | - | 80.76 | 77.22 | 96.18 | 69.33 | 80.88
✓ | ✓ | - | 84.60 | 78.88 | 95.23 | 77.17 | 83.97
- | ✓ | ✓ | 83.31 | 81.23 | 95.63 | 77.45 | 84.41
✓ | ✓ | ✓ | 85.70 | 82.55 | 96.07 | 82.60 | 86.73
TF-Cal with a Random Prototype
- | ✓ | ✓ | 71.43 | 76.32 | 85.69 | 68.95 | 75.60

Table 4:  Ablation Studies with ResNet-18 on PACS

###  4.6 Further Analysis

#### Ablation Studies of Different Components

We investigate a thorough ablation studies of different components in our
method. With only AAF, our method already outperforms most of the methods
listed in Table 1. The AAF also achieves similar performance as FACT. As
showed in Xu et al. (2021), FACT’s augmentation component obtains accuracy of
83.44 which is lower than our AAF, which empirically verifies the merit of AAF
in generating more diverse styles in features. To prove that using TF-Cal does
not serve as a way for augmenting source data, we show that only applying TF-
Cal at training does not have much difference in performance with the DeepAll.
The similar observation is also found when both AAF and TF-Cal (Train) are
activated. When we activate TF-Cal at both train-time and test-time, we see
the great improvement, which also beats most of the methods in Table 1.
Moreover, we see that AAF and TF-Cal can complement each other and further
boost the generalizability. Finally, we empirically conduct the experiment on
the importance of the source prototype. We randomly sample a prototype from
Gaussian distribution for both TF-Cal (Train) and TF-Cal (Test), and we see a
significant performance drop.

#### Visualization of Features

To qualitatively analyze our contribution in calibrating style features, we
applied t-SNE to compare the features learned by the DeepAll and our approach
on PACS dataset with sketch as the target domain. From the Fig. 2, DeepAll has
a satisfactory capacity for categorizing features, but it fails to align the
source domain Photo and the target domain Sketch with other source domains.
Specifically, AAF can well align the Sketch with the source domains, but with
unclear cluster boundaries due to the mixing of source styles. TF-Cal
precisely calibrates the target domain to the source domains, with distinct
cluster boundaries. With the TF-Cal’s advantage of aligning styles and AAF’s
advantage of diversifying styles, we observe a robust representation toward
the target domain.

#### Hyperparameter Sensitivity

We conduct the hyperparameter experiment (Fig.3a) on the calibration stength
hyperparameter $\eta$ and $\tau$ for TF-Cal with ResNet-18 and ResNet-50 on
PACS. As we define in the method section, $\eta$ and $\tau$ share the same
value. We choose $\eta$ and $\tau$ from $\\{0.1,0.3,0.5,0.75,1\\}$. The larger
$\eta$ and $\tau$ are, the more source prototype and the less target style in
the target features. Both ResNet-18 and ResNet-50 have the similar trend. We
find that TF-Cal performs the best when the $\eta$ and $\tau$ are equal to
0.5. When the $\eta$ and $\tau$ are smaller, the target style is less
calibrated. When the $\eta$ and $\tau$ are larger, class-discriminative
information in target styles might be reduced.

(a) DeepAll

(b) AAF

(c) TF-Cal

(d) TAF-Cal

Figure 2: t-SNE visualization on PACS with sketch as target domain.

(a)

(b)

Figure 3: (a) Hyperparameter Sensitivity on $\eta$ and $\tau$ for TF-Cal (b)
Where to Apply Our Method.

#### Where to Apply the Method

We explore applying TF-Cal+AAF after different residule blocks of ResNet-18 on
PACS. From the result in Fig.3b, we see that applying our method after the
second residual block performs the best, and the reason is that early layers
extract more domain information than the later layers. As we can see that
applying our method to the third and the fourth residual blocks, the
performance drops. Because the these two blocks are closer to the
classification layer, augmenting or calibrating the features will
unnecessarily affect the class-discriminative features.

##  5 Conclusion

In this paper, we propose a test-time calibration function for target styles
extracted by Fourier transformation, and a feature style augmentation
technique that complements style calibration. With TF-Cal, we alleviate the
stress of having no knowledge about target domains during training by
calibrating target styles with the source prototype at test-time.
Additionally, AAF efficiently diversifies source amplitudes using higher-level
style information extracted from features to help TF-Cal, outperforming image-
based amplitudes augmentation. While our method TAF-Cal can be easily
integrated into any CNN network, we achieve state-of-the-art performance on
two primary tasks (i.e., classification and segmentation), and we conduct
extensive experiments and qualitative analysis that validate our method’s
effectiveness and robustness.

## References

  * Albuquerque et al. [2019] Isabela Albuquerque, João Monteiro, Mohammad Javad Darvishi Bayazi, Tiago H. Falk, and Ioannis Mitliagkas.  Generalizing to unseen domains via distribution matching.  arXiv: Learning, 2019.
  * Balaji et al. [2018] Yogesh Balaji, Swami Sankaranarayanan, and Rama Chellappa.  Metareg: Towards domain generalization using meta-regularization.  In NeurIPS, 2018.
  * Ben-David et al. [2009] Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando C Pereira, and Jennifer Wortman Vaughan.  A theory of learning from different domains.  Machine Learning, 79:151–175, 2009.
  * Cha et al. [2021] Junbum Cha, Han-Cheol Cho, Kyungjae Lee, Seunghyun Park, Yunsung Lee, and Sungrae Park.  Domain generalization needs stochastic weight averaging for robustness on domain shifts.  ArXiv, abs/2102.08604, 2021.
  * Dou et al. [2019] Qi Dou, Daniel Coelho de Castro, Konstantinos Kamnitsas, and Ben Glocker.  Domain generalization via model-agnostic learning of semantic features.  In NeurIPS, 2019.
  * Ganin et al. [2016] Yaroslav Ganin, E. Ustinova, Hana Ajakan, Pascal Germain, H. Larochelle, François Laviolette, Mario Marchand, and Victor S. Lempitsky.  Domain-adversarial training of neural networks.  In J. Mach. Learn. Res., 2016.
  * Geirhos et al. [2019] Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix Wichmann, and Wieland Brendel.  Imagenet-trained cnns are biased towards texture; increasing shape bias improves accuracy and robustness.  ArXiv, abs/1811.12231, 2019.
  * He et al. [2016] Kaiming He, X. Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770–778, 2016.
  * Huang et al. [2020] Zeyi Huang, Haohan Wang, Eric P. Xing, and Dong Huang.  Self-challenging improves cross-domain generalization.  In ECCV, 2020.
  * Kuijf et al. [2019] Hugo J. Kuijf, Adrià Casamitjana, D. Louis Collins, M. Dadar, Achilleas Georgiou, and et al.  Standardized assessment of automatic segmentation of white matter hyperintensities and results of the wmh segmentation challenge.  IEEE Transactions on Medical Imaging, 38:2556–2568, 2019.
  * Li et al. [2017] Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M. Hospedales.  Deeper, broader and artier domain generalization.  2017 IEEE International Conference on Computer Vision (ICCV), pages 5543–5551, 2017.
  * Li et al. [2018a] Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M. Hospedales.  Learning to generalize: Meta-learning for domain generalization.  ArXiv, abs/1710.03463, 2018.
  * Li et al. [2018b] Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex Chichung Kot.  Domain generalization with adversarial feature learning.  2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5400–5409, 2018.
  * Li et al. [2020] Hongwei Li, Timo Loehr, and et al.  Domain adaptive medical image segmentation via adversarial learning of disease-specific spatial patterns.  arXiv: Computer Vision and Pattern Recognition, 2020.
  * Matsuura and Harada [2020] Toshihiko Matsuura and Tatsuya Harada.  Domain generalization using a mixture of multiple latent domains.  ArXiv, abs/1911.07661, 2020.
  * Nussbaumer [1982] Henri J. Nussbaumer.  The fast fourier transform.  1982\.
  * Piotrowski and Campbell [1982] Lisa Piotrowski and Fergus William Campbell.  A demonstration of the visual importance and flexibility of spatial-frequency amplitude and phase.  Perception, 11:337 – 346, 1982.
  * Ronneberger et al. [2015] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In MICCAI, 2015.
  * Sicilia et al. [2021] Anthony Sicilia, Xingchen Zhao, and Seong Jae Hwang.  Domain adversarial neural networks for domain generalization: When it works and how to improve.  ArXiv, abs/2102.03924, 2021.
  * Venkateswara et al. [2017] Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan.  Deep hashing network for unsupervised domain adaptation.  2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 5385–5394, 2017.
  * Xu et al. [2021] Qinwei Xu, Ruipeng Zhang, Ya Zhang, Yanfeng Wang, and Qi Tian.  A fourier-based framework for domain generalization.  2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 14378–14387, 2021.
  * Zhang et al. [2020] Ling Zhang, Xiaosong Wang, and et al.  Generalizing deep learning for medical image segmentation to unseen domains via deep stacked transformation.  IEEE Transactions on Medical Imaging, 39:2531–2540, 2020.
  * Zhang et al. [2021] Jian Zhang, Lei Qi, Yinghuan Shi, and Yang Gao.  More is better: A novel multi-view framework for domain generalization.  ArXiv, abs/2112.12329, 2021.
  * Zhao et al. [2021] Xingchen Zhao, Anthony Sicilia, and et al.  Robust white matter hyperintensity segmentation on unseen domain.  2021 IEEE 18th International Symposium on Biomedical Imaging (ISBI), pages 1047–1051, 2021.
  * Zhou et al. [2020] Kaiyang Zhou, Yongxin Yang, Timothy M. Hospedales, and Tao Xiang.  Deep domain-adversarial image generation for domain generalisation.  ArXiv, abs/2003.06054, 2020.
  * Zhou et al. [2021] Kaiyang Zhou, Yongxin Yang, Yu Qiao, and Tao Xiang.  Domain generalization with mixstyle.  ArXiv, abs/2104.02008, 2021.


# Sequence Learning Using Equilibrium Propagation

Malyaban Bal  Abhronil Sengupta School of Electrical Engineering and Computer
Science
The Pennsylvania State University
{mjb7906<EMAIL_ADDRESS>
###### Abstract

Equilibrium Propagation (EP) is a powerful and more bio-plausible alternative
to conventional learning frameworks such as backpropagation (BP). The
effectiveness of EP stems from the fact that it relies only on local
computations and requires solely one kind of computational unit during both of
its training phases, thereby enabling greater applicability in domains such as
bio-inspired neuromorphic computing. The dynamics of the model in EP is
governed by an energy function and the internal states of the model
consequently converge to a steady state following the state transition rules
defined by the same. However, by definition, EP requires the input to the
model (a convergent RNN) to be static in both the phases of training. Thus it
is not possible to design a model for sequence classification using EP with an
LSTM or GRU like architecture. In this paper, we leverage recent developments
in modern hopfield networks to further understand energy based models and
develop solutions for complex sequence classification tasks using EP while
satisfying its convergence criteria and maintaining its theoretical
similarities with recurrent BP. We explore the possibility of integrating
modern hopfield networks as an attention mechanism with convergent RNN models
used in EP, thereby extending its applicability for the first time on two
different sequence classification tasks in natural language processing viz.
sentiment analysis (IMDB dataset) and natural language inference (SNLI
dataset). Our implementation source code is available at
https://github.com/NeuroCompLab-psu/EqProp-SeqLearning.

##  1 Introduction

Equilibrium Propagation (EP) Scellier and Bengio (2017) is a biologically
plausible learning algorithm to train artificial neural networks. It requires
only one computational circuit and single type of network unit during the two
phases of training, whereas backpropagation (BP) requires a specialised type
of computation during the backward phase to explicitly propagate errors which
is different from the computational circuitry needed during the forward phase,
thus it is essentially considered to be biologically implausible Crick (1989).
In EP, errors are propagated implicitly in the energy based model through
local perturbations being generated at the output layer, unlike in BP.
Moreover, strong theoretical connections Scellier and Bengio (2019) between EP
and recurrent BP Almeida (1990); Pineda (1987) as well as the former’s
similarity regarding weight updates with spike-timing dependent plasticity
(STDP) Scellier and Bengio (2017); Bi and Poo (1998) (a feasible model to
understand synaptic plasticity change in neurons) makes EP a solid foundation
to further understand biological learning Lillicrap et al. (2020). Intrinsic
properties of EP also provides the opportunity to design energy efficient
implementations of the former on hardware unlike BP through time Martin et al.
(2021).

The idea that neurons collectively adjust themselves to configurations
according to the sensory input being fed into a neural network system such
that they can better predict the input data has been a popular hypothesis
Hinton (2002); Berkes et al. (2011). The collective neuron states can be
interpreted as explanations of the input data. EP also consists of this
central idea where the network, which is essentially a dynamical system
following certain dynamics, converges to lower energy states which better
explains the static input data.

In this paper, we have primarily discussed EP in a discrete time setting and
used the scalar primitive function $\phi$ Ernoult et al. (2019); Laborieux et
al. (2021) to derive the transition dynamics instead of an energy function as
used in the initial works Scellier and Bengio (2017). Primarily, algorithms
using EP have been designed for convergent RNNs Laborieux et al. (2021) which
are neural networks that take in a static input and through the recurrent
dynamics (as governed by the transition function – scalar primitive function
of the system in our case) converges to a steady state which denotes the
prediction of the network for that input. Training using EP primarily
comprises of two distinct phases. During the first phase i.e. the “free”
phase, the network converges to a steady state following only the internal
dynamics of the system. In contrast, in the second phase the output layer
(which acts as the prediction after the “free” phase has converged) is
“nudged” closer to the actual ground truth and the local perturbations
resulting from that change propagates to the other layers of the convergent
RNN, forming local error signals in time which matches the error propagation
associated with BP through time Ernoult et al. (2019). The state updates and
the consequent weight updates of the system can be done through local in space
and time Ernoult et al. (2019) computations, thus making learning algorithms
using EP highly suitable for developing low-powered energy efficient
neuromorphic implementations Martin et al. (2021).

Hopfield networks were one of the earliest energy based models which were
based on convergence to a steady state by minimizing an energy function which
is an intrinsic attribute of the system. Hopfield networks were primarily used
as associative memory, where we can store and retrieve patterns Hopfield
(1982). Classical hopfield networks were designed for storage and retrieval of
binary patterns and the capacity of storage was limited since the energy
function used was quadratic in nature. However, modern hopfield networks
Krotov and Hopfield (2016, 2018); Ramsauer et al. (2020) follow an exponential
energy function for state transitions which allows storing exponential number
of continuous patterns and allows for faster single step convergence
Demircigil et al. (2017) during retrieval. The modern hopfield layer has been
used previously for sequence-attention in conventional deep learning models
with BP as the learning framework Widrich et al. (2020). In this paper, we
integrate the capability of modern hopfield networks as a transformer-like
Vaswani et al. (2017) attention mechanism inside a convergent RNN and
consequently train the resulting model using the learning rules defined by EP.

##  2 Motivation and Primary Contributions

Until now, EP has been confined in the domain of image classification, the
primary reason being the constraint associated with convergent RNNs which
requires the input to be static. Developing an LSTM like network, which is fed
with a time-varying sequence of data instead of static data, is not possible
while maintaining the constraints of EP. Results obtained for tasks such as
sentiment analysis or inference problems by feeding the entire input sequence
as input also results in poor solutions since the dependencies between the
different parts of the sequence cannot be captured in that technique. However,
with recent developments in the domain of modern hopfield networks, we have
the ability to extend the capability of algorithms using EP to perform
efficient sequence classification.

Modern hopfield networks fit perfectly with convergent RNNs since both of them
are energy-based models which are governed by their respective state
transition dynamics. In a discrete time setting - which is primarily discussed
in this paper - both of them follow a certain state transition rule at every
time step which is governed by their defining energy function or a scalar
primitive function in case of a convergent RNN. As we will see in later
sections, the one step convergence guarantee of modern hopfield networks
allows us to seamlessly interface them with a convergent RNN allowing both the
networks to settle to their respective equilibrium states as the system
converges as a whole.

Neuromorphic Motivation for EP: The bio-plausible local learning framework
that EP provides can be best utilised in a neuromorphic system. Spiking
implementation of the proposed method following the works of Martin et al.
(2021) allows for low-powered solution for real-time scenarios such as
intrusion detection, real-time natural language processing (NLP), etc.
Implementing EP in a neuromorphic system allows for orders of magnitude more
energy savings than in GPUs and is more efficient because of its single
computational circuit and STDP like weight updates.

Algorithmic and Theoretical contributions: The primary contributions of our
paper are as follows -

  * •

We explore for the first time how hopfield networks can operate with a
convergent RNN model and how we can leverage the attention mechanism of the
former to solve complex sequence classification problems using EP.

  * •

We illustrate mathematically and empirically how to combine the state
transition dynamics of both the hopfield network and underlying convergent RNN
to converge to steady states during both phases of EP thus maintaining the
latter’s theoretical equivalence with recurrent BP w.r.t. gradient estimation.

  * •

We report for the first time the performance of EP as a learning framework on
widely known NLP tasks such as sentiment analysis (IMDB dataset) and natural
language inference (NLI) problems (SNLI dataset) and compare the results with
state-of-the-art architectures for which neuromorphic implementations can be
developed.

##  3 Methods

In this section, we will formulate the workings of EP and its theoretical
connections with backpropagation through time (BPTT). We will then delve into
the details of the modern hopfield layer and demonstrate its attention
mechanism. We will elaborate on the scalar primitive function defined for
state transition and further investigate the metastable states that arises in
modern hopfield networks and will discuss how we can leverage them to form
efficient encoding of the input sequences for our sequence classification
problems.

###  3.1 Equilibrium Propagation

EP applies to convergent RNNs whose input at each time step is static. The
state of the network eventually settles to a steady-state following a state
update rule that is derived from the scalar primitive function $\phi$.
Moreover, the weights defined between two layers are symmetric in nature i.e.
if the weight between layers $s^{i}$ and $s^{i+1}$ is $w_{i}$, then the weight
of the connection between $s^{i+1}$ and $s^{i}$ is $w_{i}^{T}$. The state
transition is defined as,

$s_{t+1}=\frac{\partial\phi}{\partial s}(x,s_{t},\theta)$ (1)

where, $s_{t}=(s_{t}^{1},s_{t}^{2},\dots,s_{t}^{n})$ is the collective state
of the convergent RNN with $n$ layers at time $t$, $x$ is the input to the
convergent RNN and $\theta$ represent the network parameters i.e. it comprises
of the weights of each of the connections between layers. We do not consider
any skip connection or self-connection in the convergent RNNs discussed but
there has been some work done in that area Gammell et al. (2021). The state
transitions result in a final convergence to a steady state $s_{*}$ after time
$T$, such that $s_{t}=s_{*}$ $\forall t\geq T$ and it satisfies the following
condition,

$s_{*}=\frac{\partial\phi}{\partial s}(x,s_{*},\theta)$ (2)

Training of convergent RNNs using EP comprises mainly of two different phases.
During the first phase or the “free” phase, the RNN follows the transition
function as shown in Eq. (1) and eventually reaches the steady state defined
as the free fixed point $s_{*}$ after $T$ time steps. We use the output layer
of the steady state $s_{*}$ i.e. $s^{n}_{*}$ to make the prediction for the
current input $x$. In the second phase or the “nudge” phase, an additional
term $-\beta\frac{\partial L}{\partial s}$ is added to the state dynamics
which immediately results in the state of the output layer being slightly
nudged in the direction to minimize the loss function $L$ (as defined between
the target $y$ and the output of the last layer of the convergent RNN). Though
the internal states of the hidden layers are initially at the free fixed point
state, they are eventually nudged to a different fixed point - weakly clamped
fixed point $s_{*}^{\beta}$ \- because of the perturbations initially
originated at the output layer. $\beta$ is a small scaling factor defined as
the “influence parameter” or “clamping factor” i.e. it controls the influence
of the loss on the actual primitive scalar function during the second phase.

Thus the initial state of the second phase is $s^{\beta}_{0}=s_{*}$ and the
transition function is defined as,

$s^{\beta}_{t+1}=\frac{\partial\phi}{\partial
s}(x,s^{\beta}_{t},\theta)-\beta\frac{\partial L}{\partial
s}(s^{\beta}_{t},y)$ (3)

Following Eq. (3), the convergent RNN settles to the steady state
$s^{\beta}_{*}$ after $K$ timesteps. After the two phases are done, the
learning rule to update the model parameters in order to minimize the loss
$L^{*}=L(s_{*},y)$ is defined as,
$\Delta\theta=\eta\nabla^{EP}_{\theta}(\beta)$, Scellier and Bengio (2017)
where $\eta$ is the learning rate and $\nabla^{EP}_{\theta}(\beta)$ can be
defined as,

$\nabla^{EP}_{\theta}(\beta)=\frac{1}{\beta}(\frac{\partial\phi}{\partial\theta}(x,s^{\beta}_{*},\theta)-\frac{\partial\phi}{\partial\theta}(x,s_{*},\theta))$
(4)

The defined convergent RNN can also be trained by BPTT Laborieux et al.
(2021). According to the property of Gradient-Descent Updates (GDU) Ernoult et
al. (2019), the gradient updates computed by the EP algorithm is approximately
equal to the gradient computed by BPTT ($\nabla^{BPTT}(t)$), according to
relative mean squared error metric (RMSE), provided the convergent RNN has
reached its steady state in $T-K$ steps during the first phase and
$\beta\rightarrow 0$. Thus for initial $K$ steps of the “nudge” phase we can
state , $\forall t=1,2,...,K$

$\nabla^{EP}_{\theta}(t,\beta)=\frac{1}{\beta}(\frac{\partial\phi}{\partial\theta}(x,s^{\beta}_{t},\theta)-\frac{\partial\phi}{\partial\theta}(x,s_{*},\theta))$
(5) $\nabla^{EP}_{\theta}(t,\beta)\xrightarrow[\beta\rightarrow
0]{}\nabla^{BPTT}(t)$ (6)

In the traditional implementations of EP, two phases are involved, one with
$\beta=0$ i.e. the “free” phase and the other with $\beta>0$ i.e. the “nudge”
phase. In order to circumnavigate the first order bias that is induced into
the system by assuming $\beta>0$, a new implementation of EP was proposed
Laborieux et al. (2021) which comprises of a second “nudge” phase with
$-\beta$ as the influence factor. Thus the algorithm comprises of three
phases. The symmetric EP gradient estimates are thus free from first order
bias and are more close to the values computed using BPTT and is defined as,

$\nabla^{EPsym}_{\theta}(\beta)=\frac{1}{2\beta}(\frac{\partial\phi}{\partial\theta}(x,s^{\beta}_{*},\theta)-\frac{\partial\phi}{\partial\theta}(x,s^{-\beta}_{*},\theta))$
(7)

###  3.2 Hopfield Network as Attention Mechanism

The recent developments in modern hopfield networks Ramsauer et al. (2020)
offers exponential storage capacity and one step retrieval of stored
continuous patterns. The increased storage capacity and the attention-like
state update rule of modern hopfield layers can be leveraged to retrieve
complex representation of the stored patterns which consists of richer
embedding similar to that of attention mechanism in transformers. In order to
allow for continuous states, the energy function of the modern hopfield
network is modified accordingly Widrich et al. (2020), and it can be
represented as,

$E=-lse(\beta,X^{T}\xi)+\frac{1}{2}(\xi^{T}\xi)+\beta^{-1}logN+\frac{1}{2}M^{2}$
(8) $lse(\beta,x)=\beta^{-1}log(\sum_{i=1}^{N}exp(\beta x_{i}))$

where, $X=(x_{1},x_{2},...x_{N})$ are $N$ continuous stored patterns, $\xi$ is
the state pattern, $M$ is the largest norm of all the stored patterns and
$\beta>0$. In general, the energy function of modern hopfield networks can be
defined as, $E=-\sum_{i=1}^{N}F(x_{i}^{T}\xi)$ Krotov and Hopfield (2016). For
example, if we use the function $F(x)=x^{2}$, we describe the classical
hopfield network which had limitations in storage capacity and also supported
only binary patterns. However, with the introduction of an exponential
interaction function like log-sum-exponential ($lse$) function, the storage
capacity can be increased exponentially while enabling continuous patterns to
be stored. Using Concave-Convex-Procedure (CCCP) Yuille and Rangarajan (2001,
2003), the state update rule for the modern hopfield network Ramsauer et al.
(2020) can be defined as,

$\xi^{new}=X\mathit{softmax}(\beta X^{T}\xi)$ (9)

where, $\xi^{new}$ is the retrieved pattern from the hopfield network.

We can define Query ($Q=RW_{Q}$) as $\xi^{T}$ and Key ($K=YW_{K}$) as $X^{T}$.
Thus the new form can be represented as,

$Q^{new}=\mathit{softmax}(\frac{1}{\sqrt{d_{k}}}RW_{Q}W_{K}^{T}Y^{T})YW_{K}$

$Q^{new}=\mathit{softmax}(\frac{1}{\sqrt{d_{k}}}QK^{T})K$ (10)

Figure 1: Simplistic scenario of application of $HopAttn$ modules in
convergent RNN for a sequence classification problem with $c$ classes has been
described. The input $X$ is directly fed as the Key $K$ to the $HopAttn$
module. During the “free” phase (as described here), the model evolves through
time following the defined state transition dynamics and converges to a steady
state configuration after $T$ time steps to settle to the prediction. The
operations inside the $HopAttn$ module are also illustrated in details.

where, $d_{k}$ is the encoding dimension of Key ($K$). The above
representation Ramsauer et al. (2020) is to show the similarity of the update
rule of modern hopfield networks and the attention mechanism in transformers
involving Query-Key pairs. $Q$ represents the state pattern whereas $Q^{new}$
represents the state pattern retrieved after the transition. The projection
matrices are defined as, $W_{Q}\in\mathbb{R}^{d_{r}\times d_{k}}$ and
$W_{K}\in\mathbb{R}^{d_{y}\times d_{k}}$.

Following the state transition as defined in Eq. (10), the Hopfield Attention
module can be defined as $HopAttn(Q,K)$, which has two critical inputs, viz.
the Query ($Q$) which can be interpreted as the state pattern and the Key
($K$) which can be considered as the stored pattern. The dimension of the
hopfield space is $d_{k}$. The underlying operations are illustrated in
details in Fig. 1. The output of the $HopAttn$ function is defined as,

$\mathit{HopAttn}(Q,K)=softmax(\frac{1}{\sqrt{d_{k}}}QK^{T})K$ (11)

$HopAttn$ module plays a critical role in our network by providing an
attention based embedding of the input sequence resulting from their
convergence to metastable states at every time step. In order to form the
Query ($Q$) for the $HopAttn$ module, we define the layers connected to the
input as projection layers (as demonstrated in Fig. 1) whose weight is
analogous to $W_{Q}$. At every timestep $t$ during the state transitions in
the convergent RNN, the recurrent dynamics of the convergent RNN computes the
Query ($Q$) to be fed to the $HopAttn$ module as illustrated in Eq. (14). The
stored sequence of patterns $K$ is fed directly into the $HopAttn$ module.

There is a theoretical guarantee that the modern hopfield network converges to
a steady state (with exponentially small separation) after a single update
step Ramsauer et al. (2020) and therefore we successfully navigate to a fixed
point at every time step in EP. We project the $HopAttn$ output through
$W_{V}\in\mathbb{R}^{d_{k}\times d_{v}}$, analogous to the weight of the
output connection, to generate the final output. The $HopAttn$ module thus
provides a rich representation of the stored patterns which helps in efficient
sequence classification.

###  3.3 Proposed Scalar Primitive function $\phi$

The input to our model is the sequence of patterns, $x\in R^{N\times D_{K}}$,
where $N$ is length of the sequence and $D_{K}$ is the encoding dimension of
the pattern. The state of the system at time $t$ is
$s_{t}=(s_{t}^{0},s_{t}^{1},s_{t}^{2},\dots,s_{t}^{n})$, where $n$ is the
number of layers in the convergent RNN. $\theta$ comprises the list of
parameters of the network. The scalar primitive function ($\phi$), which
defines the state transition rules of the convergent RNN is,

$\phi(x,{s},\theta)={s^{1}}\bullet(x\cdot w_{1})+{s^{2}}^{T}\cdot w_{2}\cdot
F(s^{1})+\sum_{i=2}^{n-1}{s^{i+1}}^{T}\cdot w_{i+1}\cdot s^{i}$ (12)

where $\bullet$ represents Euclidean scalar product of two tensors with same
dimensions, $(x\cdot w_{1})$ represents the linear projection of $x$ through
$w_{1}\in R^{D_{K}\times D_{K}}$ and $w_{i+1}$ is the weight of the connection
between $s^{i}$ and $s^{i+1}$. The flattened output of the first projection
layer ($F(s^{1})$, where $F$ is the flattening operation from $(A,B)$ to
$(1,AB)$), is then fed into the next of the $n-1$ fully connected layers. In
this section, we have primarily discussed $\phi$ for cases where we have only
one input sequence like the sentiment analysis task. For multi-sequence tasks
like NLI, we have described $\phi$ in the technical appendix.

###  3.4 Transition Dynamics and Convergence

Our methodology seamlessly interfaces modern hopfield networks with convergent
RNNs, thus allowing sequence classification using EP. The state transition
dynamics of the layers in the convergent RNN (except the last layer), where
$HopAttn$ is not applied, is given by the following,

$s^{i}_{t+1}=\sigma(\frac{\partial\phi}{\partial s^{i}}(x,s_{t},\theta))$ (13)

where, $\sigma$ is an activation function that restricts the range of the
state variables to $[0,1]$. However, for the layers where we apply the
hopfield attention mechanism, the state transition function is updated as
follows,

$s^{j}_{t+1}=\mathit{HopAttn}(\frac{\partial\phi}{\partial
s^{j}}(x,s_{t},\theta),K)$ (14)

where, $s^{j}_{t+1}$ is the final state of the $j^{th}$ layer (with $HopAttn$
applied) at time $t$+1 and $K$ is the sequence of stored patterns in the
hopfield network which is usually the input $x$ to the network.

Thus for the layers where we apply attention, we first follow the transition
function of the underlying convergent RNN as defined by
$\frac{\partial\phi}{\partial s^{j}}(x,s_{t},\theta)$ and then we follow the
update rule as derived from the modern dense hopfield network. Since
convergence to a steady state is guaranteed after a single update step, there
is no additional overhead for the hopfield network to converge and thus its
state transition rule is only applied once.

The metastable states, that the $HopAttn$ module converges to after each
timestep $t$, approach steady state w.r.t the convergent RNN during both the
phases of training of EP. The claim for such concurrent convergence can be
further substantiated empirically by analysing the convergence of the dynamics
of the model, governed by the scalar primitive $\phi$ with time. It is evident
from Fig. 2b that even after we apply the $HopAttn$ module (like in Fig. 1),
the model still converges to a steady state eventually. Moreover, even after
the integration of $HopAttn$ modules, the GDU property of the convergent RNN
still holds, thus maintaining the equivalence w.r.t gradient estimation of EP
and BPTT (Fig. 2a).

Figure 2: Results taken from experiments run on IMDB dataset. (a) Symmetric EP
gradient estimate as defined in Eq. (7) and gradient computed by BPTT, for
three randomly chosen weights in the convergent RNN integrated with $HopAttn$
module. (b) Convergence of the scalar primitive function $\phi$ with time upon
using $HopAttn$ in “free” phase taken over 50 phase transitions.

For the final layer, whose output is compared with the target label to
calculate the loss function, the state update rule is defined as,

$s^{n}_{t+1}=\sigma(\frac{\partial\phi}{\partial
s^{n}}(x,s_{t},\theta))+\beta(y-s^{n}_{t})$ (15)

where, $y$ is the target label and $s^{n}_{t}$ is the output of the final
layer at time $t$. $\beta=0$ during the “free phase” and the loss function $L$
used in this case is mean squared error.

###  3.5 Metastable States as Attention Embeddings

The factor $\beta$, as defined in Eq. (9), plays an important role in
establishing the fixed points in the modern hopfield network Ramsauer et al.
(2020). The retrieved state pattern (or $Q^{new}$ as defined above) settles to
the defined fixed points following the state transition rule (Eq. (10)).
However, if $\beta$ is very high and/or the stored patterns are well
separated, then we can easily retrieve the actual pattern stored in the
hopfield networks. On the other hand, if $\beta$ is low and the stored
patterns are not all well separated but form cluster like structures in the
encoding dimension of the patterns, then the hopfield network generates
metastable states. Thus, when we try to retrieve using a particular state
pattern i.e a Query, we might converge to a metastable point comprising of a
richer representation combining a number of patterns in that region. For
incorporating such attention like behavior, we need to keep $\beta$ small and
usually $\beta=1/\sqrt{d_{k}}$, where $d_{k}$ is the encoding dimension of the
stored patterns.

###  3.6 Local State & Parameter Updates

The state update rule, as defined earlier, is essentially local in space since
$\frac{\partial\phi}{\partial s^{i}}(x,s_{t},\theta)$ $\forall i=2,...,n-1$
can be written as,

$\frac{\partial\phi}{\partial s^{i}}(x,s_{t},\theta)=w_{i}\cdot
s_{t}^{i-1}+w_{i+1}^{T}\cdot s_{t}^{i+1}$ (16)

and for the final layer the same can be shown as,

$\frac{\partial\phi}{{\partial s^{n}}}(x,s_{t},\theta)=w_{n}\cdot s_{t}^{n-1}$
(17)

where all the associated parameters for the computation are locally connected
in space and time. The computation required w.r.t the $HopAttn$ module can
also be done locally in space and time. For the first layer (projection
layer), the local state and weight update property is still preserved:
$\frac{\partial\phi}{\partial s^{1}}(x,s_{t},\theta)=(x\cdot
w_{1})+F^{-1}(w_{2}^{T}\cdot s_{t}^{2}$).

The STDP like parameter update rule of the network, as stated in the earlier
section, is also local in space, i.e. we can compute the updated weights
directly using the state of the connecting layers. In this paper, we have used
three phase EP, thus we derive the weight update rule following Eq. (7),

$\Delta w_{i+1}=\frac{1}{2\beta}(s^{i+1,\beta}_{*}\cdot
s^{i,\beta^{T}}_{*}-s^{i+1,-\beta}_{*}\cdot s^{i,-\beta^{T}}_{*})$ (18)

where, $s^{i,\beta}_{*}$ is the state of layer $i$ after the first “nudge”
phase with influence parameter $\beta$ and $s^{i,-\beta}_{*}$ is the state of
layer $i$ after the second “nudge” phase with influence parameter $-\beta$.

Task | Model | Method | Neuromorphic Implementation | Accuracy
---|---|---|---|---
Sentiment Analysis (IMDB) | Simple Conv RNN | EP | Local opts.; No Non-diff. problem | $79.8\pm 0.4$
ReLU GRU Dey and Salem (2017) | BP | Suffers gradient Non-diff. problem; Complex circuitry | 84.8
GRU Campos et al. (2017) | 86.5
Skip GRU Campos et al. (2017) | 86.6
Skip LSTM Campos et al. (2017) | 87.0
LSTM Campos et al. (2017) | 87.2
CoRNN Rusch and Mishra (2020) | 87.4
UniCORNN Rusch and Mishra (2021) | 88.4
Our Model | EP | Local opts.; No Non-diff problem. |  88.9 $\pm$ 0.3
Natural Language Inference (SNLI) | 100D LSTM encoders Bowman et al. (2015) | BP | Suffers gradient Non-diff. problem; Complex circuitry | 77.6
Lexicalized Classifier Bowman et al. (2015) | 78.2
Parallel LMU Chilkuri and Eliasmith (2021) | 78.8
LSTM RNN encoders Bowman et al. (2016) | 80.6
DELTA - LSTM Han et al. (2019) | 80.7
300D SPINN-PI-NT Bowman et al. (2016) | 80.9
Our Model | EP | Local opts.; No Non-diff. problem |  81.4 $\pm$ 0.2
Table 1: Comparing our models with other models trained using BP on the IMDB &
SNLI datasets.

The detailed final algorithm and calculations are added in the technical
appendix. It is further validated through the experiments reported in the next
section that as we continue updating the states $s_{t}$ (following the update
rules governed by the primitive function ($\phi$) and the energy function for
the modern hopfield network), we reach steady states during both the “free”
and “nudge” phases of EP. Thus, following the weight update procedure of EP,
we can update the weights through local-in-space update rules. This enables us
to develop state-of-the-art architectures for sentiment analysis and inference
problems - potentially implementable in energy efficient neuromorphic systems.

###  3.7 Neuromorphic Viewpoint

Implementating EP in a neuromorphic setting reduces the energy consumption by
two orders of magnitude during training and three orders during inference
compared to GPUs Martin et al. (2021). Moreover, EP does not suffer from the
non-differentiability of the spiking nonlinearity that BPTT algorithms
encounter since we do not explicitly compute the error-gradient. In EP,
because of the local operations and the error-propagation being spread across
time, the memory overhead is also much lower compared to BPTT methods where we
need to store the computational graph.

##  4 Experiments

Since EP is still in a nascent stage, this is the first work to report the
performance of convergent RNNs that are trained using EP on the specified
datasets for the task of sequence classification. In the experiments reported
in this section, we focus on benchmarking with models that are trained using
BP such as LSTMs, GRUs, etc. that can be potentially implemented in a
neuromorphic setting. Unlike the used baselines, since existing attention
models (transformers) trained using BP are not directly applicable in a
neuromorphic setting; we have not compared our model with them. In the
following sub-sections, we define specific architectures for different
subtasks in NLP. For testing our proposed work on sentiment analysis problems,
we chose the IMDB Dataset and for NLI problems, we chose the Stanford Natural
Language Inference (SNLI) dataset. Though we have not tested our method on all
the tasks in GLUE Wang et al. (2018), the proposed method can perform
classification after applying simple task-specific adjustments to the model.
Details regarding coding platform and hardware used for training have been
added to the technical appendix.

###  4.1 Sentiment Analysis

We have used IMDB dataset Maas et al. (2011) to demonstrate the application of
our model on sentiment analysis tasks. IMDB dataset comprises of 50K reviews,
25K for training and 25K for testing. Each of the reviews are either
classified as positive or negative.

####  4.1.1 Architecture

300D word2vec embeddings Mikolov et al. (2013) are used for generating the
word embeddings that are then fed into the convergent RNN and the maximum
sequence length is restricted to 600. The convergent RNN used for this
experiment comprises of two fully connected hidden layers on top of the first
projection layer. We have a single attention module applied to the first layer
similar to Fig. 1.

####  4.1.2 Results

The results (Table 1) from our model are the first to report performance on
any sentiment analysis task using EP. The experimental details are shown in
Table 2.

Hyper-params & Perf. | Range | Optimal
---|---|---
Influence Factor ($\beta$) | (0.01-0.99) | 0.1
$T$ (“Free Phase”) | (40-200) | 50
$K$ (“Nudge Phase”) | (15-80) | 25
Epochs | (10-100) | 40
Layers (Linear & FC) | - | (1 & 3)
Layer-wise lr | - | 1e-4,5e-5,5e-5,5e-5
Batch Size | (8-512) | 128
Memory (GB) | - | 7 (128 batch size)
Inference Time (sec) | - | 120 (Test Set)
Table 2: Hyper-parameters & Perf. Metrics for IMDB dataset.

In order to study the advantage of using hopfield networks as an attention
mechanism, we compare the accuracy achieved using our model against a vanilla
implementation of EP on a convergent RNN with same depth but without any
$HopAttn$ modules and we see that our proposed model outperforms with a big
margin ( $>9\%$). The computational cost increases by only 17% when we use the
$HopAttn$ module. The case where hopfield attention modules were not used
seems to converge early during training, thereby resulting in over-fitting.
However, as is evident from Fig. 3, usage of modern hopfield networks results
in better generalization.

Figure 3: (a) Train Accuracy and (b) Test Accuracy comparison of two different
convergent RNN models trained with EP - one with a $HopAttn$ module applied
(as shown in Fig. 1) and one without it. Results are reported on IMDB dataset
over 5 different runs.

###  4.2 Natural Language Inference (NLI) Problems

We have used SNLI Bowman et al. (2015) dataset to evaluate our model on NLI
tasks. NLI generally deals with the classification of a pair of sentences
namely, a premise and a hypothesis. A model, given both the sentences, is able
to predict whether the relationship between the two sentences signify an
entailment, contradiction or if they are neutral. SNLI dataset comprises of
570K pairs of premises and hypotheses.

Hyper-params & Perf. | Range | Optimal
---|---|---
Influence Factor ($\beta$) | (0.01-0.99) | 0.5
$T$ (“Free Phase”) | (40-200) | 60
$K$ (“Nudge Phase”) | (15-100) | 30
Epochs | (10-100) | 50
Layers (Linear & FC) | - | (4 & 2)
Layer-wise lr | - | 4 * 5e-4,2e-4,2e-4
Batch Size | (8-1024) | 256
Memory (GB) | - | 1.9 (256 batch size)
Inference Time (sec) | - |  40 (Test Set)
Table 3: Hyper-parameters & Perf. Metrics for SNLI dataset.

####  4.2.1 Architecture

The premise and the hypothesis are encoded as a sequence of 300D word2vec word
embeddings Mikolov et al. (2013) with max sequence length of 25. The hopfield
attention modules are used as specified in Fig. 4. The architecture defined
for this problem is inspired from the decomposable attention model Parikh et
al. (2016). We denote the premise as $A=(a_{1},a_{2},...,a_{m})$ and
hypothesis as $B=(b_{1},b_{2},...,b_{n})$, where $a_{i},b_{j}\in R^{d}$, $d$
is the dimensionality of the word vector and $m$ and $n$ are the length of the
sequences of the premise and hypothesis. We compute the new vectors
$A^{{}^{\prime}}=HopAttn(s^{11},A)$ and $B^{{}^{\prime}}=HopAttn(s^{14},B)$,
where $s^{11}$ and $s^{14}$ are values of the layers as shown in Fig. 4. These
two vectors represent self-attention within $A$ and $B$ respectively. We
compute two other vectors, $\alpha=HopAttn(s^{13},A)$ which encode the soft
alignment of premise with the hypothesis $B$ and $\beta=HopAttn(s^{12},B)$
which represents the soft alignment of hypothesis with the premise $A$, where
$s^{13}$ and $s^{12}$ are values of the layers (see Fig. 4). We then
concatenate all the sequences and feed them to the next layer (see Fig. 4).

####  4.2.2 Results

The results obtained by our model (Table 1) are compared with other models
trained using BP reported in literature that can be potentially implemented in
a neuromorphic setting. The results from our model are the first to report
performance on any NLI task using EP as a learning framework. The experimental
details are shown in Table 3.

Figure 4: High-level overview of the architecture used in case of SNLI
dataset. $HopAttn$ modules are used to capture dependencies within different
parts of a text as well as cross-attention between two separate texts. The
output of each of the layers are concatenated. The network converges over time
during both phases of EP.

##  5 Conclusion

The ability to think of neural networks as a dynamical system helps to further
deepen our understanding regarding learning frameworks and intrigues us to
delve deeper into the learning processes inside the brain. In this paper, we
explore the application of EP as a learning framework in convergent RNNs
integrated with modern hopfield networks to solve sequence classification
problems. We report for the first time the performance of EP on datasets such
as IMDB and SNLI. The constraint of EP requiring static input to the
convergent RNNs makes it really difficult to train on datasets with sequence
of data, thus an attention-like-mechanism provided by modern hopfield networks
is ideal to encode the long-term dependencies in the sequence. The spatially
(and potentially temporal) local weight update feature of EP still holds even
after introducing the modern hopfield networks and therefore it can be easily
converted into a neuromorphic implementation following the works of Martin et
al. (2021).

Certain unexplored areas in the paper can be investigated in the future.
Firstly, due to limitations of EP, we employed fixed word-embeddings instead
of learnable-embeddings used in sophisticated language models. Thus
circumnavigating that challenge to achieve even better accuracy can be an
interesting problem. Secondly, another intrinsic restriction of the convergent
RNN model used in the experiments is that the weights of the connections needs
to be symmetric. In order to make the model more bio-plausible with asymmetric
connections, the Vector Field Scellier et al. (2018) can be further explored.
Finally, although storing small sequence lengths like that of SNLI is not a
big concern, the memory overhead increases with increased sequence sizes like
that of IMDB dataset. Thus a future endeavor can be made to modify EP such
that it supports time-varying inputs.

## Acknowledgments

This material is based upon work supported in part by the U.S. Department of
Energy, Office of Science, Office of Advanced Scientific Computing Research,
under Award Number #DE-SC0021562 and the National Science Foundation grant CCF
#1955815.

## References

  * Almeida (1990) Luis B Almeida.  A learning rule for asynchronous perceptrons with feedback in a combinatorial environment.  In Artificial neural networks: concept learning, pages 102–111. 1990.
  * Berkes et al. (2011) Pietro Berkes, Gergő Orbán, Máté Lengyel, and József Fiser.  Spontaneous cortical activity reveals hallmarks of an optimal internal model of the environment.  Science, 331(6013):83–87, 2011.
  * Bi and Poo (1998) Guo-qiang Bi and Mu-ming Poo.  Synaptic modifications in cultured hippocampal neurons: dependence on spike timing, synaptic strength, and postsynaptic cell type.  Journal of neuroscience, 18(24):10464–10472, 1998.
  * Bowman et al. (2015) Samuel R Bowman, Gabor Angeli, Christopher Potts, and Christopher D Manning.  A large annotated corpus for learning natural language inference.  arXiv preprint arXiv:1508.05326, 2015.
  * Bowman et al. (2016) Samuel R Bowman, Jon Gauthier, Abhinav Rastogi, Raghav Gupta, Christopher D Manning, and Christopher Potts.  A fast unified model for parsing and sentence understanding.  arXiv preprint arXiv:1603.06021, 2016.
  * Campos et al. (2017) Víctor Campos, Brendan Jou, Xavier Giró-i Nieto, Jordi Torres, and Shih-Fu Chang.  Skip rnn: Learning to skip state updates in recurrent neural networks.  arXiv preprint arXiv:1708.06834, 2017.
  * Chilkuri and Eliasmith (2021) Narsimha Reddy Chilkuri and Chris Eliasmith.  Parallelizing legendre memory unit training.  In International Conference on Machine Learning, pages 1898–1907. PMLR, 2021.
  * Crick (1989) Francis Crick.  The recent excitement about neural networks.  Nature, 337(6203):129–132, 1989.
  * Demircigil et al. (2017) Mete Demircigil, Judith Heusel, Matthias Löwe, Sven Upgang, and Franck Vermet.  On a model of associative memory with huge storage capacity.  Journal of Statistical Physics, 168(2):288–299, 2017.
  * Dey and Salem (2017) Rahul Dey and Fathi M Salem.  Gate-variants of gated recurrent unit (gru) neural networks.  In 2017 IEEE 60th international midwest symposium on circuits and systems (MWSCAS), pages 1597–1600. IEEE, 2017.
  * Ernoult et al. (2019) Maxence Ernoult, Julie Grollier, Damien Querlioz, Yoshua Bengio, and Benjamin Scellier.  Updates of equilibrium prop match gradients of backprop through time in an rnn with static input.  Advances in neural information processing systems, 32, 2019.
  * Gammell et al. (2021) Jimmy Gammell, Sonia Buckley, Sae Woo Nam, and Adam N McCaughan.  Layer-skipping connections improve the effectiveness of equilibrium propagation on layered networks.  Frontiers in computational neuroscience, 15:627357, 2021.
  * Han et al. (2019) Kun Han, Junwen Chen, Hui Zhang, Haiyang Xu, Yiping Peng, Yun Wang, Ning Ding, Hui Deng, Yonghu Gao, Tingwei Guo, et al.  Delta: A deep learning based language technology platform.  arXiv preprint arXiv:1908.01853, 2019.
  * Hinton (2002) Geoffrey E Hinton.  Training products of experts by minimizing contrastive divergence.  Neural computation, 14(8):1771–1800, 2002.
  * Hopfield (1982) John J Hopfield.  Neural networks and physical systems with emergent collective computational abilities.  Proceedings of the national academy of sciences, 79(8):2554–2558, 1982.
  * Krotov and Hopfield (2016) Dmitry Krotov and John J Hopfield.  Dense associative memory for pattern recognition.  Advances in neural information processing systems, 29, 2016.
  * Krotov and Hopfield (2018) Dmitry Krotov and John Hopfield.  Dense associative memory is robust to adversarial inputs.  Neural computation, 30(12):3151–3167, 2018.
  * Laborieux et al. (2021) Axel Laborieux, Maxence Ernoult, Benjamin Scellier, Yoshua Bengio, Julie Grollier, and Damien Querlioz.  Scaling equilibrium propagation to deep convnets by drastically reducing its gradient estimator bias.  Frontiers in neuroscience, 15:129, 2021.
  * Lillicrap et al. (2020) Timothy P Lillicrap, Adam Santoro, Luke Marris, Colin J Akerman, and Geoffrey Hinton.  Backpropagation and the brain.  Nature Reviews Neuroscience, 21(6):335–346, 2020.
  * Maas et al. (2011) Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Y Ng, and Christopher Potts.  Learning word vectors for sentiment analysis.  In Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies, pages 142–150, 2011\.
  * Martin et al. (2021) Erwann Martin, Maxence Ernoult, Jérémie Laydevant, Shuai Li, Damien Querlioz, Teodora Petrisor, and Julie Grollier.  Eqspike: spike-driven equilibrium propagation for neuromorphic implementations.  Iscience, 24(3):102222, 2021.
  * Mikolov et al. (2013) Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean.  Efficient estimation of word representations in vector space.  arXiv preprint arXiv:1301.3781, 2013.
  * Parikh et al. (2016) Ankur P Parikh, Oscar Täckström, Dipanjan Das, and Jakob Uszkoreit.  A decomposable attention model for natural language inference.  arXiv preprint arXiv:1606.01933, 2016.
  * Pineda (1987) Fernando Pineda.  Generalization of back propagation to recurrent and higher order neural networks.  In Neural information processing systems, 1987.
  * Ramsauer et al. (2020) Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, et al.  Hopfield networks is all you need.  arXiv preprint arXiv:2008.02217, 2020.
  * Rusch and Mishra (2020) T Konstantin Rusch and Siddhartha Mishra.  Coupled oscillatory recurrent neural network (cornn): An accurate and (gradient) stable architecture for learning long time dependencies.  arXiv preprint arXiv:2010.00951, 2020.
  * Rusch and Mishra (2021) T Konstantin Rusch and Siddhartha Mishra.  Unicornn: A recurrent model for learning very long time dependencies.  In International Conference on Machine Learning, pages 9168–9178. PMLR, 2021.
  * Scellier and Bengio (2017) Benjamin Scellier and Yoshua Bengio.  Equilibrium propagation: Bridging the gap between energy-based models and backpropagation.  Frontiers in computational neuroscience, 11:24, 2017.
  * Scellier and Bengio (2019) Benjamin Scellier and Yoshua Bengio.  Equivalence of equilibrium propagation and recurrent backpropagation.  Neural computation, 31(2):312–329, 2019.
  * Scellier et al. (2018) Benjamin Scellier, Anirudh Goyal, Jonathan Binas, Thomas Mesnard, and Yoshua Bengio.  Generalization of equilibrium propagation to vector field dynamics.  arXiv preprint arXiv:1808.04873, 2018.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * Wang et al. (2018) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman.  Glue: A multi-task benchmark and analysis platform for natural language understanding.  arXiv preprint arXiv:1804.07461, 2018.
  * Widrich et al. (2020) Michael Widrich, Bernhard Schäfl, Milena Pavlović, Hubert Ramsauer, Lukas Gruber, Markus Holzleitner, Johannes Brandstetter, Geir Kjetil Sandve, Victor Greiff, Sepp Hochreiter, et al.  Modern hopfield networks and attention for immune repertoire classification.  Advances in Neural Information Processing Systems, 33:18832–18845, 2020.
  * Yuille and Rangarajan (2001) Alan L Yuille and Anand Rangarajan.  The concave-convex procedure (cccp).  Advances in neural information processing systems, 14, 2001.
  * Yuille and Rangarajan (2003) Alan L Yuille and Anand Rangarajan.  The concave-convex procedure.  Neural computation, 15(4):915–936, 2003.

##  6 Appendix: Algorithm

The pseudocode for the algorithm designed to integrate modern hopfield
networks into convergent RNNs and subsequently train the latter using the
three phase (one free phase and two nudge phases) equilibrium propagation is
described in this section. We have also elaborated on the scalar primitive
function used for natural language inference task.

###  6.1 Pseudocode

Algorithm 1 ($stateUpdate$) The state update dynamics of the convergent RNN
integrated with $HopAttn$ module. Activation function $\sigma$ removed for
simplicity and we consider the input $x$ as the stored patterns in the
hopfield network.

Input: $x,y,s,$ $\theta$, $\beta$, $T_{\beta}$
Parameter: Optional list of parameters
Output: $s$

1: $s_{0}\leftarrow s$

2: for $t=0$ to $T_{\beta}$ do

3: for $i=1$ to $n-1$ do

4: $s^{i}_{t+1}\leftarrow\frac{\partial\phi}{\partial s^{i}}(x,s_{t},\theta)$

5: if $isHopAttnApplied(s^{i}_{t+1})$ is True then

6: $s^{i}_{t+1}\leftarrow HopAttn(s^{i}_{t+1},x)$

7: end if

8: end for

9: $s^{n}_{t+1}\leftarrow\frac{\partial\phi}{\partial
s^{n}}(x,s_{t},\theta)+\beta(y-s^{n}_{t})$

10: end for

11: return $s_{T_{\beta}}$

Algorithm 2 Three-phase EP applied on the proposed convergent RNN model
integrated with $HopAttn$ layer.

Input: $x,y,$ $\theta$, $\eta$
Output: $\theta$

1: $s_{0}\leftarrow 0$

2: $s_{T}\leftarrow stateUpdate(x,y,s_{0},\theta,0,T)$

3: $s_{*}\leftarrow s_{T}$

4: $s^{\beta}_{*}\leftarrow stateUpdate(x,y,s_{*},\theta,\beta,K)$

5: $s^{-\beta}_{*}\leftarrow stateUpdate(x,y,s_{*},\theta,-\beta,K)$

6:
$\nabla^{EPsym}_{\theta}\leftarrow\frac{1}{2\beta}(\frac{\partial\phi}{\partial\theta}(x,s^{\beta}_{*},\theta)-\frac{\partial\phi}{\partial\theta}(x,s^{-\beta}_{*},\theta))$

7: $\theta\leftarrow\theta+\eta\nabla^{EP}_{\theta}$

8: return $\theta$

###  6.2 Scalar Primitive Function $\phi$ \- SNLI

The scalar primitive function $\phi$ is used to derive the state transition
dynamics of the convergent RNN. In the paper, a high level overview of the
function $\phi$ is described where the task described has only one sequence as
an input. However in order to understand the sequence classification problems
with multiple inputs $x_{1},x_{2}\in R^{N\times D_{K}}$, we need to further
elaborate few underlying operations. If we consider the natural language
inference problem (SNLI dataset) with the projection layers as defined in Fig.
4, we can define the $\phi$ as,

$\phi(x,{s},\theta)=\sum_{j=1}^{2}{s^{1j}}\bullet(x_{1}\cdot
w_{1j})+\sum_{j=3}^{4}{s^{1j}}\bullet(x_{2}\cdot w_{1j})\\\ +{s^{2}}^{T}\cdot
w_{2}\cdot F(s^{1})+\sum_{i=2}^{n-1}{s^{i+1}}^{T}\cdot w_{i+1}\cdot s^{i}$
(19)

where $\bullet$ represents euclidean scalar product of two tensors with same
dimensions, $(x_{i}\cdot w_{1j})$ represents the linear projection of $x_{i}$
through $w_{1j}\in R^{D_{K}\times D_{K}}$, $\forall j=1,2,...,4$, and
$w_{i+1}$ is the weight of the connection between $s^{i}$ and $s^{i+1}$,
$\forall i=1,...,n-1$. $s^{1}$ is the concatenation of all four projection
layers $s^{1j}$, $\forall j=1,2,...,4$. The flattened output of the
concatenation of all four projection layers ($F(s^{1})$ where $F$ is the
flattening operation from $(A,B)$ to $(1,AB)$) is then fed into rest of the
$n-1$ fully connected layers.

##  7 Appendix: Experimental Details

We implemented modified EP, as given by Algorithm 2, using PyTorch 1.7.0. The
experiments were run on Nvidia GeForce RTX 2080 Ti GPUs (8). The pre-processed
IMDB dataset available in Keras is used in our experiments.

###  7.1 Sentiment Analysis (IMDB dataset)

The pre-processed IMDB dataset, available in Keras, is used in our
experiments. In the sentiment analysis problems, we have 300D word2vec
embeddings with maximum sequence length of 600 as the input to the convergent
RNN. The first layer is a linear projection layer which maintains the same
encoding dimension of 300. The hopfield attention module is integrated with
the first layer. The final output of that layer is flattened and fed into a
hidden layer of size 1000, followed by another layer of size 40 and then
followed by the output layer of size 2. Adam optimizer is used as the
optimizer to update the weights and the activation function is a modified form
of hard sigmoid function. The loss function, as discussed in the paper, is
mean squared error. GPU memory usage while training is limited to around 7GB
with batch size of 128 and running time for one experiment is around 5 hours
on the described hardware.

###  7.2 Natural Language Inference (SNLI dataset)

We have 300D word2vec embeddings with maximum sequence length of 25 for both
premise and hypothesis given as the input. The scalar primitive function can
be considered to be of a similar form as given by Eqn. 19. However, we
consider four linear projection layers connected to inputs $A$ and $B$, as
depicted in the paper. Then the output of all the projection layers are
concatenated and flattened to be fed into the remaining two fully connected
layers. The architecture is kept simple to show the effectiveness of using the
modern hopfield network. Activation function and optimizer used is the same as
the previous experiment. GPU memory usage during training is limited to around
1.9 GB with batch size of 256. Running time for one experiment is around 6
hours.

###  7.3 Implementation Details

In order to make the code simpler and intuitive, the automatic differentiation
framework of PyTorch has been used in our code to calculate the state and
weight update rules. The scalar primitive function $\phi$ has been defined for
each of our test datasets according to the underlying architecture and then
the state update dynamics and weight updates have been computed by taking the
appropriate derivatives of $\phi$. Such an implementation also makes the code
robust to further changes in architecture.