# AdaZeta: Adaptive Zeroth-Order Tensor-Train Adaption for Memory-Efficient
Large Language Models Fine-Tuning

Yifan Yang1, Kai Zhen2, Ershad Banijamal2, Athanasios Mouchtaris2 and Zheng
Zhang1
1University of California, Santa Barbara
2Amazon, Alexa AI
<EMAIL_ADDRESS>{kaizhen,ebanijam<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Fine-tuning large language models (LLMs) has achieved remarkable performance
across various natural language processing tasks, yet it demands more and more
memory as model sizes keep growing. To address this issue, the recently
proposed Memory-efficient Zeroth-order (MeZO) methods attempt to fine-tune
LLMs using only forward passes, thereby avoiding the need for a
backpropagation graph. However, significant performance drops and a high risk
of divergence have limited their widespread adoption. In this paper, we
propose the Adaptive Zeroth-order Tensor-Train Adaption (AdaZeta) framework,
specifically designed to improve the performance and convergence of the ZO
methods. To enhance dimension-dependent ZO estimation accuracy, we introduce a
fast-forward, low-parameter tensorized adapter. To tackle the frequently
observed divergence issue in large-scale ZO fine-tuning tasks, we propose an
adaptive query number schedule that guarantees convergence. Detailed
theoretical analysis and extensive experimental results on Roberta-Large and
Llama-2-7B models substantiate the efficacy of our AdaZeta framework in terms
of accuracy, memory efficiency, and convergence speed.111Code available on
GitHub https://github.com/yifanycc/AdaZeta.

AdaZeta: Adaptive Zeroth-Order Tensor-Train Adaption for Memory-Efficient
Large Language Models Fine-Tuning

Yifan Yang1, Kai Zhen2, Ershad Banijamal2, Athanasios Mouchtaris2 and Zheng
Zhang1 1University of California, Santa Barbara 2Amazon, Alexa AI
<EMAIL_ADDRESS>{kaizhen,ebanijam<EMAIL_ADDRESS><EMAIL_ADDRESS>

Figure 1: The evaluation loss curves for the SST-2, WiC, and CB tasks using
the Llama-2-7B model. The proposed AdaZeta method converges faster and
effectively addresses the divergence problem using a much smaller batch size
(BS). Both MeZO-LoRA and AdaZeta use a learning rate of 1e-4, while Sparse-
MeZO utilizes a 1e-6 learning rate.

##  1 Introduction

Fine-tuning large language models (LLMs) has demonstrated outstanding
performance in addressing numerous natural language processing applications,
such as natural language understanding Kenton and Toutanova (2019); Zhang et
al. (2023), question-answering Xu et al. ; Cheng et al. (2023b), and
summarization Zhang et al. (2024a). However, as the size of LLMs increases,
the training process consumes progressively more GPU memory. In recent years,
approaches such as quantization Dettmers et al. (2024) and parameter-efficient
fine-tuning (PEFT) Hu et al. (2021); Chen et al. (2024) have been proposed to
reduce memory costs during training by storing data with lower bit-depth or
updating only a portion of the parameters. Despite these strategies
effectively reducing memory costs, overall memory usage remains high due to
the continuous reliance on a backpropagation graph.

To further reduce the memory overhead, Malladi et al. (2023) proposed the
Memory-efficient Zeroth-order (MeZO) method for LLM fine-tuning, which shows
over 8$\times$ memory reduction compared with the first-order (FO) fine-tuning
methods like SGD Amari (1993) and AdamW Loshchilov and Hutter (2018). Unlike
FO methods, which calculate gradients via backpropagation, the MeZO method
estimates gradients based on the difference between loss values obtained from
two forward passes, thereby eliminating the need for a backpropagation graph.
However, two main challenges persist in the zeroth-order (ZO) fine-tuning of
LLMs: 1) a significant performance gap between FO and ZO approaches, and 2)
increased risk of divergence, particularly in the ZO fine-tuning of large-
scale LLMs, as observed in recent studies Gautam et al. (2024).

To improve the performance, various FO optimization techniques have been
adapted for ZO fine-tuning scenarios, like the ZO-AdaMU method Jiang et al.
(2024). However, these approaches fail to accommodate the specific needs of ZO
methods, and add significant memory overhead from the optimizer state. Given
the dimensionality-related nature of ZO convergence rates, Liu et al. (2024)
propose the Sparse-MeZO method that generates pruning masks based on the value
of the weight elements. Nevertheless, the Sparse-MeZO method yields
inconsistent performance across various tasks and hyperparameter
configurations. In contrast to this approach, we consider using the PEFT
method to reduce the number of trainable parameters. Although the ZO PEFT
method like MeZO-LoRA has been considered in Malladi et al. (2023), the
improvements are limited as the LoRA adapter fails to offer high
representational ability with an ultra-low rank. To solve this problem, we
involve tensorized adapters, which offer high performance with even lower
trainable parameters than LoRA adapters.

To address the variance-related divergence issue in large-scale ZO fine-
tuning, previous studies Malladi et al. (2023); Jiang et al. (2024) have
primarily focused on adjusting the batch size, as increasing the batch size
can reduce the noise in ZO gradient estimation. However, these approaches
introduce significant runtime overhead and fail to improve performance
significantly. To further reduce variance, Gautam et al. (2024) introduced the
MeZO-SVRG method, adapting the first-order SVRG technique to the ZO context.
Despite its success, MeZO-SVRG suffers from a slow and memory-inefficient
fine-tuning process due to the additional parameter copies and computation
process that even doubles the memory cost of the MeZO methods. In contrast to
these works, we consider reducing the ZO gradient variance with a sublinearly
increasing query222A query refers to request the gradient of the loss function
for one time in this paper Bubeck et al. (2015)[Sec. 4.1.4]. schedule that
achieves not only better accuracy but also faster convergence in terms of both
steps and time.

This paper explores task-specific PEFT training for ZO fine-tuning scenarios.
We introduce the Adaptive Zeroth-order Tensor-Train Adaption (AdaZeta)
framework, which incorporates fast-forward tensorized adapters and an adaptive
query schedule. This combination can significantly enhance the accuracy and
convergence of ZO fine-tuning, as demonstrated in Fig. 1. Our contributions
are summarized as follows:

  * •

We introduce the AdaZeta framework, outperforming other ZO fine-tuning methods
like MeZO, MeZO-LoRA, and Sparse-MeZO across different tasks with faster
convergence.

  * •

A new contraction method is introduced to accelerate the forward pass of
tensorized adapters by over 20%.

  * •

We develop an adaptive query number schedule that sub-linearly increases the
number of queries to address the persistent divergence issue in ZO fine-
tuning.

  * •

We provide both theoretical and experimental results to demonstrate the
training efficiency and performance of our method.

##  2 Background

###  2.1 Parameter-Efficient Fine-tuning

In recent years, various works related to PEFT methods have been proposed.
Beyond the most widely used methods like Adapters Houlsby et al. (2019) and
LoRA Hu et al. (2021), there are also methods exploring ultra-low trainable
parameter solutions Zaken et al. (2022); Li and Liang (2021); Liu et al.
(2022). In Malladi et al. (2023), researchers try to employ the LoRA and
prefix-tuning Li and Liang (2021) methods during the ZO fine-tuning. However,
the improvement is limited and the detailed analysis of ZO PEFT tuning is not
discussed.

In this paper, we explore tensorized adapters, an ultra-low-parameter PEFT
method that compresses the weight matrices of adapter layers using Tensor-
Train (TT) decomposition. This approach is examined in Yang et al. (2024a),
where it demonstrates strong performance in FO fine-tuning tasks. However, the
contraction process of TT format Oseledets (2011); Novikov et al. (2015)
involving a sequence of small tensor factors slows down the forward pass,
making it less suitable for ZO methods that require two forward passes per
step. To solve this problem, we propose parallel contraction methods to
improve the inference speed of tensorized adapter methods.

Figure 2: Illustration for tensorized linear layer and tensorized adapters.

###  2.2 Tensorized Adapters

As shown in Fig. 2 (a), the tensorized adapters, which are built upon
tensorized linear layers, are lightweight components injected during the fine-
tuning process to reduce the number of trainable parameters. The weight in
tensorized linear layers is represented in the TT format. Compared with a
standard weight matrix $\bm{W}\in\mathbb{R}^{m\times n}$ in a typical linear
layer, the TT format represents its reshaped $2o$-way tensor
$\mathcal{W}\in\mathbb{R}^{k_{1}\times\cdots\times k_{2o}}$ as a sequence of
tensor factors
$[\mathcal{G}_{1},\cdots,\mathcal{G}_{o},\mathcal{G}_{o+1},\cdots\mathcal{G}_{2o}]$
Oseledets (2011), where each tensor factor
$\mathcal{G}_{i}\in\mathbb{R}^{r_{i-1}\times k_{i}\times r_{i}}$ has rank
$r_{i-1}$ and $r_{i}$. The dimensions $k_{i}$ are constrainted such that
$\Pi_{i=1}^{o}k_{i}=m$ and $\Pi_{j=o+1}^{2o}k_{j}=n$. During the forward pass,
the sequence of tensor factors is contracted and reshaped back into the shape
of a weight matrix as

$\displaystyle\bm{W}=\text{Reshape}(\mathcal{G}_{1}\times\cdots\times\mathcal{G}_{2o}).$
(1)

Note that in this paper, the tensor rank is held constant, with the exception
of the first and last ranks, which are set $r_{0}=r_{2o}=1$. Also, the weights
in tensorized layers are initialized, stored, and updated in TT-format instead
of the matrix form in a traditional linear layer.

The structure of tensorized adapters is shown in Fig. 2 (b). Each tensorized
adapter contains two tensorized layers and a non-linear layer in between. For
each encoder/decoder block, the tensorized adapters are attached after the
attention and feed-forward layer. Different from Yang et al. (2024a) that
makes both tensorized adapters and layer norm trainable, we freeze the layer
norm during the ZO fine-tuning, as noisy gradient estimation of the scaling
factor in layer normalization can seriously degrade model performance. The
tensorized adapters reduce trainable parameters by over $80\times$, making
them a better fit for ZO fine-tuning.

##  3 Methods

In this section, we first introduce some basic knowledge of the ZO gradient
estimator. Then, we present our AdaZeta method, a powerful framework designed
to improve the performance of ZO LLM fine-tuning with two main components: 1)
the fast-forward tensorized adapters, and 2) an adaptive query number
schedule. Finally, we provide a theoretical analysis of the convergence rate
of the AdaZeta method, demonstrating the improved convergence rate
theoretically.

###  3.1 Zeroth-order Estimation

Traditional ZO estimation has been widely studied in both convex and non-
convex optimization setups Ghadimi and Lan (2013); Malladi et al. (2023); Chen
et al. (2019). In our problem, considering a supervised dataset $\mathcal{D}$,
mini-batch $\mathcal{B}$ with the size of $D$ and $B$ respectively, we set the
loss function for our fine-tuning problem to be $\ell(\bm{w};\mathcal{B})$,
where the trainable parameter in the tensorized adapters
$\bm{w}\in\mathbb{R}^{d}$ has a size of $d$. Then, the Randomized Zeroth-order
Gradient Estimation (RGE) at training step $k$ is given as:

$\nabla\hat{\ell}(\bm{w}_{k})=\sum_{q=1}^{Q_{k}}\frac{\ell_{\mathcal{B}}(\bm{w}_{k}+\epsilon\bm{z}_{q})-\ell_{\mathcal{B}}(\bm{w}_{k}-\epsilon\bm{z}_{q})}{2\epsilon}\bm{z}_{q}$

where $Q_{k}$ is the query number at the training step $k$,
$\bm{z}_{q}\sim\mathcal{N}(0,\mathbb{I}_{d})$ is the vector-wise random
perturbation for each query $q$, and $\epsilon$ is a scaling factor for the
perturbation.

Unlike FO fine-tuning, which relies on backpropagation, RGE requires only two
forward passes with perturbations added to the weights of tensorized adapters,
eliminating the need for a backpropagation graph. Additionally, by sublinearly
increasing the number of queries at the beginning of each epoch, we
effectively reduce the variance of the ZO gradient estimation by involving
distinct perturbations $\bm{z}_{q}$ at each time of query. Details of the
setup will be discussed in the following section.

Input: Parameters $\bm{w}$, loss function $\ell(\cdot)$, random seed $s_{q}$,
scaling factor $\epsilon$, Query-realted constant $\alpha,\beta$, maximum
query $Q_{max}$, learning rate $\eta$.

1: for $k=1,\cdots,K$ do

2: Calculating query number at epoch $e_{k}$ start:  $Q_{k}:=\min(\alpha
e_{k}^{\beta},Q_{max})$

3: for $q=1,\cdots,Q_{k}$ do

4:
$\bm{w}\leftarrow\bm{w}+\epsilon\bm{z}_{q},\quad\bm{z}_{q}\sim\mathcal{N}(0,\mathbb{I}_{d},s_{q})$

5: $\ell_{+}^{q}\leftarrow\ell(\bm{w},\mathcal{B})$

6:
$\bm{w}\leftarrow\bm{w}-2\epsilon\bm{z}_{q},\quad\bm{z}_{q}\sim\mathcal{N}(0,\mathbb{I}_{d},s_{q})$

7: $\ell_{-}^{q}\leftarrow\ell(\bm{w},\mathcal{B})$

8:
$\bm{w}\leftarrow\bm{w}+\epsilon\bm{z}_{q},\quad\bm{z}_{q}\sim\mathcal{N}(0,\mathbb{I}_{d},s_{q})$

9: Reset random seed $s_{q}$ for generating $\bm{z}_{q}$

10: end for

11:
$\nabla_{\bm{w}}\hat{\ell}(\bm{w})=\frac{1}{Q_{k}}\sum_{q=1}^{Q_{k}}\left[\frac{\ell_{+}^{q}-\ell_{-}^{q}}{2\epsilon}\bm{z}_{q}\right]$

12: $\bm{w}\leftarrow\bm{w}-\eta*\nabla_{\bm{w}}\hat{\ell}(\bm{w})$

13: end for

Algorithm 1 AdaZeta Algorithm

###  3.2 The AdaZeta Framework

Previous ZO fine-tuning methods, such as MeZO, typically estimate the gradient
for a large number of trainable parameters simultaneously using RGE. This
approach results in high variance due to the dimension-related nature of the
RGE method. Although techniques like LoRA and prefix tuning have been
considered, few works consider the tasks-specific PEFT adapters for the ZO
LLMs fine-tuning. Additionally, as shown in Fig. 1, we have observed an
increased risk of divergence when using the MeZO-LoRA method during fine-
tuning. To address these issues, we propose our AdaZeta framework to improve
performance and solve the instability problem of the vanilla MeZO method. Our
framework includes the following components:

Fast Forward Tensorized Adapters. The Parameter-efficient issue has been
widely studied in the FO cases, where people often freeze the pre-trained
model parameters and fine-tune the LLMs by adding trainable adapters along
with the frozen pretrain weights. Since the ZO estimation accuracy is
dimension-dependent, reducing dimensionality can significantly help improve
the gradient estimation quality. Thus, we consider injecting the ultra-low
parameter tensorized adapters in our AdaZeta framework to reduce the number of
trainable parameters while retaining the performance.

As we have mentioned, ZO fine-tuning mainly relies on gradient estimation with
two forward passes at each step. Thus, the speed of the forward pass is a
crucial factor for the overall speed of ZO fine-tuning. Instead of using the
sequential contraction method during the forward pass as in previous work, we
propose a new parallel contraction method to speed up the forward passes. This
method divides the sequence of tensor factors into several groups to enable
parallel processing and avoid the presence of high-dimensional tensors. Taking
a bipartite case as an example, the contraction process in eq. (1) is replaced
by:

$\displaystyle\bm{W}=\text{R}(\prod_{i=1}^{o}\mathcal{G}_{i}\prod_{j=o+1}^{2o}\mathcal{G}_{j}),$

where $\mathcal{G}_{i}$ represents the $i$-th tensor factor, $\text{R}(\cdot)$
represents the reshape operation. For larger models, the tensor factors can be
organized into tripartite or quadripartite structures to accelerate the
inference speed of the tensorized methods.

Adaptive Query Adjustment for ZO estimation. As previously noted, the training
process for existing ZO methods often exhibits instability, particularly with
large-size models where divergence issues frequently occur. Previous studies
Chen et al. (2019); Jiang et al. (2024) have explored using a fixed multiple
queries scheme to improve the estimation accuracy in the optimization
community. However, utilizing a fixed number of queries may significantly
hinder the training efficiency of large-scale ZO fine-tuning tasks, as naively
increasing the number of perturbations greatly escalates training durations.
To solve this problem, we consider a simple but effective sublinear increasing
query number adjustment schedule, where the number of queries is updated at
the beginning of each epoch $e_{k}$. By expressing the epoch in terms of the
global training steps as $e_{k}=\lfloor k/\lceil\frac{D}{B}\rceil\rfloor$, we
have:

$Q_{k}:=\min(\alpha e_{k}^{\beta},Q_{max})$ (2)

with a fixed scaling factor $\alpha\in(0,1)$, a sublinear increasing factor
$\beta\in(0,1)$ and a max query threshold $Q_{max}$. Then, the query number is
fixed for all training steps within each epoch. This adjustment solves all
divergence problems we observed with theoretical guarantee and performs even
faster than the traditional way to solve the divergence problem for ZO LLMs
fine-tuning by increasing the batch size.

The corresponding optimization algorithm used in the AdaZeta framework is
shown in Alg. 1. We adjust the query number at the beginning of each epoch.
Different from the MeZO algorithm, we obtain the gradient used for the model
update by taking the average over multiple query results. Note that we fix the
query number to be 1 when fine-tuning medium-size models like Roberta-Large
since the noise of ZO estimation is relatively low when the number of
trainable parameters is small. Later, we will show that a sublinear increasing
query number benefits the convergence of the problem when the model size is
large, both theoretically and experimentally.

Table 1: Comparative analysis of various ZO fine-tuning methods on the
Roberta-Large models.

Methods | RTE | SST-2 | SST-5 | QNLI | MNLI | SNLI | MR
---|---|---|---|---|---|---|---
FT | 66.4 | 91.9 | 47.5 | 63.4 | 70.0 | 77.5 | 88.2
Zero-Shot | 51.4 | 79.0 | 35.5 | 50.9 | 48.8 | 50.2 | 80.2
LP | 59.4 | 76.0 | 40.3 | 57.6 | 56.5 | 66.0 | 86.6
BS=16
MeZO | 52.7 | 90.5 | 31.1 | 59.9 | 60.5 | 63.5 | 85.5
MeZO(LoRA) | 52.7 | 84.2 | 44.8 | 60.3 | 58.5 | 65.6 | 85.7
AdaZeta | 66.8 | 91.4 | 48.3 | 61.3 | 58.1 | 69.1 | 87.0
BS=64
MeZO | 64.0 | 90.5 | 45.5 | 60.5 | 58.7 | 68.5 | 85.0
MeZO(LoRA) | 63.9 | 91.3 | 43.0 | 59.0 | 64.0 | 69.7 | 87.4
AdaZeta | 64.3 | 91.5 | 49.6 | 60.7 | 68.1 | 68.7 | 87.5

###  3.3 Theoretical Analysis

In this subsection, we give the theoretical analysis for the AdaZeta
framework. Our theoretical analysis highlights why the tensorized adapter and
adaptive query schedule can significantly help to improve the ZO convergence
rate. Unlike the theoretical analysis in the MeZO paper, which focuses on the
”effective rank” for the Hessian of loss, we focus on the dimension of the
optimized models $d$ (number of trainable parameters) instead. As the
trainable parameters with PEFT adapters are much smaller than the model size,
the theoretical analysis based on the exact dimension of the optimization
problem can better help us to explore the behavior of different PEFT methods.

To align our analysis with LLM fine-tuning, we consider a non-convex
optimization setup and study the convergence behavior regarding the training
steps $k$. It is important to note that the ZO estimated gradient
$\nabla\hat{\ell}$ by the RGE, is an unbiased estimation of the true gradient
$\nabla\ell$ when $\epsilon\rightarrow 0$, which gives the fact
$\mathbb{E}_{\bm{z}}[\nabla\hat{\ell}]=\nabla\ell$ Nesterov and Spokoiny
(2017). First, we list the following assumptions for our analysis:

A1: The loss function $\ell$ has an L-Lipschitz continuous gradient, where for
$L>0$ we have:

$\|\nabla\ell(\bm{w}_{i})-\nabla\ell(\bm{w}_{j})\|\leq
L\|\bm{w}_{i}-\bm{w}_{j}\|,\forall\bm{w}_{i},\bm{w}_{j}$

A2: At each step $k$, the gradient of loss function $\ell$ is upper bounded as
$\|\nabla\ell\|\leq\delta,\forall k$.

Then, we offer the global convergence rate for our AdaZeta algorithm:

######  Theorem 1.

Under A1 and A2, randomly pick $\bm{w}_{T}$ from history with probability
$P(T=k)=\frac{1}{K}$, the convergence of the AdaZeta algorithm can be bounded
by:

$\displaystyle\mathbb{E}[\|\nabla\ell(\bm{w}_{T})\|^{2}]\leq\mathcal{O}(\frac{R+\epsilon^{2}L+C(d,\epsilon)\sum_{k}\frac{1}{Q_{k}}}{K\epsilon}),$

where $R$ is defined by the distance between the start point and the optimal
solution $\ell(\bm{w}_{1})-\ell^{*}$, the ZO perturbation scaling factor is
represented as $\epsilon$, and $C(d,\epsilon)$ is a constant related to the
model parameter size $d$, which is defined at the end of the proof in Appendix
C.

###### Proof.

Details can be found in Appendix C. ∎

According to Theorem 1, we can observe that the bound is related to the query
schedule. For convenience, take a simplified case with $\alpha=\beta=0.5$ and
ignore the minimum in eq. (2), we have $Q_{k}=\frac{1}{2}\sqrt{\lfloor
k/\lceil\frac{D}{B}\rceil\rfloor}$, gives $\sum_{k=1}^{K}\frac{1}{Q_{k}}\leq
2\left\lceil\frac{D}{B}\right\rceil\sqrt{\left\lfloor\frac{K}{\lceil\frac{D}{B}\rceil}\right\rfloor}$,
which guarantees the true gradient approaches zero when $K\rightarrow\infty$.
In contrast, using a small constant such as $Q=1$ results in an upper bound of
$\mathcal{O}(C(d,\epsilon)/K\epsilon)$, which becomes challenging to minimize
due to the term $C(d,\epsilon)$ is directly proportional to the model size
$d$. Additionally, we observe that the convergence rate is significantly
influenced by the model dimension $d$. Consequently, in this paper, we also
try to reduce the number of trainable parameters with the tensorized adapters.

Table 2: Comparative analysis of various ZO fine-tuning methods on the
Llama-2-7B model.

Methods | RTE | CB | BoolQ | WSC | WIC | SST2 | MultiRC | COPA | ReCoRD | SQuAD
---|---|---|---|---|---|---|---|---|---|---
FT | 61.7 | 66.1 | 84.6 | 63.4 | 65.9 | 94.0 | 45.4 | 86.0 | 81.1 | 90.7
LoRA | 85.5 | 67.8 | 84.8 | 62.5 | 73.9 | 94.8 | 85.0 | 81.0 | 79.4 | 90.5
Zero-Shot | 49.5 | 32.1 | 65.1 | 36.5 | 50.6 | 79.7 | 55.8 | 59.7 | 80.9 | 54.7
ICL | 54.5 | 58.9 | 67.4 | 65.4 | 52.7 | 81.2 | 58.7 | 84.4 | 80.1 | 67.1
MeZO | 54.6 | 73.0 | 68.6 | 52.8 | 57.8 | 85.8 | 62.6 | 86.0 | 70.8 | 72.5
MeZO-LoRA | 59.6 | 74.0 | 71.6 | 53.0 | 55.2 | 86.8 | 67.2 | 89.0 | 72.0 | 80.0
Sparse-MeZO | 58.6 | 76.0 | 67.8 | 53.0 | 56.8 | 85.2 | 61.2 | 86.0 | 70.6 | 64.4
AdaZeta | 74.0 | 75.0 | 79.4 | 52.2 | 58.0 | 91.0 | 68.2 | 94.0 | 71.2 | 80.0

##  4 Experiemnts

In this section, we conduct comprehensive experiments to evaluate the
performance of our proposed AdaZeta framework across several LLMs with
different scales on a variety of natural language understanding and generation
tasks Socher et al. (2013); Williams et al. (2017); Rajpurkar et al. (2016).
We demonstrate that our methods surpass a comprehensive array of memory-
efficient baselines, including inference-only methods such as Zero-shot Brown
et al. (2020), In-Context Learning (ICL), and Linear Probing (LP) Kumar et al.
(2021), as well as ZO fine-tuning methods like MeZO, MeZO-LoRA Malladi et al.
(2023), and Sparse-MeZO Liu et al. (2024). Also, the first-order fine-tuning
(FT) baseline is also provided as a reference.

Initially, we present experimental evidence using Roberta-Large models Liu et
al. (2019), illustrating that the integration of tensorized adapters can
significantly enhance the efficiency of ZO fine-tuning by reducing the number
of trainable parameters. Subsequently, we enabled our proposed adaptive query
schedule method to show the effectiveness of the AdaZeta framework on large-
scale Llama-2-7B models Touvron et al. (2023), which not only enhances
performance but also ensures robust convergence. All experiments are conducted
on NVIDIA Tesla A100-40GB GPUs, with further details about the experimental
setup available in Appendix A.

###  4.1 Medium-size Roberta-Large Models

We initially evaluated the effectiveness of using tensorized adapters on
RoBERTa-large models across various tasks, including single-sentence tasks
like SST-2 and SST-5, natural language inference tasks such as QNLI, MNLI,
SNLI, RTE, and the sentiment analysis dataset Movie Reviews (MR). The results
are summarized in Table 1. Experiments were conducted under a 16-shot setup,
with 16 data samples in each class of the datasets. We monitored the best test
accuracy every 500 steps, using a test pool of 1,000 data samples. Note that,
similar to previous ZO fine-tuning studies, we fixed the number of queries to
1 in this subsection. This decision is based on the observation that gradient
noise is relatively small in medium-sized Bert-based models. The following
conclusions have been reached:

AdaZeta Shows Higher Accuracy than Other ZO Fine-Tuning Methods. According to
our observations in Table 1, AdaZeta outperforms other ZO fine-tuning
approaches in term of evaluation accuracy. Compared with MeZO-LoRA, which also
involves PEFT adapters, AdaZeta outperforms in 6 out of 7 tests under both 16
and 64 batch size (BS) settings. This advantage shows the effectiveness of
improving ZO estimation accuracy by further reducing the number of trainable
parameters with the tensorized adapter. This is supported by the dimension-
related convergence rate proved in Section 3.3.

AdaZeta Demonstrates Improved Convergence. Compared to the MeZO-LoRA method,
the AdaZeta method exhibits superior convergence when the batch size is 16.
Given our 16-shot training setup, it is reasonable to expect that the 16 batch
size scenario would outperform the 64 batch size scenario if the fine-tuning
process converges effectively. However, a performance decline is observed with
the MeZO-LoRA method, indicating that it is adversely affected by ZO gradient
noise. Comparatively, the AdaZeta method achieves consistent results across
both setups by reducing such noise with less trainable parameters, effectively
showcasing its ability to aid in convergence.

###  4.2 Large-scale Llama-2 Models

In the previous section, we demonstrated how utilizing the tensorized adapter
method enhances ZO fine-tuning performance by reducing gradient noise through
a decrease in trainable parameters. In this section, we assess the
effectiveness of the AdaZeta framework with the large-scale Llama-2-7B model.
Differing from the experiments on the Roberta-Large models, we enabled the
adaptive query schedule method proposed in our AdaZeta framework to mitigate
the commonly observed divergence issues in large-scale ZO fine-tuning.

To hilight the challenge of our experiments, we adopt a low-data resource
approach using datasets from SuperGLUE Wang et al. (2019) and generative tasks
such as SQuAD Rajpurkar et al. (2016) and DROP Dua et al. (2019). Our
experimental protocol follows the prompted-based fine-tuning strategy outlined
in the MeZO paper Malladi et al. (2023). The quantitive results are summarized
in Table 2 and the training curves have been shown in Fig. 1. The following
conclusions are drawn:

AdaZeta Method Demonstrates Superior Performance Over Traditional ZO Fine-
Tuning. The AdaZeta framework delivers exceptional accuracy results across a
variety of tasks, outperforming all ZO baseline methods such as MeZO and MeZO-
LoRA in 8 out of 10 tasks. Compared with traditional inference-only methods
like ICL and Zero-shot, AdaZeta significantly surpasses them with respect to
test accuracy. Moreover, the AdaZeta method even outperforms the FO AdamW
methods over several tasks like RTE, CB, and COPA, which require 8$\times$
more GPU memory.

AdaZeta Method Effectively Addresses Divergence Issues in ZO Fine-Tuning. We
can observe from the table that the MeZO and MeZO-LoRA methods achieve
unsatisfied results in some tasks like SST2, RTE, and BoolQ compared with our
proposed method, which is led by the convergence issue. Also, we have shown
that the AdaZeta method achieves lower evaluation loss much faster than the
MeZO-LoRA and Sparse-MeZO methods across all tasks in Fig. 1. For example, the
MeZO-LoRA method requires nearly 6K steps to achieve a loss of 0.4, whereas
the AdaZeta method achieves the same degree of loss minimization in less than
1K steps, which represents a 6$\times$ speed-up with the same 1e-4 learning
rate. Traditional ways to solve such divergence issues through increasing the
batch size are hard to follow in the large-scale LLMs fine-tuning tasks. In
contrast, the adaptive query schedule in the AdaZeta framework successfully
mitigates this issue without increasing the training memory, thereby improving
training outcomes.

Table 3: Required GPU hours (GPU numbers $\times$ Training hours) to achieve
each training loss for different ZO fine-tuning methods on Llama-2-7B model.
AdaZeta${}_{\text{seq}}$ represents the AdaZeta method utilizing traditional
sequential contracted tensorized adapters.

Methods | SST2 | WIC | CB | MultiRC
---|---|---|---|---
MeZO-LoRA(BS=64) | 3.0 | 4.8 | 8,6 | 30.0
MeZO-LoRA(BS=16) | 0.6 | 1.1 | 3.1 | 10.8
Sparse-MeZO | 4.1 | 3.6 | 4.3 | 6.4
AdaZeta${}_{\text{seq}}$ | 1.1 | 1.0 | 0.9 | 12.1
AdaZeta | 0.9 | 0.9 | 0.8 | 10.6

Figure 3: Trade-off between the accuracy and memory cost for different fine-
tuning methods. We can observe that the AdaZeta method achieves the best
accuracy among the memory-efficient methods.

###  4.3 Memory Training Time Efficiency

In this section, we evaluate the memory and time efficiency of the AdaZeta
method. Specifically, we test the peak memory cost of different fine-tuning
methods over the Llama-2-7B model and study the trade-off between memory,
accuracy, and training time. The result is summarized in Fig. 3 and further
discussion about training memory can be referred to Appendix B.1.

According to Fig. 3 (refer to Appendix B.1 for numberial results), the AdaZeta
method requires only 14GB of memory to fine-tune the SST2 tasks on the
Llama-2-7B model, which achieves over 8$\times$ Memory Reduction Relative to
the FT Method. Also, compared with other ZO fine-tuning methods like MeZO,
MeZO-LoRA, and Sparse-MeZO, the AdaZeta method utilizes similar or even less
memory to achieve variance reduction. Traditional ways to reduce the ZO
gradient estimation noise like increasing the batch size, consume
significantly more memory than the AdaZeta method as shown in Fig. 3.

In Table 3, we measure the total GPU hours required to achieve a certain
threshold of training loss across four tasks (SST2, WIC, CB, MultiRC). For the
applicability of the experiments, we established an evaluation loss threshold
that all methods could achieve. According to the results, it is evident that
the AdaZeta method converges on-par or faster than other ZO fine-tuning method
with even better results than the MeZO-LoRA and Sparse-MeZO methods under the
large-batch size case. Note that we did not utilize the gradient accumulation
technique for the 64 batch size case, which may significantly increase the
training time. Furthermore, we observe that the use of the parallel
contraction method can speed up the training time by 20% compared with the
traditional sequential comparison, as demonstrated by the method
AdaZeta${}_{\text{seq}}$.

##  5 Conclusion

In this paper, we propose an adaptive zeroth-order fine-tuning framework with
tensor-train decomposition, named AdaZeta. Compared with previous ZO fine-
tuning works, the AdaZeta method achieves significantly better fine-tuning
results across various tasks and models. Theoretical analysis has confirmed
that our proposed methods enjoy better convergence, which is consistency with
our experimental results on both Roberta-Large and Llama-2 models across
various fine-tuning tasks.

## Limitations

The primary limitation of this work is related to accelerating the proposed
method. Currently, multiple queries at each training step are executed
sequentially in a for-loop, which restricts further speed enhancements. This
process can potentially be optimized by implementing parallel or distributed
optimization techniques on GPUs, allowing for the simultaneous execution of
multiple queries, as these queries are independent of each other with
different random seeds. Also, Also, other gradient-free methods can be
considered for this problem Hinton (2022); Mei et al. (2023) in different
applications Zhang et al. (2024b); Cheng et al. (2023a)

Additionally, while we have introduced a fast-forward tensorized adapter in
this study, further optimization is possible in the contraction process of
these adapters during the forward pass. Advances in CUDA programming could
significantly enhance the performance of our method. However, our current
expertise in the domain of GPU optimization is limited, suggesting a potential
area for collaboration or future research.

## Potential Risks

This paper provides a cost-effective solution that operates with a minimal
memory footprint. Event though we need to fine-tune large-scale models, the
proposed method can alleviate the burden on data centers and reduce $CO_{2}$
emissions. However, we acknowledge that prolonged training times, especially
with multiple GPUs, can pose environmental challenges. Consequently, our
ongoing research endeavors are focused on developing more efficient training
methods and preserving computational power with ecological considerations in
mind.

## References

  * Amari (1993) Shun-ichi Amari. 1993.  Backpropagation and stochastic gradient descent method.  _Neurocomputing_ , 5(4-5):185–196.
  * Bowman et al. (2015) Samuel Bowman, Gabor Angeli, Christopher Potts, and Christopher D Manning. 2015.  A large annotated corpus for learning natural language inference.  In _Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing_ , pages 632–642.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901.
  * Bubeck et al. (2015) Sébastien Bubeck et al. 2015.  Convex optimization: Algorithms and complexity.  _Foundations and Trends® in Machine Learning_ , 8(3-4):231–357.
  * Chen et al. (2019) Xiangyi Chen, Sijia Liu, Kaidi Xu, Xingguo Li, Xue Lin, Mingyi Hong, and David Cox. 2019.  Zo-adamm: Zeroth-order adaptive momentum method for black-box optimization.  _Advances in neural information processing systems_ , 32.
  * Chen et al. (2024) Zhuotong Chen, Zihu Wang, Yifan Yang, Qianxiao Li, and Zheng Zhang. 2024.  Pid control-based self-healing to improve the robustness of large language models.  _Transactions on Machine Learning Research_.
  * Cheng et al. (2023a) Xuxin Cheng, Qianqian Dong, Fengpeng Yue, Tom Ko, Mingxuan Wang, and Yuexian Zou. 2023a.  M 3 st: Mix at three levels for speech translation.  In _ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pages 1–5. IEEE.
  * Cheng et al. (2023b) Xuxin Cheng, Zhihong Zhu, Ziyu Yao, Hongxiang Li, Yaowei Li, and Yuexian Zou. 2023b.  Ghostt5: generate more features with cheap operations to improve textless spoken question answering.  In _Proc. INTERSPEECH_ , pages 1134–1138.
  * Dettmers et al. (2024) Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. 2024.  Qlora: Efficient finetuning of quantized llms.  _Advances in Neural Information Processing Systems_ , 36.
  * Dua et al. (2019) Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, and Matt Gardner. 2019.  Drop: A reading comprehension benchmark requiring discrete reasoning over paragraphs.  _arXiv preprint arXiv:1903.00161_.
  * Gautam et al. (2024) Tanmay Gautam, Youngsuk Park, Hao Zhou, Parameswaran Raman, and Wooseok Ha. 2024.  Variance-reduced zeroth-order methods for fine-tuning language models.  _arXiv preprint arXiv:2404.08080_.
  * Ghadimi and Lan (2013) Saeed Ghadimi and Guanghui Lan. 2013.  Stochastic first-and zeroth-order methods for nonconvex stochastic programming.  _SIAM journal on optimization_ , 23(4):2341–2368.
  * Hinton (2022) Geoffrey Hinton. 2022.  The forward-forward algorithm: Some preliminary investigations.  _arXiv preprint arXiv:2212.13345_.
  * Houlsby et al. (2019) Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. 2019.  Parameter-efficient transfer learning for nlp.  In _International Conference on Machine Learning_ , pages 2790–2799. PMLR.
  * Hu et al. (2021) Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2021.  Lora: Low-rank adaptation of large language models.  _arXiv preprint arXiv:2106.09685_.
  * Jiang et al. (2024) Shuoran Jiang, Qingcai Chen, Youcheng Pan, Yang Xiang, Yukang Lin, Xiangping Wu, Chuanyi Liu, and Xiaobao Song. 2024.  Zo-adamu optimizer: Adapting perturbation by the momentum and uncertainty in zeroth-order optimization.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 38, pages 18363–18371.
  * Kenton and Toutanova (2019) Jacob Devlin Ming-Wei Chang Kenton and Lee Kristina Toutanova. 2019.  Bert: Pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of NAACL-HLT_ , pages 4171–4186.
  * Kumar et al. (2021) Ananya Kumar, Aditi Raghunathan, Robbie Matthew Jones, Tengyu Ma, and Percy Liang. 2021.  Fine-tuning can distort pretrained features and underperform out-of-distribution.  In _International Conference on Learning Representations_.
  * Li and Liang (2021) Xiang Lisa Li and Percy Liang. 2021.  Prefix-tuning: Optimizing continuous prompts for generation.  In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_ , pages 4582–4597.
  * Liu et al. (2022) Haokun Liu, Derek Tam, Mohammed Muqeeth, Jay Mohta, Tenghao Huang, Mohit Bansal, and Colin A Raffel. 2022.  Few-shot parameter-efficient fine-tuning is better and cheaper than in-context learning.  _Advances in Neural Information Processing Systems_ , 35:1950–1965.
  * Liu et al. (2018) Sijia Liu, Bhavya Kailkhura, Pin-Yu Chen, Paishun Ting, Shiyu Chang, and Lisa Amini. 2018.  Zeroth-order stochastic variance reduction for nonconvex optimization.  _Advances in Neural Information Processing Systems_ , 31.
  * Liu et al. (2019) Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019.  Roberta: A robustly optimized bert pretraining approach.  _arXiv preprint arXiv:1907.11692_.
  * Liu et al. (2024) Yong Liu, Zirui Zhu, Chaoyu Gong, Minhao Cheng, Cho-Jui Hsieh, and Yang You. 2024.  Sparse mezo: Less parameters for better performance in zeroth-order llm fine-tuning.  _arXiv preprint arXiv:2402.15751_.
  * Lohr (2009) Sharon L Lohr. 2009.  _Sampling: design and analysis_.  Nelson Education.
  * Loshchilov and Hutter (2018) Ilya Loshchilov and Frank Hutter. 2018.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_.
  * Malladi et al. (2023) Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D Lee, Danqi Chen, and Sanjeev Arora. 2023.  Fine-tuning language models with just forward passes.  _arXiv preprint arXiv:2305.17333_.
  * Mei et al. (2023) Yongsheng Mei, Tian Lan, Mahdi Imani, and Suresh Subramaniam. 2023.  A bayesian optimization framework for finding local optima in expensive multimodal functions.  In _ECAI 2023_ , pages 1704–1711. IOS Press.
  * Nesterov and Spokoiny (2017) Yurii Nesterov and Vladimir Spokoiny. 2017.  Random gradient-free minimization of convex functions.  _Foundations of Computational Mathematics_ , 17(2):527–566.
  * Novikov et al. (2015) Alexander Novikov, Dmitrii Podoprikhin, Anton Osokin, and Dmitry P Vetrov. 2015.  Tensorizing neural networks.  _Advances in neural information processing systems_ , 28.
  * Oseledets (2011) Ivan V Oseledets. 2011.  Tensor-train decomposition.  _SIAM Journal on Scientific Computing_ , 33(5):2295–2317.
  * Pang et al. (2002) Bo Pang, Lillian Lee, and Shivakumar Vaithyanathan. 2002.  Thumbs up? sentiment classification using machine learning techniques.  In _Proceedings of the ACL-02 conference on Empirical methods in natural language processing-Volume 10_ , pages 79–86.
  * Rajpurkar et al. (2016) Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. 2016.  Squad: 100,000+ questions for machine comprehension of text.  _arXiv preprint arXiv:1606.05250_.
  * Socher et al. (2013) Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng, and Christopher Potts. 2013.  Recursive deep models for semantic compositionality over a sentiment treebank.  In _Proceedings of the 2013 conference on empirical methods in natural language processing_ , pages 1631–1642.
  * Touvron et al. (2023) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. 2023.  Llama: Open and efficient foundation language models.  _arXiv preprint arXiv:2302.13971_.
  * Wang et al. (2019) Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. 2019.  Superglue: A stickier benchmark for general-purpose language understanding systems.  _Advances in neural information processing systems_ , 32.
  * Wang et al. (2018) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. 2018.  Glue: A multi-task benchmark and analysis platform for natural language understanding.  _arXiv preprint arXiv:1804.07461_.
  * Williams et al. (2017) Adina Williams, Nikita Nangia, and Samuel R Bowman. 2017.  A broad-coverage challenge corpus for sentence understanding through inference.  _arXiv preprint arXiv:1704.05426_.
  * Williams et al. (2018) Adina Williams, Nikita Nangia, and Samuel R Bowman. 2018.  A broad-coverage challenge corpus for sentence understanding through inference.  In _Proceedings of NAACL-HLT_ , pages 1112–1122.
  * (39) Han Xu, Jingyang Ye, Yutong Li, and Haipeng Chen.  Can speculative sampling accelerate react without compromising reasoning quality?  In _The Second Tiny Papers Track at ICLR 2024_.
  * Yang et al. (2024a) Yifan Yang, Jiajun Zhou, Ngai Wong, and Zheng Zhang. 2024a.  Loretta: Low-rank economic tensor-train adaptation for ultra-low-parameter fine-tuning of large language models.  _arXiv preprint arXiv:2402.11417_.
  * Yang et al. (2024b) Zi Yang, Samridhi Choudhary, Xinfeng Xie, Cao Gao, Siegfried Kunzmann, and Zheng Zhang. 2024b.  Comera: Computing-and memory-efficient training via rank-adaptive tensor optimization.  _arXiv preprint arXiv:2405.14377_.
  * Zaken et al. (2022) Elad Ben Zaken, Yoav Goldberg, and Shauli Ravfogel. 2022.  Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers)_ , pages 1–9.
  * Zhang et al. (2024a) Tianyi Zhang, Faisal Ladhak, Esin Durmus, Percy Liang, Kathleen McKeown, and Tatsunori B Hashimoto. 2024a.  Benchmarking large language models for news summarization.  _Transactions of the Association for Computational Linguistics_ , 12:39–57.
  * Zhang et al. (2024b) Yigeng Zhang, Fabio A González, and Thamar Solorio. 2024b.  Interpreting themes from educational stories.  _arXiv preprint arXiv:2404.05250_.
  * Zhang et al. (2023) Yigeng Zhang, Mahsa Shafaei, Fabio Gonzalez, and Thamar Solorio. 2023.  Positive and risky message assessment for music products.  _arXiv preprint arXiv:2309.10182_.

##  Appendix A Detail of Experiment Setup

###  A.1 Dataset Setup

Table 4: Metrics that we use to evaluate the benchmark for Roberta-Large
Model.

Task Name | Metric
---|---
SST-2 | Accuracy
SST-5 | Accuracy
QNLI | Accuracy
MNLI | Matched Acc.
SNLI | Accuracy
RTE | Accuracy

Our research utilized a variety of tasks to measure the performance of the
Roberta-Large model, including sentiment analysis (SST-2, SST-5 Socher et al.
(2013), MR Pang et al. (2002)), and natural language inference (MNLI Wang et
al. (2018), QNLI Williams et al. (2018), SNLI Bowman et al. (2015), RTE Wang
et al. (2018)) tasks. Table 4 summarizes the evaluation metrics used for these
tasks.

Further, we extended our experiments on a large-scale Llama-2-7B model to
include tasks from the SuperGLUE benchmark Wang et al. (2019), which involves
both classification (CB, BoolQ, WSC) and reasoning tasks (COPA and ReCoRD), as
well as additional generation tasks, SQuAD Rajpurkar et al. (2016). For these
tests, we introduced a challenging low-resource data condition, limiting our
samples to 1,000 for training, 500 for validation, and 1,000 for testing, as
detailed in the prompt-based task settings from Appendix D of Malladi et al.
(2023). The metrics for these evaluations are outlined in Table 5.

Table 5: Metrics that we use to evaluate SuperGLUE and generations tasks.

Task Name | Metric
---|---
CB | F1
BoolQ | Accuracy
WSC | F1
COPA | Accuracy
ReCoRD | F1
SQuAD | F1

###  A.2 Baselines

In this section, we provide a detailed introduction to the baseline method
considered in our experiments, which are listed as follows:

Full-model First-Order Fine-Tuning (FT) is the most widely used method for
fine-tuning LLMs. In this process, the model is initialized with pre-trained
weights, and all model parameters are updated by the first-order optimizer. In
this paper, the AdamW optimizer Loshchilov and Hutter (2018) is used to
conduct the first-order experiments.

Zero-shot/In-context-learning (ICL) is the most widely used method for fine-
tuning large language models (LLMs). In this process, the model is initialized
with pre-trained weights, and all model parameters are updated by the first-
order (FO) optimizer. In this paper, the AdamW optimizer Loshchilov and Hutter
(2018) is used to conduct the first-order experiments.

Linear-probing (LP) method involves freezing the pretrained weights of the
model and adding a final linear classifier layer, implemented using the scipy
package. By fine-tuning this layer with the first-order method, we only need
to construct a small backpropagation graph. However, this method is not
suitable for generative tasks. Therefore, we only apply the LP method in the
Roberta-Large experiments.

Memory-Efficient Zeroth-Order (MeZO) was first proposed in Malladi et al.
(2023), which fine-tunes LLMs using only the forward pass. The MeZO method
significantly reduces memory costs by eliminating the need for a
backpropagation graph and has demonstrated superior performance compared to
inference-only methods like Zero-shot, ICT, and LP methods across various
downstream tasks.

Memory-Efficient Zeroth-Order with LoRA adapters (MeZO-LoRA) is a derivative
method introduced in Malladi et al. (2023), which freezes the pretrained
weights and fine-tunes only the injected LoRA adapters Hu et al. (2021). The
MeZO-LoRA method is the most relevant baseline in this field compared to our
work. However, its performance improvement over the MeZO method is limited,
and the mechanisms behind zeroth-order parameter-efficient fine-tuning are not
extensively discussed.

Sparse Memory-efficient Zeroth-Order (Sparse-MeZO) is a recently proposed
method aiming to enhance the performance and convergence speed of the MeZO
method Liu et al. (2024). However, as the code and detailed layer-wise
hyperparameter setup have not been released, we have reproduced the method
using a fixed sparsity ratio for each layer. This ratio is selected based on
the best overall outcome as presented in Fig. 6 of their paper.

###  A.3 Hyperparameters

In this section, we outline the detailed setup of hyperparameters utilized in
our study. The specific choices of hyperparameters, such as learning rate,
training steps, and batch size, are summarized in Table 6. In our experiments,
we strive to maintain a consistent learning rate across different methods for
the same tasks. However, for approaches like full-model fine-tuning, we opt
for a lower learning rate to ensure convergence. This principle is also
applied in our large-scale experiments on the Llama-2-7B model, details of
which are summarized in Table 7.

In addition to the standard hyperparameter configuration, we also consider the
shape of tensor factors in our methods. To represent a layer with input and
output dimensions of $o$ and $p$, respectively, we employ a list of $m$ tensor
factors $\mathcal{G}_{i}\in\mathbb{R}^{r\times k_{i}r}$, where the product
$\Pi k_{1}\cdots k_{m}=o\cdot p$. The specific shapes of $k_{i}$ corresponding
to different values of $o$ and $p$, given a bottleneck size of 8 or 64 for the
tensorized methods, are detailed in Table 8. Note that the optimal factors
shape and tensor rank for the tensor-train method can only be determined by
the experiments’ trail. However, previous work also explores the possibility
of utilizing the adaptive rank to improve the performance Yang et al. (2024b),
which may further improve the performance of our AdaZeta method.

Table 6: The hyperparameter grids used for Roberta-Large experiments are
detailed as follows. We fine-tune each task for 80K steps, except for the FT
method, which is conducted over 20 epochs. We record the best model checkpoint
based on the validation loss every 200 training steps.

Experiment | Hyperparameters | Values
---|---|---
FT | Batch size | {8, 16, 64}
| Learning rate | {1e-6, 5e-7}
MeZO | Batch size | {16, 64}
| Learning rate | {1e-6, 5e-7}
| $\epsilon$ | 1e-3
MeZO-LoRA | Batch size | {16, 64}
| Learning rate | {1e-4, 5e-5}
| LoRA rank | 8
| $\epsilon$ | 1e-3
Sparse-MeZO | Batch size | {16, 64}
| Learning rate | {1e-5, 1e-6}
| sparse ratio | 0.75
| $\epsilon$ | 1e-3
AdaZeta | Batch size | {16, 64}
| Learning rate | {1e-4, 5e-5}
| Bottleneck dimension | 64
| Tensor Rank | 5
| $\epsilon$ | 1e-3

Table 7: The hyperparameter grids used for Llama-2-7B experiments are outlined
as follows. We fine-tune each task for 5K steps using our AdaZeta method, 10K
steps for other ZO fine-tuning methods (MeZO, MeZO-LoRA, Sparse-MeZO), and 5
epochs for the first-order Full-model Fine-Tuning (FT) method. We record the
best model checkpoint based on the validation loss every 200 training steps.

Experiment | Hyperparameters | Values
---|---|---
FT | Batch size | {8, 16, 64}
| Learning rate | {1e-6, 5e-7}
MeZO | Batch size | {16, 64}
| Learning rate | {1e-6, 5e-7}
| $\epsilon$ | 1e-3
MeZO-LoRA | Batch size | {16, 64}
| Learning rate | {1e-4, 5e-5}
| LoRA rank | {5, 8, 16}
| $\epsilon$ | 1e-3
Sparse-MeZO | Batch size | {16, 64}
| Learning rate | {1e-5, 1e-6}
| sparse ratio | 0.75
| $\epsilon$ | 1e-3
AdaZeta | Batch size | {16, 64}
| Learning rate | {1e-4, 5e-5}
| Bottleneck dimension | {8, 64}
| Tensor Rank | {5, 8, 16}
| Query Constants | $\alpha=0.85,\beta=0.45$
| Maximum Query | $Q_{max}=20$
| $\epsilon$ | 1e-3

Table 8: The shape settings of the tensorized adapters in AdaZeta Method

Bottleneck size | Matrix Shape | Tensor Shape
---|---|---
8 | $768\times 64$ | [8, 8, 12, 4, 4, 4]
| $4096\times 64$ | [16, 16, 16, 4, 4, 4]
| $64\times 768$ | [4, 4, 4, 12, 8, 8]
| $64\times 4096$ | [4, 4, 4, 16, 16, 16]
64 | $768\times 8$ | [8, 8, 12, 2, 2, 2]
| $4096\times 8$ | [16, 16, 16, 2, 2, 4]
| $8\times 768$ | [2, 2, 2, 12, 8, 8]
| $8\times 4096$ | [2, 2, 2, 16, 16, 16]

##  Appendix B Additional Experiments

###  B.1 Additional Momeory Comparison results

In this section, we provide more quantitative results about the training
memory comparison between the FO and ZO fine-tuning methods. In addition to
the training memory on SST2 tasks we measure in Section 4.3, we further
profile the memory cost on WIC, CB, and MultiRC tasks. The results are shown
in Table 9.

We can observe from the table that the AdaZeta method achieves 5-8$\times$
memory reduction on different tasks. Also, the AdaZeta method utilizes similar
or even less memory than the other MeZO, MeZo-LoRA, and Sparse-MeZO methods
with an additional variance reduction feature, which largely improves the ZO
fine-tuning accuracy.

Table 9: Quantitive results for the memory profiling over SST2 and MultiRC
tasks.

Methods | SST2 | WIC | CB | MultiRC
---|---|---|---|---
FT | 118.65 | 115.3 | 151.97 | 191.97
MeZO | 15.08 | 15.22 | 23.01 | 41.17
MeZO-LoRA | 14.75 | 15.23 | 23.02 | 41.18
MeZO-LoRA(BS=64) | 21.07 | 25.3 | 71.70 | 84.3
Sparse-MeZO | 14.35 | 15.21 | 23.01 | 42.13
AdaZeta | 14.73 | 15.22 | 23.01 | 41.17

##  Appendix C Proof of Theorem 1

To retain the readability of the proof, we use a single-column format in the
following. To provide the proof of Theorem 1, we first present a Lemma
regarding the bound of gradient noise. Recall from the gradient estimation
rule that:

$\displaystyle\nabla\hat{\ell}(\bm{w}_{k})$
$\displaystyle=\frac{1}{B}\sum_{b_{i}\in\mathcal{B}}\hat{g}(\bm{w}_{k};b_{i})$
(3) $\displaystyle\hat{g}(\bm{w}_{k};b_{i})$
$\displaystyle=\frac{1}{Q_{k}}\sum_{j=1}^{Q_{k}}\hat{g}(\bm{w}_{k};b_{i},u_{i,j}),$
(4)

where there are two sources of randomness: a) The randomness leads by the
mini-batch sampling and b) The randomness leads by the presence of ZO gradient
estimation. Based on these two randomnesses, we define two gradient noises as
$h_{k}$ and $e_{k}$, respectively.

$\displaystyle h_{k}$
$\displaystyle:=\nabla\hat{\ell}(\bm{w}_{k})-\nabla\ell(\bm{w}_{k})=\frac{1}{B}\sum_{b_{i}\in\mathcal{B}}\hat{g}(\bm{w}_{k};b_{i})-\nabla\ell(\bm{w}_{k})$
(5) $\displaystyle e_{k}$
$\displaystyle:=\hat{g}(\bm{w}_{k};b_{i})-\nabla\ell(\bm{w}_{k})=\frac{1}{Q_{k}}\sum_{j=1}^{Q_{k}}\hat{g}(\bm{w}_{k};b_{i},u_{i,j})-\nabla\ell(\bm{w}_{k})$
(6)

Here, we first bound the gradient noise $h_{k}$ with a fact given in
stochastic gradient descent theory. We consider the noise concerning the mean
of the ZO estimated gradient $\nabla\ell(\bm{w}_{k})$, where the loss function
$\ell$ is a randomized smoothing version of $\ell$.

######  Lemma 1.

Based on the definition in eq. (5) and the Assumption A2, we can bound the
L2-norm of the gradient noise $h_{k}$ by taking expectation:

$\displaystyle\mathbb{E}[\|h_{k}\|^{2}]\leq\frac{N-B}{NB(B-1)Q_{k}}\sum_{i}(2d\delta^{2}+\frac{\epsilon^{2}L^{2}d^{2}}{2}+2\delta^{2})$
(7)

###### Proof.

For convenience, we consider a general case that the mini-batch $\mathcal{B}$
is formed by uniform sampling without replacement and follows the i.i.d.
fashion. Then, according to Lohr (2009)[Section 2.8, Page 48], the following
holds for a random sampling noise:

$\displaystyle\mathbb{E}[\|h_{k}\|^{2}]=\frac{N-B}{NB}\Lambda^{2},$ (8)

where $\Lambda^{2}$ is the sample variance of the gradient
$\hat{g}(\bm{w}_{k};b_{i})$, which is defined as:

$\displaystyle\Lambda^{2}$
$\displaystyle=\frac{1}{B-1}\sum_{i=1}^{B}\|\hat{g}(\bm{w}_{k};b_{i})-\nabla\ell(\bm{w}_{k})\|^{2}$
(9)
$\displaystyle=\frac{1}{B-1}\sum_{i=1}^{B}\|\nabla\ell(\bm{w}_{k})+e_{k}-\nabla\ell(\bm{w}_{k})\|^{2}$
(10) $\displaystyle=\frac{1}{B-1}\sum_{i=1}^{B}\|e_{k}\|^{2},$ (11)

where $e_{k}$ is defined as the gradient noise leads by the ZO estimation in
eq. (5).

Finally, we need to bound the variance $\Lambda^{2}$, related to the ZO
gradient estimation noise. Taking expectation with respect to the i.i.d.
random perturbation vector $\bm{u}$, we have:

$\displaystyle\mathbb{E}_{\bm{u}}[\Lambda^{2}]$
$\displaystyle\leq\mathbb{E}_{\bm{u}}[\frac{1}{B-1}\sum_{i=1}^{B}\|e_{k}^{i}\|^{2}]$
(12)
$\displaystyle\leq\frac{1}{(B-1)Q_{k}^{2}}\sum_{i}\mathbb{E}_{\bm{u}}[\|\sum_{j=1}^{Q_{k}}(\hat{g}(\bm{w}_{k};b_{i},u_{i,j})-\nabla\ell(\bm{w}_{k}))\|]$
(13)
$\displaystyle\overset{(a)}{=}\frac{1}{(B-1)Q_{k}}\sum_{i}\mathbb{E}_{\bm{u}}[\|\hat{g}(\bm{w}_{k};b_{i},u_{i,1})-\nabla\ell(\bm{w}_{k})\|],$
(14)

where (a) is given under the case that $u_{i,j}$ is i.i.d, which obtain:

$\displaystyle\mathbb{E}_{\bm{u}}[\|\hat{g}(\bm{w}_{k};b_{i},u_{i,j})-\nabla\ell(\bm{w}_{k})\|]=\mathbb{E}_{\bm{u}}[\|\hat{g}(\bm{w}_{k};b_{i},u_{i,1})-\nabla\ell(\bm{w}_{k})\|]$
(15)

Finally, we need to bound the term
$\mathbb{E}_{\bm{u}}[\|\hat{g}(\bm{w}_{k};b_{i},u_{i,1})-\nabla\ell(\bm{w}_{k})\|]$,
which gives:

$\displaystyle\mathbb{E}_{\bm{u}}[\|\hat{g}(\bm{w}_{k};b_{i},u_{i,1})-\nabla\ell(\bm{w}_{k})\|]]$
(16)
$\displaystyle\leq\mathbb{E}_{\bm{u}}[\|\hat{g}(\bm{w}_{k};b_{i},u_{i,1})-\nabla\ell(\bm{w}_{k};b_{i})\|]+\mathbb{E}_{\bm{u}}[\|\nabla\ell(\bm{w}_{k};b_{i})-\nabla\ell(\bm{w}_{k})\|]$
(17)
$\displaystyle\overset{(a)}{\leq}2d\|\hat{g}(\bm{w}_{k};b_{i},u_{i,1})\|+\frac{\epsilon^{2}L^{2}d^{2}}{2}+\mathbb{E}_{\bm{u}}[\|\nabla\ell(\bm{w}_{k})\|]+\mathbb{E}_{\bm{u}}[\|\nabla\ell(\bm{w}_{k};b_{i})\|]$
(18)
$\displaystyle\overset{(b)}{\leq}2d\delta^{2}+\frac{\epsilon^{2}L^{2}d^{2}}{2}+2\delta^{2},$
(19)

where (a) follows a similar idea of the proof in Ghadimi and Lan (2013)[eq.
(3.21)] and (b) is given by using the bound of the gradient in Assumption A2.

Putting it all together we can obtain the upper bound for the gradient noise
$\|h_{k}\|$. ∎

Now we begin to present the proof of Theorem 1:

We start from the gradient updating rule in the AdaZeta algorithm, which gives
$\bm{w}_{t+1}=\bm{w}_{t}-\eta\nabla\hat{\ell}(\bm{w}_{k})$. By using Taylor’s
theorem on the exact smoothed loss $\ell(\bm{w}_{k})$, we have:

$\displaystyle\ell(\bm{w}_{k+1})$
$\displaystyle=\ell(\bm{w}_{k}-\eta\nabla\hat{\ell}(\bm{w}_{k}))$ (20)
$\displaystyle=\ell(\bm{w}_{k})-\eta\nabla\hat{\ell}_{k}(\bm{w}_{k})^{\top}\nabla\ell(\bm{w}_{k})+\frac{\eta^{2}}{2}\nabla\hat{\ell}(\bm{w}_{k})^{\top}\nabla\ell(\bm{w}_{k})^{2}\nabla\hat{\ell}(\bm{w}_{k})$
(21)

Taking expectations on both sides gives:

$\displaystyle\mathbb{E}_{\bm{w}_{t}}[\ell(\bm{w}_{k+1})]$
$\displaystyle=\mathbb{E}_{\bm{w}_{t}}[\ell(\bm{w}_{k})]-\eta\mathbb{E}_{\bm{w}_{t}}[\nabla\hat{\ell}(\bm{w}_{k})^{\top}\nabla\ell(\bm{w}_{k})]+\frac{\eta^{2}}{2}\mathbb{E}_{\bm{w}_{t}}[\nabla\hat{\ell}(\bm{w}_{k})^{\top}\nabla\ell(\bm{w}_{k})^{2}\nabla\hat{\ell}(\bm{w}_{k})]$
$\displaystyle\overset{\text{(a)}}{\leq}\mathbb{E}_{\bm{w}_{t}}[\ell(\bm{w}_{k})]-\eta\mathbb{E}_{\bm{w}_{t}}[\nabla\ell(\bm{w}_{k})^{2}]+\frac{\eta^{2}L}{2}\mathbb{E}_{\bm{w}_{t}}[\nabla\hat{\ell}(\bm{w}_{k})^{2}],$

where (a) can be proved with the use of the Lipschitz smoothness gradient
denied in Assumption A1 that gives $x$ and $y$, we have
$\|\nabla\ell(x)-\nabla\ell(y)\|\leq L\|x-y\|$. Additionaly, by the mean value
theorem for vector-valued functions, there exists for any point $c$ on the
line segment between $x$ and $y$ such that:

$\displaystyle\nabla f(y)-\nabla f(x)=\nabla^{2}f(c)(y-x).$ (22)

Taking the norms on both sides and using the Lipschitz condition, we have:

$\displaystyle\left\|\nabla^{2}f(c)(y-x)\right\|=\|\nabla f(y)-\nabla
f(x)\|\leq L\|y-x\|.$ (23)

Finally, since this must hold for any $y$ and $x$, and since the norm of the
Hessian matrix is the supremum of $\left\|V^{2}f(c)(y-x)\right\|/\|y-x\|$ for
non-zero $y-x$, it follows that:

$\displaystyle\left\|\nabla^{2}f(c)\right\|\leq L$ (24)

Rearrange and we obtain:

$\displaystyle\eta\mathbb{E}[\|\nabla\ell(\bm{w}_{k})\|^{2}]\leq\mathbb{E}[\ell(\bm{w}_{k})]-\mathbb{E}[\ell(\bm{w}_{k+1})]+\frac{\eta^{2}L}{2}\mathbb{E}[\nabla\hat{\ell}(\bm{w}_{k})^{2}]$
(25)

Taking summation over steps $k=1,\cdots,K$ gives:

$\displaystyle\sum_{k=1}^{K}\eta\mathbb{E}[\|\nabla\ell(\bm{w}_{k})\|^{2}]$
$\displaystyle\leq\mathbb{E}[\ell(\bm{w}_{0})-\ell(\bm{w}_{K})]+\sum_{k=1}^{K}\frac{\eta^{2}L}{2}\mathbb{E}[\nabla\hat{\ell}(\bm{w}_{k})^{2}]$
(26)
$\displaystyle\overset{\text{(a)}}{\leq}\mathbb{E}[\ell_{0}-\ell^{*}]+\epsilon^{2}L+\sum_{k=1}^{K}\frac{\eta^{2}L}{2}\mathbb{E}[\nabla\hat{\ell}(\bm{w}_{k})^{2}]$
(27)
$\displaystyle\overset{\text{(b)}}{\leq}R+\epsilon^{2}L+\sum_{k=1}^{K}\frac{\eta^{2}L}{2}\mathbb{E}[\nabla\hat{\ell}(\bm{w}_{k})^{2}],$
(28)

where (a) is using the Lemma 1 in Liu et al. (2018) that
$\ell(\bm{w}_{0})-\ell(\bm{w}_{T})\leq\ell(\bm{w}_{0})-\ell(\bm{w}_{0})+\ell^{*}-\ell^{*}\leq(\ell(\bm{w}_{0})-\ell^{*})+\epsilon^{2}L$
and (b) is given by setting $R:=\ell(\bm{w}_{1})-\ell^{*}$. Now, the key to
the bound comes from the last term in the right of the inequation.

To bound the last term, we first represent the noise gradient
$\nabla\hat{\ell}_{k}(\bm{w}_{k})^{2}$ as a combination of the true gradient
and the gradient noise introduced in eq. (5), which gives:

$\displaystyle\nabla\hat{\ell}(\bm{w}_{k}):=\nabla\ell(\bm{w}_{k})+h_{k}$ (29)

Taking eq. (29) back into eq. (26), using the results from Lemma 1, taking the
expectation over all randomness and average over the maximum steps $K$, we
obtain:

$\displaystyle\frac{1}{K}\sum_{k=1}^{K}\eta\mathbb{E}[\|\nabla\ell(\bm{w}_{k})\|^{2}]$
$\displaystyle\leq\frac{R}{K}+\frac{\epsilon^{2}L}{K}+\frac{1}{K}\sum_{k=1}^{K}\frac{\eta^{2}L}{2}\mathbb{E}[\nabla\hat{\ell}(\bm{w}_{k})^{2}]$
$\displaystyle\leq\frac{R}{K}+\frac{\epsilon^{2}L}{K}+\frac{1}{K}\sum_{k=1}^{K}\frac{\eta^{2}L}{2}\mathbb{E}[\|\nabla\ell(\bm{w}_{k})\|+\|h_{k}\|]$
$\displaystyle=\frac{R}{K}+\frac{\epsilon^{2}L}{K}+\frac{\eta^{2}L\delta}{2}+\frac{1}{K}\sum_{k=1}^{K}\frac{\eta^{2}L}{2}\frac{N-B}{NB}\Lambda^{2}$
$\displaystyle\leq\frac{R}{K}+\frac{\epsilon^{2}L}{K}+\frac{\eta^{2}L\delta}{2}+\frac{1}{K}\sum_{k=1}^{K}\frac{\eta^{2}L}{2}\frac{N-B}{NB}(\frac{1}{(B-1)Q_{k}}\sum_{i}\mathbb{E}_{\bm{u}}[\|e^{i,1}\|]]+\frac{B\epsilon^{2}L}{2(B-1)})$
$\displaystyle\leq\frac{R}{K}+\frac{\epsilon^{2}L}{K}+\frac{\eta^{2}L\delta}{2}+\frac{1}{K}\sum_{k=1}^{K}\frac{\eta^{2}L}{2}\frac{N-B}{NB}(\frac{\sum_{i}(2d\delta^{2}+\frac{\epsilon^{2}L^{2}d^{2}}{2}+2\delta)}{(B-1)Q_{k}}+\frac{B\epsilon^{2}L}{2(B-1)})$
$\displaystyle=\frac{R+\epsilon^{2}L+C(d,\epsilon)\sum_{k}\frac{1}{Q_{k}}}{K}+\frac{\eta^{2}L\delta}{2}+\frac{B\epsilon^{2}L}{2(B-1)}$
$\displaystyle=\mathcal{O}(\frac{R+\epsilon^{2}L+C(d,\epsilon)\sum_{k}\frac{1}{Q_{k}}}{K}),$

where $C(d,\epsilon)$ is a constant defined as
$C(d,\epsilon):=\sum_{k=1}^{K}\frac{\eta^{2}L}{2}\frac{N-B}{NB}(\frac{\sum_{i}(2d\delta^{2}+\frac{\epsilon^{2}L^{2}d^{2}}{2}+2\delta)}{(B-1)})$

Divide both side with $\eta$ and use the trick to introduce some randomly
chosen $\bm{w}_{T}$ from the histroy with probability $P(T=k)=\frac{1}{K}$, we
finish the proof as:

$\displaystyle\mathbb{E}[\|\nabla\ell(\bm{w}_{T})\|^{2}]=\frac{1}{K}\sum_{k=1}^{K}\mathbb{E}[\|\nabla\ell(\bm{w}_{k})\|^{2}]\leq\mathcal{O}(\frac{R+\epsilon^{2}L+C(d,\epsilon)\sum_{k}\frac{1}{Q_{k}}}{K\epsilon})$
(30)


# Multi-modal fusion with gating using audio, lexical and disfluency features
for Alzheimer’s Dementia recognition from spontaneous speech

###### Abstract

This paper is a submission to the Alzheimer’s Dementia Recognition through
Spontaneous Speech (ADReSS) challenge, which aims to develop methods that can
assist in the automated prediction of severity of Alzheimer’s Disease from
speech data. We focus on acoustic and natural language features for cognitive
impairment detection in spontaneous speech in the context of Alzheimer’s
Disease Diagnosis and the mini-mental state examination (MMSE) score
prediction. We proposed a model that obtains unimodal decisions from different
LSTMs, one for each modality of text and audio, and then combines them using a
gating mechanism for the final prediction. We focused on sequential modelling
of text and audio and investigated whether the disfluencies present in
individuals’ speech relate to the extent of their cognitive impairment. Our
results show that the proposed classification and regression schemes obtain
very promising results on both development and test sets. This suggests
Alzheimer’s Disease can be detected successfully with sequence modeling of the
speech data of medical sessions.

Index Terms: Cognitive Decline Detection, Affective Computing, Computational
Paralinguistics

##  1 Introduction

Alzheimer’s Disease (AD) is a chronic neurodegenerative condition and the most
common form of dementia. AD gradually affects the memory, language and
cognitive skills and ultimately the ability to perform basic tasks in the
everyday lives of patients. Early diagnosis of AD has become essential in
disease management as it has not been possible to reverse the degenerative
process, even with significant efforts focused on therapies [1].

Discrepancies in speech comprehension, speech production and memory functions
are closely tied in with AD as suggested by a decrease in global vocabulary
and a loss in evocative memory [2]. Patients with AD have difficulty
performing tasks that leverage semantic information; they exhibit problems
with verbal fluency and identification of objects [3]. The semantics and
pragmatics of their language appear affected throughout the entire span of the
disease more than syntax [4]. AD Patients talk more gradually with longer
pauses and invest extra time seeking the right word, which contributes to
disfluency of speech [3].

AD diagnosis demands the existence of cognitive dysfunction to be validated by
neuropsychological assessments like the mini mental state examination (MMSE)
performed in medical clinics [5]. Diagnosis is typically based on the clinical
analysis of patients’ history and the presence of typical neurological and
neuropsychological features. It is costly and not accessible to all patients
who have concerns about their memory functions.

Recent experimental research has looked at AD’s automated analysis from
multimodal data as alternative, less invasive tools for diagnostics. Studying
behaviours of individuals could also help detect AD earlier. There has been
research on building systems which use a broad range of multimodal features to
identify AD severity. A meaningful association between MMSE scores and
language measures such as articulation and disfluency has been found [6].

Much of the work to date has looked separately at the properties of the
language of an individual: acoustic and lexical characteristics of speech, or
syntax, fluency, and content of information. Usually these are studied within
language tasks in specific domains or in conversational dialogue. Several
studies have suggested various forms of speech analysis to identify AD.
Researchers found that the number of pauses, pause proportion, phonation time,
phonation–to-time ratio, speach rate, articulation rate, and noise-to-harmonic
ratio correlate with the severity of AD [7]. Weiner et al. [8] developed a
Linear Discriminant Analysis (LDA) classifier with a set of acoustic features
such as the mean of silent segments, speech and silence durations and silence
to speech ratio to distinguish subjects with AD from the control group and
achieved a classification accuracy of 85.7 percent. Ambrosini et al. [9]
showed an accuracy of 73 percent when using selected acoustic features (pitch,
voice breaks, shimmer, speech rate, syllable duration) to detect mild
cognitive impairment from a spontaneous speech task.

In terms of the features which aid AD detection, lexical features from
spontaneous speech are shown to be informative. Jarrold et al. [10] extracted
the frequency occurrence of 14 different part of speech features and combined
them with acoustic features. Abel et al. [11] modeled patient speech errors
(naming and repetition disorders) to the problem of AD diagnosis.

There has also been work on modelling multimodal input for AD detection.
Gosztolya et al. [12] examined the fusion of two SVM models with separate
feature sets. The first model used a set of acoustic features, and the second
model was developed using linguistic features extracted from manually
annotated transcripts. Their work showed the complementary information that
audio and lexical features may contain about a subject with AD.

Among other similar tasks, using multimodal fusion to predict a cognitive
state, research has been done on integrating temporal information from two or
more modalities in a recurrent approaches to classify emotions or detecting
different mental states, such as depression [13]. One key challenge these
models have is addressing the various predictive capacity of each modality and
their different levels of noise. The application of a gating mechanism in
various multimodal tasks has been shown to be successful in controlling the
level of contribution of each modality to the eventual prediction.

This paper addresses AD classification and MMSE score regression tasks, which
are part of the Alzheimer’s Dementia Recognition through Spontaneous Speech
(ADReSS) challenge [14]. In ADReSS, participants are required to assess the AD
severity of different subjects, where the target severity is based on their
MMSE scores.

We performed a binary classification of samples of speech into AD and non-AD
classes and create regression models to predict MMSE scores. Using the ADReSS
Challenge data which consists of speech recordings and transcripts of spoken
picture descriptions, we explored various features as diagnostically relevant
tools. We focused in particular on sequential modelling of sessions and
whether the disfluencies and self-repairs present in individuals’ speech can
help predict the level of cognitive impairment.

Our approach is motivated by [13] that developed the ability to learn
difficult decision boundaries which other models with different methods of
fusion have trouble managing, and maximise the use and combination of each
modality. We employed data of individuals under controlled conditions, and
modeled the sessions with audio and text features in a Long-Short Term Memory
(LSTM) neural network to detect AD. Our findings indicate that AD can be
detected with minimal information available on the structure of the
description tasks by pure sequential modelling of a session. We also found
that disfluency markers have predictive power for AD recognition.

##  2 Proposed Approach

Our approach is to model the speech of individuals giving picture descriptions
as a sequence to predict whether they have AD or not, and if so, to what
degree. To predict AD, we performed three sets of experiments using features
from the audio and text data:

  * 1

LSTM models utilising unimodal audio and text features.

  * 2

LSTM model with gating to test the effect of using multimodality.

  * 3

A multimodal LSTM model using acoustic and lexical information, including
disfluency tagging.

The details of the three experiments are outlined below in the following sub-
sections. In line with the standard assumption in deep learning, we take the
approach that for a model to be genuinely data-driven, minimal feature
engineering is required. The model’s power is in its capacity to represent
information through non-linear transforms, at varying spatial and temporal
units, and from different modalities. Since we were interested in modelling
temporal session changes, we used a bi-directional Long Short-Term Memory
(LSTM) neural network as it has the added benefit of sequential data
modelling. For each of the audio and text modalities we trained an LSTM model
separately, using the audio and text features.

###  2.1 Multimodal Features

Lexical Features from Text A pre-trained GloVe model [15] was used to extract
the lexical feature representations from the picture description transcript
and convert the utterance sequences into word vectors. We selected the
hyperparameter values, which optimised the output of the model on the training
set. The optimal dimension of the embedding was found to be 100.

Audio Features A set of 79 audio features were extracted using the COVAREP
acoustic analysis framework software, a package used for automatic extraction
of features from speech [16]. We sampled the audio features at 100Hz and used
the higher-order statistics (mean, maximum, minimum, median, standard
deviation, skew, and kurtosis) of COVAREP features. The features include
prosodic features (fundamental frequency and voicing), voice quality features
(normalized amplitude quotient, quasi open quotient, the difference in
amplitude of the first two harmonics of the differentiated glottal source
spectrum, maxima dispersion quotient, parabolic spectral parameter, spectral
tilt/slope of wavelet responses, and shape parameter of the Liljencrants-Fant
model of the glottal pulse dynamics) and spectral features (Mel cepstral
coefficients 0-24, Harmonic Model and Phase Distortion mean 0-24 and
deviations 0-12). Segments without audio data were set to zero. A standard
zero-mean and variance normalization was applied to features. We omitted all
features with no statistically significant univariate correlation with the
results of training set.

###  2.2 Sequence Modeling

The potential of neural networks lies in the power to derive representations
of features by non-linear input data transformations, providing greater power
than traditional models. As we were interested in modelling temporal nature of
speech recordings and transcripts, we used a bi-directional LSTM. For each of
the audio and text modalities we trained a separate unimodal LSTM model, using
different sets of features. For the input data we explored different timesteps
and strides. After exploring different hyper-parameters, the model using audio
data has a timestep of 20 and stride of 1 with 4 bi-directional LSTM layers
with 256 hidden nodes. The model using text input has an input with a timestep
of 10 and stride of 2 and has 2 LSTM layers with 16 hidden nodes. The code
used in the experiments are publicly available in an online
repository.111https://github.com/mortezaro/ad-recognition-from-speech

###  2.3 Multimodal Fusion with Gating

Audio and text features can include not only discriminative and temporarily
changing information about the current state of a subject, but supporting
information as well.

The model consists of two branches of the LSTM, one for each of the
modalities, with their outputs combined into final feed-forward highway
layers. The branches are made up of different hyperparameters and configured
with respect to each modality’s properties. Their outputs are concatenated and
passed through N highway layers (where the best value N was determined from
optimizing on heldout data). We pad the size of the training examples in the
text set (which was the smaller set) to meet the audio set by mapping together
instances that occurred in the same session, as the audio and text inputs for
each branch of the LSTM had different timesteps and strides.

Gating Mechanism Data from two modalities affect the final output differently,
and it is important to consider the amount of noise when aggregating them into
a single representation. Since learned representation for the text can be
undermined by corresponding audio representation, during multimodal fusion we
need to minimise the effects of noise and overlaps. We use feed-forward
highway layers [17], with gating units that learn by weighing text and audio
inputs at each time step to regulate information flow through network work.

Each highway layer consists of two non-linear transformations: a Carry ($Cr$)
and a Transform ($Tr$) gate which determine the degree to which the output is
generated by transforming and carrying the input. Each layer uses the gates
and feed-forward layer $H$ to regulate its input vector at timestep $t$,
$D_{t}$, to generate output $y$:

$y=Tr\cdot H+Cr\cdot D_{t}$ (1)

where $Cr$ is simply defined as $1-Tr$, giving:

$y=Tr\cdot H+(1-Tr)\cdot D_{t}$ (2)

The transform gate $Tr$ is defined as $\sigma(W_{Tr}D_{t}+b_{Tr})$, where
$W_{Tr}$ is the weight matrix and $b_{Tr}$ the bias vector for the gates.
Based on the transform gates outputs, highway layers adjusts their performance
from multiple-unit layers to layers that only pass through their inputs. As
inspired by [17] and to help resolve long-term learning dependencies faster we
initialise $b_{Tr}$ with a negative value (biased towards the Carry gate). We
use a block of 3 stacked highway layers. The overall architecture of the LSTM
with Gating model is shown in Figure 1.

Figure 1: Multimodal fusion with gating.

###  2.4 Multi-modal Model with Disfluency Markers

Disfluencies like self-repairs, pauses and fillers are widespread in everyday
speech [18]. Disfluencies are usually seen as indicative of communication
problems, caused by production or self-monitoring issues [19]. Individuals
with AD are likely to deal with troubles in language and cognitive skills.
Patients with AD speak more slowly and with longer breaks, and invest extra
time seeking the right word, which in effect contributes to disfluency [3].
The present research explores the disfluencies present in the speech of AD
patients as they contribute to severity of symptoms.

Self-repair disfluencies are typically assumed to have a reparandum-
interregnum-repair structure, in their fullest form as speech repairs [20]. A
reparandum is a speech error subsequently fixed by the speaker; the corrected
expression is a repair. An interregnum word is a filler or a reference
expression between the words of repair and reparandum, often a halting step as
the speaker produces the repair, giving the structure as in (3)

$\texttt{ John }\underbrace{\texttt{ [ likes
}}_{\texttt{reparandum}}+\underbrace{\texttt{
\\{~{}uh~{}\\}}}_{\texttt{interregnum}}\underbrace{\texttt{ loves
]}}_{\texttt{repair}}\texttt{ Mary}$ (3)

In the absence of reparandum and repair, the disfluency reduces to an isolated
edit term. A marked, lexicalized edit term such as a filled pause (“uh” or
“um”) or more phrasal terms like “I mean” and “you know” can occur.
Recognizing these elements and their structure is then the task of disfluency
detection.

We automatically annotated self-repairs using a deep-learning-driven model of
incremental detection of disfluency developed by Hough and Schlangen [21, 22].
It consists of deep learning sequence models that use word embeddings of
incoming words, part-of-speech annotations, and other features in a left-to-
right, word-by-word manner to predict disfluency tags. Here each word is
either tagged as a repair onset tag (marking first word of the repair phase)
edit term, or fluent word by the disfluency detector- we concatenate the
disfluency tags with the word vectors to create the input for text-based LSTM.

##  3 Experiments

###  3.1 Data

The ADReSS challenge’s data consists of speech recordings and transcripts of
spoken picture descriptions gathered from participants via the Boston
Diagnostic Aphasia Exam’s Cookie Theft picture [14]. The training set includes
108 subjects, and the state of the subjects is assessed on the basis of the
MMSE score. MMSE is a commonly used cognitive function test for older people.
It involves orientation, memory, language, and visual-spatial skills tests.
Scores of 25-30 out of 30 are considered as normal, 21-24 as mild, 10-20 as
moderate and $<$10 as severe impairment.

The total number of speech segments each participant had generated was 24.86
on average. The annotations for the test set were not included in the public
release of the ADReSS Challenge, so all models were tested on both the
development and test set. The data is pre-processed acoustically and is
balanced in terms of age and gender.

###  3.2 Implementation and Metrics

We set up our model to learn the most useful information from modalities for
predicting AD. All experiments are carried out without being conditioned on
the identity of the speaker. The sizes of layers and the learning rates are
calculated by grid search on validation test. The LSTM models were trained
using ADAM [23] with a learning rate of 0.0001. For the loss function we used
Binary Cross-Entropy to model binary outcomes, and Mean Square Error (MSE) to
model regression outcomes. For binary classification of AD and non-AD, we
report accuracy, precision, recall, and F1 scores and for the MMSE prediction
task we report the Root Mean Square Error (RMSE).

###  3.3 Baseline Models

We compare the performance of our models to the ADReSS Challenge baseline [14]
with an ensemble of audio features which was provided with the dataset. The
baseline classification experiments were different methods of linear
discriminant analysis (LDA), decision trees (DT), and support vector machines
(SVM). The baseline regression experiments were different methods of DT,
gaussian process regression (GPR), and SVM.

Table 1: Result of the AD classification and regression experiments with our
models in cross validation

Models | Features | Accuracy | RMSE
---|---|---|---
LSTM | Acoustic | 0.64 | 6.01
LSTM | Lexical | 0.69 | 5.42
LSTM | Lexical+ Dis | 0.73 | 5.08
LSTM with Gating | Acoustic + Lexical | 0.76 | 5.01
LSTM with Gating | Acoustic + Lexical + Dis | 0.77 | 4.98

Table 2: Result of the AD classification and regression experiments with our
models against baseline models on test set

Models | Features | Accuracy | RMSE
---|---|---|---
Baseline ([14])
LDA | Acoustic | 0.625 | -
DT | Acoustic | 0.625 | 6.14
SVM | Acoustic | 0.563 | 6.12
GPR | Acoustic | - | 6.33
Our Models
LSTM | Acoustic | 0.666 | 5.93
LSTM | Lexical | 0.708 | 5.45
LSTM | Lexical + Dis | 0.729 | 4.88
LSTM with Gating | Acoustic + Lexical | 0.771 | 4.57
LSTM with Gating | Acoustic + Lexical + Dis | 0.792 | 4.54

##  4 Results

In Table 1, we present our proposed model’s performance in a cross-validation
setting and in Table 2 against that of baselines models on AD detection on the
provided test set. For AD detection, our proposed LSTM model with gating and
disfluency features achieves an accuracy of 0.792 and RMSE of 4.54,
outperforming all the baselines. The overall findings confirm our assumption
that a model with a gating structure can more efficiently minimise the errors
and noise of the individual modalities.

Effect of disfluency features We found that disfluency tags help as features
in AD detection. Adding disfluency features to the lexical features lead to
improvement in both unimodal (ACC 0.70 vs. 0.72; RMSE 5.45 vs. 4.88) and
multimodal models (ACC 0.77 vs. 0.79; RMSE 4.57 vs. 4.54).

Effect of multimodality The multimodal LSTM with gating model outperforms the
single modality AD detection models in both the classification and regression
tasks. A performance increase is obtained by combining textual and audio
modalities with gating over single modality models (ACC 0.72 vs. 0.79; RMSE
4.88 vs. 4.54). Adding audio features improves performance despite having
different steps and timesteps inputs for each LSTM branch. In terms of our
competitor baselines (without the information from the manual transcripts),
multimodal classifiers performed better than all the baseline models,
indicating the potential benefits of multimodal fusion in AD detection. We
found that while the baseline audio-based models have some discriminative
capacity, sequence modelling is more accurate (ACC scores 0.67 vs. 0.63) and
has lower (better) RMSE (5.93 vs. 6.12) for predicting AD.

For AD classification, the text features alone are more informative than the
audio features, as using only the text modality gives a better AD prediction
than utilizing unimodal audio modality sequentially (Acc scores 0.73 vs. 0.67;
RMSE 4.88 vs. 5.93).

We can see that all models provide more accurate results on the test set than
in cross validation. LSTM with gating models accuracy improved more than other
models on the test set (RMSE 4.54 and 4.57 vs. 4.98 and 5.01).

Error analysis The results in Table 3 show that the LSTM model with gating and
disfluency features obtains the highest precision and recall for both AD and
non-AD classes. The model achieves F1 scores of 0.7826 for AD and 0.8000 for
non-AD. The addition of gating particularly improves the recall of AD class:
the LSTM model with lexical and disfluency features without gating has a
recall 0.6667 for the AD class compared to the 0.7500 achieved with gating,
while its 0.7910 recall for the non-AD class is not as far beneath the 0.8333
achieved by the full gating model. Depending on the application the model is
used for, false negatives or false positives for AD detection will be more or
less desirable, but as it stands our full gating model considerably reduces
the false negatives of diagnosis whilst still marginally reducing the false
positives.

Table 3: Results of AD classification task on test set

Models | Class | Precision | Recall | F1 Score | Accuracy
---|---|---|---|---|---
LSTM (Lexical+ Dis) | AD | 0.7619 | 0.6667 | 0.7111 | 0.7292
non-AD | 0.7037 | 0.7910 | 0.7451
LSTM with Gating (Acoustic + Lexical) | AD | 0.7826 | 0.7500 | 0.7660 | 0.7708
non-AD | 0.7600 | 0.7917 | 0.7755
LSTM with Gating (Acoustic + Lexical+ Dis) | AD | 0.8182 | 0.7500 | 0.7826 | 0.7917
non-AD | 0.7692 | 0.8333 | 0.8000

##  5 Conclusions

We have presented a deep multi-modal fusion model that learns the AD
indicators from audio and text modalities as well as disfluency features. We
trained and tested the model on audio and transcript data from individuals
doing a description task under controlled conditions, and modeled the sessions
with an LSTM and feed-forward highway layers as gating mechanism for AD
detection. Our findings indicate that AD can be identified by pure sequential
modelling of a session, with limited information available on the structure of
the description tasks. We also found that markers of disfluency hold
predictive power for identification of AD.

In future work we intend to study a series of language markers associated with
AD severity, as well as interactions between them. In particular, we want to
undertake a more principled approach to lexical markers, disfluency markers in
terms of a study of self-repair and structural markers with a look at
grammatical fluency. Furthermore, we want to find acoustic features that
contribute more to the prediction of AD and have higher correlation with
linguistic information.

##  6 Acknowledgments

Purver is partially supported by the EPSRC under grant EP/S033564/1, and by
the European Union’s Horizon 2020 programme under grant agreements 769661
(SAAM, Supporting Active Ageing through Multimodal coaching) and 825153
(EMBEDDIA, Cross-Lingual Embeddings for Less-Represented Languages in European
News Media). The results of this publication reflect only the authors’ views
and the Commission is not responsible for any use that may be made of the
information it contains.

## References

  * [1] A. Burns and S. Iliffe, “Alzheimer’s disease,” _B M J_ , vol. 338, no. 7692, pp. 467–471, 2 2009.
  * [2] D. Kempler, _Neurocognitive disorders in aging_. Sage, 2005.
  * [3] K. López-de Ipiña, J.-B. Alonso, C. M. Travieso, J. Solé-Casals, H. Egiraun, M. Faundez-Zanuy, A. Ezeiza, N. Barroso, M. Ecay-Torres, P. Martinez-Lage _et al._ , “On the selection of non-invasive methods based on speech analysis oriented to automatic alzheimer disease diagnosis,” _Sensors_ , vol. 13, no. 5, pp. 6730–6745, 2013.
  * [4] K. A. Bayles and D. R. Boone, “The potential of language tasks for identifying senile dementia,” _Journal of Speech and Hearing Disorders_ , vol. 47, no. 2, pp. 210–217, 1982.
  * [5] M. F. Folstein, S. E. Folstein, and P. R. McHugh, ““mini-mental state”: a practical method for grading the cognitive state of patients for the clinician,” _Journal of psychiatric research_ , vol. 12, no. 3, pp. 189–198, 1975.
  * [6] M. F. Weiner, K. E. Neubecker, M. E. Bret, and L. S. Hynan, “Language in alzheimer’s disease,” _The Journal of clinical psychiatry_ , vol. 69, no. 8, p. 1223, 2008.
  * [7] G. Szatloczki, I. Hoffmann, V. Vincze, J. Kalman, and M. Pakaski, “Speaking in alzheimer’s disease, is that an early sign? importance of changes in language abilities in alzheimer’s disease,” _Frontiers in aging neuroscience_ , vol. 7, p. 195, 2015.
  * [8] J. Weiner, C. Herff, and T. Schultz, “Speech-based detection of alzheimer’s disease in conversational german.” in _INTERSPEECH_ , 2016, pp. 1938–1942.
  * [9] E. Ambrosini, M. Caielli, M. Milis, C. Loizou, D. Azzolino, S. Damanti, L. Bertagnoli, M. Cesari, S. Moccia, M. Cid _et al._ , “Automatic speech analysis to early detect functional cognitive decline in elderly population,” in _2019 41st Annual International Conference of the IEEE Engineering in Medicine and Biology Society (EMBC)_. IEEE, 2019, pp. 212–216.
  * [10] W. Jarrold, B. Peintner, D. Wilkins, D. Vergryi, C. Richey, M. L. Gorno-Tempini, and J. Ogar, “Aided diagnosis of dementia type through computer-based analysis of spontaneous speech,” in _Proceedings of the Workshop on Computational Linguistics and Clinical Psychology: From Linguistic Signal to Clinical Reality_ , 2014, pp. 27–37.
  * [11] S. Abel, W. Huber, and G. S. Dell, “Connectionist diagnosis of lexical disorders in aphasia,” _Aphasiology_ , vol. 23, no. 11, pp. 1353–1378, 2009\.
  * [12] G. Gosztolya, V. Vincze, L. Tóth, M. Pákáski, J. Kálmán, and I. Hoffmann, “Identifying mild cognitive impairment and mild alzheimer’s disease based on spontaneous speech using asr and linguistic features,” _Computer Speech & Language_, vol. 53, pp. 181–197, 2019.
  * [13] M. Rohanian, J. Hough, M. Purver _et al._ , “Detecting depression with word-level multimodal fusion,” _Proc. Interspeech 2019_ , pp. 1443–1447, 2019.
  * [14] S. Luz, F. Haider, S. de la Fuente, D. Fromm, and B. MacWhinney, “Alzheimer’s dementia recognition through spontaneous speech: The ADReSS Challenge,” 2020. [Online]. Available: https://arxiv.org/abs/2004.06833
  * [15] J. Pennington, R. Socher, and C. D. Manning, “Glove: Global vectors for word representation,” in _Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP)_ , 2014, pp. 1532–1543.
  * [16] G. Degottex, J. Kane, T. Drugman, T. Raitio, and S. Scherer, “Covarep—a collaborative voice analysis repository for speech technologies,” in _2014 ieee international conference on acoustics, speech and signal processing (icassp)_. IEEE, 2014, pp. 960–964.
  * [17] R. K. Srivastava, K. Greff, and J. Schmidhuber, “Highway networks,” _arXiv preprint arXiv:1505.00387_ , 2015.
  * [18] E. A. Schegloff, G. Jefferson, and H. Sacks, “The preference for self-correction in the organization of repair in conversation,” _Language_ , vol. 53, no. 2, pp. 361–382, 1977.
  * [19] W. J. Levelt, “Monitoring and self-repair in speech,” _Cognition_ , vol. 14, no. 1, pp. 41–104, 1983.
  * [20] E. E. Shriberg, “Preliminaries to a theory of speech disfluencies,” Ph.D. dissertation, Citeseer, 1994.
  * [21] J. Hough and D. Schlangen, “Recurrent neural networks for incremental disfluency detection,” ser. Interspeech 2015, 2015, pp. 849–853.
  * [22] ——, “Joint, incremental disfluency detection and utterance segmentation from speech,” in _Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 1, Long Papers_ , 2017, pp. 326–336.
  * [23] D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” _arXiv preprint arXiv:1412.6980_ , 2014.