TRANSAVS: END-TO-END AUDIO-VISUAL SEGMENTATION WITH TRANSFORMER

For the AVS task, the input data comprises a sequence of video frames $\mathcal{V}=\{v_{i}\}_{i=1}^{T}$, where $v_{i}$ $\in\mathbb{R}^{3\times H_{v}\times W_{v}}$, and $T$-second audio stream $\mathcal{A}$. The goal of AVS is to segment all sounding objects in each frame $v_{i}$ under the acoustic guidance $\mathcal{A}$. The segmentation results are binary masks $\mathcal{M}=\{m_{i}\}_{i=1}^{T}$, where $m_{i}\in\{0,1\}^{H_{v}\times W_{v}}$, with ‘1’ indicates sounding objects while ‘0’ corresponds to background or silent objects.

Our TransAVS framework consists of 3 modules: (1) a feature extractor which is responsible for extracting multi-scale image features $F_{v}$ and audio features $F_{a}$; (2) an audio-visual transformer-based fusion module, which disentangles $F_{a}$ into audio queries $A_{q}$ and fuses with $F_{v}$ in a transformer fashion; (3) a mask generation module predicting binary masks with probability of sounding objects.

As mentioned before, objects of the same category tend to produce similar sound frequencies, resulting in a significant degree of homogeneity that can hinder the model’s performance. Toward this, we propose two loss functions at query and mask level, namely the Audio Query Distance Loss (AQDL) and the Audio Query Mask Loss (AQML), with the goal of increasing heterogeneity and thus enhancing segmentation accuracy.

To be specific, AQDL, denoted as $\mathcal{L}_{AQDL}$, is a penalty on $q_{i}$ that predicts sounding objects but getting too close to each other, indicating that they have a high similarity with less clear guidance:

where $h$ represents a projection head in $\mathbb{R}^{d\times d}$, $q_{i}$ and $q_{j}$ are elements of the set $S_{1}=\{z_{k}|p_{k}>\delta_{1}\}$, $\delta_{1}$ is the confidence threshold of $\mathcal{L}_{AQDL}$, $n_{1}$ is the cardinality of $S_{1}$, $||$·$||_{2}$ is the $\mathcal{L}_{2}$ norm, and $d_{0}$ is the threshold for $d(a_{i},a_{j})$. AQDL promotes heterogeneity by restricting queries from getting too close.

On the other hand, AQML encourages $q_{i}$ to predict exclusive sounding masks $m_{i}$ as much as possible. It focuses on the intersecting pixels between different sounding-object masks:

where $m_{i}$ and $m_{j}$ belong to the set $S_{2}=\{z_{k}|$ $p_{k}>\delta_{2}\}$, $\delta_{2}$ is the confidence threshold of $\mathcal{L}_{AQML}$, $n_{2}$ is the cardinality of $S_{2}$, ‘Bin’ denotes the binary operation with threshold set $0.5$ and $\odot$ is the Hadamard product. AQML forces queries to attend to different parts of images, thus reducing their heterogeneity.

Besides AQDL and AQML, we also use 2 supervised segmentation losses, including focal classification loss $\mathcal{L}_{class}$[29] and dice loss $\mathcal{L}_{dice}$[30]. All losses are linearly combined as optimization objective during our end-to-end training process:

During the inference, TransAVS predicts each pixel at location $(x,y)$ based on $\mathcal{Z}$:

that is, when and only when both the class probability $p_{i}(c_{i})$ and the mask prediction probability $m_{i}(x,y)$ are high enough will a pixel $(x,y)$ be assigned to $z_{i}$ .

All videos in AVSBench dataset are trimmed into 5 seconds and separated into 2 subsets based on the number of sound source: single-source sound segmentation (S4) and multiple-sound source segmentation (MS3). Then each video is divided into 5 non-overlapping 1-second clips, each clip is sampled one frame. As shown in Table 2, only the first sampled frame in the training split of S4 is annotated while all frames in other split are annotated. S4 contains 23 classes (Cls.), covering sounds from humans, animals, vehicles, and musical instruments. Each video in the MS3 subset includes two or more categories from the S4 subset.

For the visual backbone, we choose 2 representative ones: the standard CNN-based ResNet[31] backbone R101 and Transformer-based Swin-Transformer[32] backbone swin-base. R101 is pretrained on ImageNet-1K[33] while swin-base on ImageNet-22K. For the loss weights, we set $\lambda_{1}=\lambda_{2}=2.0$ and $\lambda_{3}=\lambda_{4}=5.0$. For the optimizer, we use AdamW [34] with an initial learning rate of 0.0001 for both R101 and swin-base backbones. A learning rate multiplier of $0.1$ is also applied. All models are trained with 8 3090 GPUs for 90k iterations with a batch size of $8$.

Since AVS is a newly proposed problem, we compare our network with baseline in [16] and methods from three related tasks, including sound source localization (SSL), video object segmentation (VOS), and salient object detection (SOD). For each task, we report the results of two SOTA methods on AVSBench dataset, i.e., LVS[21] and MSSL[22] for SSL, 3DC[23] and SST[24] for VOS, iGAN[25] and LGVT[26] for SOD. To ensure fairness, all backbones of these methods were pretrained on the ImageNet-1K[33] dataset.

In Table 3, we verify the effectiveness of each key design in the proposed method with ResNet backbone on both S4 and MS3 subsets. Based on the first 2 rows and the last row, our results show that the optimal performance obtained when the number of queries $N$ was set to 300. In the third and fourth rows, $\delta_{1}$ and $\delta_{2}$ are set to $0.6$, respectively. Both fixed mode show a notable decrease compared with the increasing mode adapted in the last row:

where $a=0.55$, $b=0.65$, $n_{iter}=5000$. We hypothesize that this phenomenon can be explained as follows: as TransAVS becomes increasingly confident about its mask prediction, a fixed threshold may not penalize audio queries with less confidence at earlier stages while pushing too many audio queries away at later epochs. This leads to poorer performance compared to using an increasing threshold.