Improving Text-Audio Retrieval by Text-aware Attention Pooling and Prior Matrix Revised Loss

In existing TAR works, the aggregation function $p$ does not directly consider the input text and is merely a function of the audio frames such as max-pooling, mean-pooling schemes [6]. However, a text is most semantically related to the sub-segments of an audio clip. What’s more, there could be multiple texts matching a specific audio clip, but the frames that are most semantically similar would vary. As such, text-agnostic aggregation functions would capture superfluous and distracting information not stated in the text, which impairs the TAR performance. Therefore, it is important to match a given text with its most semantically similar audio frames.

To that end, we present a learnable text-aware attention pooling (TAP) module $\psi$ for TAR to perform the cross-modal reasoning on the audio frames that are most semantically related to a given text. The core mechanism is a scaled dot product attention [3, 16] between the text $t$ and the frames of an audio clip $a$. By conditioning $\psi$ on $t$, we can generate the audio aggregated embedding that learns to capture the most semantically similar audio frames as described in $t$. The resulting aggregated audio embedding is denoted as $z_{a|t}$, and our similarity function $s(t,a)$ is defined as:

$$z_{a|t}=\psi(c_{a}|t),\quad s(t,a)=\frac{z_{t}\cdot z_{a|t}}{\left\|z_{t}\right\|\left\|z_{a|t}\right\|}.$$

(4)

To elaborate, as shown in Fig. 1, we first project a text embedding $c_{t}\in\mathbb{R}^{D}$ output by the text encoder into a query $Q_{t}\in\mathbb{R}^{1\times D_{p}}$ and an audio embedding $c_{a}\in\mathbb{R}^{T\times D}$ generated by the audio encoder into key $K_{a}\in\mathbb{R}^{T\times D_{p}}$ and value $V_{a}\in\mathbb{R}^{T\times D_{p}}$ matrices, where $D_{p}$ is the size of the projection dimension. The projections are defined as:

$$Q_{t}={\rm LN}(c_{t}^{T})W_{Q},$$

(5)

$$K_{a}={\rm LN}(c_{a})W_{K},$$

(6)

$$V_{a}={\rm LN}(c_{a})W_{V},$$

(7)

where $LN$ represents a layer normalization layer [17] and $W_{Q}$, $W_{K}$ and $W_{V}$ are projection matrices in $\mathbb{R}^{D\times D_{p}}$. In order to flexibly learn the relevance between the given text and the audio frames, we then adapt the scaled dot product attention from the query-projected text embedding to the key-projected frame embedding. The dot product attention provides relevance weights from a text to each audio frame, which we adopt to aggregate the value-projected frame embedding:

$${\rm Attention}(Q_{t},K_{a},V_{a})={\rm softmax}(\frac{Q_{t}K_{a}^{T}}{\sqrt{D_{p}}})V_{a}.$$

(8)

Specifically, the query-projected text embedding is utilized to search frames with high relevance from the key-projected frame embedding. The value-projected embedding represents the audio’s context, from which we aggregate frames conditioned on the given text. To embed an audio clip into a shared space with a text, we project the aggregated audio feature from the attention module back into $\mathbb{R}^{D}$ by leveraging a weight $W_{O}\in\mathbb{R}^{D_{p}\times D}$ to obtain:

$$z_{a|t}={\rm LN}({\rm Attention}(Q_{t},K_{a},V_{a})W_{O}),$$

(9)

where the resulting output $z_{a|t}$ is an aggregated audio embedding depending on the text $t$. By introducing the text-aware attention pooling, the model can concentrate on the most pertinent audio frames as described in a given text, thus effectively mitigating the negative impacts of the heterogeneity of contents between text and audio. Next, we will introduce a prior matrix revised loss to further revise the results for both $t2a$ and $a2t$ tasks.

Previous TAR loss functions [13, 18, 19] (e.g., the NTXent loss) only conduct the softmax for every single-side retrieval, which is just inferred with the similarity score for each row in the original similarity matrix, thus ignoring the potential cross-retrieval information.

Based on the introduced dual optimal match hypothesis from previous extensive TAR experiments [6, 7, 13] that when a $t2a$ (or $a2t$) pair reaches the single-side match, the symmetric $a2t$ (or $t2a$) score should also be the highest, a prior probability matrix is introduced to be calculated in the cross direction for $t2a$ and $a2t$ to fully exploit the cross-retrieval information. Specifically, the prior matrix is obtained by calculating the softmax score along each column of the original similarity matrix and then we multiply the prior matrix with the original similarity matrix to revise the similarity score, that is, for the $t2a$ (or $a2t$) task, we first calculate the softmax score of $a2t$ (or $t2a$), and then incorporate it into the loss calculation of $t2a$ (or $a2t$). Lastly, we conduct the softmax operation along each row of the revised similarity matrix to get the final probability result. Our prior matrix revised (PMR) loss is formulated as below:

$\displaystyle L_{t2a}=-\frac{1}{B}\sum_{i}^{B}{\rm log}\frac{{\rm exp}(s(t_{i},a_{i})\cdot Pr_{i,i}^{t2a}/\tau)}{\sum_{j}^{B}{\rm exp}(\bm{s(t_{i},a_{j})}\cdot Pr_{i,j}^{t2a}/\tau)},$

(10)

$\displaystyle L_{a2t}=-\frac{1}{B}\sum_{i}^{B}{\rm log}\frac{{\rm exp}(s(t_{i},a_{i})\cdot Pr_{i,i}^{a2t}/\tau)}{\sum_{j}^{B}{\rm exp}(\bm{s(t_{j},a_{i})}\cdot Pr_{j,i}^{a2t}/\tau)},$

(11)

$\displaystyle L=L_{t2a}+L_{a2t},$

(12)

where $Pr^{t2a}$, $Pr^{a2t}$ denote the prior matrix for text-to-audio and audio-to-text tasks, respectively:

$\displaystyle Pr_{i,j}^{t2a}=\frac{{\rm exp}(\omega\cdot s(t_{i},a_{i}))}{\sum_{j}^{B}{\rm exp}(\omega\cdot\bm{s(t_{j},a_{i})})},$

(13)

$\displaystyle Pr_{j,i}^{a2t}=\frac{{\rm exp}(\omega\cdot s(t_{i},a_{i}))}{\sum_{j}^{B}{\rm exp}(\omega\cdot\bm{s(t_{i},a_{j})})},$

(14)

where $\omega$ represents a logit scaling parameter to smooth the gradients. In this way, $t2a$ and $a2t$ can coordinately revise each other’s similarity scores, which provides prior knowledge of each other to filter the outliers and sharpen the more convincing points, thus achieving the dual optimal result.

We evaluate our methods on two publicly available datasets: AudioCaps [20] and Clotho [21] datasets. AudioCaps contains about 50K audio samples, which are all 10-second long. The training set consists of 49274 audio clips, each with one corresponding human-annotated caption. The validation and test sets contain 494 and 957 audio clips, each with five human-annotated captions. The Clotho v2 dataset contains 6974 audio samples between 15 and 30 seconds in length. Each audio sample is annotated with 5 sentences. The numbers of training, validation, and test samples are 3839, 1045, and 1045, respectively.

In our work, we follow the same pipeline in [13] to train our network. We adopt BERT [9] as the text encoder, while employing the ResNet-38 in Pre-trained audio neural networks (PANNs) [10] as the audio encoder if not otherwise specified. We conduct experiments by fine-tuning the pre-trained models. The query, key and value projection dimension size is set as $D_{p}=512$. Recall at rank k (R@k) is utilized as the evaluation metric, which is a popular cross-modal retrieval evaluation protocol. R@k measures the proportion of targets retrieved within the top-k ranked results, so a higher score means better performance. The results of R@1, R@5, and R@10 are reported.

In this section, we first compare the performance of our TAP with various text-agnostic pooling functions on the AudioCaps and Clotho datasets. To demonstrate the superiority of our TAP, we compare it with previously popular aggregation strategies that achieve SOTA results, where Mean denotes the average pooling function, MeanMax [13] denotes we both use an average and max pooling layer to aggregate the frame-level feature, NetRVLAD [6, 22] is a descriptor-based pooling method that enables back-propagation by adopting soft assignment to clusters, and TAP represents our text-aware attention pooling method. The experiments are repeated three times with different training seeds.

As shown in Table 1, our TAP outperforms all other works that use text-agnostic pooling on all datasets and across all metrics, thereby highlighting the importance of our text-aware aggregation scheme that can learn to match a text with its most relevant frames while suppressing distracting information from irrelevant audio frames. In addition to adopting ResNet-38 as the backbone of our audio encoder, we also provide the results of using CNN14 [10] as our audio encoding model. It can be seen that our method also achieves performance boosts by a large margin, which strongly demonstrates the effectiveness and generalization of our TAP module. Notably, although the NetRVLAD aggregation method achieves relatively good results, it needs to manually select the number of clusters for different datasets and its tuning of the hyper-parameter is task and data specific. In contrast, our TAP can adaptively learn the optimal amount of information to extract for each text-audio pair, which thus removes the need to manually specify the hyper-parameter and can be more robust to different tasks and instances.

To evaluate our PMR loss, we compare it with the previous SOTA loss for TAR: NTXent, which has be shown in [13] to outperform the popular triplet-based losses. As can be seen in Table 2, our PMR loss shows stable performance boosts on both AudioCaps and Clotho datasets with different aggregation schemes. Besides, our PMR also achieves consistent performance gains on different baseline models, which further demonstrates the effectiveness of our PMR loss.