# Fader Networks for domain adaptation on fMRI: ABIDE-II study

Marina Pominova  Skolkovo University of Science and Technology  Ekaterina
Kondrateva  Skolkovo University of Science and Technology  Maxim Sharaev
Skolkovo University of Science and Technology  Alexander Bernstein  Skolkovo
University of Science and Technology  Evgeny Burnaev  Skolkovo University of
Science and Technology

###### Abstract

ABIDE is the largest open-source autism spectrum disorder database with both
fMRI data and full phenotype description. These data were extensively studied
based on functional connectivity analysis as well as with deep learning on raw
data, with top models accuracy close to 75% for separate scanning sites. Yet
there is still a problem of models transferability between different scanning
sites within ABIDE. In the current paper, we for the first time perform domain
adaptation for brain pathology classification problem on raw neuroimaging
data. We use 3D convolutional autoencoders to build the domain irrelevant
latent space image representation and demonstrate this method to outperform
existing approaches on ABIDE data.

###### keywords:

ABIDE, functional MRI, Domain Adaptation, Fader Networks, Invariant
Representation

##  1 INTRODUCTION

###  1.1 Autism spectrum disorder classification

Autism spectrum disorder (ASD) is wide diagnosis capturing different
developmental disorders [1]. This disorder is being widely studied now by
separate laboratories and consortiums resulting in several publicly available
ASD neuroimaging datasets.

Modern machine and deep learning models are one of the most powerful
techniques to extract information from complex and high dimensional fMRI
imaging data. However, due to the difficulty of collecting medical data, the
typical size of such datasets usually varies within a few tens or hundreds of
observations, which is too small for most statistical instruments and, most
importantly, proper training of machine learning and deep learning models. The
results of neuroimaging studies of ASD have often been inconsistent due to
small sample sizes in relation to subtle effects [1].

One of the possible solutions is the meta-analysis of several datasets
collected for a single research task. Unfortunately, in this case, another
problem arises: due to differences across studies (such as differences in
methodology, equipment, scanner settings, processing pipelines, etc.), the
obtained scan data vary in imaging properties, which significantly affects the
generalization ability of the resulting models and might nullify the meta-
analysis benefits. Moreover, the applicability of the derived models to the
new test data from other sources is also usually hurt due to inter-site
variability effects.

Autism Brain Imaging Data Exchange (ABIDE) initiative is the largest open-
source data collection, which has two releases: ABIDE-I (2014) and ABIDE-II
(2017). To date, ABIDE-II has collected over 1000 additional studies, while
the total number of observations in the two releases has exceeded 2000
subjects from 29 different sites. The difference between sites is not only in
processing pipeline or patient cohort but also in scanning equipment ($1.5$T
or $3$T MR machines). Analysis of fMRI data collected in the scope of ABIDE
initiative may help to shed light on the causes and mechanisms of ASD, as well
as to identify objective biomarkers of this disorder which could facilitate
diagnostics, therapy and outcome predictions [2]. However, to get the most out
of the analysis of this large and diverse dataset, the difficulties associated
with inter-site data variation and resulting low robustness and generalization
ability of analytical models have to be overcome.

###  1.2 ABIDE: previous studies

To the moment, there exists a large scope of research focused on the study of
the first phase ABIDE - ABIDE-1 dataset. However, in most cases, the authors
analyzed small subsets from separate sites, separately or as part of a joint
sample, but without trying to exclude the cross-site variability effects.
Several ASD detection models trained on small samples achieved classification
accuracies up to 75% [2]. At the same time, the majority of studies conducted
on the whole multi-site dataset showed significantly lower accuracy around 65%
[3] [4].

The studies also differ methodologically: more frequently they use functional
connectivity analysis [5], while others apply deep learning approaches
directly to raw or preprocessed data, see [6] [4]. While working with
vectorized features from fMRI functional connectivity matrix is not new, some
authors explore graph-theoretic and topological features from connectivity
graphs and report comparable or better performance (around 65-80% accuracy
depending on patient cohort) than traditional approaches [5].

The first application of networks with LSTM is demonstrated in [7] to process
raw fMRI time-series data from a multi-site heterogeneous ABIDE-I dataset with
an accuracy of 68.5% on cross-validation. It is important that the authors
also made one more step towards revealing autism biomarkers from fMRI data,
trying to interpret the best model and identify brain regions important for
distinguishing subjects with ASD from typical controls

With the ability to automatically learn the task-specific features from
unstructured image or sequence data, deep learning techniques also have the
potential to be applied to raw full-size fMRI data. Such architectures already
demonstrated a possibility to detect various psychoneurological conditions,
for example, epilepsy, schizophrenia, and major depressive disorder from full-
size fMRI without any manual feature extraction [Dakka, Pominova]. Several
modifications of the neural network architecture are made to capture the one-
dimensional temporal and three-dimensional spatial properties of 4D fMRI data,
for example, in [4]. Here, in the experiment with several sites, the accuracy
is significantly (64% versus 78% for the best model) lower than in the
experiment with a single site, which is consistent with observations in other
works. An interesting combination of biologically plausible feature extraction
and deep learning approach is demonstrated in [8]. Here the input to the 3D
CNN is formed by concatenating channel-wise several three-dimensional voxel-
level maps of functional connectivity, calculated from raw fMRI time-series.
Furthermore, the authors applied the saliency map approach to their 3D CNN
model to visualize discriminative features and assumed that unsupervised
learning on these maps can provide insights into the clinical origins of
disease [8].

The small number of exceptions in which there have been attempts to explicitly
account for site-related differences and data heterogeneity include, for
example, work [1], in which the authors fitted linear mixed-effects models for
each calculated feature with respect to the data from 54 data sets (data sets
were coded as factor variables). Federated learning and domain adaptation
methods were proposed in [9] to train the models locally on separate sites
data and to account for differences of fMRI signal distributions from
different sites. Specifically, a mixture of experts (MoE, adaptation near the
output layer), and adversarial domain alignment (adaptation on the data
knowledge representation level), significantly increased the model accuracy,
see [9] for details. It is reasonable to assume that reliable biomarkers could
be detected from a reliable (robust across sites) model. Thus, federated
learning might be helpful to detect replicable and robust biomarkers across
different sites.

In the present study, we focus on the problem of eliminating site-related
differences from fMRI data and propose an approach to train robust and
transferable between sites neural network models on a multi-site dataset. We
conduct experiments on the data collected in the scope of the ABIDE
initiative. Since the first phase of ABIDE (ABIDE-I) is already a well-studied
dataset, here we focus on the recently released second phase (ABIDE-II).
First, for the four data sites with largest contribution in the ABIDE-II
dataset, we provide ASD recognition baselines obtained with two widely used
approaches — conventional functional connectivity matrices analysis and
recently proposed full-size MRI series analysis with a conv-temporal neural
network [10][11]. We present classification results on the merged dataset of
all 4 sites, as well as on each site data separately. Second, we propose a
novel approach based on 3D convolutional autoencoders, fader networks and
domain-adaptation techniques to remove site-related variability and reduce the
heterogeneity of the data from different sites. To validate if the proposed
method allows enhancing the generalization ability of ASD classification
models, we compare it with baseline conventional methods for connectivity
matrices and full-size fMRI on the same data sample.

##  2 Methods

###  2.1 Data and preprocessing

We perform the study on the publicly available ABIDE-II
dataset111http://fcon_1000.projects.nitrc.org/indi/abide/abide_II.html which
contains resting-state fMRI sequences of patients with ASD and healthy
controls [12]. We use the whole data sample from 4 collection sites, totaling
352 subjects. The distribution of ASD subjects per-site represented in the
Table 1.

Table 1: ABIDEII distribution of ASD subjects per-site  | Description
---
| NYU
---
(NYU_1, NYU_2)
| OHSU
---
(OHSU_1)
| KKI
---
(KKI_1)
| GU
---
(GU_1)
All sites
Total Subject | 105 | 93 | 79 | 75 | 352
ASD Subject | 30 | 56 | 62 | 36 | 184
Mean age (STD) | 10.01 (2.4) | 9.60 (4.61) | 10.20 (4.02) | 10.64 (1.72) | 10.15 (2.98)

The sequences of images were preprocessed according to Functional Connectomes
Project / INDI [13] protocols which include slice timing correction, motion
realignment, intensity normalization, and rigid registration to a template. We
use the denoised dataset, which is as for now the only available option with
samples from all collecting sites.

For Functional Connectivity (FC) analysis we applied the standard Nilearn
pipeline222https://nilearn.github.io/modules/generated/nilearn.input_data.NiftiLabelsMasker.html
with band-pass filtering lowpass $=0.1$ and highpass $=0.01$ Hz, signal
standardization, and de-trending. For brain parcellation, we use an Automated
Anatomical Labeling (AAL) atlas consisting of 116 regions. For each region,
the corresponding time series is taken as the first component of Singular
Value Decomposition with high variance confounds function of all the region
voxels. Next, the functional connectivity matrix is computed as a $116\times
116$ matrix of pairwise correlations of the obtained time series for all
regions.

###  2.2 Fader network for elimination of site-related information

Various domain adaptation techniques are proven candidates to perform training
of domain-invariant models with enhanced generalization ability. One of the
most widely used methods is a DANN approach proposed by [14]. This allows
adapting a pretrained model to a new domain by fine-tuning it with an
additional network, the domain discriminator. The model is trained for a
target task only on a sample from “source domain”. At the same time, the
discriminator takes features extracted by the model from both ”source” and
”target” domain data and tries to distinguish which domain these features come
from. Usually, these features are taken from the last hidden layer of the
model. The main model is trained with gradient descent in a standard way.
However, when the discriminator is optimized, the gradients backwarded through
the weights of the main model are reversed. As a result, the model learns to
extract the features that confuse the discriminator most and complicate
distinguishing between domains.

However, a potential drawback of this approach in the case of complex and
high-dimensional functional MRI data is the difficulty of stabilizing
adversarial training. Our preliminary experiments revealed that when the
domain adaptation approach is applied during model training for the target
task, the main classification network converges too fast and begins to
overfit. Thus, the domain discriminator does not have enough time to properly
guide the learning of domain-invariant features.

For this reason, we opted for applying a similar technique and train a fader
network [15] \- an autoencoder-like model for preliminary extraction of site-
independent features, which can be further used for training an ASD-
recognizing classifier.

Fader networks were suggested as a method for training a conditional
autoencoder for generation images with predefined values of some attributes,
such as age, hair color, and so on. The trained model is able to produce
different versions of input image conditioned on the chosen attribute values.
This is achieved by disentangling the salient information of the image and the
values of attributes directly in the latent space. The decoder takes desirable
attribute values as an additional input along with the latent code and
generates an image with the corresponding characteristics. At the same time,
the latent representation itself remains independent and does not contain any
attribute-related information.

The elimination of the information on attributes from the latent space of the
autoencoder is achieved in a manner similar to the domain adaptation
approaches. During training the encoder-decoder model, an additional
discriminator network tries to predict attribute values of each training image
from its latent representations. The latent discriminator is optimized
adversarially with the encoder, forcing it to filter out the information
related to the attributes.

Thus, the entire model consists of three main components: — the encoder $E$, a
convolutional network that maps the input image $x$ to its latent
representation $z(x)$, — the decoder $D$, a deconvolutional network that takes
as input latent representation $z$ of the image and a vector of attribute
values y, and generates a new version of the image with visual characteristics
corresponding to attributes - $x^{\prime}(z,y)$, — the latent discriminator
$L$, which tries to identify true attribute values of the input image x based
on its latent representation $z$. The discriminator is optimized adversarially
with the encoder-decoder architecture, and to make it unable to predict the
right attributes, the encoder has to learn attribute-invariant
representations.

In this study, we apply the Fader networks approach to obtain an encoder
capable of extracting site-independent latent representations from full-size
fMRI data. We train the encoder-decoder model on 3D brain images sampled from
fMRI time series at random time steps and consider site labels as a single
attribute variable that the discriminator tries to predict from latent
representations.

---
Figure 1:  Fader Network architecture used. An fMRI image - $x$ and ABIDE-II
site id (one-hot-encoded) - $y$ are given as input. The encoder maps $x$ to
it’s latent representation; the discriminator is trained to predict $y$ given
$z$ whereas the encoder is trained to make it impossible for the discriminator
to predict $y$ given vector of latent features only. The decoder should
reconstruct $x$ given the latent representation and $y$.

###  2.3 Classification of latent representation sequences

After the Fader Network has converged, we use the encoder to transform the
fMRI sequence into a sequence of one-dimensional vectors of latent features.
These features are not trained for any specific task and are expected to
contain maximum information from the original fMRI series in a compressed
form, except for the site-related information, which is excluded by the
training procedure. They can be further used to train any predictive model,
for example, a classification recurrent network for ASD detection.

It may also be noted that combining the encoder that extracts feature
descriptions from brain images in the fMRI sequence with a recurrent
classification network that processes the temporal information is
architecturally equivalent to the ConvTempNet, a proven architecture for full-
size fMRI analysis proposed by [Dakka, Pominova]. This architecture consists
of convolutional and recurrent components aimed to take into account both
spatial and temporal patterns in the data. During the training process, the
convolutional part takes as input 3D brain images at each time step of the
fMRI sequence and transforms them into the one-dimensional feature vectors.
Thus, the whole fMRI time series is compressed into a sequence of 1D vectors,
which is then fed into the recurrent part of the network to extract temporal
information.

Therefore, it is natural to compare the described model with our proposed
fader network-based architecture. We also consider a comparison with a model
consisting of the pretrained encoder of the plain autoencoder and recurrent
classification network to directly assess the effect of site information
elimination.

###  2.4 Experiment details

To estimate the influence of site-related information on the performance of
ASD detection from full-size fMRI data, we examine the following three
architectures:

— an original ConvTempNet, consisting of the convolutional and recurrent
parts,

— a pretrained encoder of the plain autoencoder connected with a recurrent
neural network,

— our proposed approach with a pretrained site-independent fader network
encoder connected with a recurrent neural network.

We also consider two degrees of elimination of the information associated with
the data site - the slight adjusting of latent vectors to reduce the share of
site-related variability and almost complete its removal of it. We denote the
former option as FaderNetwork-slight - FaderNetwok with slight site-related
information elimination, and the latter as FaderNetwork-strong - FaderNetwok
with strong site-related information elimination.

For a uniform and fair comparison, we use an identical convolutional
architecture for the convolutional part of ConvTempNet and encoders of the
conventional autoencoder and the fader network. Specifically, for the fader
net, we adapted a 3D-version of the architecture proposed in [15]. The encoder
part consists of 6 blocks of convolutional layers with kernel size $4\times
4\times 4$, a stride of $2$, and a padding of $1$. The number of channels
starts from $32$ and multiplies by $2$ with each next block, while the spatial
resolution of the activations reduces twice along each axis. We also apply
leaky-ReLU with a slope of $0.2$ in the encoder, as well as batch
normalization after all convolutional layers except for the first one. Each
time point of an fMRI sequence is represented by a three-dimensional image of
size $64\times 64\times 64$. After going through the encoder, it is
transformed into a one-dimensional latent representation with $1024$ feature
maps. The decoder part architecture accurately mirrors that of the encoder.
The only difference is that here we apply 3D transposed convolutions with the
same parameters, and thus each of 6 blocks of layers doubles the spatial
resolution of activation maps back to $[64,64,64]$. Also, convolutional layers
are followed by ReLU activation and batch normalization. The plain autoencoder
model has the same architecture and the convolutional part of ConvTempNet is
analogous to the encoder.

For the latent discriminator, we construct a small classification network with
3 fully-connected layers of $1024$, $256$, and $64$ units respectively, ReLU
activations, batch normalization, and dropout of $0.3$ after each layer. In
our experiments, both dropout and batch normalization turned out to greatly
influence the success of an adversarial training procedure.

Similarly, we also fix a single recurrent architecture for the temporal part
of ConvTempNet and for processing sequences of one-dimensional feature
vectors, extracted by autoencoder’s and fader network’s encoders. Our
recurrent network consists of $2$ layers of GRU units with $512$ memory units
each. After they processed an input sequence of one-dimensional feature
vectors extracted from the fMRI time series, the hidden states of the second
GRU layer memory cells are averaged over time and used to predict the final
probability of ASD presence.

Training of all three models is performed on the same training set and with a
fixed validation strategy. We assess the classification performance of each
model with the ROC AUC metric computed on $10$-fold cross-validation. Cross-
validation splits for all models were identical. We also perform model
evaluation on the leave-one-site-out validation. In this case, we alternately
select one site for testing, train the model on the data from all remaining
sites and then assess its performance on the data from the selected site that
it has not seen before. It allows us to estimate model robustness to the site-
related variations in a training set and transferability to the data from new
sources.

###  2.5 Classification of functional connectivity matrices

We also performed a comparison of the proposed models for full-size fMRI data
with the classical approach to fMRI analysis based on functional connectivity
matrices.

Feature descriptions for each object were obtained by flattening the upper-
triangular part of the symmetric functional correlation matrix into a vector
of $6670$ connectivity features. Since their number exceeds the number of
objects, feature selection is an important step before training the model. We
select the $200$ most important variables according to their weights in the
logistic regression model on the training set. Then, these features are used
to train the final classification mode on the same training. As a final
classifier, we used the SVM algorithm with hyper-parameters optimized with
grid search.

##  3 Results

Table 2: Results of using different training strategies on data from different ABIDE sites, on 10-fold cross-validation and leave-one-site-out validation, ROC AUC(Std) Model | GU | KKI | NYU | OHSU | All sites | LOSO
---|---|---|---|---|---|---
FC, SVC (kernel = rbf) |  | 0.737
---
(0.120)
| 0.687
---
(0.152)
| 0.656
---
(0.120)
| 0.539
---
(0.119)
| 0.635
---
(0.051)
| 0.579
---
(0.064)
ConvGRU on full-size fMRI |  | 0.543
---
(0.098)
| 0.586
---
(0.214)
| 0.771
---
(0.063)
| 0.664
---
(0.072)
| 0.665
---
(0.049)
| 0.541
---
(0.077)
GRU on latent vectors (Autoencoder) |  | 0.624
---
(0.101)
| 0.519
---
(0.225)
| 0.760
---
(0.059)
| 0.646
---
(0.103)
| 0.667
---
(0.075)
| 0.534
---
(0.082)
GRU on site-invariant latent vectors (FaderNetwok-slight) |  | 0.611
---
(0.089)
| 0.579
---
(0.201)
| 0.763
---
(0.065)
| 0.678
---
(0.091)
| 0.731
---
(0.078)
| 0.652
---
(0.080)
GRU on site-invariant latent vectors (FaderNetwok-strong) |  | 0.575
---
(0.093)
| 0.535
---
(0.170)
| 0.729
---
(0.049)
| 0.637
---
(0.161)
| 0.692
---
(0.094)
| 0.597
---
(0.091)

In this section, we demonstrate the results of our experiments. The Table 2
shows ROC AUC scores of the three models for full-size data and the baseline
model for functional connectivity features. For each model, we present its
classification performance on small datasets from individual sites, on the
joint dataset from all 4 sites, and on leave-one-site-out validation.

When analyzing the results, we primarily focus on a) the impact of the
information associated with the data sites on the model generalization, b) the
effect of removing the site-related information on the generalization ability
and transferability of the model.

a) The impact of the information associated with the data sites.

First, for the standard ConvTempNet and the plain encoder connected with a
recurrent neural network, we observe that the classification performance of
the same model on the data from a single site may noticeably vary depending on
the site. This can be explained by the differences in data quality, degree of
noise removal, and characteristics of the study participants, such as age and
comorbidities. Furthermore, in both cases, the model trained on the joint
multi-site dataset underperforms two out of four single-site models, which
also indicates the negative effect of site-related variability on the model
generalization ability. Even if several data sources were presented in a
training sample, the model has difficulties in revealing robust and site-
independent biomarkers of autism. This is additionally supported by the fact
that model performance on the leave-one-site-out validation turns out to be
extremely low and close to the random guessing. Thus, almost all the
distinctive patterns learned by the model partly rely on the properties of a
particular data site and are of no value in the analysis of data from a new
source.

b) The effect of removing site-related information.

In the case of the models with fader network encoders connected with recurrent
neural networks we observe that the performance on individual sites is in
general slightly lower, especially for the model with strong elimination of
site-related information. It is potentially caused by a slight distortion of
latent features during adversarial training. However, on the multi-site data,
these models demonstrate considerably better performance. Most significantly,
the use of site-independent latent vectors improves the results in leave-one-
site-out validation. ASD recognition on the completely new data from the
previously unseen site is still less accurate than on the new data from sites
presented in the training set. However, now the model noticeably outperforms
the random guessing, which indicates that the disease biomarkers learned by it
are more general and probably rely on the properties of the brain itself
rather than on the properties of the data source.

It’s also worth noting that a more thorough removal of site-specific
variability does not improve the quality of the final model. Apparently, part
of this variability turns out to be difficult to separate from the patient-
specific patterns, and the attempt to completely eliminate it leads to the
corruption of other information in latent vectors and the decrease in quality
of the disease detection.

Figure 2:  tSNE embedding of fMRI latent representations learned by (a)
conventional autoencoder, (b) fader network with slight site-related
information elimination, (c) fader network with strong site-related
information elimination

###  3.1 Comparison of latent representations learned by conventional
Autoencoder and Fader Networks

Since the models with fader network encoders proved to be the most successful
in the discussed above problem of autism recognition from multi-site full-size
fMRI, here we visually inspect the latent representations that they extract
from the brain images. We also compare them with the representations learned
by the standard autoencoder in order to figure out how introducing the site-
discriminator in the encoder training affects the features learned.

For each subject in the considered dataset of samples from 4 largest sites in
ABIDE-II, we encode it’s 3D brain image for a single time-step, randomly
sampled from fMRI sequence, into a latent vector using 1) the encoder of a
plain autoencoder, 2) the encoder of fader network with slight elimination of
site-related information, 3) the encoder of fader network with strong
elimination of site-related information. After that, we reduce the
dimensionality and plot the tSNE 2d-embedding for samples from all sites.

The results, presented in Figure 2.a, indicate that for the plain autoencoder,
latent representations contain information which allows clearly separating
samples from different sites into distinct clusters. Figure 2.b demonstrates
that after we force the encoder to avoid preserving site - related knowledge
by introducing discriminator into a latent space, the embeddings learned are
tending to mix across the sites. In Figure 2.c we observe a complete mixing of
objects from different sites, which, however, is achieved at the cost of a
partial loss of valuable information.

##  4 Discussion

We performed the first to our knowledge attempt toward the generalizable fMRI
deep classification models. As we can see from the obtained results, the model
performance can be significantly affected by the presence site-related
variability in training data. However, we prove the ability of fader networks
to be applied for construction site-invariant latent representations, which
allow obtaining notable improvement of autism recognition, especially on the
data from the previously unseen site. At the same time, the complete removal
of any source-related part of the information is still undesirable, since it
inevitably affects other patterns in the data that are important for pathology
recognition.

##  5 Acknowledgements

The reported study was funded by RFBR according to the research project
20-37-90149.

## References

  * [1] Postema, M. C., Van Rooij, D., Anagnostou, E., Arango, C., Auzias, G., Behrmann, M., Busatto Filho, G., Calderoni, S., Calvo, R., Daly, E., et al., “Altered structural brain asymmetry in autism spectrum disorder in a study of 54 datasets,” Nature communications 10(1), 1–12 (2019).
  * [2] Yahata, N., Morimoto, J., Hashimoto, R., Lisi, G., Shibata, K., Kawakubo, Y., Kuwabara, H., Kuroda, M., Yamada, T., Megumi, F., et al., “A small number of abnormal brain connections predicts adult autism spectrum disorder,” Nature communications 7(1), 1–12 (2016).
  * [3] Abraham, A., Milham, M. P., Di Martino, A., Craddock, R. C., Samaras, D., Thirion, B., and Varoquaux, G., “Deriving reproducible biomarkers from multi-site resting-state data: An autism-based example,” NeuroImage 147, 736–745 (2017).
  * [4] El-Gazzar, A., Quaak, M., Cerliani, L., Bloem, P., van Wingen, G., and Thomas, R. M., “A hybrid 3dcnn and 3dc-lstm based model for 4d spatio-temporal fmri data: An abide autism classification study,” in [OR 2.0 Context-Aware Operating Theaters and Machine Learning in Clinical Neuroimaging ], 95–102, Springer (2019).
  * [5] Kazeminejad, A. and Sotero, R. C., “Topological properties of resting-state fmri functional networks improve machine learning-based autism classification,” Frontiers in neuroscience 12, 1018 (2019).
  * [6] Price, T., Wee, C.-Y., Gao, W., and Shen, D., “Multiple-network classification of childhood autism using functional connectivity dynamics,” in [International Conference on Medical Image Computing and Computer-Assisted Intervention ], 177–184, Springer (2014).
  * [7] Cheng, D. and Liu, M., “Classification of alzheimer’s disease by cascaded convolutional neural networks using pet images,” in [International Workshop on Machine Learning in Medical Imaging ], 106–113, Springer (2017).
  * [8] Khosla, M., Jamison, K., Kuceyeski, A., and Sabuncu, M. R., “Ensemble learning with 3d convolutional neural networks for functional connectome-based prediction,” Neuroimage 199, 651–662 (2019).
  * [9] Li, X., Gu, Y., Dvornek, N., Staib, L., Ventola, P., and Duncan, J. S., “Multi-site fmri analysis using privacy-preserving federated learning and domain adaptation: Abide results,” arXiv preprint arXiv:2001.05647 (2020).
  * [10] Dakka, J., Bashivan, P., Gheiratmand, M., Rish, I., Jha, S., and Greiner, R., “Learning neural markers of schizophrenia disorder using recurrent neural networks,” arXiv preprint arXiv:1712.00512 (2017).
  * [11] Pominova, M., Artemov, A., Sharaev, M., Kondrateva, E., Bernstein, A., and Burnaev, E., “Voxelwise 3d convolutional and recurrent neural networks for epilepsy and depression diagnostics from structural and functional mri data,” in [2018 IEEE International Conference on Data Mining Workshops (ICDMW) ], 299–307, IEEE (2018).
  * [12] Di Martino, A., O’connor, D., Chen, B., Alaerts, K., Anderson, J. S., Assaf, M., Balsters, J. H., Baxter, L., Beggiato, A., Bernaerts, S., et al., “Enhancing studies of the connectome in autism using the autism brain imaging data exchange ii,” Scientific data 4(1), 1–15 (2017).
  * [13] Craddock, C., Benhajali, Y., Chu, C., Chouinard, F., Evans, A., Jakab, A., Khundrakpam, B. S., Lewis, J. D., Li, Q., Milham, M., et al., “The neuro bureau preprocessing initiative: open sharing of preprocessed neuroimaging data and derivatives,” Neuroinformatics 4 (2013).
  * [14] Ganin, Y. and Lempitsky, V., “Unsupervised domain adaptation by backpropagation,” in [International conference on machine learning ], 1180–1189, PMLR (2015).
  * [15] Lample, G., Zeghidour, N., Usunier, N., Bordes, A., Denoyer, L., and Ranzato, M., “Fader networks: Manipulating images by sliding attributes,” in [Advances in neural information processing systems ], 5967–5976 (2017).


11institutetext:  Cornell University, Ithaca NY 14850, USA
11email<EMAIL_ADDRESS>
# Adaptive Learning Rate and Momentum for Training Deep Neural Networks

Zhiyong Hao 11 Yixuan Jiang 11 Huihua Yu 11 Hsiao-Dong Chiang 11

###### Abstract

Recent progress on deep learning relies heavily on the quality and efficiency
of training algorithms. In this paper, we develop a fast training method
motivated by the nonlinear Conjugate Gradient (CG) framework. We propose the
Conjugate Gradient with Quadratic line-search (CGQ) method. On the one hand, a
quadratic line-search determines the step size according to current loss
landscape. On the other hand, the momentum factor is dynamically updated in
computing the conjugate gradient parameter (like Polak-Ribiere). Theoretical
results to ensure the convergence of our method in strong convex settings is
developed. And experiments in image classification datasets show that our
method yields faster convergence than other local solvers and has better
generalization capability (test set accuracy). One major advantage of the
paper method is that tedious hand tuning of hyperparameters like the learning
rate and momentum is avoided.

###### Keywords:

Optimization Algorithm Line Search Deep Learning.

##  1 Introduction

We consider the minimization problem commonly used in the machine learning
setting with the following finite-sum structure:

$\underset{\theta}{\min}\leavevmode\nobreak\
f(\theta)=\frac{1}{n}\sum_{i=1}^{n}f_{i}(\theta)$ (1)

where $f:D\subset\mathds{R}^{n}\to\mathds{R}$ is continuously differentiable
and $D$ is the search space. A point $\theta^{*}\in D$ is called a local
minimum if $f(\theta^{*})\leq f(\theta)$ for all $\theta\in D$ with
$\left\lVert\theta-\theta^{*}\right\rVert<\sigma$ for $\sigma>0$. A general
training procedure is in Algorithm 1, where key steps are determining the
update direction $p_{t}$ and step size $\alpha_{t}$ in each iteration.

Algorithm 1 Iterative Approach for Optimizing (1)

1:procedure Optimize($f(\theta),maxiter$)

2: Initialize $\theta_{0}$

3: for $t\leftarrow 1$ to $maxiter$ do

4: Choose a search direction $p_{t}$;

5: Determine a update step size $\alpha_{t}$;

6: $\theta_{t}\leftarrow\theta_{t-1}+\alpha_{t}p_{t}$;

7: return $\theta_{t}$

Several methods propose various choices of search directions. The Gradient
Descent (GD) method directly takes the negative gradient
$-\nabla_{\theta}\leavevmode\nobreak\ f(\theta)$. Newton’s method selects
$p_{t}$ as $-(\nabla^{2}_{\theta}\leavevmode\nobreak\
f(\theta))^{-1}\nabla_{\theta}f(\theta)$ (where $\nabla^{2}_{\theta}f(\theta)$
is the Hessian matrix). On the other hand, momentum [23, 19] is used to
stabilize $p_{t}$ in stochastic optimization. Adam [12] formalizes $p_{t}$ as
a moving average of first moments divided by that of second moments. Other
methods focus on the choice of $\alpha_{t}$. Various classic line-search
methods were proposed, like the Newton-Raphson, Armijo, and Wolfe line search.
Recently, [31] applied the Armijo rule in SGD. In this paper, we propose the
Conjugate Gradient with Quadratic line-search method that jointly optimizes
both $p_{t}$ and $\alpha_{t}$ for fast convergence.

The performance of SGD is well-known to be sensitive to step size $\alpha_{t}$
(or the learning rate called in machine learning). Algorithms, such as AdaGrad
[7] and Adam [12], ease this limitation by tuning the learning rate
dynamically using past gradient information. Line-search strategy is another
option to adaptively tune step size $\alpha_{t}$. Exact line-search requires
solving a simpler sub-optimization problem, such as applying quadratic
interpolation [17]. An inexact line-search method can also be applied, such as
the Armijo rule [1], to guarantee a sufficient decrease at each step. A
stochastic variant of Armijo rule was proposed in [31] to set the step size
for SGD.

Conjugate Gradient method (CG) is also popular for solving nonlinear
optimization problems. In CG, $p_{t}$ is chosen as $(-\nabla
f(\theta)+\beta_{t}p_{t-1})$, where $\beta_{t}$ is the conjugate parameter
computed according to various formulas. $\alpha_{t}$ is determined by a line
search ($\arg\min_{\alpha}f(\theta_{t}+\alpha p_{t})$). In fact, CG can be
understood as a Gradient Descent with an adaptive step size and dynamically
updated momentum. For the classic CG method, step size is determined by the
Newton-Raphson method or the Secant method, both of which need to exactly
compute or approximate the Hessian matrix. This makes it less appealing to
apply classic CG in deep learning.

In this paper, we propose the Conjugate Gradient with the Quadratic line-
search method (CGQ), which uses quadratic interpolation line-search to choose
the step size and adaptively tunes momentum term according to the Conjugate
Gradient formula. CGQ requires no Hessian information and can be viewed as a
variant of SGD with adaptive learning rate and momentum. For quadratic line-
search, we propose two variants, i.e., the 2-point method and the least
squares method, which are introduced in detail in Sections 3.2.1 and 3.2.3.
And we illustrate how we dynamically determine the momentum term during
training in Section 3.3. Our analyses indicate that SGD with 2-point quadratic
line-search converges under convex and smooth settings. sCGQ is further
proposed to improve efficiency and generalization on large models. Through
various experiments, CGQ exhibits fast convergence, and is free from manually
tuning the learning rate and momentum term.

Our major contributions are summarized as follows:

  * •

A quadratic line-search method is proposed to dynamically adjust the learning
rate of SGD.

  * •

An adaptive scheme for computing the effective momentum factor of SGD is
designed.

  * •

Theoretical results are developed to prove the convergence of SGD with
quadratic line-search.

  * •

Line search on subset of batches is proposed to improve efficiency and
generalization capability.

##  2 Related Work

In modern machine learning, SGD and its variants [7, 12, 16, 24], are the most
prevalent optimization methods used to train large-scale neural networks
despite their simplicity. Theoretically, the performance of SGD with a
constant learning rate has been analyzed that under some assumptions, such as,
the Strong Growth Condition, has linear convergence if the learning rate is
small enough [26]. However, there are still two non-negligible limitations:
(1) the theoretical convergence rate is not as satisfying as the full gradient
descent method; and (2) the performance is very sensitive to hyperparameters,
especially the learning rate.

In SGD, random sampling introduces randomness and degrades the theoretical
convergence of SGD. One direction of research focuses on reducing the variance
of random sampling so that larger learning rates are allowed. Based on this
idea, a family of SGD variants called Variance Reduction methods were proposed
[25, 11, 27, 6]. These methods have been proved to possess impressive
theoretical convergence on both convex and non-convex optimization problems.
However, the applications of these methods in machine learning are limited to
optimization of simple logistic regression models because the additional
memory and computational cost of these methods are unacceptable when training
large-scale deep neural networks.

On the other hand, significant efforts seek to ameliorate SGD by designing an
effective method to adaptively tune the learning rate as training. [21]
divides the optimization process of SGD into transient phase and stationary
phase, while the learning rate will be reduced after entering the stationary
phase. In [29], scheduling learning rate cyclically for each iteration
empirically boosts the convergence of SGD. Meta-learning [2] techniques that
optimize the learning rate according to the hyper-gradient also have an
outstanding performance in experiments. However, none of these techniques
gives solid theoretical proof of fast convergence. In pursuit of an optimal
learning rate schedule with a theoretical convergence guarantee, Vaswani et
al.hybrid SGD with classic optimization techniques like Armijo line-search
[31] and Polyak step size [15, 3], and the developed methods not only have
sound theory foundation but also empirically show much faster convergence than
many widely-used optimization methods. Our work is inspired by Quadratic
Interpolation (QI) technique used in line-search [30]. We developed two
versions of quadratic line-search method to assist SGD to automatically
determine an ideal learning rate schedule.

To accelerate convergence of SGD, heavy ball algorithm (SGD with fixed
momentum) [23] is proposed . In practical neural network training, Nesterov
accelerated momentum [19] is often used. As discussed in [4], heavy ball
method is highly related to conjugate gradient (CG) method, and many efforts
have been put into improving SGD by injecting the spirit of CG method. CG
method presents promising performance in training an auto-encoder comparing
with SGD and LBFGS [14]. A CG-based ADAM [13] was proposed and showed
comparative performance to Adam, AdaGrad in solving text and image
classification problems, but no generalization performance was provided. Jin
et al. [10] proposed a stochastic CG with Wolfe line-search that was tested on
low-dimensional datasets. All of these related works motivate us to build a
bridge between SGD and CG method and create a dynamic acceleration scheme for
SGD. Instead of tuning momentum term manually, we dynamically adjust it during
training according to Polak-Ribiere formula [22] which shows the best based on
our experiments.

Most recently, [18] independently proposed a line search method based on
parabolic approximation (PAL). Our major differences are: 1). CGQ introduces
feedback to check the estimation quality that improves reliability. 2). CGQ
eliminates the momentum hyperparameter, as discussed in Section 3.3; 3). For
stochastic settings, CGQ does line search on a subset of batches to improve
the runtime and generalization; 4). Convergence analysis of CGQ is conducted
with no quadratic objective assumption.

##  3 The CGQ Method

###  3.1 Overview

The Conjugate Gradient with Quadratic Line-Search (CGQ) method follows the
iterative training framework in Algorithm 1, and eliminates hand-tuning two
major hyperparameters in SGD: the learning rate and momentum factor.

In each training iteration, CGQ optimizes the learning rate $\alpha_{t}$ by
estimating the one-dimensional loss landscape as a parabola. Two variants, the
2-point quadratic interpolation and the Least Squares estimation, are
illustrated in Section 3.2. For direction $p_{t}$, methods from Conjugate
Gradient are adopted to dynamically adjust the momentum factor in SGD instead
of using a hand-tuned fixed momentum, which is described in section 3.3.

###  3.2 Dynamically Adjusting the Learning Rate

####  3.2.1 Two-point Quadratic Interpolation

(a)

(b)

Figure 1: Comparison of the two quadratic estimation algorithms (2-point
method and least squares method). (a) When the landscape is smooth, both
estimate well. (b) When the landscape is rough, least squares tend to capture
more details. It is also observed in experiments that roughness only occurs at
the beginning of a training, after which the landscapes are almost always like
(a).

Given a twice differenciable loss function $f(\theta)$
($\theta\in\mathds{R}^{m}$ is the parameter vector), its second order Taylor
expansion around $\theta_{0}$ is expressed as:

$\begin{split}f(\theta)&=f(\theta_{0})+\nabla
f(\theta_{0})^{T}(\theta-\theta_{0})\\\ &\leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\
+\frac{1}{2}(\theta-\theta_{0})^{T}\nabla^{2}f(\theta_{0})(\theta-\theta_{0})\\\
&\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\
\leavevmode\nobreak\ +O((\theta-\theta_{0})^{2}).\end{split}$ (2)

Along direction $p_{t}\in\mathds{R}^{m}$, the 1D loss landscape $\phi(\alpha)$
w.r.t. the step size $\alpha\in\mathds{R}$ can be described as:

$\begin{split}\phi(\alpha)&=f(\theta-\alpha p_{t})\\\
&=f(\theta_{0})-\alpha\cdot\leavevmode\nobreak\ \nabla
f(\theta_{0})^{T}p_{t}\\\ &\leavevmode\nobreak\ \leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\ +\frac{1}{2}\alpha^{2}\cdot
p_{t}^{T}\nabla^{2}f(\theta_{0})p_{t}+O(\alpha^{2}).\\\ \end{split}$ (3)

When the residual term $O(\alpha^{2})$ is small enough, the line-search
function can be expressed as:

$\phi(\alpha)\approx q(\alpha)\triangleq\leavevmode\nobreak\
A\alpha^{2}+B\alpha+C$ (4)

where $A=\frac{1}{2}p_{t}^{T}\nabla^{2}f(\theta_{0})p_{t}$, $B=-\nabla
f(\theta_{0})^{T}p_{t}$, and $C=f(\theta_{0})$. When $A>0$, $q(\alpha)$
reaches minimum at $\alpha^{*}=-B\leavevmode\nobreak\ /\leavevmode\nobreak\
(2A)$.

To get $A,B,C$ at each training iteration without evaluating the Hessian
matrix $\nabla^{2}f(\theta_{0})$, we propose to apply the 2-point quadratic
interpolation method. Given two sample points $(0,q(0))$,
$(\alpha_{1},q(\alpha_{1}))$ and the slope $q^{\prime}(0)$, parabola
parameters can be calculated as:

$\begin{split}A&=\frac{q(\alpha_{1})-q(0)-q^{\prime}(0)\cdot\alpha_{1}}{\alpha_{1}^{2}}\\\
B&=q^{\prime}(0)\\\ C&=q(0)\\\ \end{split}$ (5)

The stopping criteria and details are shown in Algorithm 2.

Algorithm 2 Quadratic Line-Search (2-point interpolation)

1:procedure QLS2($f(\cdot),\theta_{t},p_{t},slope$)

2: Initialize: $\gamma,K_{max},\alpha_{max},\sigma$

3: $(x_{0},y_{0})\leftarrow(0,f(\theta_{t}))$

4: for $k\leftarrow 1$ to $K_{max}$ do

5: $(x_{1},y_{1})\leftarrow(\gamma,f(\theta_{t}+\gamma\cdot p_{t}))$

6:
$q_{A,B,C}(\cdot)\leftarrow\texttt{Interpolate}(x_{0},y_{0},x_{1},y_{1},slope)$

7: $\alpha^{*}=\min(-B\leavevmode\nobreak\ /\leavevmode\nobreak\
2A,\alpha_{max})$

8: if $A<=0$ then

9: return $\gamma$

10: else if $q_{ABC}(\alpha^{*})>0\leavevmode\nobreak\ \&\leavevmode\nobreak\
f(\theta_{t}+\alpha^{*}p_{t})<f(\theta_{t})$ then

11: return $\alpha^{*}$

12: else

13: $\gamma\leftarrow\alpha^{*}$

14: return $\alpha^{*}$

One of the major advantages of our method is that when moving closer to a
local optimum as training proceeds, $\alpha^{*}$ will decay automatically, due
to the nature of a convex parabola. This is also verified and visualized in
our experiment section. Another issue is when the initial $\gamma$ is too
large compared to the estimated $\alpha^{*}$. In such cases, $q(\alpha)$ is
not a good surrogate to represent the local landscape of $\phi(\alpha)$ near
the origin. This is handled by decaying the initial value of $\gamma$ when it
is consistently larger than the recent average of $\alpha^{*}$. Through
experiments, we observe that $A$ is usually estimated positive in training
neural networks, while the $A<0$ case is also properly dealt with in Algorithm
2.

####  3.2.2 Convergence Analysis

Our method is majorly motivated by empirical results. To give some insights on
the theoretical side, we provide a convergence analysis of quadratic line-
search under SGD update rule. To start with, we use the following assumptions
on the loss function $f(\theta)=\sum_{i}f_{i}(\theta)$, which are commonly
used for convergence analysis in deep learning [31, 15]. (a) $f$ is lower-
bounded by a finite value $f^{*}$; (b) interpolation: This requires that
$\nabla f(\theta^{*})=0$ implies $\nabla f_{i}(\theta^{*})=0$ for all
components $f_{i}$; (c) $f$ is $L$-smooth, i.e., the gradient $\nabla f$ is
$L$-Lipschitz continuous; and (d) $f$ is $\mu$ strong-convex. Proofs are given
in the appendix.

The following lemma provides the bound of the step size returned by quadratic
line-search at each iteration $k$.

###### Lemma 1

Assuming (1) $f_{i}$’s are $L_{i}$-smooth, (2) $\mu_{i}$ strong-convexity of
$f_{i}$, the step size $\alpha_{k}$ returned by quadratic line-search in
algorithm 2 that is constrained within $(0,\alpha_{max}]$ satisfies:

$\alpha_{k}\in[\min\\{\frac{1}{L_{ik}}\text{, }\alpha_{max}\\}\text{,
}\min\\{\frac{1}{\mu_{ik}}\text{, }\alpha_{max}\\}].$ (6)

Next, we provide the convergence of quadratic line-search under strong-convex
and smooth assumptions.

###### Theorem 3.1

Assuming (1) interpolation, (2) $f_{i}$’s are $L_{i}$-smooth, and (3)
$\mu_{i}$ strong-convexity of $f_{i}$’s. When setting
$\alpha_{max}\in(0,\leavevmode\nobreak\
\min_{i}\\{\frac{1}{L_{i}}+\frac{\mu_{i}}{L_{i}^{2}}\\})$, the parameters of a
neural network trained with quadratic line-search satisfies:

$\mathbb{E}[\left\lVert\theta_{k}-\theta^{*}\right\rVert^{2}]\leq
M^{k}\left\lVert\theta_{0}-\theta^{*}\right\rVert^{2}$ (7)

where $M=\mathbb{E}_{i}[\max\\{g_{i}(\alpha_{max})\text{,
}1-\frac{\mu_{i}}{L_{i}}\\}]\in(0,1)$ is the convergence rate, and
$g_{i}(\alpha_{k})=1-(\mu_{i}+L_{i})\alpha_{k}+L_{i}^{2}\alpha_{k}^{2}$.

####  3.2.3 Least Squares Estimation

Notice that theorem 3.1 assumes the loss landscape to be strongly convex and
$L$-smooth. In reality, we find that such conditions are satisfied easier in
the later than earlier phase of training. In other words, the loss landscape
is more rough at the beginning, and the 2-point method sometimes counters
difficulties getting a good estimation $q(\alpha)$. For instance, in rough
surfaces like fig. 1(b), the 2-point method may focus too much on the local
information at the origin (both zero’s order and first order moments.), which
may lead to poor estimation. This motivates us to further develop a more
robust version of the quadratic line-search using least squares estimation.

Instead of only using two samples (plus gradient), the Least Squares estimator
takes more samples. And instead of finding a perfect match on the samples,
Least Squares minimizes the mean squared error among all samples. In our case,
Least Squares minimizes the following objective:

$\min_{\omega}\left\lVert
X\cdot\omega-Y\right\rVert_{2}^{2}=\frac{1}{n}\sum_{i}(A\alpha_{i}^{2}+B\alpha_{i}+C-\phi(\alpha_{i}))^{2}$
(8)

where

$X=\begin{bmatrix}\alpha_{1}^{2}&\alpha_{1}&1\\\
\alpha_{2}^{2}&\alpha_{2}&1\\\ &...\end{bmatrix}\in\mathds{R}^{n\times 3}$ (9)

is the sample matrix formulated by each step $\alpha_{i}$,
$\omega=[A,B,C]^{T}$ are the coefficients to be optimized, and
$Y=[\phi(\alpha_{1}),\phi(\alpha_{2}),...]^{T}\in\mathds{R}^{n}$ is the
observation vector. The Least Squares estimator is:
$\omega=(X^{T}X)^{-1}X^{T}Y$, where $(X^{T}X)^{-1}X^{T}$ is known as the
Moore-Penrose inverse.

Algorithm 3 Quadratic line-search (Least Squares)

1:procedure
QLS2($f(\cdot),\vec{\theta}_{t},\vec{p_{t}},K_{max},\alpha_{max},slope$)

2: Initialize: $\gamma$(sample distance); $\alpha_{0}$;

3: $X=\\{(0,1,0),(0,0,1)\\}$; $Y=\\{slope,f(\theta_{t})\\}$

4: for $k\leftarrow 1$ to $K_{max}$ do

5: $x_{k}\leftarrow(\gamma^{2},\gamma,1)$

6: $X\leftarrow X\bigcup\\{x_{k}\\}$

7: $Y\leftarrow Y\bigcup\\{f(\theta_{t}+\gamma\cdot\leavevmode\nobreak\
p_{t})\\}$

8: $q_{A,B,C}(\cdot)\leftarrow\texttt{LeastSquare}(X,Y)$

9: if $A<=0$ then

10: return $\gamma$

11: else if $q_{ABC}(\alpha^{*})>0\leavevmode\nobreak\ \&\leavevmode\nobreak\
f(\theta_{t}+\alpha^{*}p_{t})<f(\theta_{t})$ then

12: return $\alpha^{*}$

13: else

14: $\gamma\leftarrow\alpha^{*}$

15: return $\alpha^{*}$

The algorithm is described in Algorithm 3. When the loss landscape
$\phi(\alpha)$ is rough, as shown in Figure 1(b), the slope of one point is
not useful in order to capture the landscape trend. By using least squares
estimation on multiple sample points without using slope, the quadratic
interpolation becomes insensitive to gradient information near the origin, and
seeks to accommodate more global information. Through experiments in Section
4, we show that using least squares is more robust.

###  3.3 Dynamically Adjusting Momentum

In this section, we propose an automatic computation of the momentum factor.
In SGD or its variants, this is manually set and fixed in value. We build a
bridge between the momentum in SGD and the conjugate gradient parameter
$\beta_{t}$. (Recall that in the Nonlinear Conjugate Gradient (CG) method, the
conjugate direction $s_{t}$ at iteration $t$ is updated as $s_{t}=-\nabla
f+\beta_{t}\cdot s_{t-1}$, and parameters are updated by
$x_{t+1}=x_{t}+\alpha_{t}\cdot s_{t}$.) $\beta_{t}$ has a deterministic
optimal value (s.t. $s_{t}$ is conjugate to $s_{t-1}$) using the Hessian
matrix in a quadratic objective. But for a general nonlinear case, various
heuristic formulas are proposed [9, 8, 22, 5, 28]. $\alpha_{t}$ can be
determined using either the Newton–Raphson method or a line-search such as the
Armijo rule [1] and Wolfe conditions [32, 33]. The Nonlinear CG method
generalizes various gradient-based optimization methods. For instance, when
$\beta_{t}=0$ and $\alpha_{t}$ is constant, CG degenerates as the Gradient
Descent,and when $\beta_{t}$ is a non-zero constant, it instantiates the
Gradient Descent with momentum, which is also discussed in [4]. Moreover, the
CG framework can also be used in the stochastic settings, where gradients
$\nabla_{\theta}c(\theta)$ are replaced by partial gradients
$\nabla_{\theta}\leavevmode\nobreak\ c_{batch}(\theta)$ in each iteration.

With our proposed Quadratic line-search, the local solver is able to
adaptively change the learning rate for performance. We next show that using
the CG framework, the momentum term in SGD ($\beta_{t}$ in CG), which is
usually constant in a SGD solver, can also be adaptively adjusted.

For a quadratic objective of the following form:

$\min f(\theta)=\frac{1}{2}\theta^{T}Q\theta+b^{T}\theta$ (10)

where $\theta\in\mathds{R}^{n}$ is the model parameters, and
$Q\in\mathds{R}^{n\times n}$ and $b\in\mathds{R}^{n}$ are problem-dependent
matrices. The optimal $\beta_{t}$ can be chosen as

$\beta_{t}=\frac{\nabla f(\theta_{t})^{T}\cdot Q\cdot
p_{t-1}}{p_{t-1}^{T}Qp_{t-1}}$ (11)

where $p_{t-1}$ is the update direction at time $t-1$. Then the next update
direction $p_{t}=-\nabla f(\theta_{t})+\beta_{t}\cdot p_{t-1}$ is guaranteed
to be $Q$-conjugate to $p_{t-1}$, which accelerates training by eliminating
influences on previous progress. For general nonlinear objective functions
with no explicit $Q$ and the Hessian is difficult to compute, different
heuristics were proposed in [9, 8, 22, 5, 28]. Since the classic Gradient
Descent with momentum method is a stationary version of the Conjugate Gradient
[4], a dynamic choice of $\beta_{t}$, instead of fixed, can potentially speed
up the training process. It is also notable that to compute $\beta_{t}$ using
the aforementioned formula, no Hessian is required. For example, the Polak-
Ribiere [22] formula:

$\beta^{PR}_{t}=\frac{\nabla f(\theta_{t})^{T}(\nabla f(\theta_{t})-\nabla
f(\theta_{t-1}))}{\nabla f(\theta_{t-1})^{T}\nabla f(\theta_{t-1})}$ (12)

only requires the recent two gradients.

A bounded version of the Polak-Ribiere momentum is used, i.e., to compute the
momentum as
$\beta^{\prime}_{t}=\min\\{\max\\{0,\beta_{t}^{PR}\\},\beta_{max}\\}$. Through
experiments, it is shown that by applying such dynamic momentum
$\beta^{\prime}_{t}$, the performance is comparable to or even better than the
best manual momentum. The major advantage is that another crucial
hyperparameter is eliminated from hand tuning.

###  3.4 Optimizations for Large Datasets

Figure 2: Running time comparison. sCGQ: stochastic line search with
probability 0.1. fix_m: fixed momentum; dyn_m: Polak-Ribiere momentum. 2-point
interpolation and Least Squares takes similar time.

Conceptually, CGQ performs a line search at every iteration. This
unsurprisingly introduces more computation. To alleviate this issue on large
systems, it is reasonable to use partially observed curvature (from batches)
to estimate that of the whole dataset. This motivates us to implement the sCGQ
that performs line search stochastically: For each training batch, the
optimizer will perform a line search (that optimizes learning rate) with a
certain probability $p$. For batches without line searches, the moving average
of past line search results will be applied. fig. 2 show that this makes the
run time of sCGQ with fixed momentum comparative to naive SGD, and sCGQ with
dynamic momentum comparative to other line serach methods. Experiments in
section 4.1.2 and the runtime comparison in fig. 2 show that sCGQ benefits
from both fast execution and good estimation of 1-D loss landscape (in terms
of learning rate).

##  4 Experiments

To empirically evaluate the performance of the proposed method, we design a
thorough experiment consisting of two parts described in the following
sections. We focus on comparing our method with different configurations in
Section 4.1. In Section 4.2, we benchmark our method with six popular
optimization methods: SGD+momentum, Adam [12], COCOB [20], Lookahead optimizer
[34], SGD+Armijo Line-Search [31], and PAL [18]. We manually tune their
hyperparameters and the results summarized in the following sections are their
best performances in the experiments to our best knowledge. Specifically for
SGD and ADAM, the learning rates are fixed (SGD: 0.02 for VGG, 0.01 on MNIST,
and 0.1 elsewhere; Adam: 0.001). There is no learning rate for COCOB. For our
own method, we tested four varients: CGQ(2pt)/CGQ(LS): complete line search
with quadratic interpolation or least squares estimation; sCGQ(2pt)/sCGQ(LS):
stochastic line search with quadratic interpolation/least squares estimation.
We showcase the fast convergence of our method on training multi-class
classifiers for widely used benchmark image datasets including MNIST, SVHN,
CIFAR-10, and CIFAR-100. For fairness, model parameters are initialized the
same among all solvers for each test case, and we run multiple times and show
the average performance for each case. All experiments are conducted on a
NVIDIA GTX-1080Ti GPU, and the program is developped under PyTorch 1.7.0
framework. Code is available in https://github.com/zh272/CGQ-solver.

Method | Train Loss / Test Accuracy
---|---
MNIST | SVHN
MLP | VGG-16 | ResNet-110 | DenseNet-100
SGD | .076 / 97.24 | .080 / 93.16 | .217 / 92.68 | .158 / 94.34
Adam [12] | .014 / 98.02 | .026 / 93.14 | .134 / 95.34 | .093 / 95.35
Lookahead [34] | .131 / 95.99 | .039 / 93.95 | .157 / 93.14 | .112 / 94.12
COCOB [20] | .267 / 92.72 | .036 / 91.81 | .069 / 90.34 | .035 / 91.30
Armijo [31] | .023 / 97.85 | .008 / 92.42 | – | –
PAL [18] | .021 / 97.60 | – | .016 / 94.73 | .017 / 95.10
CGQ(2pt) | .007/98.24 | .011 / 94.70 |  .009/ 95.31 | .004/96.11
CGQ(LS) | .007/98.24 |  .006/ 94.79 | .012 / 95.81 | .008 / 95.99
sCGQ(2pt) |  .008 /98.24 | .020 / 93.89 |  .043 / 95.84 | .029 / 94.23
sCGQ(LS) | .008 / 98.12 |  .016 / 94.93 | .039 / 95.63 | .032 / 95.43
Table 1: Performance comparison of different optimization algorithms on the
MNIST and SVHN datasets. Training budget: 20 epochs (MNIST) / 40 epochs
(SVHN). “-” means not converging.

###  4.1 Ablation Test

####  4.1.1 Dynamical Momentum

Applying SGD with momentum to reach faster convergence has become common in
deep neural network training. Heavy ball [23] and Nesterov momentum [19] are
two classic static acceleration scheme that consider the momentum term as a
fixed hyperparameter. In this section, we compare the performances of SGD with
quadratic line-search plus (1) Heavy ball momentum[23] with coefficients from
0.1 to 0.9, (2) Fletcher-Reeves momentum [8], (3) Polak-Ribiere momentum [22],
(4) Hestenes-Stiefel momentum [9], and (5) Dai-Yuan momentum [5].

(a)

(b)

(c)

Figure 3: Training loss (One-layer MLP trained on MNIST dataset) w.r.t.: (a)
Different momentum settings. (b) Line search probability. “0”: naive SGD with
no line search; “1”: line search on every mini-batch. (c) $\alpha_{max}$ and
$\beta_{max}$.

The benchmark experiment is run on the MNIST dataset. The comparison results
are summarized in Figure 3(a). On the basis of the results of the ablation
test, it is promising to use Polak-Ribiere momentum as a dynamic acceleration
scheme for our method so that the momentum term is no longer a hyperparameter
that needs tuning manually.

####  4.1.2 Line Search Probability

To evaluate the impact of the randomness on line search, we test our method
with probability thresholds $p$ varying from 0(no line search) to 1(complete
line search). Results in fig. 3(b) show that line search improves the final
model performance even with small probability $p$, and it does not yield
further improvements as $p$ increases.

####  4.1.3 $\alpha_{max}$ and $\beta_{max}$

$\alpha_{max}$ (upper bound of learning rate) and $\beta_{max}$ (upper bound
of momentum), are two important hyperparameters in CGQ. In this section we
perform an ablation test to show the feasible regions. We perform a grid
search on the 2D space by $\alpha_{max}$ (0.01 to 1.1) and $\beta_{max}$ (0 to
1). It is observed from Figure 3(c) that the model performs well when
$\alpha_{max}$ and $\beta_{max}$ are not both large. And the model will
diverge when $\beta_{max}$ is larger than 1 (thus not shown on figure).
Empirically good choices are: $\beta_{max}=0.8$; $\alpha_{max}=0.3$ (ResNet) /
0.1 (DenseNet) / 0.05 (VGG).

Method | Train Loss / Test Accuracy
---|---
VGG-16 | ResNet-164 | DenseNet-100
CIFAR-10 | CIFAR-100 | CIFAR-10 | CIFAR-100 | CIFAR-10 | CIFAR-100
SGD |  .116 / 89.64 | .685 / 64.37 | .271 / 86.06 | .860 / 63.55 | .255 / 86.71 | .901 / 62.06
Adam | .044 / 88.06 | .408 / 55.21 | .144 / 90.06 | .419 / 68.07 | .127 / 90.42 | .364 / 68.23
Lookahead | .069 / 89.21 | .369 / 60.97 | .155 / 88.33 | .526 / 59.21 | .153 / 88.50 | .594 / 61.62
COCOB | .087 / 84.78 | .772 / 44.83 | .330 / 78.96 | 1.450/ 48.00 | .205 / 80.63 | 1.103/ 51.56
Armijo | .298 / 81.18 | .025 / 62.50 |  .003/ 93.64 |  .004/ 73.26 | .009 / 92.72 | .086 / 65.55
PAL | – | – | .017 / 92.04 | .452 / 64.01 | .023 / 90.99 | .072 / 65.90
CGQ(2pt) |  .008/ 88.21 | .025 / 65.46 | .007 / 93.04 | .014 / 75.26 |  .009/ 91.49 |  .020/ 75.01
CGQ(LS) | .010 / 88.64 | .019/66.49 | .010 / 93.52 | .019 / 73.43 | .010 / 92.23 | .029 / 74.35
sCGQ(2pt) | .035 /89.16 | .038 / 63.12 |  .017 /93.75 |  .016 /75.94 |  .020 / 93.71 |  .035 /76.15
sCGQ(LS) | .041 / 88.88 | .044 / 62.93 | .016 / 93.56 | .016 / 75.60 | .028 /93.13 | .032 / 76.04
Table 2: Performance comparison for the CIFAR-10 and CIFAR-100 datasets.
Training budget: 200 epochs. “-” means the method does not converge.

###  4.2 Multi-class classification on image datasets

Recent evidence has shown that deep convolutional neural networks lead to
impressive breakthroughs in multi-class image classification. We mainly select
the three popular families of architectures of deep convolutional neural
networks: ResNet, DenseNet and VGG to benchmark our method with the other
powerful optimizers. The comparison results will be presented in the following
two sections, according to different datasets used in experiments.

####  4.2.1 MNIST and SVHN

The MNIST dataset contains 60,000 hand-written digit images for training, and
10,000 images for testing. Images in MNIST are grayscale images with size
$28\times 28$. SVHN dataset consists of 73,257 images in the training set,
26032 images in the testing set and another 531,131 additional images provided
for training. All the images in SVHN are colored digit images with a size of
$32\times 32$. For MNIST, we construct a fully connected network with one
hidden layer of 1000 neurons as a toy benchmark example. For SVHN, we
experiment with three more expressive architectures: VGG-16, ResNet-110, and
DenseNet-100.

The main results are summarized in Table 1. We observed that with the same
training budget, CGQ presents better performances in terms of both training
loss and test accuracy. The stochastic version of We omit SGD with Armijo-
line-search in the SVHN experiment since it appeared to be unstable and easy
to diverge.

####  4.2.2 CIFAR-10 and CIFAR-100

(a)

(b)

(c)

(d)

Figure 4: Dynamics of training a ResNet-164 with various optimization methods
on the CIFAR100. (a) Training Loss. (b) Test Accuracy. (c) Learning Rate. (d)
Momentum

The CIFAR-10 and CIFAR-100 datasets both consist of 50,000 images for training
and 10,000 images for testing, while images are in 10 and 100 different
classes, respectively. We run the experiments using VGG-16, ResNet-164, and
DenseNet-100. Each of model-algorithm pair is trained for 200 epochs with same
initialization points.

In Table 2, we record the training loss and test accuracy of each model and
optimizer pair in the CIFAR-10 and CIFAR-100 experiments. In most of the
cases, CGQ and sCGQ not only reach the best performance among all line search
methods, but also converges to a local minimum faster. In the the more
challenging CIFAR-100 case, the margin is more significant. fig. 4 shows the
dynamics for training ResNet on the CIFAR100 dataset. Each line plots the mean
over multiple runs, while corresponding standard deviations are shown as
shaded area around the mean. (a) and (b) tells that CGQ variants performs
stably better in training loss and testing accuracy compared to other compared
methods. Moreover, (c) and (d) demonstrates the automatic decaying effect on
learning rate and momentum of the CGQ method.

It is observed from the tables that CGQ performs better than sCGQ on the
training set, but sCGQ has better test time performance. In this sense, the
sCGQ is preferable to CGQ for larger datasets, not only because of faster
execution time, but also because of better generalization. It is also observed
that 2-point interpolation method performs better than its Lease Squares
variant on larger datasets (CIFAR). This could due to an underestimation of
the smoothness for deep neural networks. This is also validated through a
visualization of the loss landscape along the line search direction in fig. 1.

##  5 Conclusion

In this paper, we propose CGQ method to replace hand tuning two crucial
hyperparameters in SGD, i.e., learning rate and momentum. The learning rate is
determined by our quadratic line-search method, and momentum is adaptively
computed using the bounded Polak-Ribiere formula.

Experiments on modern convolutional neural networks show that CGQ converges
faster due to its adaptiveness to the loss landscape. Theoretical results are
also provided towards the convergence of quadratic line search under convex
and smooth assumptions. To improve efficiency on larger models, we further
propose the sCGQ method that performs line search on fewer iterations. This
modification improves not only the run time, but also the generalization
capability in terms of test accuracy. In most cases, the CGQ method
outperforms other local methods.

## References

  * [1] Armijo, L.: Minimization of functions having lipschitz continuous first partial derivatives. Pacific J. Math. 16(1), 1–3 (1966)
  * [2] Baydin, A.G., Cornish, R., Rubio, D.M., Schmidt, M., Wood, F.: Online learning rate adaptation with hypergradient descent. ICLR (2018)
  * [3] Berrada, L., Zisserman, A., Kumar, M.P.: Training Neural Networks for and by Interpolation. International Conference on Machine Learning (2020)
  * [4] Bhaya, A., Kaszkurewicz, E.: Steepest descent with momentum for quadratic functions is a version of the conjugate gradient method. Neural Networks 17, 65–71 (2004). https://doi.org/10.1016/S0893-6080(03)00170-9
  * [5] Dai, Y.H., Yuan, Y.X.: A nonlinear conjugate gradient method with a strong global convergence property. SIAM J. Optimization 10(1), 177–182 (1999)
  * [6] Defazio, A., Bach, F., Lacoste-Julien, S.: Saga: A fast incremental gradient method with support for non-strongly convex composite objectives. In: Ghahramani, Z., Welling, M., Cortes, C., Lawrence, N., Weinberger, K.Q. (eds.) Advances in Neural Information Processing Systems. vol. 27, pp. 1646–1654. Curran Associates, Inc. (2014)
  * [7] Duchi, J., Hazan, E., Singer, Y.: Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research 12(61), 2121–2159 (2011)
  * [8] Fletcher, R., Reeves, C.M.: Function minimization by conjugate gradients. The Computer Journal 7(2), 149–154 (01 1964). https://doi.org/10.1093/comjnl/7.2.149
  * [9] Hestenes, M.R., Stiefel, E.: Methods of conjugate gradients for solving linear systems. Journal of research of the National Bureau of Standards 49, 409–436 (1952)
  * [10] Jin, X.B., Zhang, X.Y., Huang, K., Geng, G.G.: Stochastic Conjugate Gradient Algorithm With Variance Reduction. IEEE Transactions on Neural Networks and Learning Systems 30(5), 1360–1369 (2019). https://doi.org/10.1109/TNNLS.2018.2868835
  * [11] Johnson, R., Zhang, T.: Accelerating stochastic gradient descent using predictive variance reduction. In: Burges, C.J.C., Bottou, L., Welling, M., Ghahramani, Z., Weinberger, K.Q. (eds.) Advances in Neural Information Processing Systems. vol. 26, pp. 315–323. Curran Associates, Inc. (2013)
  * [12] Kingma, D.P., Ba, J.L.: Adam: A method for stochastic optimization. International Conference on Learning Representations (ICLR) (2015)
  * [13] Kobayashi, Y., Iiduka, H.: Conjugate-gradient-based Adam for stochastic optimization and its application to deep learning (2020), http://arxiv.org/abs/2003.00231
  * [14] Le, Q.V., Ngiam, J., Coates, A., Lahiri, A., Prochnow, B., Ng, A.Y.: On optimization methods for deep learning. In: ICML (2011)
  * [15] Loizou, N., Vaswani, S., Laradji, I., Lacoste-Julien, S.: Stochastic polyak step-size for sgd: An adaptive learning rate for fast convergence. In: Proceedings of the 24th International Conference on Artificial Intelligence and Statistics (AISTATS). pp. 1–33 (2021)
  * [16] Loshchilov, I., Hutter, F.: Fixing weight decay regularization in adam. arXiv:1711.05101 (2017)
  * [17] M. J. D. Powell: An efficient method for finding the minimum of a function of several variables without calculating derivatives. Computer Journal 7(2), 155–162 (1964)
  * [18] Mutschler, M., Zell, A.: Parabolic approximation line search for dnns. In: NeurIPS (2020)
  * [19] Nesterov, Y.: A method of solving a convex programming problem with convergence rate $o(1/k^{2})$. Soviet Mathematics Doklady 27, 372–376 (1983)
  * [20] Orabona, F., Pal, D.: Coin betting and parameter-free online learning. In: Lee, D., Sugiyama, M., Luxburg, U., Guyon, I., Garnett, R. (eds.) Advances in Neural Information Processing Systems. vol. 29, pp. 577–585. Curran Associates, Inc. (2016)
  * [21] Pesme, S., Dieuleveut, A., Flammarion, N.: On convergence-diagnostic based step sizes for stochastic gradient descent. In: Proceedings of the International Conference on Machine Learning 1 pre-proceedings (ICML 2020) (2020)
  * [22] Polak, E., Ribiere, G.: Note sur la convergence de méthodes de directions conjuguées. ESAIM: Mathematical Modelling and Numerical Analysis - Modélisation Mathématique et Analyse Numérique 3(R1), 35–43 (1969)
  * [23] Polyak, B.T.: Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics 4(5), 1–17 (1964)
  * [24] Reddi, S.J., Kale, S., Kumar, S.: On the convergence of adam and beyond. In: The 35th International Conference on Machine Learning (ICML) (2018)
  * [25] Schmidt, M., Le Roux, N., Bach, F.: Minimizing finite sums with the stochastic average gradient. Mathematical Programming 162(1), 83–112 (2017)
  * [26] Schmidt, M., Roux, N.L.: Fast convergence of stochastic gradient descent under a strong growth condition (2013)
  * [27] Shalev-Shwartz, S., Zhang, T.: Stochastic dual coordinate ascent methods for regularized loss minimization. Journal of Machine Learning Research 14(1), 567599 (2013)
  * [28] Shewchuk, J.R.: An introduction to the conjugate gradient method without the agonizing pain (August 1994), http://www.cs.cmu.edu/~quake-papers/painless-conjugate-gradient.pdf
  * [29] Smith, L.N.: Cyclical learning rates for training neural networks. In: IEEE Winter Conference on Applications of Computer Vision (WACV) (2017)
  * [30] Vandebogert, K.: Method of quadratic interpolation (September 2017), https://people.math.sc.edu/kellerlv/Quadratic˙Interpolation.pdf
  * [31] Vaswani, S., Mishkin, A., Laradji, I., Schmidt, M., Gidel, G., Lacoste-Julien, S.: Painless stochastic gradient: Interpolation, line-search, and convergence rates. In: Advances in Neural Information Processing Systems. pp. 3727–3740 (2019)
  * [32] Wolfe, P.: Convergence conditions for ascent methods. SIAM Review 11(2), 226–000 (1969)
  * [33] Wolfe, P.: Convergence conditions for ascent methods. ii: Some corrections. SIAM Review 13(2), 185–000 (1969)
  * [34] Zhang, M., Lucas, J., Ba, J., Hinton, G.E.: Lookahead optimizer: k steps forward, 1 step back. In: Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., Garnett, R. (eds.) Advances in Neural Information Processing Systems 32, pp. 9597–9608 (2019)

##  Appendix 0.A Proof of Lemma 1

Suppose $h_{ik}(\alpha)=f_{ik}(\theta_{k}-\alpha\nabla f_{ik}(\theta_{k}))$.
Besides $\alpha=0$, sample the other point $\alpha=\alpha_{1}$ and get the
quadratic interpolant of $f_{ik}$ as follows,

$q(\alpha)=\frac{\left\lVert\nabla
f_{ik}(\theta_{k})\right\rVert^{2}\alpha_{1}+h_{ik}(\alpha_{1})-h_{ik}(0)}{\alpha_{1}^{2}}\alpha^{2}-\left\lVert\nabla
f_{ik}(\theta_{k})\right\rVert^{2}\alpha+h_{ik}(0).$

Thus, the step size $\alpha_{k}$ returned by a quadratic line-search is

$\alpha_{k}=\frac{\left\lVert\nabla
f_{ik}(\theta_{k})\right\rVert^{2}\alpha_{1}^{2}}{2(\left\lVert\nabla
f_{ik}(\theta_{k})\right\rVert^{2}\alpha_{1}+h_{ik}(\alpha_{1})-h_{ik}(0))}.$

According to the assumption on the smoothness of $f_{ik}$ and the stochastic
gradient descent update rule, the following inequality holds:

$h_{ik}(\alpha_{1})\leq h_{ik}(0)+\nabla f_{ik}(\theta_{k})^{\prime}(-\nabla
f_{ik}(\theta_{k})\alpha_{1})+\frac{L_{ik}}{2}\left\lVert-
f_{ik}(\theta_{k})\alpha_{1}\right\rVert^{2}.$

Thus,

$h_{ik}(\alpha_{1})-h_{ik}(0)\leq(-\alpha_{1}+\frac{L_{ik}}{2}\alpha_{1}^{2})\left\lVert\nabla
f_{ik}(\theta_{k})\right\rVert^{2}.$

As a result,

$\alpha_{k}\geq\frac{\left\lVert\nabla
f_{ik}(\theta_{k})\right\rVert^{2}\alpha_{1}^{2}}{2(\left\lVert\nabla
f_{ik}(\theta_{k})\right\rVert^{2}\alpha_{1}+(-\alpha_{1}+\frac{L_{ik}}{2}\alpha_{1}^{2})\left\lVert\nabla
f_{ik}(\theta_{k})\right\rVert^{2})}=\frac{1}{L_{ik}}.$

Similarly, based on the assumption on the strong convexity of $f_{ik}$ and the
stochastic gradient descent update rule,

$h_{ik}(\alpha_{1})-h_{ik}(0)\geq(-\alpha_{1}+\frac{\mu_{ik}}{2}\alpha_{1}^{2})\left\lVert\nabla
f_{ik}(\theta_{k})\right\rVert^{2}.$

Therefore,

$\alpha_{k}\leq\frac{1}{\mu_{ik}}.$

By definition, $\alpha_{k}\in(0,\alpha_{max}]$, and thus, the step size
returned by quadratic line-search satisfies the following relation:

$\alpha_{k}\in[min\\{\frac{1}{L_{ik}}\text{, }\alpha_{max}\\}\text{,
}min\\{\frac{1}{\mu_{ik}}\text{, }\alpha_{max}\\}].$

∎

##  Appendix 0.B Proof of Theorem 1

Suppose the loss function $f(\theta)$ is minimized at $\theta^{*}$. According
to the stochastic gradient descent update rule,

$\left\lVert\theta_{k+1}-\theta^{*}\right\rVert^{2}=\left\lVert\theta_{k}-\alpha_{k}\nabla
f_{ik}(\theta_{k})-\theta^{*}\right\rVert^{2}=\left\lVert\theta_{k}-\theta^{*}\right\rVert^{2}-2\alpha_{k}\langle\nabla
f_{ik}(\theta_{k})\text{,
}\theta_{k}-\theta^{*}\rangle+\alpha_{k}^{2}\left\lVert\nabla
f_{ik}(\theta_{k})\right\rVert^{2}.$

Based on the assumption on the strong convexity of $f_{ik}(\cdot)$,

$-\langle\nabla f_{ik}(\theta_{k})\text{, }\theta_{k}-\theta^{*}\rangle\leq
f_{ik}(\theta^{*})-f_{ik}(\theta_{k})-\frac{\mu_{ik}}{2}\left\lVert\theta_{k}-\theta^{*}\right\rVert^{2}.$

Thus,

$\left\lVert\theta_{k+1}-\theta^{*}\right\rVert^{2}\leq(1-\mu_{ik}\alpha_{k})\left\lVert\theta_{k}-\theta^{*}\right\rVert^{2}+2\alpha_{k}[f_{ik}(\theta^{*})-f_{ik}(\theta_{k})]+\alpha_{k}^{2}\left\lVert\nabla
f_{ik}(\theta_{k})\right\rVert^{2}.$

Based on the assumption on the smoothness of $f_{ik}(\cdot)$,

$f_{ik}(\theta^{*})-f_{ik}(\theta_{k})\leq-\frac{1}{2L_{ik}}\left\lVert\nabla
f_{ik}(\theta_{k})\right\rVert^{2}$ $\left\lVert\nabla
f_{ik}(\theta_{k})-\nabla f_{ik}(\theta^{*})\right\rVert^{2}\leq
L_{ik}^{2}\left\lVert\theta_{k}-\theta^{*}\right\rVert^{2}.$

In the interpolation setting, $\nabla f_{ik}(\theta^{*})=0$; therefore,

$\left\lVert\nabla f_{ik}(\theta_{k})\right\rVert^{2}\leq
L_{ik}^{2}\left\lVert\theta_{k}-\theta^{*}\right\rVert^{2}.$

As a result,

$\left\lVert\theta_{k+1}-\theta^{*}\right\rVert^{2}\leq[1-(\mu_{ik}+L_{ik})\alpha_{k}+L_{ik}^{2}\alpha_{k}^{2}]\left\lVert\theta_{k}-\theta^{*}\right\rVert^{2}$

Let
$g_{ik}(\alpha_{k})=1-(\mu_{ik}+L_{ik})\alpha_{k}+L_{ik}^{2}\alpha_{k}^{2}$.
According to the assumptions, it is trivial to get $L_{ik}>\mu_{ik}>0$, and
thus,

$g_{ik}(\alpha_{k})>1-2L_{ik}\alpha_{k}+L_{ik}^{2}\alpha_{k}^{2}\geq 0.$

For each $f_{i}$, define a constant
$C_{i}=\frac{1}{L_{i}}+\frac{\mu_{i}}{L_{i}^{2}}$. If
$\alpha_{max}\in(0,C_{min})$, $g_{ik}(\alpha_{max})<1$. Based on Lemma 1,

$0\leq g_{ik}(\alpha_{k})\leq max\\{g_{ik}(\alpha_{max})\text{,
}g_{ik}(\frac{1}{L_{ik}})\\}=max\\{g_{ik}(\alpha_{max})\text{,
}1-\frac{\mu_{ik}}{L_{ik}}\\}<1.$

Therefore,

$\left\lVert\theta_{k+1}-\theta^{*}\right\rVert^{2}\leq
max\\{g_{ik}(\alpha_{max})\text{,
}1-\frac{\mu_{ik}}{L_{ik}}\\}\left\lVert\theta_{k}-\theta^{*}\right\rVert^{2}.$

Take the expectation with respect to $i_{k}$:

$\mathbb{E}[\left\lVert\theta_{k+1}-\theta^{*}\right\rVert^{2}]\leq\mathbb{E}_{ik}[max\\{g_{ik}(\alpha_{max})\text{,
}1-\frac{\mu_{ik}}{L_{ik}}\\}]\left\lVert\theta_{k}-\theta^{*}\right\rVert^{2}$

where $\mathbb{E}_{ik}[\max\\{g_{ik}(\alpha_{max})\text{,
}1-\frac{\mu_{ik}}{L_{ik}}\\}]\in(0,\leavevmode\nobreak\ 1)$. ∎