# Formatting Instructions for ICLR 2024
Conference Submissions

Manh Luong1, Khai Nguyen2, Nhat Ho2, Dinh Phung1, Gholamreza Haffari1, Lizhen
Qu1
1Monash University, Australia, 2 University of Texas at Austin, USA
<EMAIL_ADDRESS>
<EMAIL_ADDRESS>
# Revisiting Deep Audio-Text Retrieval through the Lens of Transportation

Manh Luong1, Khai Nguyen2, Nhat Ho2, Dinh Phung1, Gholamreza Haffari1, Lizhen
Qu1
1Monash University, Australia, 2 University of Texas at Austin, USA
<EMAIL_ADDRESS>
<EMAIL_ADDRESS>
###### Abstract

The Learning-to-match (LTM) framework proves to be an effective inverse
optimal transport approach for learning the underlying ground metric between
two sources of data, facilitating subsequent matching. However, the
conventional LTM framework faces scalability challenges, necessitating the use
of the entire dataset each time the parameters of the ground metric are
updated. In adapting LTM to the deep learning context, we introduce the mini-
batch Learning-to-match (m-LTM) framework for audio-text retrieval problems.
This framework leverages mini-batch subsampling and Mahalanobis-enhanced
family of ground metrics. Moreover, to cope with misaligned training data in
practice, we propose a variant using partial optimal transport to mitigate the
harm of misaligned data pairs in training data. We conduct extensive
experiments on audio-text matching problems using three datasets: AudioCaps,
Clotho, and ESC-50. Results demonstrate that our proposed method is capable of
learning rich and expressive joint embedding space, which achieves SOTA
performance. Beyond this, the proposed m-LTM framework is able to close the
modality gap across audio and text embedding, which surpasses both triplet and
contrastive loss in the zero-shot sound event detection task on the ESC-50
dataset. Notably, our strategy of employing partial optimal transport with
m-LTM demonstrates greater noise tolerance than contrastive loss, especially
under varying noise ratios in training data on the AudioCaps dataset. Our code
is available at https://github.com/v-manhlt3/m-LTM-Audio-Text-Retrieval

##  1 Introduction

Audio-text matching is a crucial task that has tremendous applications, such
as audio-text searching (Elizalde et al., 2019; Koepke et al., 2022), audio
captioning (Kim et al., 2019), and text-to-audio generation (Liu et al.,
2023). The mainstream approaches for audio-text matching involve the
construction of a shared embedding space that encompasses both audio and text
modalities. The shared embedding space should ideally be endowed with richness
and expressiveness. The primary objective of learning a joint audio-text
embedding space is to establish a seamless alignment between all text captions
and their corresponding audio counterparts. Moreover, it is crucial that this
embedding space should minimize the modality gap between audio and text
embeddings, thereby enhancing its transferability to downstream tasks (Liang
et al., 2022).

The widely adopted approaches for learning the joint embedding space are
contrastive learning (Jia et al., 2021; Radford et al., 2021; Mei et al.,
2022; Wu et al., 2022b) , and triplet loss (Wei et al., 2021). Both methods
have shown their success in learning the joint embedding space cross
modalities. Even though they are effective in training cross-modal matching
models, these methods only aim to minimize the ground metric, either Euclidean
distance or cosine similarity, for an anchor point at a time and push away all
negative samples. They treat all negative samples equally, therefore, they
might ignore the geometric properties of the joint embedding space. As a
consequence, a suboptimal metric may make the relative distances between a set
of positively aligned embeddings similar to those between misaligned
embeddings. Furthermore, (Liang et al., 2022) showed that contrastive learning
preserves the modality gap in the joint embedding space. Therefore,
contrastive learning hinders the transferability of the joint embedding space
to downstream tasks.

An additional hurdle encountered in the process of acquiring a joint embedding
space across modalities is the presence of noisy data, which is prevalent in
audio. The noisy data for multimodal matching is misaligned data pairs due to
the data collection process. It is desirable that the examples in a training
dataset for multimodal matching are precisely aligned across modalities. If
the multimodal training data is collected from the Internet (Jia et al., 2021;
Sharma et al., 2018), it is unavoidable to include substantial noisy
correspondence pairs as ground truth. To address this issue, (Huang et al.,
2021) introduces the Noise Correspondence Rectifier approach for image-text
matching. However, this technique heavily depends on a particular choice of
neural network architectures, therefore, it is less amenable to
straightforward adaptation with alternative tasks.

To deal with the aforementioned issues, we propose the minibatch Learning-to-
Match (m-LTM)framework to learn the joint embedding space across audio and
text modalities through the lens of optimal transport. Specifically, we
formulate the cross-modal matching problem as solving an optimal transport
problem with entropic regularization in the embedding space to match the point
sets between audios and texts. We derive a scalable solution to the problem
based on mini-batches to ensure that it scales up to large audio-text
datasets. The resulting optimal plan is an implicit ranking matrix in which
each row points out the best match between two modalities and ranks all
negative points from one modality to the other. To further reduce the modality
gap between audio and text embeddings, we replace the widely used Euclidean
distance with Mahalanobis distance (Dhouib et al., 2020) as the ground metric,
which is introduced for the first time for audio-text matching. To mitigate
the prevalent alignment noises in the training data for the first time for
audio-text matching, we address the problem by using Partial Optimal Transport
(POT) (Chapel & Alaya, 2020). To sum up, our key contributions are three-fold:

  1. 1.

We adapt the learning-to-match framework (Li et al., 2018) to the deep
learning setting based on a subsampling training procedure for audio-text
matching, which leads to the mini-batch learning-to-match framework. We
leverage the soft-matching of the entropic optimal plan and Mahalanobis
distance as a ground metric to learn a shared embedding space. Our approach
facilitates the acquisition of a shared embedding space that is rich and
expressive, thus we acquire SOTA performance in audio-text retrieval tasks on
AudioCaps and Clotho datasets.

  2. 2.

The mini-batch m-LTM framework is capable of bridging the modality gap,
therefore, it is more transferable to downstream tasks than triplet and
contrastive learning. We evaluate the modality gap across audio and text
embedding on three datasets: AudioCaps, Clotho, and ESC50. We also carry out
zeroshot sound event detection on ESC50 dataset to illustrate the
transferability of the m-LTM framework.

  3. 3.

We propose a variant of the m-LTM framework which utilizes Partial Optimal
Transport (POT) to deal with noisy training data circumstances. Our
experiments demonstrate that the minibatch LTM with POT is more noise-
tolerance than triplet and contrastive loss on AudioCaps dataset.

##  2 Preliminary

###  2.1 Deep Audio-text Retrieval

Audio-text retrieval is to learn the cross-modality alignment between audio
and text captions. The alignment can be achieved by learning a joint embedding
space across audio and text modalities. Recently, contrastive learning has
been the most effective learning method for learning expressive cross-modal
embedding space (Jia et al., 2021; Wu et al., 2022b; Radford et al., 2021).

Training. Denote $D=\\{(x_{i},y_{i})\\}_{i=1}^{n}$ as all pairs of audio and
caption in the training dataset. The goal is to learn a joint embedding space
between audio and text embedding. Given an audio encoder $f_{\theta}(.)$ and a
text encoder $g_{\phi}(.)$, the objective function for learning the embedding
space is:

$\displaystyle\underset{\theta,\phi}{\min}\;{\mathbb{E}_{(X^{b},Y^{b})\sim
D}}\bigg{[}-\sum_{i=1}^{b}\log{\bigg{(}\frac{\exp(s(f_{\theta}(x_{i}),g_{\phi}(y_{i}))/\tau)}{\sum_{k=1}^{b}\exp(s(f_{\theta}(x_{i}),g_{\phi}(y_{k}))/\tau)}\bigg{)}}$
(1)
$\displaystyle-\sum_{i=1}^{b}\log{\bigg{(}\frac{\exp(s(f_{\theta}(x_{i}),g_{\phi}(y_{i}))/\tau)}{\sum_{k=1}^{b}\exp(s(f_{\theta}(x_{k}),g_{\phi}(y_{i}))/\tau)}\bigg{)}}\bigg{]}$

, where $(X^{b},Y^{b})$ is a minibatch of $b$ audio-caption pairs drawn from
training data $D$, $\tau$ is the temperature hyperparameter, and $s(.,.)$ is
the cosine similarity between audio and text embedding

Retrieval. Given $n^{\prime}$ audio $X_{test}=\\{x_{i}\\}_{i=1}^{n^{\prime}}$
and $m^{\prime}$ caption $Y_{test}=\\{y_{j}\\}_{j=1}^{m^{\prime}}$, the
ranking matrix $R(X_{test},Y_{test})$ between two sets of audio and caption
can be retrieved by computing pairwise cosine similarity between audio and
text embedding. For example, the correspondence caption of a given audio
$x_{i}$ is

$\hat{y}=\underset{y_{j}\in
Y_{test}}{\operatorname*{arg\,max}}\frac{\exp(s(f_{\theta}(x_{i}),g_{\phi}(y_{j})))}{\sum_{k=1}^{m^{\prime}}\exp(s(f_{\theta}(x_{i}),g_{\phi}(y_{k})))}$
(2)

###  2.2 Learning to Match

Entropic optimal transport. Given two empirical probability measures
$P_{X}=\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}}$ and
$P_{Y}=\frac{1}{m}\sum_{j=1}^{m}\delta_{y_{j}}$ where $x_{i}\in\mathcal{X}$
for $i=1,\ldots,n$ and $y_{j}\in\mathcal{Y}$ for $j=1,\ldots,m$, the entropic
optimal transport between $P_{X}$ and $P_{Y}$ (Cuturi, 2013) with the ground
metric $c:\mathcal{X}\times\mathcal{Y}\to\mathbb{R}^{+}$ is defined as:

$\displaystyle\pi_{\epsilon,c}^{X,Y}=\underset{\pi\in\Pi(P_{X},P_{Y})}{\operatorname*{arg\,min}}\sum_{i=1}^{n}\sum_{j=1}^{n}\pi_{ij}c(x_{i},y_{j})-\epsilon\sum_{i=1}^{n}\sum_{j=1}^{m}\pi_{ij}\log\pi_{ij},$
(3)

where $\Pi(P_{X},P_{Y})=\\{\pi\in\mathbb{R}^{n\times
m}|\pi\mathbf{1}_{m}=\mathbf{1}_{n}/n,\pi^{T}\mathbf{1}_{n}=\mathbf{1}_{m}/m\\}$
denotes the set of transportation plans or couplings between $P_{X}$ and
$P_{Y}$. When $\epsilon>0$, the optimization problem Equation 3 can be solved
by using the differentiable Sinkorn algorithm (Marshall & Olkin, 1968).

Learning to match. In practice, we are given $\\{(x_{i},y_{i})\\}_{i=1}^{n}$
which is all pairs of audio and caption. Therefore, we would like to infer the
underlying ground cost metric $c_{\phi}(x_{i},y_{j})$ (parameterized by
$\phi$). Learning-to-match (LTM) (Li et al., 2019) is one of the most
effective solutions for doing such inference. The LTM framework considers the
following optimization:

$\displaystyle\inf_{c\in\mathcal{C}}\text{KL}(\hat{\pi}||\pi_{\epsilon,c}^{X,Y})),$
(4)

where $\mathcal{C}$ is a set of ground metrics, $X$ denotes the set of
$\\{x_{1},\ldots,x_{n}\\}$ and $Y$ denotes the set of
$\\{y_{1},\ldots,y_{n}\\}$, and $\hat{\pi}$ is a matrix of size $n\times n$
such that $\hat{\pi}_{ii}=1/n$ and $\hat{\pi}_{ij}=0$ ($i\neq j$) for
$i=1,\ldots,n$ and $j=1,\ldots,n$.

Retrieval. Given $n^{\prime}$ audio $X_{test}=\\{x_{i}\\}_{i=1}^{n^{\prime}}$,
$m^{\prime}$ caption $Y_{test}=\\{y_{j}\\}_{j=1}^{m^{\prime}}$, and the
optimized ground metric $c$ from the Equation. 4, the ranking matrix
$R(X_{test},Y_{test})$ between two sets of audio and text caption is the
optimal solution of the Equation. 3 between
$P_{X_{test}}=\frac{1}{n^{\prime}}\sum_{i=1}^{n^{\prime}}{\delta_{x_{i}}}$,
where $x_{i}\in X_{test}$ and
$P_{Y_{test}}=\frac{1}{m^{\prime}}\sum_{j=1}^{m^{\prime}}{\delta_{y_{j}}}$,
where $y_{j}\in Y_{test}$. For example, the correspondence caption of a given
audio $x_{i}$ is

$\hat{y}=\underset{y_{j}\in
Y_{test}}{\operatorname*{arg\,max}}\;\pi_{\epsilon,c}^{x_{i},y_{j}}$ (5)

, where $\pi_{\epsilon,c}^{x_{i},y_{j}}$ is the coupling value of i-th audio
and j-th text from the optimal plan. The same retrieval procedure can be done
for text-to-audio retrieval.

##  3 Deep Audio-text Matching via Mini-batch Learning to Match

This section introduces our proposed method, the m-LTM framework, for the
audio-text matching task and a variant to deal with noisy correspondence
training data. Our model consists of the audio encoder $f_{\theta}(.)$, the
text encoder $g_{\phi}(.)$, and the interaction matrix $M$. Our proposed
framework utilizes the soft-matching from entropic optimal transport and
incorporates the Mahanalobis distance as a ground metric to learn a flexible
ground metric to solve audio-text matching problems.

###  3.1 Mini-batch Learning to Match

By abuse of notation, we denote $(X^{b},Y^{b})\sim D$ as sampling uniformly a
mini-batch which contains $b$ pairs of audio and caption. The mini-batch
learning-to-match framework is defined as follows:

######  Definition 1.

Given a paired training data $D=\\{(x_{i},y_{i})\\}_{i=1}^{n}$, the mini-batch
learning-to-match (m-LTM) framework solves the following optimization problem:

$\displaystyle\inf_{c\in\mathcal{C}}\mathbb{E}_{(X^{b},Y^{b})\sim
D}[\text{KL}(\hat{\pi}^{b}||\pi_{\epsilon,c}^{X^{b},Y^{b}})],$ (6)

where $b\geq 2$ is the mini-batch size, $\pi_{\epsilon,c}^{X^{b},Y^{b}}$ is
the optimal transportation plan from Equation 3 with the corresponding
empirical measures over mini-batches $X^{b}$ and $Y^{b}$, and $\hat{\pi}_{b}$
is a matrix of size $b\times b$ such that $\hat{\pi}_{ii}=1/b$ and
$\hat{\pi}_{ij}=0$ ($i\neq j$) for $i=1,\ldots,b$ and $j=1,\ldots,b$.

Neural network parameterization. As in the preliminary section, we
parameterize the ground metric
$c_{\theta,\phi}(x,y)=d(f_{\theta}(x_{i}),g_{\phi}(y_{j})))$ where
$f_{\theta}:\mathcal{X}\to\mathcal{Z}$ ($\theta\in\Theta$) and
$g_{\phi}\to\mathcal{Z}$ ($\phi\in\Phi$) are two neural networks parameterized
by $\theta$ and $\phi$ respectively, and
$d:\mathcal{Z}\times\mathcal{Z}\to\mathbb{R}^{+}$ is a chosen metric e.g.,
$\mathbb{L}_{2}$ norm. Therefore, the set of ground metric $\mathcal{C}$
becomes the set of all possible weights of two neural networks i.e.,
$\Theta\times\phi$. As a result, Equation 6 is equivalent to:

$\displaystyle\min_{(\theta,\phi)\in\Theta\times\Phi}\mathbb{E}_{(X^{b},Y^{b})\sim
D}[\text{KL}(\hat{\pi}^{b}||\pi_{\epsilon,c_{\theta,\phi}}^{X^{b},Y^{b}})].$
(7)

Stochastic gradient estimation. To solve the optimization in Equation 7, we
perform stochastic gradient descent algorithm to update $\theta$ and $\phi$.
In particular, we sample $B\geq 1$ mini-batches
$(X^{b},Y^{b})_{1},\ldots,(X^{b},Y^{b})_{B}\overset{i.i.d}{\sim}D$, then form
the following stochastic estimation of gradients:

$\displaystyle\nabla_{(\theta,\phi)}\mathbb{E}_{(X^{b},Y^{b})\sim
D}[\text{KL}(\hat{\pi}^{b}||\pi_{\epsilon,c_{\theta,\phi}}^{X^{b},Y^{b}})]\approx\frac{1}{B}\sum_{i=1}^{B}\nabla_{(\theta,\phi)}\text{KL}(\hat{\pi}^{b}||\pi_{\epsilon,c_{\theta,\phi}}^{(X^{b},Y^{b})_{i}})].$
(8)

Since this estimation is unbiased, the stochastic optimization procedure can
converge to a local minima. In practice, we often set $B=1$ for having a
computationally fast estimation.

Retrieval. During the inference phase, it is noteworthy that the quantity of
audio and caption data is considerably smaller compared to the extensive
training dataset. Moreover, the conventional evaluation protocol for retrieval
necessitates that each query must be compared to the entirety of available
items in the test set to compute the evaluation metric. Consequently,
employing the minibatch configuration during the retrieval phase does not
yield discernible advantages. For more detailed information on retrieval,
please refer to Section 2.2.

###  3.2 Mahalanobis-Enhanced Ground Metric

Although the parameterization of the ground metric in Section 3.1 i.e.,
$c_{\theta,\phi}(x,y)=d(f_{\theta}(x_{i}),g_{\phi}(y_{j})))$ has been widely
used in practice, it still has a limitation. In particular, the embedding from
the two encoders $f_{\theta}(.)$ and $g_{\phi}(.)$ might not be well aligned
in terms of scaling in each dimension. To partially address the issues, we
propose a more generalized family of ground metric which is based on the
Mahalanobis distance.

######  Definition 2.

Given two encoder functions $f_{\theta}:\mathcal{X}\to\mathcal{Z}$ and
$g_{\phi}:\mathcal{Y}\to\mathcal{Z}$, a metric
$d:\mathcal{Z}\times\mathcal{Z}\to\mathbb{R}^{+}$, the Mahalanobis enhanced
ground metric is defined as:

$\displaystyle
c_{\theta,\phi,M}(x,y)=\sqrt{(f_{\theta}(x_{i})-g_{\phi}(y_{j}))^{\top}M(f_{\theta}(x_{i})-g_{\phi}(y_{j}))},$
(9)

for $\theta\in\Theta$ and $\phi\in\Phi$ which are spaces of parameters and M
is a positive definite matrix.

Mini-batch learning to match with Mahalanobis-Enhanced Ground Metric. By using
the family of Mahalanobis-Enhanced ground metrics in Definition 2, the m-LTM
objective in Definition 1 becomes:

$\displaystyle\min_{(\theta,\phi,M)\in\Theta\times\Phi\times\mathcal{M}}\mathbb{E}_{(X^{b},Y^{b})\sim
D}[\text{KL}(\hat{\pi}^{b}||\pi_{\epsilon,c_{\theta,\phi,M}}^{X^{b},Y^{b}})],$
(10)

where $\mathcal{M}$ is the set of all possible positive definite matrices
e.g., $x^{\top}Mx>0$ for all $x\in\mathcal{Z}$.

Hybrid stochastic gradient descent. the optimization problem in Equation 10
consists of three parameters $\theta,\phi,$ and $M$. In contrast to $\theta$
and $\phi$ which are unconstrained, $M$ is a constrained parameter. Therefore,
we propose to use a hybrid stochastic gradient descent algorithm. In
particular, we still update $\theta,\phi$ using the estimated gradients in
Equation 8. However, we update $M$ using the projected gradient descent update
(Rockafellar, 1997). We first estimate the stochastic gradient with respect to
$M$:

$\displaystyle\nabla_{M}\mathbb{E}_{(X^{b},Y^{b})\sim
D}[\text{KL}(\hat{\pi}^{b}||\pi_{\epsilon,c_{\theta,\phi,M}}^{X^{b},Y^{b}})]\approx\frac{1}{B}\sum_{i=1}^{B}\nabla_{M}\text{KL}(\hat{\pi}^{b}||\pi_{\epsilon,c_{\theta,\phi,M}}^{(X^{b},Y^{b})_{i}})].$
(11)

After that, we update $M=\text{Proj}(F(M,\nabla M))$ where $F(M,\nabla M)$
denotes the one-step update from a chosen optimization scheme e.g., Adam,
gradient descent, and so on, and $\text{Proj}(\cdot)$ denotes the projecting
function that maps a unconstrained matrix to the space of positive definite
matrix $\mathcal{M}$. In greater detail,
$\text{Proj}(A)=U\text{diag}(\bar{S})V^{T}$ with $U,S,V=SVD(A)$ and
$\bar{S}=\text{diag}(\max\\{0,\sigma_{1}\\},...,\max\\{0,\sigma_{n}\\})$. We
refer the reader to the training algorithm in the Algorithm. 1 in the Appendix
A.2.

###  3.3 Partial OT for Noisy Correspondence

Noisy correspondence setting. We first introduce the noisy correspondence
setting for the audio-text matching task. That is, given the training data
$D=\\{(x_{i},y_{i})\\}_{i=1}^{N}$ where $N$ is the number of training samples,
a proportion of training data $N_{cor}$, $N_{cor}<N$, is corrupted, for
instance, due to the data collection process. The amount of data corruption is
unknown, and we do not know which samples are corrupt, thus presenting a
challenging learning scenario. Due to the noisy correspondence phenomenon, the
empirical matching $\hat{\pi}$ is now corrupted to an incomplete matching
$\bar{\pi}$, in which there is no matching in some rows. Assuming an absence
of noise in the original training data, we introduce a noisy correspondence
training dataset by randomly shuffling a portion of audio-text pairs within
the dataset. We denote a random variable $z\in\\{0,1\\}$ which is sampled from
a binomial distribution $Binomial(N,\frac{N_{cor}}{N})$, if $z=1$ indicates
the audio-text pair is shuffled. The training data is now
$\tilde{D}=\\{(z_{i},x_{i},y_{i})\\}_{i=1}^{N}$, where $z_{i}\sim
Binomial(N,\frac{N_{cor}}{N})$.

Given the absence of knowledge regarding the nature of data corruption, the
training process relies solely on the empirical matching $\hat{\pi}$. The
training objective is to infer the incomplete matching $\bar{\pi}$ from the
noisy training dataset $\tilde{D}$. However, it is important to note that the
m-LTM framework may encounter challenges when attempting to recover the
incomplete matching $\bar{\pi}$. This difficulty arises due to the constraint
imposed by transportation preservation, which dictates that each row must be
associated with at least one column.

To solve the aforementioned issue, we propose to use Partial OT, which relaxes
the transportation preservation constraint, to mitigate the harmfulness of
noisy empirical matching for approximating the incomplete matching
$\bar{\pi}$. The objective function 10 is rewritten as

$\displaystyle\min_{(\theta,\gamma,M)\in\Theta\times\Phi\times\mathcal{M}}\mathbb{E}_{(\tilde{X}^{b},\tilde{Y}^{b})\sim\tilde{D}}[\text{KL}(\hat{\pi}^{b}||\pi_{s,\epsilon,c_{\theta,\phi,M}}^{\tilde{X}^{b},\tilde{Y}^{b}})],,$
(12)

, where $(\tilde{X}^{b},\tilde{Y}^{b})$ is a minibatch sampled from noisy
training data $\tilde{D}$, and
$\pi_{\epsilon,s}^{\tilde{X}^{b},\tilde{Y}^{b}}$ is the optimal solution of
the equation

$\displaystyle\pi_{s,\epsilon,c_{\theta,\gamma,M}}^{\tilde{X}^{b},\tilde{Y}^{b}}=\underset{\pi\in\Pi_{s}(P_{\tilde{X}^{b}},P_{\tilde{Y}^{b}})}{\operatorname*{arg\,min}}\sum_{i=1}^{b}\sum_{j=1}^{b}\pi_{ij}c(x_{i},y_{j})-\epsilon\sum_{i=1}^{b}\sum_{j=1}^{b}\pi_{ij}\log\pi_{ij},$
(13)

where
$\Pi_{s}(P_{\tilde{X}^{b}},P_{\tilde{Y}^{b}})=\\{\pi\in\mathbb{R}_{+}^{b\times
b}|\pi\mathds{1}\leq P_{\tilde{X}^{b}},\pi^{\top}\mathds{1}\leq
P_{\tilde{Y}^{b}},\mathds{1}\mathbf{\pi}^{\top}\mathds{1}=s\\}$. The partial
OT optimization can be solved by using Bregman projection (Benamou et al.,
2015). Intuitively, the partial transportation plan should be such that true
correspondence audio-text pairs within a minibatch are mapped to each other,
while noisy correspondence audio-text pairs are discarded from the
transportation plan. We conducted an experiment to examine the relationship
between the transportation mass and noise ratio in training data in the
Section. 5.2

##  4 Related Works

Cross-modal Matching. Cross-modal matching is to learn a shared representation
space (Wang et al., 2013) between two modalities like image-text (Jia et al.,
2021; Radford et al., 2021; Wei et al., 2020) or audio-text (Wu et al., 2022b;
Deshmukh et al., 2022; Oncescu et al., 2021; Mei et al., 2022). The most
popular approach for cross-modal matching is metric learning, such as triplet
loss (Wei et al., 2020) and contrastive loss (Radford et al., 2021; Yang et
al., 2022). The cross-modal training procedure exposes a great benefit of the
transferability to downstream tasks (Zeng et al., 2022) . (Xu et al., 2021)
shows that contrastive cross-modal pretraining procedure can be used to train
a single model for both text and video understanding tasks without fine-
tuning. However, cross-modal training is beneficial for downstream tasks, this
training procedure requires a huge amount of data with high alignment quality.
(Huang et al., 2021) introduces noisy correspondence training circumstances
for image-text matching. This work demonstrates that their proposed method is
capable of dealing with high-level image-caption mismatching, but it heavily
depends on neural network architecture. To handle the noisy correspondence
training data, we introduce a new loss function that is not only capable of
mitigating the harmfulness of label mismatch for cross-modal matching but also
easy to incorporate into existing backbones for cross-modal matching.
Furthermore, we apply a new ground metric, Mahalanobis distance, which is able
to learn a rich and expressive joint embedding space across audio-text
modalities and close the gap between audio and text embedding space.

Optimal transport. Optimal transport (OT) is a powerful mathematical tool to
measure the discrepancy between two probability distributions. Therefore, it
has a wide range of applications, such as domain adaption (Courty et al.,
2017), generative models (Genevay et al., 2018; Nguyen & Ho, 2024; Nguyen et
al., 2024), and representation learning (Chen et al., 2020). Our work is close
to the Inverse Optimal Transport (IOT) problem, which has been well-studied
(Dupuy et al., 2016; Li et al., 2018; Stuart & Wolfram, 2019; Chiu et al.,
2022) by the research community aiming to learn matching between two sources
of interest. (Wu et al., 2022a) proposed the OTTER framework which leverages
entropic optimal plan as soft matching labels to learn many-to-many
relationships in the image-text matching task based on contrastive learning
procedure. (Shi et al., 2023) proposed a unifying understanding for
contrastive learning from an inverse optimal transport view. All previous
works either adapt cosine similarity as the ground metric (Shi et al., 2023)
or study the whole empirical plan of tabular data for learning-to-matching
(Dupuy et al., 2016; Li et al., 2018; Stuart & Wolfram, 2019). Mini-batch
optimal transport has been investigated in Fatras et al. (2020); Nguyen et al.
(2022b; a). In our paper, we adapt the learning-to-match framework (Li et al.,
2018) for the minibatch setting with a more flexible ground cost metric,
Mahalanobis distance, to model the expressive shared embedding space between
audio and text modalities. Mahalanobis distance represents an infinite set of
cost functions in the theoretical analysis of robust optimal transport
optimization (Paty & Cuturi, 2019; Dhouib et al., 2020). However, there is a
gap between theory and application of Mahalanobis distance as a ground metric
for optimal transport optimization due to the difficult constraint of
Mahalanobis distance. We close the gap by successfully applying Mahalanobis
distance for the audio-text matching application. The strict constraint for
learning Mahalanobis distance is relaxed by Projection Gradient Descent.

##  5 Experiments

In this section, we design experiments to examine the effectiveness of our
framework in order to learn a joint embedding space across audio and text
modalities. We conduct the audio-text retrieval task on two datasets:
AudioCaps (Kim et al., 2019) and Clotho (Drossos et al., 2019) to answer the
question of whether our framework is able to learn a rich and expressive joint
embedding space. In addition, we evaluate the modality gap metric (Liang et
al., 2022) in the joint embedding space on three datasets: AudioCaps, Clotho,
and ESC-50 (Piczak, 2015) to illustrate that the m-LTM framework is capable of
bridging the modality gap between audio and text embedding space. We perform
zero-short sound event detection on ESC-50 test set to demonstrate the
transferability of our method. We verify the noise-tolerance capability of the
m-LTM with POT framework in a variant of noise ratio training data on the
AudioCaps dataset. Finally, we conduct ablation studies for hyperparameters
tuning. The implementation details of our framework are provided in Appendix.
A.1

Baselines. We compare against all state-of-the-art models (Mei et al., 2022;
Deshmukh et al., 2022; Oncescu et al., 2021; Wu et al., 2022b) for audio-text
retrieval tasks to illustrate the remarkable performance of our framework. All
the results for these baselines are copied from the reference papers.
Furthermore, we examine our framework with diverse learning objectives, such
as contrastive and triplet loss.

Evaluation metrics. To evaluate our proposed framework, we use the recall at
rank $k(R@k)$ metric, which is the standard evaluation metric for cross-modal
retrieval tasks. $R@k$ reports the percentage of the ground-truth retrieval
within the top-k ranks, thus, higher is better. $R@1,R@5$, and $R@10$ are
shown to compare performance among baselines and our method. We also report
the modality gap metric (Liang et al., 2022) to examine the shared embedding
space. The lower modality gap is better for downstream tasks.

###  5.1 Expressiveness and Transferability of The Joint Embedding Space

Table 1: The comparison of m-LTM framework with baselines on audio-text retrieval task on two benchmark datasets, AudioCaps and Clotho dataset. Dataset | Method | Text->Audio | Audio->Text
---|---|---|---
R@1 | R@5 | R@10 | R@1 | R@5 | R@10
Audiocaps | (Oncescu et al., 2021) | 28.1 | - | 79.0 | 33.7 | - | 83.7
(Mei et al., 2022) | 33.9 | 69.7 | 82.6 | 39.4 | 72 | 83.9
(Deshmukh et al., 2022) | 33.07 | 67.30 | 80.3 | 39.76 | 73.72 | 84.64
(Wu et al., 2022b) | 36.7 | 70.9 | 83.2 | 45.3 | 78 | 87.7
m-LTM(our) | 39.10 | 74.06 | 85.78 | 49.94 | 80.77 | 90.49
Clotho | (Oncescu et al., 2021) | 9.6 | - | 40.1 | 10.7 | - | 40.8
(Mei et al., 2022) | 14.4 | 36.6 | 49.9 | 16.2 | 37.5 | 50.2
(Deshmukh et al., 2022) | 15.79 | 36.78 | 49.93 | 17.42 | 40.57 | 54.26
(Wu et al., 2022b) | 12.0 | 31.6 | 43.9 | 15.7 | 36.9 | 51.3
m-LTM(our) | 16.65 | 39.78 | 52.84 | 22.1 | 44.4 | 56.74
Table 2: Experiments of loss functions on the AudioCaps dataset. There are two audio encoders, ResNet38 and HTSAT, used to extract audio embedding, and the text encoder is the BERT text encoder. | Loss | Text->Audio | Audio->Text
---|---|---|---
R@1 | R@5 | R@10 | R@1 | R@5 | R@10
ResNet38 Audio Encoder | Triplet | 32.2 | 68.2 | 81.6 | 36.1 | 69.2 | 81.4
Contrastive | 33.9 | 69.7 | 82.6 | 39.4 | 72.0 | 83.9
m-LTM(our) | 39.10 | 74.06 | 85.78 | 49.94 | 80.77 | 90.49
HTSAT Audio Encoder | Triplet | 33.87 | 69.91 | 83.23 | 33.95 | 70.03 | 82.46
Contrastive | 34.38 | 71.16 | 84.64 | 35.84 | 71.06 | 83.39
m-LTM(our) | 40.02 | 75.86 | 86.42 | 40.13 | 74.09 | 85.89

We compare our proposed framework with state-of-the-art models in the audio-
caption retrieval task to verify the expressiveness of our proposed framework.
To have a fair comparison, we train all the models and objective functions
with the same amount of training data. As shown in Table. 1, our minibatch
learning-to-match framework outperforms four baseline methods in terms of
$R@1$, $R@5$, and $R@10$ evaluation metrics on two datasets in two retrieval
tasks, audio-to-text and text-to-audio retrievals. Furthermore, we carry out
the experiment to compare the m-LTM objective function with contrastive and
triplet loss on the AudioCaps dataset shown in Table. 2. Our method
illustrates the highest performance for cross-modal retrieval tasks, there is
a significant gap regarding the $R@1$ score compared with the best loss
function, contrastive loss. For the ResNet audio encoder, the utilization of
m-LTM results in a marked and notable augmentation of the $R@1$ score,
showcasing an improvement from $33.9\%$ to $37.97\%$ for the text-to-audio
matching task, and a distinct escalation from $39.4\%$ to $47.43\%$ for the
audio-to-text matching task. This performance increase is also discernible in
the case of the HTSAT audio encoder. We also provide some qualitative results
for audio-text retrieval tasks in Appendix A.3.

Table 3: The zero-shot sound event detection on the ESC50 test set, the $R@1$ score is equivalent to accuracy. Loss | Audio->Sound
---|---
R@1 | R@5 | R@10 | mAP
Triplet | 71.25 | 91.75 | 95.75 | 80.09
Contrastive | 72.25 | 93 | 96.75 | 80.84
m-LTM | 81.0 | 97.0 | 99.25 | 87.57

Table 4: The modality gap between audio and text embedding in the shared embedding space. Lower is better for downstream tasks. Loss | Modality gap($\lVert\vec{\Delta}_{gap}\rVert$)
---|---
AudioCaps | Clotho | ESC50
Triplet | 0.149 | 0.283 | 0.937
Contrastive | 0.181 | 0.266 | 0.922
m-LTM | 0.117 | 0.142 | 0.224

We conduct a zero-shot sound event detection experiment on the ESC50 test set
to examine the transferability of the learned embedding space. We replace the
task of sound event detection with audio-text retrieval, all the classes in
the test set are converted to the template caption as ”This is a sound of
{class}”. We pretrain all models on the AudioCaps dataset on the audio-caption
matching task. We report $R@1$, $R@5$, $R@10$, and mAP scores for triplet
loss, contrastive loss, and the m-LTM loss in Table. 4. The m-LTM loss
acquires the highest $R@1$ score, equivalent to achieving the highest accuracy
for the sound event detection task. Also, there is a large performance gap in
terms of $R@5$, $R@10$, and mAP metrics between m-LTM loss and contrastive
loss. As depicted in Figure. LABEL:fig:tSNE, audio embeddings learned from the
m-LTM loss are clearly clustered by the text sound event classes, which means
the learned embedding space is more expressive for the zero-shot sound event
detection setting. In contrast, audio embeddings learned from contrastive loss
are spread uniformly in the joint embedding space. This observation
consolidates to the result in Table. 4 which is the m-LTM is more transferable
than contrastive loss for zero-shot setting. Furthermore, we study the
modality gap of our proposed loss in Table. 4, the $\vec{\Delta}_{gap}$
modality gap metric between two modalities is computed as the mean of
embedding between two modalities
$\vec{\Delta}_{gap}=\frac{1}{n}\sum_{i=1}^{n}f(x_{i})-\frac{1}{n}\sum_{i=1}^{n}g(y_{i})$.
Our proposed loss achieves the smallest gap between audio and text embedding,
which means it is more transferable to downstream tasks than triplet and
contrastive loss.

###  5.2 Noisy Correspondence Tolerance

Table 5: The performance of learning-to-match and metric learning methods for audio-text retrieval task under the variant ratio of noisy training data.  |  | Text->Audio | Audio->Text
---|---|---|---
Noise | Method | R@1 | R@5 | R@10 | R@1 | R@5 | R@10
| Triplet loss | 23.01 | 54.98 | 69.98 | 28.52 | 58.09 | 70.11
| Contrastive loss | 31.34 | 67.73 | 81.27 | 40.12 | 70.84 | 82.54
| m-LTM | 35.51 | 71.32 | 84.01 | 46.64 | 78.68 | 87.87
20% | m-LTM with POT | 35.92 | 72.28 | 84.11 | 47.12 | 79.2 | 88.19
| Triplet loss | 0.1 | 1.19 | 2.75 | 1.25 | 5.43 | 9.4
| Contrastive loss | 26.68 | 62.98 | 78.18 | 34.69 | 66.66 | 78.99
| m-LTM | 32.58 | 67.75 | 80.89 | 40.31 | 71.16 | 84.57
40% | m-LTM with POT | 33.64 | 69.23 | 82.27 | 42.63 | 73.35 | 86.1
| Triplet loss | 0.1 | 0.52 | 1.06 | 0.1 | 0.52 | 1.46
| Contrastive loss | 20.58 | 53.96 | 70.72 | 27.37 | 58.72 | 75.21
| m-LTM | 25.26 | 59.72 | 75.03 | 34.08 | 66.77 | 79.62
60% | m-LTM with POT | 27.73 | 62.61 | 76.17 | 35.42 | 68.65 | 80.56

To study the robustness of our proposed framework with noisy correspondence,
we carry out the experiment under varying noisy correspondence training data.
As defined in the section 3.3, we use three different noise ratios
$\xi=\\{0.2,0.4,0.6\\}$ of training data. We perform the experiment on the
AudioCaps dataset for four losses: triplet loss, contrastive loss, m-LTM loss,
and m-LTM with POT loss described in section 3.3. First, the binary value
$z_{i}$ is sampled from a $Binomial(N,\xi)$ distribution for each audio-text
pair in the training data, so we have N tuples of training data
$\\{z_{i},x_{i},y_{i}\\}_{i=1}^{N}$. If $z_{i}=1$ we replace the text $y_{i}$
with a random text $\tilde{y}_{i}$ in training data, the noisy training tuple
is $\\{z_{i}=1,x_{i},\tilde{y}_{i}\\}$. The random variable $z\sim
Binomial(N,\xi)$ is only used for data sampling, and it is not used during
training. Table. 5 shows the audio-text retrieval performance under the
variant noise ratio of four loss functions: triplet loss, contrastive loss,
and our proposed methods. The triplet loss is less robust to high levels of
noise, its performance plunges from $28.52\%$ to $1.25\%$ and from $23.01\%$
to $0.1\%$ in terms of $R@1$ score when the noise ratio increases from $20\%$
to $40\%$ for audio-to-text and text-to-audio retrieval, respectively.
Contrastive loss and two variants of our proposed method are robust to a wide
range of noise ratios, however, our proposed losses surpass contrastive loss
for all levels of noise in terms of recall at rank $k$ metrics. There is a
significant gap between contrastive loss and the m-LTM with POT loss on the
$R@1$ metric for two retrieval tasks. At $60\%$ noise of training data, the
m-LTM with POT loss boosts the performance from $20.58\%$ and $27.37\%$ of
contrastive loss to $27.73\%$ and $35.42\%$ in terms of $R@1$ metric for text-
to-audio and audio-to-text, respectively. Furthermore, it is noteworthy that
in both of the two retrieval tasks, the m-LTM with POT consistently
outperforms the conventional m-LTM across all levels of noise. As a result, it
can be inferred that the m-LTM with POT exhibits state-of-the-art (SOTA)
performance in effectively addressing challenging scenarios characterized by
noisy correspondences.

###  5.3 Ablation Study

Table 6: Ablation studies on AudioCaps dataset for choosing the ground metric and hyper-parameter for minibatch learning-to-match loss function. | $\epsilon$ | Text->Audio | Audio->Text
---|---|---|---
R@1 | R@5 | R@10 | R@1 | R@5 | R@10
Maha | 0.05 | 39.10 | 74.06 | 85.78 | 49.94 | 80.77 | 90.49
0.1 | 29.86 | 65.51 | 79.64 | 39.81 | 70.53 | 83.07
0.5 | 9.59 | 31.99 | 48.19 | 10.55 | 34.16 | 46.91
Eucli | 0.05 | 38.13 | 73.47 | 85.43 | 48.15 | 79.10 | 89.23
0.1 | 29.55 | 65.60 | 79.43 | 38.45 | 70.84 | 82.23
0.5 | 10.17 | 33.06 | 48.50 | 13.06 | 34.69 | 48.27

Figure 2: Ablation study on the AudioCaps dataset for transportation mass of
m-LTM with POT.

In this section, we conduct two ablation studies. As illustrated in Table. 6,
Mahalanobis distance as the ground metric for learning-to-match framework
surpasses the Euclidean distance for audio-text retrieval task. Also, the
$\epsilon$ hyper-parameter is crucial in our framework. If $\epsilon$ is too
high, the optimal transportation plan will be a uniform distribution and lead
to a decrease in performance. The value of $\epsilon=0.05$ acquires the
highest performance, thus, all the reported results in this paper using
$\epsilon=0.05$. The second ablation study is the effect of transportation
mass for noisy correspondence training data demonstrated in Figure. 2, we
report the average $R@1$ score for both audio-to-text and text-to-audio
retrieval tasks. For different noise ratios, there is a best transportation
mass value, which corresponds to noise level, to acquire the highest
performance. This observation supports our hypothesis in section. 3.3 which is
choosing the appropriate transportation mass for noisy correspondence
circumstances is essential.

##  6 Conclusion

In this paper, we introduce the mini-batch learning-to-match (m-LTM) framework
for audio-text matching, the m-LTM utilizes the Mahalanobis distance and soft-
matching from the entropic optimal plan to facilitate learning of a shared
embedding space for audio and text modalities. By leveraging soft matching
from the optimal plan, our framework is capable of learning a rich and
expressive joint embedding space across audio and text modality, therefore,
the m-LTM acquires state-of-the-art performance for audio-text retrieval on
AudioCaps and Clotho datasets. Furthermore, we adopt the Mahalobis distance as
the ground metric to reduce the modality gap, thus, our framework is more
transferable to downstream tasks compared to conventional methods relying on
triplet and contrastive loss functions. Finally, we propose a variant of the
m-LTM with POT to alleviate the adverse effects of noisy correspondence in
training data under various noise conditions.

## Acknowledgements

NH acknowledges support from the NSF IFML 2019844 and the NSF AI Institute for
Foundations of Machine Learning.

## References

  * Benamou et al. (2015) Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré.  Iterative bregman projections for regularized transportation problems.  _SIAM Journal on Scientific Computing_ , 37(2):A1111–A1138, 2015.
  * Chapel & Alaya (2020) Laetitia Chapel and Mokhtar Z. Alaya.  Partial optimal tranport with applications on positive-unlabeled learning.  _arXiv: Machine Learning_ , 2020.
  * Chen et al. (2022) Ke Chen, Xingjian Du, Bilei Zhu, Zejun Ma, Taylor Berg-Kirkpatrick, and Shlomo Dubnov.  Hts-at: A hierarchical token-semantic audio transformer for sound classification and detection.  _ICASSP 2022 - 2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pp. 646–650, 2022.
  * Chen et al. (2020) Yen-Chun Chen, Linjie Li, Licheng Yu, Ahmed El Kholy, Faisal Ahmed, Zhe Gan, Yu Cheng, and Jingjing Liu.  Uniter: Universal image-text representation learning.  In _European conference on computer vision_ , pp. 104–120. Springer, 2020.
  * Chiu et al. (2022) Wei-Ting Chiu, Pei Wang, and Patrick Shafto.  Discrete probabilistic inverse optimal transport.  In _International Conference on Machine Learning_ , pp. 3925–3946. PMLR, 2022.
  * Courty et al. (2017) Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamonjy.  Joint distribution optimal transportation for domain adaptation.  _Advances in neural information processing systems_ , 30, 2017.
  * Cuturi (2013) Marco Cuturi.  Sinkhorn distances: Lightspeed computation of optimal transport.  In _NIPS_ , 2013.
  * Deshmukh et al. (2022) Soham Deshmukh, Benjamin Elizalde, and Huaming Wang.  Audio retrieval with wavtext5k and clap training.  _ArXiv_ , abs/2209.14275, 2022.
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Dhouib et al. (2020) Sofien Dhouib, Ievgen Redko, Tanguy Kerdoncuff, Rémi Emonet, and Marc Sebban.  A swiss army knife for minimax optimal transport.  In _International Conference on Machine Learning_ , pp. 2504–2513. PMLR, 2020.
  * Drossos et al. (2019) Konstantinos Drossos, Samuel Lipping, and Tuomas Virtanen.  Clotho: an audio captioning dataset.  _ICASSP 2020 - 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pp. 736–740, 2019.
  * Dupuy et al. (2016) Arnaud Dupuy, Alfred Galichon, and Yifei Sun.  Estimating matching affinity matrix under low-rank constraints.  _Econometric Modeling: Theoretical Issues in Microeconometrics eJournal_ , 2016.
  * Elizalde et al. (2019) Benjamin Elizalde, Shuayb Zarar, and Bhiksha Raj.  Cross modal audio search and retrieval with joint embeddings based on text and audio.  In _ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pp. 4095–4099. IEEE, 2019.
  * Fatras et al. (2020) Kilian Fatras, Younes Zine, Rémi Flamary, Remi Gribonval, and Nicolas Courty.  Learning with minibatch wasserstein: asymptotic and gradient properties.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 2131–2141. PMLR, 2020.
  * Gemmeke et al. (2017) Jort F. Gemmeke, Daniel P. W. Ellis, Dylan Freedman, Aren Jansen, Wade Lawrence, R. Channing Moore, Manoj Plakal, and Marvin Ritter.  Audio set: An ontology and human-labeled dataset for audio events.  _2017 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pp. 776–780, 2017.
  * Genevay et al. (2018) Aude Genevay, Gabriel Peyré, and Marco Cuturi.  Learning generative models with sinkhorn divergences.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 1608–1617. PMLR, 2018.
  * Huang et al. (2021) Zhenyu Huang, Guocheng Niu, Xiao Liu, Wenbiao Ding, Xinyan Xiao, Hua Wu, and Xi Peng.  Learning with noisy correspondence for cross-modal matching.  _Advances in Neural Information Processing Systems_ , 34:29406–29419, 2021.
  * Jia et al. (2021) Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoc Le, Yun-Hsuan Sung, Zhen Li, and Tom Duerig.  Scaling up visual and vision-language representation learning with noisy text supervision.  In _International conference on machine learning_ , pp. 4904–4916. PMLR, 2021.
  * Kim et al. (2019) Chris Dongjoo Kim, Byeongchang Kim, Hyunmin Lee, and Gunhee Kim.  Audiocaps: Generating captions for audios in the wild.  In _North American Chapter of the Association for Computational Linguistics_ , 2019.
  * Kingma & Ba (2014) Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _CoRR_ , abs/1412.6980, 2014.
  * Koepke et al. (2022) A Sophia Koepke, Andreea-Maria Oncescu, Joao Henriques, Zeynep Akata, and Samuel Albanie.  Audio retrieval with natural language queries: A benchmark study.  _IEEE Transactions on Multimedia_ , 2022.
  * Kong et al. (2019) Qiuqiang Kong, Yin Cao, Turab Iqbal, Yuxuan Wang, Wenwu Wang, and Mark D. Plumbley.  Panns: Large-scale pretrained audio neural networks for audio pattern recognition.  _IEEE/ACM Transactions on Audio, Speech, and Language Processing_ , 28:2880–2894, 2019.
  * Li et al. (2018) Ruilin Li, Xiaojing Ye, Haomin Zhou, and Hongyuan Zha.  Learning to match via inverse optimal transport.  _J. Mach. Learn. Res._ , 20:80:1–80:37, 2018.
  * Li et al. (2019) Ruilin Li, Xiaojing Ye, Haomin Zhou, and Hongyuan Zha.  Learning to match via inverse optimal transport.  _Journal of machine learning research_ , 20, 2019.
  * Liang et al. (2022) Weixin Liang, Yuhui Zhang, Yongchan Kwon, Serena Yeung, and James Y. Zou.  Mind the gap: Understanding the modality gap in multi-modal contrastive representation learning.  _ArXiv_ , abs/2203.02053, 2022.
  * Liu et al. (2023) Haohe Liu, Zehua Chen, Yi Yuan, Xinhao Mei, Xubo Liu, Danilo Mandic, Wenwu Wang, and Mark D Plumbley.  Audioldm: Text-to-audio generation with latent diffusion models.  _arXiv preprint arXiv:2301.12503_ , 2023.
  * Marshall & Olkin (1968) Albert W. Marshall and Ingram Olkin.  Scaling of matrices to achieve specified row and column sums.  _Numerische Mathematik_ , 12:83–90, 1968.
  * Mei et al. (2022) Xinhao Mei, Xubo Liu, Jianyuan Sun, MarkD . Plumbley, and Wenwu Wang.  On metric learning for audio-text cross-modal retrieval.  In _Interspeech_ , 2022.
  * Nguyen & Ho (2024) Khai Nguyen and Nhat Ho.  Sliced Wasserstein estimation with control variates.  In _The Twelfth International Conference on Learning Representations_ , 2024.
  * Nguyen et al. (2022a) Khai Nguyen, Dang Nguyen, Quoc Dinh Nguyen, Tung Pham, Hung Bui, Dinh Phung, Trung Le, and Nhat Ho.  On transportation of mini-batches: A hierarchical approach.  In _International Conference on Machine Learning_ , pp. 16622–16655. PMLR, 2022a.
  * Nguyen et al. (2022b) Khai Nguyen, Dang Nguyen, Tung Pham, Nhat Ho, et al.  Improving mini-batch optimal transport via partial transportation.  In _International Conference on Machine Learning_ , pp. 16656–16690. PMLR, 2022b.
  * Nguyen et al. (2024) Khai Nguyen, Shujian Zhang, Tam Le, and Nhat Ho.  Sliced wasserstein with random-path projecting directions.  _arXiv preprint arXiv:2401.15889_ , 2024.
  * Oncescu et al. (2021) Andreea-Maria Oncescu, A. Sophia Koepke, João F. Henriques, Zeynep Akata, and Samuel Albanie.  Audio retrieval with natural language queries.  In _Interspeech_ , 2021.
  * Paty & Cuturi (2019) François-Pierre Paty and Marco Cuturi.  Subspace robust wasserstein distances.  In _International conference on machine learning_ , pp. 5072–5081. PMLR, 2019.
  * Piczak (2015) Karol J. Piczak.  Esc: Dataset for environmental sound classification.  _Proceedings of the 23rd ACM international conference on Multimedia_ , 2015.
  * Radford et al. (2021) Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In _International conference on machine learning_ , pp. 8748–8763. PMLR, 2021.
  * Rockafellar (1997) R Tyrrell Rockafellar.  _Convex analysis_ , volume 11.  Princeton university press, 1997.
  * Sharma et al. (2018) Piyush Sharma, Nan Ding, Sebastian Goodman, and Radu Soricut.  Conceptual captions: A cleaned, hypernymed, image alt-text dataset for automatic image captioning.  In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pp. 2556–2565, 2018.
  * Shi et al. (2023) Liangliang Shi, Gu Zhang, Haoyu Zhen, Jintao Fan, and Junchi Yan.  Understanding and generalizing contrastive learning from the inverse optimal transport perspective.  2023\.
  * Stuart & Wolfram (2019) Andrew M. Stuart and Marie-Therese Wolfram.  Inverse optimal transport.  _SIAM J. Appl. Math._ , 80:599–619, 2019.
  * Wang et al. (2013) Kaiye Wang, Ran He, Wei Wang, Liang Wang, and Tieniu Tan.  Learning coupled feature spaces for cross-modal matching.  In _Proceedings of the IEEE international conference on computer vision_ , pp. 2088–2095, 2013.
  * Wei et al. (2020) Jiwei Wei, Xing Xu, Yang Yang, Yanli Ji, Zheng Wang, and Heng Tao Shen.  Universal weighting metric learning for cross-modal matching.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 13005–13014, 2020.
  * Wei et al. (2021) Jiwei Wei, Yang Yang, Xing Xu, Xiaofeng Zhu, and Heng Tao Shen.  Universal weighting metric learning for cross-modal retrieval.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 44(10):6534–6545, 2021.
  * Wu et al. (2022a) Bichen Wu, Ruizhe Cheng, Peizhao Zhang, Tianren Gao, Joseph E. Gonzalez, and Peter Vajda.  Data efficient language-supervised zero-shot recognition with optimal transport distillation.  In _The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022_. OpenReview.net, 2022a.
  * Wu et al. (2022b) Yusong Wu, K. Chen, Tianyu Zhang, Yuchen Hui, Taylor Berg-Kirkpatrick, and Shlomo Dubnov.  Large-scale contrastive language-audio pretraining with feature fusion and keyword-to-caption augmentation.  _ArXiv_ , abs/2211.06687, 2022b.
  * Xu et al. (2021) Hu Xu, Gargi Ghosh, Po-Yao Huang, Dmytro Okhonko, Armen Aghajanyan, Florian Metze, Luke Zettlemoyer, and Christoph Feichtenhofer.  Videoclip: Contrastive pre-training for zero-shot video-text understanding.  _arXiv preprint arXiv:2109.14084_ , 2021.
  * Yang et al. (2022) Jinyu Yang, Jiali Duan, Son Tran, Yi Xu, Sampath Chanda, Liqun Chen, Belinda Zeng, Trishul Chilimbi, and Junzhou Huang.  Vision-language pre-training with triple contrastive learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 15671–15680, 2022.
  * Zeng et al. (2022) Andy Zeng, Maria Attarian, Brian Ichter, Krzysztof Choromanski, Adrian Wong, Stefan Welker, Federico Tombari, Aveek Purohit, Michael Ryoo, Vikas Sindhwani, et al.  Socratic models: Composing zero-shot multimodal reasoning with language.  _arXiv preprint arXiv:2204.00598_ , 2022.

##  Appendix A Appendix

###  A.1 Implementation details

AudioCaps is the biggest audio captioning dataset that consists of around 50k
audio-caption pairs. All audio clips are extracted from AudioSet (Gemmeke et
al., 2017), a large-scale dataset for audio tagging. There are a total of
40,582 audio clips in training data, and all audio clips are 10 seconds long.
Each audio clip has a single human-annotated caption. The validation and test
sets have 494 and 957 audio clips, respectively, and each audio clip has five
ground-truth captions.

Clotho is an audio captioning dataset collected from the Freesound platform.
The audio clips’ length varies from 15 to 30 seconds. We use the second
version of the Clotho dataset to carry out the experiment. There are 3839
audio clips in the training set and over 1,000 audio clips in the validation
and test set. All captions in the dataset consist of 5 human-annotated
captions.

ESC50 is an audio dataset for sound event detection. The dataset incorporates
2,000 labeled environmental recordings with 50 classes. We only use the test
set of the ESC50 dataset to verify the transferable ability of our proposed
method. There are 400 audio clips in the test set used in our experiment.

Implementation details. We follow the implementation details in (Mei et al.,
2022). Log mel-spectrograms are used as the audio features for the audio
encoder, and the settings used to extract Log mel-spectrograms are window size
of 1024, Hanning window, hop-size of 320, and 64 first mel bins. The audio
encoder is the ResNet-38 model (Kong et al., 2019) that is pretrained on the
AudioSet dataset with the audio tagging task. After pertaining, we discard the
last two layers and then apply an average-pooling layer to accumulate features
along the frequency dimension on the extracted feature map. We also conducted
an experiment on the HTSAT audio encoder (Chen et al., 2022), which is a
transformer-based encoder. The text encoder is the BERT model (Devlin et al.,
2018), which is pretrained on various NLP tasks to be able to extract
contextual-aware embeddings. Due to the dimension mismatch between output of
the audio encoder and the text encoder, we utilize two linear project blocks
to project two encoder outputs to the shared embedding space. The shared
embedding space is a 1024-dimensional space. The interaction matrix $M$ is
initialized as follows $M=\frac{1}{2}(A+A^{T})+\mathds{1}$, where $A$ is a
random initialized matrix of size 1024 by 1024. All the models and the matrix
$M$ are trained for 30 epochs with Adam optimizer (Kingma & Ba, 2014). The
hyperparameters for training are learning rate $lr=1\times 10^{-4}$, batch
size $b=256$, and dropout ratio $p=0.2$. All experiments are performed on a
single A100 GPU.

###  A.2 Algorithm

Input : Initialize audio encoder $f_{\theta}$, text encoder $g_{\phi}$,
interaction matrix $M$, and training data $p(x,y)$

Output : Learned audio encoder $f_{\theta}$, text encoder $g_{\phi}$ and
interaction matrix $M$

while  _until converged_ do

$(x^{i},y^{i})_{i=1}^{b}\sim p(x,y)$;

$Z_{1}=\\{f_{\theta}(x^{i})\\}_{i=1}^{b}$;

$Z_{2}=\\{g_{\phi}(y^{i})\\}_{i=1}^{b}$;

$\hat{\pi}=$ Eye($\frac{1}{b}$);

$C_{M}^{Z_{1},Z_{2}}=\sqrt{(Z_{1}-Z_{2})^{T}M(Z_{1}-Z_{2})}$;

$\pi^{X^{b},Y^{b}}_{\epsilon,c_{\theta,\phi,M}}=$
Sinkhorn($C_{M}^{Z_{1},Z_{2}},\epsilon$);

$\mathcal{L}(\theta,\phi,M)=KL(\hat{\pi}||\pi^{X^{b},Y^{b}}_{\epsilon,c_{\theta,\phi,M}})$;

$\theta=\theta-\eta\nabla_{\theta}{\mathcal{L}}$;

$\phi=\phi-\eta\nabla_{\phi}{\mathcal{L}}$;

$M=Proj(M-\eta\nabla_{M}{\mathcal{L}})$

return _$f_{\theta},g_{\phi}$ , and $M$_

Algorithm 1 Learning ground cost metric using m-LTM framework and Mahanalobis
distance

Input : $C\in\mathbb{R}^{n\times n}$ and $\epsilon$

Output : $\pi_{\epsilon}$

$K=\exp{(-C/\epsilon)}$, $v\leftarrow\frac{1_{n}}{n}$;

while  _until converged_ do

$u=\frac{1_{n}}{n}\oslash Kv$; $v=\frac{1_{n}}{n}\oslash K^{T}u$;

$\pi_{\epsilon}=$ diag($u$)$K$diag($v$);

return _$\pi_{\epsilon}$_

Algorithm 2 Sinkhorn Algorithm

###  A.3 Qualitative Experiments

Figure 5: Qualitative results for text-to-audio retrieval task. top-1, top-2,
and top-3 retrieved audio results are from left to right in the figure. The
ground-truth audio for the caption is marked in red border. Figure 6:
Qualitative results for text-to-audio retrieval task. top-1, top-2, and top-3
retrieved audio results are from left to right in the figure. The ground-truth
audio for the caption is marked in red border.

###  A.4 Additional Experiments and Ablation Study

We compare the m-LTM loss with contrastive loss and a variant of contrastive
loss with entropic regularization in Table. 7. The training objective of
contrastive loss with entropic regularization is as follows

$\displaystyle\mathcal{L}(\theta,\phi)=\mathcal{L}_{CL}+\epsilon*\mathcal{L}_{ent}(R)$
(14)

, where $\mathcal{L}_{CL}$ is contrastive loss described in the Equation. 1,
$R$ is the ranking matrix between two sets of audio and caption, and
$\epsilon$ is the coefficient of the entropic regularization term. The
experiment shows that entropic regularization can help to boost the
performance of contrastive loss from $33.9\%$ to $34.04\%$ and from $39.4\%$
to $44.11\%$ for text-to-audio and audio-to-text retrievals, respectively.
However, there is a remarkable gap in performance between the m-LTM and the
best performance acquired by contrastive loss with entropic regularization at
$\epsilon=0.05$.

Table 7: The additional baseline experiment of loss functions, the m-LTM, contrastive loss, and contrastive loss with entropy regularization, on the AudioCaps dataset. Audio Encoder | Method | Epsilon | Text->Audio | Audio->Text
---|---|---|---|---
R@1 | R@5 | R@10 | R@1 | R@5 | R@10
ResNet38 Audio Encoder | CL | - | 33.9 | 69.7 | 82.6 | 39.4 | 72.0 | 83.9
CL with entropic regularization | 0.05 | 34.04 | 69.11 | 82.51 | 44.11 | 74.71 | 85.37
0.1 | 34.27 | 70.21 | 82.92 | 44.41 | 73.87 | 86.41
0.5 | 33.79 | 69.51 | 82.90 | 42.52 | 72.10 | 85.78
m-LTM | 0.05 | 39.10 | 74.06 | 85.78 | 49.94 | 80.77 | 90.49
HTSAT Audio Encoder | CL | - | 34.38 | 71.16 | 84.64 | 35.84 | 71.06 | 83.39
CL with entropic regularization | 0.05 | 35.71 | 72.3 | 85.19 | 37.45 | 73.91 | 84.44
0.1 | 35.12 | 72.05 | 85.01 | 36.96 | 73.55 | 84.34
0.5 | 34.45 | 71.33 | 85.11 | 35.98 | 72.53 | 83.78
m-LTM | 0.05 | 40.02 | 75.86 | 86.42 | 40.13 | 74.09 | 85.89

We also provide an additional ablation study for choosing the best
hyperparameter $\epsilon$ for the m-LTM framework in Table. 8 and conducted
across datasets retrieval experiment to illustrate the robustness of the m-LTM
framework compared with triplet and contrastive loss in Table. 9.

Table 8: Further ablation study on AudioCaps dataset for choosing the hyperparameter $\epsilon$ for the minibatch learning-to-match framework utilizing Mahalanobis distance. Epsilon | Text->Audio | Audio->Text
---|---|---
R@1 | R@5 | R@10 | R@1 | R@5 | R@10
0.01 | 0.10 | 0.52 | 1.04 | 0.10 | 0.31 | 0.41
0.02 | 0.17 | 0.77 | 2.06 | 0.45 | 1.31 | 3.78
0.03 | 38.28 | 74.85 | 85.13 | 49.88 | 79.77 | 89.86
0.04 | 37.22 | 73.58 | 85.37 | 48.38 | 80.04 | 90.17
0.05 | 39.10 | 74.06 | 85.78 | 49.94 | 80.77 | 90.49
Table 9: The robustness of minibatch learning-to-match, triplet, and contrastive loss for audio-text retrieval across datasets. Train-Test | Method | Text->Audio | Audio->Text
---|---|---|---
R@1 | R@5 | R@10 | R@1 | R@5 | R@10
AudioCaps->Clotho | Triplet | 10.41 | 27.38 | 38.04 | 11.39 | 28.42 | 39.80
CL | 12.31 | 30.93 | 43.44 | 13.87 | 33.30 | 43.54
m-LTM | 15.01 | 35.42 | 47.71 | 19.42 | 37.03 | 48.61
Clotho->AudioCaps | Triplet | 14.33 | 40 | 54.42 | 17.34 | 43.05 | 55.69
CL | 14.64 | 38.64 | 53.58 | 17.76 | 42.11 | 55.59
m-LTM | 17.38 | 43.78 | 58.87 | 22.04 | 47.85 | 61.65


# Unsupervised Data Splitting Scheme for Federated Edge Learning in IoT
Networks

Boubakr Nour Soumaya Cherkaoui Email<EMAIL_ADDRESS><EMAIL_ADDRESS>Department of Electrical and Computer
Science Engineering, Université de Sherbrooke, Canada

###### Abstract

Federated Edge Learning (FEEL) is a promising distributed learning technique
that aims to train a shared global model while reducing communication costs
and promoting users’ privacy. However, the training process might
significantly occupy a long time due to the nature of the used data for
training, which leads to higher energy consumption and therefore impacts the
model convergence. To tackle this issue, we propose a data-driven federated
edge learning scheme that tends to select suitable participating nodes based
on quality data and energy. First, we design an unsupervised data-aware
splitting scheme that partitions the node’s local data into diverse samples
used for training. We incorporate a similarity index to select quality data
that enhances the training performance. Then, we propose a heuristic
participating nodes selection scheme to minimize the communication and
computation energy consumption, as well as the amount of communication rounds.
The obtained results show that the proposed scheme substantially outperforms
the vanilla FEEL in terms of energy consumption and the number of
communication rounds.

###### Index Terms:

federated learning, edge computing, data-driven node selection

##  I Introduction

Nowadays, the world is witnessing a substantial development in the field of
artificial intelligence (AI), which is sweeping major fields. This development
is made because of the tremendous advances in machine learning (ML)
techniques. Novel ML applications emerge every day, ranging from autonomous
driving and finance to marketing and healthcare, potential applications are
boundless [1]. In parallel, wireless mobile technologies, such as the fifth-
generation (5G), promise to connect billions of heterogeneous devices to the
network edge [2]. This connection opens various opportunities to verticals
toward the rise of new applications and business models under the banner of
the Internet of Things (IoT) [3]. Edge nodes seek to collect a massive amount
of data and open new avenues for ML applications. The prevalent approach of ML
implementations on edge devices is to amass all the relevant data at
cloud/edge servers and train powerful ML models using the available data and
processing capacities [4, 5]. However, such a centralized fashion is not
fitting many use cases’ requirements, which might also violate the latency
requirements of the underlying application, scatter communication resources,
and infringe user/data privacy [6]. For example, an autonomous car may
generate from 5 to 20 TB of data per day. This large volume of data contains
sensitive information and needs to be transmitted over the bottleneck core
network to be processed at distant cloud servers. The low-latency
requirements, information density, and data privacy become challenging to be
addressed in such circumstances.

To fulfill the requirements of most IoT applications and promote data privacy,
federated edge learning (FEEL) [7] has been proposed an advanced machine
learning technique to train models in a distributed fashion, using mobile edge
computing. In FEEL,  (i) a centralized edge server shares the model with a set
of nodes (i.e., participating nodes),  (ii) participating nodes train the
model locally without sharing their data with the edge server,  (iii) the
upload only the model parameters to the edge server,  (iv) the latter
aggregates the received updates to improve the global model, and  (v) re-sends
the model to participating nodes for another training round.  The process is
repeated multiple rounds or until the model reaches a pre-defined threshold.

For each training round, the edge node performs local computation for model
training as well as communication with the edge server for model updating.
Multiple communication and computation rounds significantly drain the energy
of devices that are energy-constrained in nature, and hence restrict their
ability to conduct more training rounds. This circumstance negatively impacts
the training performance and consequently leads to a slower convergence [8,
9]. An efficient way to overcome this issue is to reduce the number of
training rounds, which reflects to decreasing the computation and
communication overhead, and hence, sustain the node’s energy. Using fewer data
samples for training can help in reaching less communication overhead. Yet,
this is a double-edged problem since the model might not converge with fewer
data and the accuracy is controversial. Notwithstanding, work in [10] reveals
that the model convergence is highly relying on the quality of data used
during the training. ML models can be trained with fewer data samples, as long
as this data has high quality (i.e., variety). Determining the data used in
the model training can substantially improve the training performance [11].

To this end, in this paper, we design a data-driven federated edge learning
scheme for IoT networks. Instead of using the entire dataset for training,
edge nodes split their data into a set of sub-datasets by minimizing the data
similarity (i.e., increase the diversity within each sub-dataset). This
approach is motivated by the fact that less yet quality data samples
significantly contribute to model learning [10]. Thereafter, the edge server
selects a subset of nodes to participate in the training based on their
similarity score and communication/computation energy. Overall, the major
contributions of this work can be outlined as follows:

  * •

we formulate data splitting and node selection problems with the aim to
maximize data diversity and minimize energy consumption.

  * •

we design an unsupervised data-aware splitting scheme that tends to partition
local data into a set of sub-datasets with diverse and quality samples.

  * •

we propose a heuristic node selection scheme that chooses participating nodes
with minimum similarity scores while reducing the node’s energy consumption.

  * •

we validate the proposed scheme through simulation using MNIST, a non-
independent and identically distributed (non-IID) and unbalanced dataset. The
designed scheme helps in model convergence with fewer communication rounds
while minimizing energy consumption.

The rest of this paper is organized as follows. Section II presents an
overview of federated learning and review on related work. Section III
introduces the system model and the problem formulation. Section IV details
the proposed data-driven federated edge learning scheme. The performance and
effectiveness of the proposed scheme are evaluated in Section V. Finally, we
conclude the paper in Section VI.

##  II Related Work

The integration of artificial intelligence and edge computing drives a mutual
benefit that helps to overcome various issues and challenges in today’s
networks [12]. For instance, edge computing strengthens from the intelligent
decisions made by AI mechanisms, while AI models will be further improved by
availing the distributed nature of edge computing. Yet, the edge server
requires to collect all data generated by users for training and inference
purposes, which might violate the user’s privacy [13].

Federated learning [14] aims to overcome the privacy issues in conventional
machine learning techniques. In FEEL, multiple edge devices collaboratively
train a global model without sharing their data. The process is orchestrated
by a centralized edge server that sends the model to multiple nodes. Each node
trains the model locally using its own data and sends the model parameters to
the server for aggregation. The model is refined and updated for multiple
rounds until its convergence [15]. There have been various efforts in the
literature that tend to enhance network resources while ensuring fast model
convergence [16]. For instance, authors in [17] propose a client selection
scheme for federated learning, taking into account the nature of heterogeneous
mobile devices. The edge server sets a deadline to be met by clients, which
includes model download, train, update, and upload. The edge server, based on
the clients’ resources, selects which nodes can participate in the learning
phase and hence lower the training time. Mohammed et al. [18] tend to minimize
the number of IoT candidates contributing to the training process. Inspired by
the solution of the secretary problem, the authors design an online heuristic
scheme to find the best candidate clients. The objective is to overcome
communication and computation constraints while ensuring a predefined
accuracy. Albaseer et al. [19] propose a client selection approach that
leverages devices’ heterogeneity. The approach aims to accelerate the
convergence rate by using clients’ round latency and exploiting the bandwidth
reuse for clients that consume more time to update the model. The server
performs model averaging and then clusters the clients based on a predefined
threshold. Clients with less latency are selected to update the model. Works
in [20] and [21] propose a data exclusion approach to minimize energy
consumption during federated learning rounds. Toward this, the authors design
a customized local training algorithm that selects only the samples that
improve the model’s quality based on a predefined threshold probability.

Although the aforementioned solutions have contributed to enhancing federated
learning architecture, some of them did not consider the data in their
studies. Indeed, data is a pillar component when selecting participating
clients since it has a direct impact on model accuracy. This work emphasizes
data as the pivotal element for model training without ignoring energy
consumption.

##  III System Model & Problem Formulation

###  III-A System Model

Figure 1 depicts the used system model in this work. Similar to [21], we
consider a network that has a set $\mathcal{K}$ edge devices connected to a
centralized edge server. The edge server serves as a coordinator for
participating edge nodes to train a global model.

Figure 1: System model for centralized federated edge learning.

Learning Model. Each edge device $k\in\mathcal{K}$ has its own data
$\mathcal{D}_{k}$ used to train the model locally and sends the update
$\omega_{k}$ back to the edge server. Let $|\mathcal{D}_{k}|$ denote the
number of local data samples for edge device $k$. The edge server selects a
subset of edge devices, denoted $\mathcal{S}\subset\mathcal{K}$ to participate
in model training. Those devices receive the initialization parameters
$\omega_{0}$ from the edge server (Step 1), train the model locally using
local data $\mathcal{D}_{k}$ (Step 2), and send the updates to the edge server
(Step 3). The edge server, upon receiving the updates, aggregates the received
parameters ((Step 4), generates a new global model, and sends the updated
model parameters $\omega_{r}$ to the participating edge devices for another
training round. The process is repeated $r_{max}$ rounds until the model
converges or reaches a predefined threshold.

Each participating edge device tends to minimize the local loss function. Let
$F_{k}^{r}$, defined in Eq. (1), denote the local loss function for the $k$-th
participating device at the $r$-th round:

$F_{k}^{r}(\omega)\triangleq\frac{1}{|\mathcal{D}_{k}|}\sum_{s\in\mathcal{D}_{k}}f_{s}(\omega).$
(1)

where $f_{s}$ is the error for each local sample in local data. Let
$D\triangleq\sum_{k=1}^{|\mathcal{S}|}|\mathcal{D}_{k}|$ denote the total
amount of data samples used in edge network, while
$\delta_{k}=\frac{|\mathcal{D}_{k}|}{D}$ denote the weight of the local data
samples at the $k$-th device.

Edge devices use stochastic gradient descent (SGD) [22] to minimize the loss
function, Eq. (1), for $\epsilon$ epochs. The local model parameters
$\omega_{k}$ are updated accordingly as shown in Eq. (2).

$\omega_{k}(n)=\omega_{k}(n-1)-\eta\nabla F_{k}^{r}(\omega_{k}(n)).$ (2)

where $n=1,2,\dots,\epsilon$ is the number of local updates performed by the
$k$-th device and $\eta$ is the learning rate at each round. $\omega_{k}(0)$
is the initial global parameters sent by the edge server and received by the
$k$-th device, and $\omega_{k}(\epsilon)$ is the last local parameters update
sent by the $k$-th device to the edge server after $\epsilon$ rounds.

After receiving all updates from all participating edge devices, the weighted
global loss function at the $r$-th round is calculated as shown in Eq. (3).

$F_{r}(\omega)\triangleq\sum_{k=1}^{|\mathcal{K}|}\delta_{k}F_{k}^{r}(\omega).$
(3)

Accordingly, the global model parameters are calculated as shown in Eq. (4).

$\omega_{r}=\sum_{k=1}^{|\mathcal{K}|}\delta_{k}~{}\omega_{k},$ (4)

Then, both $F_{r}(\omega)$ and $\omega_{r}$ are sent to all participating
devices to be employed for the next round ($r+1$), aiming to minimize
$F(\omega)$.

$\overset{*}{\omega}\triangleq\arg\min F(\omega).$ (5)

Local Computation Model. The participating device randomly partitions local
data $\mathcal{D}_{k}$ into multiple batches of size $b$ and trains the model
for $\epsilon$ epochs [23]. Assuming that the device uses the entire data for
training, the local computation time $T_{k}^{\operatorname{cmp}}$ is defined
as shown in Eq. (6).

$T_{k}^{\operatorname{cmp}}=\epsilon\frac{|\mathcal{D}_{k}|\Phi}{f_{k}^{\operatorname{cmp}}},$
(6)

where $f_{k}^{\operatorname{cmp}}$ denotes the local CPU frequency and $\Phi$
is the number of cycles required to process one data sample.

In order to provide an efficient synchronization of all updates and avoid long
waiting time, the edge server sets a deadline $\mathcal{T}$ to send the update
back to him. Edge devices have to accomplish the computation and computation
within $\mathcal{T}$. Hence, the local computation time is defined as shown in
Eq. (7).

$T_{k}^{\operatorname{cmp}}=\mathcal{T}-T_{k}^{\operatorname{up}}.$ (7)

where $T_{k}^{\operatorname{up}}$ is the required time to upload the update to
the edge server.

Energy Consumption Model. The energy consumption for the $k$-th device is
based on the training process and defined as shown in Eq. (8).

$E_{k}^{\operatorname{cmp}}=\frac{\alpha_{k}}{2}(f_{k}^{\operatorname{cmp}})^{3}T_{k}^{\operatorname{cmp}},$
(8)

where $\frac{\alpha_{k}}{2}$ is the device’s energy capacitance coefficient.
From Eq. (6) and Eq. (8), we find the local energy consumption as:

$E_{k}^{\operatorname{cmp}}=\frac{\alpha_{k}}{2}(\varepsilon(f_{k}^{\operatorname{cmp}})^{2}|D_{k}|\Phi).$
(9)

Considering a Time Division Multiple Access (TDMA), the model upload latency
is defined as in Eq. (10).

$T_{k}^{\operatorname{up}}=\frac{\xi}{\Gamma_{k}}.$ (10)

where $\xi$ denotes the model size and $\Gamma_{k}$ the uplink data rate
achieved by the $k$-th device. Therefore, the transmission energy consumption
of the $k$-th device is defined based on the upload latency
$T_{k}^{\operatorname{up}}$ and the transmit power
$P_{k}^{\operatorname{up}}$, as shown in Eq. (11).

$E_{k}^{\operatorname{up}}=T_{k}^{\operatorname{up}}P_{k}^{\operatorname{up}}.$
(11)

###  III-B Problem Formulation

Maximize the data diversity: the first objective is to accelerate the training
phase [10]. In doing so, we tend to build sub-datasets with diverse yet
quality samples. Therefore, the objective of Eq. (12) aims at minimizing the
similarity between selected samples within the same sub-dataset.

$\textbf{P1}:\operatorname{minimize}\sum\limits_{\begin{subarray}{c}\forall
d\in\mathcal{D}\\\ \forall u\in d\\\ \forall u\neq v\in
D\end{subarray}}x_{u}^{d}x_{v}^{d}\operatorname{Sim}(u,v),$ (12)

subject to,

$\sum_{u\in d}x_{u}^{d}=1,\qquad\forall d\in\mathcal{D},$ (13a)
$\sum_{d\in\mathcal{D}}x_{u}^{d},\geq 2\qquad\forall u\in d,$ (13b)
$\sum_{d\in\mathcal{D}}\sum_{u\in d}x_{u}^{d}=|D|,$ (13c)
$x_{u}^{d}\in\\{0,1\\},\quad\forall u\in d,\forall d\in\mathcal{D}.$ (13d)

where $x_{i}^{d}$ is a variable indicating the assignment of the
$i^{\text{th}}$ sample to the sub-dataset $d$. Constraints (13a) and (13b)
setup up the outliers-free model where each data sample should be assigned to
only one sub-dataset and each sub-dataset should contain at least 2 samples.
Constraint (13c) imposes that all data samples should be used to construct the
sub-dataset, which implies the fairness in using all data samples in the
learning phase. Finally, the non-negativity constraint is shown in (13d).

Nodes selection: the second objective (14) tends to select the appropriate
nodes that have qualified data in the training phase while minimizing the
energy consumption.

$\textbf{P2}:\operatorname{minimize}\sum\limits_{\begin{subarray}{c}\forall
k\in\mathcal{K}\end{subarray}}\gamma_{k}(E_{k}^{\operatorname{cmp}}+E_{k}^{\operatorname{up}}),$
(14)

subject to,

$\gamma_{k}(T_{k}^{\operatorname{cmp}}+T_{k}^{\operatorname{up}})\leq\mathcal{T},\quad\forall
k\in\mathcal{K},$ (15a)
$\gamma_{k}(E_{k}^{\operatorname{cmp}}+E_{k}^{\operatorname{up}})\leq
E_{k},\quad\forall k\in\mathcal{K},$ (15b)
$\gamma_{k}\in\\{0,1\\},\quad\forall k\in\mathcal{K}.$ (15c)

where $\gamma_{k}$ indicates the node selection variable. $\gamma_{k}=1$ if
the node $k$ has been selected in the training phase, otherwise
$\gamma_{k}=0$. Constraint (15a) imposes that the computation and computation
should be accomplished before the fixed deadline by the edge server.
Similarly, constraint (15b) enforces that computation and transmission energy
should not exceed the node available energy. This constraint helps in ensuring
that the selected node is able to provide training and successfully send the
update to the edge service without its power running off. Finally, constraint
(15c) expresses the non-negativity.

Input:

$D$: local dataset;

$k$: number of sub-datasets;

Output:

$\\{\mathcal{D}_{i},\dots,\mathcal{D}_{k}\\}$: sub-datasets;

1

2$R$ $\leftarrow D$.random($k$);

3 $i$ $\leftarrow$ 1;

4 for  _$r\in R$_ do

5 $\mathcal{D}_{i}$.append($r$);

6 $r$.marked $\leftarrow$ True;

7 $i\leftarrow i+1$;

8

9 end for

10$Q\leftarrow D-R$;

11

12while  _( $|Q|$ > 0)_ do

13 $v\leftarrow Q$.random();

14

15 if  _( $v$.marked $\leftarrow$ False)_ then

16 for  _$i\in[1,~{}k]$_ do

17 while  _( $e\leq|\mathcal{D}|$)_ do

18
$\mathcal{S}(\mathcal{D}_{i})\leftarrow\operatorname{Sim}(\mathcal{D}(e),v$);

19 $e\leftarrow e+1$;

20

21 end while

22
$j\leftarrow\operatorname{min}(\mathcal{S}(\\{\mathcal{D}_{i},\dots,\mathcal{D}_{k}\\})$);

23 $\mathcal{D}_{j}$.append($v$);

24 $v$.marked $\leftarrow$ True;

25

26 end for

27

28 end if

29

30 end while

31return $\mathcal{D}$

Algorithm 1 Unsupervised Data-Aware Splitting Scheme.

##  IV Data-Driven Federated Edge Learning

In order to solve the aforementioned problems, we design unsupervised data-
aware splitting and heuristic node selection schemes. The unsupervised data-
aware splitting scheme runs at the edge device and enables the node to split
its local dataset into a set of sub-dataset that include quality samples by
maximizing the data diversity. The node selection scheme, on the other hand,
is executed at the edge server to select participating nodes, that own quality
data, in the training phase with minimum energy consumption. The joint goal is
to significantly enhance the global model training and improve the overall
network performance.

###  IV-A Data-Aware Splitting Scheme

Algorithm 1 shows the process of a data-aware splitting scheme that tends to
solve P1. The algorithm allows the edge device to efficiently split its large-
size dataset into a set of sub-datasets. The splitting process needs to
maximize the data diversity so that the model can converge quickly with fewer
data samples.

In doing so, each node independently selects a number ($k$) of desirable sub-
datasets to be created, randomly picks $k$ samples from the original data, and
then assigns each one to each sub-dataset (lines 1-7). Afterward, and for each
sample $v$ in the remaining data, if $v$ has not been assigned to any sub-
dataset, the scheme calculates the similarity index of the said sample versus
each sub-dataset (lines 13-16). At the end of the process, the sample is
assigned to the sub-dataset with the minimum similarity score. The algorithm
returns a list of $k$ sub-datasets that has minimum similarity scores.

In this work, we apply Jaccard Index to calculate the similarity between data
samples which is measured as shown in Eq. 16, where $d_{i}$ and $d_{j}$ are
two data samples. An index of 0 means that samples are completely dissimilar,
an index of 1 means that they are identical, and an index between 0 and 1
represents a degree of similarity.

$J(d_{i},d_{j})=\frac{d_{i}\cap d_{j}}{d_{i}\cup d_{j}}.$ (16)

###  IV-B Participating Node Selection Scheme

To solve P2, we design a heuristic nodes selection scheme. After the edge
server sends participation requests, nodes willing to participate in the
training phase send information, such as the number of created sub-datasets
with their size and similarity scores, available energy, and communication and
computation energy, etc., back to the edge server. It is important to
highlight that this information does not carry any sensible input that
violates the data/user privacy. Algorithm 2 shows the overall process executed
by the edge server.

Input:

$N$: list of edge nodes;

$k$: number of required participating nodes;

Output:

$\mathcal{S}$: list of selected participatig nodes;

1

2Send participation initialization;

3 Nodes execute Algorithm 1;

4 Nodes reply with available information;

5 Assign coefficients to nodes;

6 Order nodes based on similarities and energy;

7 Pick the first $k$ nodes;

8

9for  _$r\leq r_{max}$_ do

10 for  _$c\in\mathcal{K}$_ do

11 Send the global model;

12 Train model locally;

13 Send updated model with available energy;

14

15 end for

16 Aggregate the model;

17 Update nodes information; Order nodes based on similarities and energy;

18 Pick the first $k$ nodes;

19 $r\leftarrow r+1$;

20

21 end for

22

Algorithm 2 Participating nodes selection scheme.

For each node, the edge server assigns a coefficient based on the number of
available sub-dataset, similarity index, and energy. Smaller coefficients are
assigned to nodes with a minimum similarity index and larger energy. The edge
server sorts the information and picks the first $n$ nodes to participate in
the learning. The node’s coefficient helps the edge server to pick a minimum
number of nodes with diverse data samples instead of picking multiple nodes.
The smaller number of selected nodes is, the least communication overhead
performed. Since the energy of nodes can be changed over time, the node
selection process is repeated for each round. New nodes can be selected for
the next training round while already selected nodes could be dropped out.

##  V Numerical Results

In the experiment part, we consider a realistic federated edge learning
environment, as depicted in Figure 1, that consists of 100 nodes. We used the
MNIST database [24], an imbalanced and non-IID data distribution with
different CCN models. Edge devices use mini-batch stochastic gradient descent
as a local solver [22]. The global model evaluation is performed every round
with a learning rate of $0.001$. The original dataset is randomly partitioned
into $N$ pieces where each node is assigned one part. The database is split
into 80% for training and 20% for testing. We compared the proposed scheme
against vanilla federated edge learning (i.e., standard FEEL scheme) and
random nodes selection (i.e., selecting random participating nodes). Results
are collected and shown in an average of 10 trials.

Figure 2: Average accuracy. Figure 3: Average loss.

Figure 4: Average cumulative energy consumption.

Figure 5: Average instantaneous energy consumption.

Figures 2 and 3 show the evaluation of the accuracy and the loss function for
200 training rounds, respectively. The loss function used is Mean Squared
Error. We can see that the proposed scheme is not affected by reducing the
number of participating nodes or the number of samples used in the learning.
Indeed, the obtained accuracy and loss are better than the vanilla FEEL and
the random scheme. With a fewer number of participating nodes and a small size
of data, the proposed scheme was able to converge quickly.

Figures 5 and 5 present the cumulative and instantaneous energy consumption
results, respectively. We can observe that the proposed scheme shows a
significant reduction in energy consumption compared with vanilla FEEL and
random node selection scheme. The reduction gain is proportional to the number
of training rounds. This is due to the fact that the proposed scheme selects a
minimum number of nodes to participate in the training with diverse and
quality data samples. The diversity helps in fast model convergence with fewer
data and training rounds. Both vanilla FEEL and random schemes need to
incorporate the entire dataset that requires more time for training,
computation, and communication and by consequence more energy consumption.

##  VI Conclusion

In this paper, we designed a data-driven federated edge learning scheme for
IoT networks. Our design goal is to enhance the model convergence with the
minimum energy consumption. In doing so, we proposed an unsupervised data-
splitting scheme that partitions the node’s local data into diverse and
quality sub-datasets to be used in the training. We applied the Jaccard index
to ensure the diversity of data samples. We also proposed a heuristic node
selection scheme to select participating nodes with minimum energy
consumption. The experimental evaluation demonstrates the effectiveness and
efficiency of the proposed scheme in terms of model convergence and energy
consumption. As a future step, we will focus on designing a distributed data-
aware exclusion and node scheduling architecture for federated learning.

## Acknowledgments

The authors would like to thank the Natural Sciences and Engineering Research
Council of Canada (NSERC) for the financial support of this research.

## References

  * [1] F. Tang, B. Mao _et al._ , “Survey on Machine Learning for Intelligent End-to-End Communication towards 6G: From Network Access, Routing to Traffic Control and Streaming Adaption,” _IEEE Communications Surveys & Tutorials_, 2021.
  * [2] A. Filali, A. Abouaomar _et al._ , “Multi-access edge computing: A survey,” _IEEE Access_ , vol. 8, pp. 197 017–197 046, 2020.
  * [3] M. Laroui, B. Nour _et al._ , “Edge and Fog Computing for IoT: A Survey on Current Research Activities & Future Directions,” _Computer Communications_ , vol. 180, pp. 210–231, 2021.
  * [4] B. Nour and S. Cherkaoui, “A Network-based Compute Reuse Architecture for IoT Applications,” _arXiv preprint arXiv:2104.03818_ , 2021.
  * [5] Y. Liu, M. Peng _et al._ , “Toward Edge Intelligence: Multiaccess edge computing for 5G and Internet of Things,” _IEEE Internet of Things Journal_ , vol. 7, no. 8, pp. 6722–6747, 2020.
  * [6] L. U. Khan, W. Saad _et al._ , “Federated learning for internet of things: Recent advances, taxonomy, and open challenges,” _IEEE Communications Surveys & Tutorials_, 2021.
  * [7] F. Malandrino and C. F. Chiasserini, “Federated learning at the network edge: When not all nodes are created equal,” _arXiv preprint arXiv:2101.01995_ , 2021.
  * [8] L. U. Khan, S. R. Pandey _et al._ , “Federated learning for edge networks: Resource optimization and incentive mechanism,” _IEEE Communications Magazine_ , vol. 58, no. 10, pp. 88–93, 2020.
  * [9] H. Wang, Z. Kaplan _et al._ , “Optimizing federated learning on non-iid data with reinforcement learning,” in _IEEE Conference on Computer Communications_. IEEE, 2020, pp. 1698–1707.
  * [10] A. A. Abdellatif, C. F. Chiasserini _et al._ , “Active learning-based classification in automated connected vehicles,” in _IEEE Conference on Computer Communications Workshops_. IEEE, 2020, pp. 598–603.
  * [11] Z. Gong, P. Zhong _et al._ , “Diversity in machine learning,” _IEEE Access_ , vol. 7, pp. 64 323–64 350, 2019.
  * [12] X. Wang, Y. Han _et al._ , “Convergence of edge computing and deep learning: A comprehensive survey,” _IEEE Communications Surveys & Tutorials_, vol. 22, no. 2, pp. 869–904, 2020.
  * [13] B. S. Gu, L. Gao _et al._ , “Privacy on the edge: Customizable privacy-preserving context sharing in hierarchical edge computing,” _IEEE Transactions on Network Science and Engineering_ , 2019.
  * [14] A. Taïk, B. Nour _et al._ , “Empowering Prosumer Communities in Smart Grid with Wireless Communications and Federated Edge Learning,” _IEEE Wireless Communications_ , vol. 28, no. 6, pp. 26–33, 2021.
  * [15] B. Nour, S. Cherkaoui _et al._ , “Federated Learning and Proactive Computation Reuse at the Edge of Smart Homes,” _IEEE Transactions on Network Science and Engineering_ , 2021.
  * [16] C. Yang, Q. Wang _et al._ , “Characterizing Impacts of Heterogeneity in Federated Learning upon Large-Scale Smartphone Data,” in _Proceedings of the Web Conference_ , 2021, pp. 935–946.
  * [17] T. Nishio and R. Yonetani, “Client selection for federated learning with heterogeneous resources in mobile edge,” in _IEEE International Conference on Communications_. IEEE, 2019, pp. 1–7.
  * [18] I. Mohammed, S. Tabatabai _et al._ , “Budgeted Online Selection of Candidate IoT Clients to Participate in Federated Learning,” _IEEE Internet of Things Journal_ , vol. 8, no. 7, pp. 5938–5952, 2020.
  * [19] A. Albaseer, M. Abdallah _et al._ , “Client Selection Approach in Support of Clustered Federated Learning over Wireless Edge Networks,” _arXiv preprint arXiv:2108.08768_ , 2021.
  * [20] A. M. Albaseer, M. Abdallah _et al._ , “Fine-Grained Data Selection for Improved Energy Efficiency of Federated Edge Learning,” _IEEE Transactions on Network Science and Engineering_ , 2021.
  * [21] A. Albaseer, M. Abdallah _et al._ , “Threshold-Based Data Exclusion Approach for Energy-Efficient Federated Edge Learning,” _arXiv preprint arXiv:2104.05509_ , 2021.
  * [22] Z. Wu, Q. Ling _et al._ , “Federated variance-reduced stochastic gradient descent with robustness to byzantine attacks,” _IEEE Transactions on Signal Processing_ , vol. 68, pp. 4583–4596, 2020.
  * [23] M. Bakhtiari, R. Nasirigerdeh _et al._ , “Federated Multi-Mini-Batch: An Efficient Training Approach to Federated Learning in Non-IID Environments,” _arXiv preprint arXiv:2011.07006_ , 2020.
  * [24] Y. LeCun, L. Bottou _et al._ , “Gradient-based learning applied to document recognition,” _Proceedings of the IEEE_ , pp. 2278–2324, 1998\.