# TMGAN-PLC: Audio Packet Loss Concealment using Temporal Memory Generative
Adversarial Network

###### Abstract

Real-time communications in packet-switched networks have become widely used
in daily communication, while they inevitably suffer from network delays and
data losses in constrained real-time conditions. To solve these problems,
audio packet loss concealment (PLC) algorithms have been developed to mitigate
voice transmission failures by reconstructing the lost information. Limited by
the transmission latency and device memory, it is still intractable for PLC to
accomplish high-quality voice reconstruction using a relatively small packet
buffer. In this paper, we propose a temporal memory generative adversarial
network for audio PLC, dubbed TMGAN-PLC, which is comprised of a novel nested-
UNet generator and the time-domain/frequency-domain discriminators.
Specifically, a combination of the nested-UNet and temporal feature-wise
linear modulation is elaborately devised in the generator to finely adjust the
intra-frame information and establish inter-frame temporal dependencies. To
complement the missing speech content caused by longer loss bursts, we employ
multi-stage gated vector quantizers to capture the correct content and
reconstruct the near-real smooth audio. Extensive experiments on the PLC
Challenge dataset demonstrate that the proposed method yields promising
performance in terms of speech quality, intelligibility, and PLCMOS.

Index Terms: audio packet loss concealment, generative adversarial network,
temporal memory, vector quantization

##  1 Introduction

With the end-to-end digital packet switched telephone system becoming
increasingly important in daily communication, the degradation of transmitted
audio quality has become an endogenous problem to be solved, which is
typically caused by packet loss, delay, and unrecoverable bit errors during
transmission [1]. The receiver outputs silence when packets are not received
properly, and audible distortion can be identified even at very low packet-
loss rates. In this respect, ”Packet Loss Concealment” (PLC) can be applied to
restore the missing content of packet-loss audio, so as to improve audio
quality in real-time packetized audio communication applications.

With the renaissance of deep neural networks (DNNs), several approaches based
on deep generative models have thrived in the PLC area, namely Deep PLC. By
conducting supervised training on numerous clean-lossy audio pairs, a
multitude of Deep PLC algorithms have demonstrated their superior capability
in handling packet losses and predicting the future frames in a post-
processing manner. In [2], by leveraging the spectral features of previous
audio frames, a feed-forward neural network was proposed to predict the
spectral features and then reconstruct the time-domain signal for each missing
frame. Following the naive neural networks, a novel PLC framework based on
recurrent neural network (RNN) showed the effectiveness of capturing more
audio-related information in estimating lossy audio [3]. Furthermore, by
combining the usage of the long short-term memory networks (LSTM) and
convolutional networks (CNNs), the convolutional recurrent network (CRN) for
PLC [4] achieved remarkable performance in improving speech quality using the
time-domain waveform mapping. More recently, on the account of its superiority
in the audio generation, a generative adversarial network (GAN) was introduced
to restore the missing part of the speech [5] in PLC tasks. For example, a
GAN-based approach [6] utilized mel-spectrogram of the lossy audio as input,
and reconstructed the complete audio in the same way with the vocoders. Our
preliminary work [7] considered both time-domain waveform and frequency-domain
spectrum optimization, in other words, time-domain and frequency-domain
discriminators were integrated into the classical GAN framework, leading to
the improved PLC performance.

Motivated by the aforementioned Deep PLC studies, we aim to establish a low-
complexity PLC framework, which can compensate for short-term losses totally
and smooth over longer loss bursts. In this paper, following the typical GAN
framework, we propose a novel generator inserted with gated vector quantizers,
which learns to map the audio information to a high-level abstract space. On
the basis of audio representation learning, we attach a gating mechanism to
the vector quantizers to perform temporal collation of the encoding features
representing the important content of speech. These modified vector quantizers
enable the proposed network to ”remember” the correlation of speech contexts
like human memory, so as to reduce the impact of lost perturbations. Besides,
an interleaved structure with the nested-UNet and the temporal feature-wise
linear modulation (TFiLM) is utilized to assign temporal features for joint
exploitation of long-term and short-term information correlation. Inspired by
the preliminary study [7], we integrate multi-resolution temporal
discriminators and a complex-valued-spectrum discriminator to collaboratively
provide external guidance in terms of the multi-scale waveform, spectral
magnitude and phase, thus alleviating the time-frequency distortion.
Comprehensive experimental results on the public dataset, namely INTERSPEECH
2022 Deep PLC challenge [8], show that the proposed method achieves
competitive performance and outperforms many state-of-the-art PLC baselines.

Figure 1: The overall diagram of the proposed system. Different modules are
indicated with different colors for better visualization. ”skip” refers to the
skip connection and $\sigma(\lambda)$ is the sigmoid function.
$\left(\bullet\right)$ after each convolutional layer represents the number of
output channels for this convolution, while $\left(\bullet\right)$ after each
sub-UNet indicates the number of (de)encoding layers $M$ within this block.

##  2 Methodlogy

Following with the traditional PLC approach [9], the audio sequence decoded
from a single packet is assumed to contain only one frame without loss of
generality. For audio transmission, the continuous audio ${\bf
x}\in\mathbb{R}^{1\times L}$ is often split into short-time frames with the
frame length $N$, and thus the $j$th audio frame ${\bf
x}_{j}\in\mathbb{R}^{1\times N}$ can be defined as:

$x_{j}(n)=x\left(N\cdot\left(j-1\right)+n\right),n=0,\ldots,N-1,$ (1)

where $n$ is the time index of the $j$th frame, $L$ is the length of the input
mixture and the number of frames $F$ is given by $F=\lfloor L/N\rfloor$. When
the $j$th audio frame is lost, this frame becomes silent, and all observations
in this frame are set to zero, i.e., ${\bf x}_{j}={\bf 0}^{1\times N}$. We
further define a received frame ${\bf\widetilde{x}}_{j}$ which can represent
the clean state and lossy state of the $j$th frame simultaneously. With this
definition, we have ${\widetilde{bfx}}_{j}=I_{j}{bfx}_{j}$, where $I_{j}$
becomes 1 when the $j$th frame is correctly received, otherwise $I_{j}=0$. In
a nutshell, the task for deep PLC is to map the received signal
$\widetilde{x}(n)$ to the original high-quality audio $x(n)$. For real-time
applications, the deep PLC needs to map the $j$th audio frame with the $j$th
received audio frame and its previous received/mapped audio frames for
causality.

As illustrated in Figure 1(a), the proposed system is mainly comprised of two
sub-networks, namely a nested-UNet generator and the collaborative time-domain
and frequency-domain discriminators. The nested-UNet generator aims at
exporting the reconstructed frames ${\bf\widehat{x}}_{j}$ that can deceive the
time-frequency integrated discriminators, while these discriminators attempt
to find the best decision boundary between the reconstructed frames and the
correctly received ones. By using adversarial training, the generated output
is forced to be indistinguishable from the natural audio, and meanwhile
multiple objective metrics can be optimized with different discriminators,
such as the spectral details and voiced$/$unvoiced energy distribution. The
detailed architectures of each module are illustrated below.

###  2.1 Nested-UNet generator

Although the previous generative modeling studies [10] have verified the
effectiveness of UNet-style architectures in the audio reconstruction task,
the construction of deeper UNet architectures usually sacrifices high-
resolution feature maps, which may cause the temporal information loss [11]
and degrade the quality of reconstructed audio. To tackle this shortcoming, we
follow the architecture design patterns from the nested-UNet structure [11],
and utilize a combination of the gated linear unit (GLU) and a sub-UNet in the
proposed generator to mitigate the information loss and enhance the temporal
dependencies, as presented in Figure 1(b).

Specifically, we employ a novel sub-UNet instead of the normal convolutional
layers in original UNet to accomplish the multi-scale audio compression and
restoration, which is a classical symmetric encoder-decoder structure. Then,
before feeding the input features into each sub-Net, we employ GLUs to reduce
the impact of missing information, which have been demonstrated its temporal
modeling performance [12]. Input feature map is gradually downsampled by
multiple encoding layers, and the mirrored decoding layers reconstruct the
encoded features to the original size. Compared with non-nested UNet, the
stacking of proposed sub-UNets can leverage several contextual collations,
progressive sampling and more filters with various receptive fields. More
frequent sampling operations and the spatial resolution adjustments allow the
proposed model to simulate finer reconstruction of lost information, and
effectively extenuate the loss of waveform details caused by direct large-
scale upsampling.

In [13], it has been revealed that the reconstruction of lost frames requires
not only nearby contextual information but also the long-term information to
improve the overall consistency of the reconstructed audio. In this regard, we
employ the temporal feature-wise linear modulation (TFiLM) [14] as the
normalization layer of sub-UNet. By combining the pooling operations and LSTM,
TFiLM can assist in capturing inter-frame dependencies along with frame-level
inputs $Y\in\mathbb{R}^{C\times F\times N}$, which can be formulated as:

$\displaystyle O_{\rm pool}=MaxPooling(Y),$ (2) $\displaystyle O_{\rm
TFiLM}=LSTM(O_{\rm pool})\cdot Y,$ (3)

where the pooling output $O_{\rm pool}\in\mathbb{R}^{C\times F\times 1}$ is
obtained by the pooling operation along the temporal dimension $N$, allowing
LSTM to focus on the information correlation refinement of the frame dimension
$F$. Note that on the account of LSTMs invoking the pooling features, the
computational efficiency of TFiLM layers still meets the requirement of real-
time systems.

Each encoding and decoding layer in a sub-UNet is comprised of plain
convolution, TFiLM normalization and PReLU activation [15]. The downsampling
and upsampling operations are based on max-pooling and sub-pixel convolution
[16], respectively. The number of (de)encoding layers $M$ in the sub-UNets are
set to $\left\\{4,3,2,1\right\\}$ for the encoder and mirrored
$\left\\{1,2,3,4\right\\}$ for the decoder.

###  2.2 Gated vector quantizers

In real-world communication scenarios, the potential large-scale loss
bursts(_e.g._ up to 1000 milliseconds) cause the difficulty to output normal
speech with a small amount of useful context. Motivated by recent studies on
Vector Quantization (VQ) [17], we insert the vector quantizers into the
proposed generator to capture the correct speech content of processing frames.
To be specific, the vector quantizer maps the encoded features
$enc({\bf\widetilde{x}})$ to discrete latent variables, and the decoder learns
to reconstruct the smooth speech ${\bf\widehat{x}}$ with the learnable
discrete codebook ${\bf Q}^{V}=\left\\{q_{1},q_{2},\ldots,q_{V}\right\\}$,
which can be formulated as:

$\displaystyle Quantize(enc({\bf\widetilde{x}}_{j}))=q_{j},$ (4)
$\displaystyle q_{j}=\mathop{\arg\min}\limits_{q\in{\bf
Q}^{V}}(\left\|enc({\bf\widetilde{x}}_{j})-q\right\|_{2}^{2}),$

where $V$ is the codebook size and the function $Quantize$ selects the vector
$q_{j}$ to be closest to $enc({\bf\widetilde{X}}^{l})$ based on Euclidean
distances. The subscript $j$ still denotes the frame index.

Furthermore, owing to the fact that the recording of burst losses in PLC tasks
may be detrimental to the codebook, we incorporate the gating mechanism into
the original VQ architecture. As shown in Figure 1(c), an additional gated
branch is attached to the input layer of the modified vector quantizer, which
consists of a plain convolutional layer and a sigmoid function. The gating
mechanism enhances the contextual relevance of the latent codebook and picks
out prioritized speech content for quantitative learning. In addition, we
cascade multiple vector quantizers in a residual fashion [18] to learn more
fine-grained ambiguous speech content in lossy speech. The unquantized input
vector is passed through the first modified vector quantizer and the error
between the original and quantized vector is then encoded with a sequence of
additional quantizers, resulting in a progressively refined quantization of
the input vector.

The modified vector quantizer is dubbed as gated vector quantizer (GVQ), in
which the kernel size and stride are set to (1, 3) and (1, 1), respectively.
The number of gated vector quantizers is set to 3 and each quantizer uses a
discrete codebook of size 512. To simplify the training process, we optimize
each vector quantizer with exponential moving average k-means, similar to
[19].

###  2.3 Time-frequency discriminators

Inspired by the classification learning in multiple domains, two
discriminators are adopted to separately learn the classification of the
waveform and complex spectrum, which encourage the generator to produce
reconstructed signals that are indistinguishable from correctly received
audio. The diagrams of the proposed discriminators are shown in Figure 1(d)
and (e).

For the waveform discriminator, we utilize the multi-resolution temporal
discriminator [20], in which three structurally identical discriminators are
applied to the input audio at different resolutions. As described in [20],
each single-scalar discriminator is mainly composed of an initial
convolutional layer and four grouped convolutional layers, followed by batch
normalization (BN) [21] and PReLU. For the complex spectrum discriminator, it
follows the design of our previous study [7], in which complex-valued 2D
convolutions are used for the classification of complex-valued spectra and
promote, resulting in fine-grained prediction of both the magnitude and phase
of the lost frames. The complex instance normalization (IN) and complex PReLU
[22] operate on both real and imaginary values. Finally, a $1\times 1$
convolutional layer is employed in both waveform and complex spectrum
discriminators to aggregate the corresponding discrimination results.

###  2.4 Loss function

During the adversarial training process, we adopt a combination of the
adversarial loss and the feature loss. The adversarial loss is used to
establish the competitive relationship between discriminators and the
generator. Similar to [23], we adopt the relativistic average least-square
loss to stabilize the adversarial training, in which the relativistic logit
aggregates the means of multiple discriminative results and reduces the
discrepancy between reconstructed audio ${\bf\widehat{x}}$ and correctly
received audio ${\bf x}$, which can be formulated as:

$\bar{D}_{{\bf x}}=\frac{1}{K}\sum_{k}D_{k}({\bf
x})-\mathbb{E}_{{\bf\widehat{x}}}\left[\frac{1}{K}\sum_{k}D_{k}({\bf\widehat{x}})\right],$
(5)

$\bar{D}_{{\bf\widehat{x}}}=\frac{1}{K}\sum_{k}D_{k}({\bf\widehat{x}})-\mathbb{E}_{{\bf
x}}\left[\frac{1}{K}\sum_{k}D_{k}({\bf x})\right],$ (6)

where $D_{k}({\bf x})$ and $D_{k}({\bf\widehat{x}})$ represent the output
vector of $k$th discriminator with the correctly received frames ${\bf x}$ and
the reconstructed frames ${\bf\widehat{x}}$ as input, respectively. Note that
$k\in\left\\{0,1,\dots,K\right\\}$ denotes the index over the individual
discriminators, where $k=0$ presents the complex spectrum discriminator and
$k\in\left\\{1,\dots,K\right\\}$ presents the different resolutions of the
waveform discriminators ($K=3$ in this paper). Hence, the loss function of the
discriminator and its corresponding generator can be expressed as below:

$\mathcal{L}_{D}=\mathbb{E}_{{\bf x}}[(\bar{D}_{{\bf
x}}-1)^{2}]+\mathbb{E}_{{\bf\widehat{x}}}[(\bar{D}_{{\bf\widehat{x}}}+1)^{2}].$
(7)

$\mathcal{L}_{G}^{\rm
adv}=\mathbb{E}_{{\bf\widehat{x}}}[(\bar{D}_{{\bf\widehat{x}}}-1)^{2}]+\mathbb{E}_{{\bf
x}}[(\bar{D}_{{\bf x}}+1)^{2}],$ (8)

For the feature loss, we incorporate the Mean Squared Error (MSE) loss and the
multi-scale spectral reconstruction loss [24] to facilitate the audio fidelity
and perceptual quality. Given the inpainted time-domain signal
${\bf\widehat{x}}$ and the target one ${\bf x}$ as well as their corresponding
short-time Fourier transforms (STFT) $\widehat{S}$ and $S$, the feature loss
can be given by:

$\displaystyle\mathcal{L}_{G}^{\rm mse}=\left\|{\bf
x}-{\bf\widehat{x}}\right\|^{2}_{F},$ (9) $\displaystyle\mathcal{L}_{G}^{\rm
spec}=\frac{\left\|\left|S(l,k)\right|-\left|\widehat{S}(l,k)\right|\right\|_{F}}{\left\|\left|S(l,k)\right|\right\|_{F}}+\frac{1}{L}\left\|log\frac{\left|S(l,k)\right|}{\left|\widehat{S}(l,k)\right|}\right\|_{1},$
(10)

where $l$ and $k$ denote the frame index and the frequency index,
respectively. The multi-resolution STFT loss is the summation of the STFT
losses with different lengths of window to improve the frequency robustness of
the output speech.

Therefore, the full loss for the generator is given by the weighted sum of all
the above mentioned loss components:

$\mathcal{L}_{G}=\mathcal{L}_{G}^{\rm adv}+\lambda_{\rm
mse}\mathcal{L}_{G}^{\rm mse}+\lambda_{\rm spec}\mathcal{L}_{G}^{\rm spec},$
(11)

where $\lambda_{\rm mse}$ and $\lambda_{\rm spec}$ are the weighting hyper-
parameters, which are initially set to be 1.0 and 0.5, respectively, aiming to
balance the learning rate of each loss.

Table 1: Comparison with other state-of-art systems on the PLC Challenge non-
blind test set. BOLD indicates the highest performance

| Models | Domain | Param.(M) | MACs(G/s) | RTF | WB-PESQ | STOI($\%$) | PLCMOS
---|---|---|---|---|---|---|---|---
| Burst subset | - | - | - | - | (0, 120] | (120,320] | (320,1000] | (0, 120] | (120,320] | (320,1000] | (0, 120] | (120,320] | (320,1000]
| Lossy speech | - | - | - | - | 2.59 | 1.75 | 1.72 | 89.42 | 70.83 | 65.02 | 3.13 | 2.47 | 2.82
| NSNet2 | - | - | - | - | 2.39 | 1.71 | 1.67 | 84.82 | 64.83 | 60.96 | 3.08 | 2.39 | 2.61
| CRN | T | 4.92 | 0.52 | 0.57 | 2.63 | 1.78 | 1.74 | 89.64 | 72.70 | 67.27 | 3.26 | 2.55 | 2.49
| TFNet | T-F | 2.55 | 1.59 | 0.88 | 2.90 | 1.93 | 1.81 | 91.92 | 78.02 | 69.93 | 3.88 | 2.80 | 2.59
| TFGAN | T-F | 1.85 | 0.71 | 0.57 | 3.03 | 2.14 | 1.86 | 92.32 | 78.27 | 70.16 | 3.99 | 3.07 | 2.66
| TMGAN | T-F | 2.85 | 2.12 | 0.74 | 3.37 | 2.33 | 2.08 | 93.37 | 79.50 | 71.56 | 4.57 | 4.28 | 4.21

##  3 Experiments

###  3.1 Datasets

To evaluate the performance of our framework, we conduct extensive experiments
on the INTERSPEECH 2022 Audio Deep Packet Loss Concealment Challenge dataset
[8], which consists of 23,184 clean-lossy pairs for training. The Microsoft
team created this dataset, where each clip of audio was sampled from a base
public domain podcast dataset with the duration about 10 seconds. All the
lossy audio clips were obtained by zeroing out the corresponding regions
according to the losses from the sampled traces [8]. Specifically, the lossy
utterances can be divided into three subsets according to the maximum burst
loss length in the trace:
$\left\\{(0,120]ms,(120,320]ms,(320,1000]ms\right\\}$. For the test set, we
utilize the test dataset of PLC-Challenge, which consists of 966 real
recordings in the English language.

###  3.2 Implementation setup

To leverage the gated vector quantizer to learn more speech content, we
incorporate additional clean utterances from the VCTK corpus [25] in the
training process. All clips of audio are sampled at 16 kHz and the length of
each frame is set to 20 ms. The proposed model and baselines are trained for
100 epochs using the RAdam optimizer [26] with a learning rate of 0.0001, and
the corresponding attenuation rates are $\beta_{1}=0.9$ and $\beta_{2}=0.999$.
To stabilize the training process, the learning rates of the discriminators
are set to 0.00005. The generator directly extracts the features from the
original waveform, while the complex spectrum discriminator needs to perform
320-point FFT on waveform beforehand. The training batch size is set to 8 at
the utterance level. To further smooth the transition between reconstructed
frames and correctly received frames, we introduce a speech enhancement post-
processing module according to our preliminary studies. [27].

##  4 Experimental results and discussion

In this study, we use the wide-band perceptual evaluation of speech quality
(WB-PESQ) [28], short-time objective intelligibility (STOI) [29], word error
rate (WER) and PLCMOS [8] to evaluate PLC performance of the proposed model.
Higher values indicate better performance.

###  4.1 Ablation study

For ablation studies on the proposed model, we reassemble the above modules
and calculate the average evaluation scores of each scheme on the test set. We
first train two UNet-style architectures with the guidance of mixed
time/frequency-domain discriminators, named traditional UNet and nested-UNet.
As shown in Table 2, the nested structure achieves significant improvements
than traditional UNet in terms of all metrics, especially in improving STOI
and PLCMOS scores. Subsequently, the combination of the nested-UNet and TFiLM
achieves consistently better performance, indicating that the memory mechanism
of LSTM can effectively improve the reconstruction quality of lossy speech.
Finally, when incorporating the single raw vector quantizer$/$multi-stage
gated vector quantizers into the nested-UNet, one can observe that the naive
VQ affected by the burst losses even degrades speech quality, while GVQ
demonstrates its superiority in capturing correct speech content.

###  4.2 Comparison with SOTA methods

In different burst-loss length subsets, we further compare our model with
several state-of-the-art PLC methods, including the PLC-Challenge baseline
NSNet2 [8], CRN [4], TFNet [30] and our earlier proposed model TFGAN [7]. From
Table 1, several observations can be obtained. First, by introducing the
spectral features, all the mixed time/frequency methods outperform CRN by a
large margin, which only relies on the time-domain information. Second, when
comparing TFNet with TFGAN, one can observe that transferring the time-
frequency mapping task to discriminators does not degrade the reconstructed
speech quality. This operation simplifies the inference procedure of the
generator, leading to less real-time processing time (RTF). Third, compared
with other advanced PLC methods, the proposed method achieves relatively
better performance in terms of all metrics. Furthermore, in the large-scale
burst subset of $(320,1000]ms$, the proposed method significantly surpasses
other baselines in terms of the PLCMOS score, demonstrating the remarkable
reconstruction performance of TMGAN. In Table 2, we present the subjective
results of the proposed method in terms of PLCMOS [8], DNSMOS [31] and Crowd-
Sourced Mean Opinion Score (CMOS) [8] on the PLC Challenge blind test set.
Compared with the zero-filling baseline, TMGAN provides significant
improvements in terms of overall speech quality. The number of trainable
parameters of the proposed framework is 2.85 million, and the number of
multiply-accumulate operations (MACs) per second is 2.12 G as shown in Table
3. The one-frame processing time of our system with the PyTorch implementation
is around 14.8 ms on an Intel i5-4300U PC. With two frames looking ahead
(40ms) and the strade time (20 ms), the algorithm latency of proposed model is
about 80 ms.

Table 2: Ablation study on the proposed method. “$+$TFiLM” denotes the
introduction of temporal feature-wise linear modulations. “$+$VQ” denotes that
a raw vector quantizer is incorporated, while “$+$GVQ” represents the multi-
stage gated vector quantizers.

| Models | WB-PESQ | STOI($\%$) | PLCMOS
---|---|---|---|---
| Traditional UNet | 2.44 | 82.24 | 3.25
| Nested-UNet | 2.59 | 84.37 | 3.54
|  \+ TFiLM | 2.73 | 84.91 | 4.04
|  \+ VQ | 2.65 | 84.02 | 3.86
|  \+ GVQ | 2.83 | 85.47 | 4.14

Table 3: Subjective evaluation on PLC Challenge blind test set.

| models | PLCMOS | DNSMOS | CMOS
---|---|---|---|---
| Zero-filling baseline | 2.904 | 3.444 | -1.231
| TMGAN | 4.406 | 3.958 | -0.279

##  5 Conclusions

This paper proposes a novel packet loss concealment model based on temporal
memory generative adversarial network (TMGAN-PLC). Taking a nested-UNet and
the temporal feature-wise linear modulation as the backbone, the proposed
generator can aggregate local and global features by leveraging the memory
function in LSTMs, thus reducing the effect of buffer size on the
reconstruction quality. Furthermore, we design the multi-stage vector
quantizers with a gating mechanism to deal with burst losses. In the gated
vector quantizers, the encoding learning can autonomously avoid missing the
large regions, and the decoding reconstruction incorporates the codebook and
quantization features of previous frames, which can recover the high-quality
audio even for long burst losses. Experimental results demonstrate the
effectiveness of each temporal memory module, and show the competitive
performance of the proposed method in terms of speech quality, intelligibility
and PLCMOS.

## References

  * [1] A. Takahashi, H. Yoshino, and N. Kitawaki, “Perceptual QoS assessment technologies for VoIP,” _IEEE Communications Magazine_ , vol. 42, no. 7, pp. 28–34, 2004.

  * [2] B. K. Lee and J. H. Chang, “Packet loss concealment based on deep neural networks for digital speech transmission,” _IEEE/ACM Transactions on Audio Speech and Language Processing (TASLP)_ , vol. 24, no. 2, pp. 378–387, 2016\.
  * [3] R. Lotfidereshgi and P. Gournay, “Speech prediction using an adaptive recurrent neural network with application to packet loss concealment,” in _2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_. IEEE, 2018, pp. 5394–5398.
  * [4] J. Lin, Y. Wang, K. Kalgaonkar, G. Keren, D. Zhang, and C. Fuegen, “A time-domain convolutional recurrent network for packet loss concealment,” in _2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_. IEEE, 2021, pp. 7148–7152.
  * [5] Y. Shi, N. Zheng, Y. Kang, and W. Rong, “Speech loss compensation by generative adversarial networks,” in _2019 Asia-Pacific Signal and Information Processing Association Annual Summit and Conference (APSIPA ASC)_. IEEE, 2019, pp. 347–351.
  * [6] S. Pascual, J. Serrà, and J. Pons, “Adversarial auto-encoding for packet loss concealment,” in _2021 IEEE Workshop on Applications of Signal Processing to Audio and Acoustics (WASPAA)_. IEEE, 2021, pp. 71–75.
  * [7] J. Wang, Y. Guan, C. Zheng, R. Peng, and X. Li, “A temporal-spectral generative adversarial network based end-to-end packet loss concealment for wideband speech transmission,” _The Journal of the Acoustical Society of America_ , vol. 150, no. 4, pp. 2577–2588, 2021.
  * [8] L. Diener, S. Sootla, S. Branets, A. Saabas, R. Aichner, and R. Cutler, “Interspeech 2022 audio deep packet loss concealment challenge,” in _INTERSPEECH 2022 - 23rd Annual Conference of the International Speech Communication Association_ , 2022 (submitted).
  * [9] E. Gunduzhan and K. Momtahan, “Linear prediction based packet loss concealment algorithm for PCM coded speech,” _IEEE Transactions on Speech and Audio Processing_ , vol. 9, no. 8, pp. 778–785, 2001.
  * [10] V. Kuleshov, S. Z. Enam, and S. Ermon, “Audio super-resolution using neural nets,” in _the Fifth International Conference on Learning Representations (ICLR 2017-Workshop Track)_ , 2017.
  * [11] X. Qin, Z. Zhang, C. Huang, M. Dehghan, O. R. Zaiane, and M. Jagersand, “U2-net: Going deeper with nested u-structure for salient object detection,” _Pattern Recognition_ , vol. 106, p. 107404, 2020.
  * [12] A. Li, C. Zheng, L. Cheng, R. Peng, and X. Li, “A time-domain monaural speech enhancement with feedback learning,” in _2020 Asia-Pacific Signal and Information Processing Association Annual Summit and Conference (APSIPA ASC)_. IEEE, 2020, pp. 769–774.
  * [13] A. Marafioti, P. Majdak, N. Holighaus, and N. Perraudin, “Gacela: A generative adversarial context encoder for long audio inpainting of music,” _IEEE Journal of Selected Topics in Signal Processing_ , vol. 15, no. 1, pp. 120–131, 2020.
  * [14] S. Birnbaum, V. Kuleshov, Z. Enam, P. W. W. Koh, and S. Ermon, “Temporal film: Capturing long-range sequence dependencies with feature-wise modulations.” _Advances in Neural Information Processing Systems (NeurIPS 2019)_ , vol. 32, pp. 10 287–10 298, 2019.
  * [15] K. He, X. Zhang, S. Ren, and J. Sun, “Delving deep into rectifiers: Surpassing human-level performance on imagenet classification,” in _Proceedings of the 2015 IEEE International Conference on Computer Vision (ICCV)_ , 2015, pp. 1026–1034.
  * [16] W. Shi, J. Caballero, F. Huszár, J. Totz, A. P. Aitken, R. Bishop, D. Rueckert, and Z. Wang, “Real-time single image and video super-resolution using an efficient sub-pixel convolutional neural network,” in _Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR)_ , 2016, pp. 1874–1883.
  * [17] A. van den Oord, O. Vinyals, and K. Kavukcuoglu, “Neural discrete representation learning,” in _Proceedings of the 31st International Conference on Neural Information Processing Systems (NIPS’17)_ , California, USA, 2017, pp. 6309–6318.
  * [18] A. Vasuki and P. Vanathi, “A review of vector quantization techniques,” _IEEE Potentials_ , vol. 25, no. 4, pp. 39–47, 2006.
  * [19] A. Razavi, A. van den Oord, and O. Vinyals, “Generating diverse high-fidelity images with vq-vae-2,” vol. 32, Vancouver, Canada, 2019, pp. 14 866–14 876.
  * [20] K. Kumar, R. Kumar, T. de Boissiere, L. Gestin, W. Z. Teoh, J. Sotelo, A. de Brébisson, Y. Bengio, and A. C. Courville, “Melgan: Generative adversarial networks for conditional waveform synthesis,” _Advances in Neural Information Processing Systems (NeurIPS 2019)_ , vol. 32, 2019.
  * [21] S. Ioffe and C. Szegedy, “Batch normalization: Accelerating deep network training by reducing internal covariate shift,” in _Proceedings of the 32nd International Conference on Machine Learning_. PMLR, 2015, pp. 448–456.
  * [22] Y. Hu, Y. Liu, S. Lv, M. Xing, S. Zhang, Y. Fu, J. Wu, B. Zhang, and L. Xie, “DCCRN: Deep complex convolution recurrent network for phase-aware speech enhancement,” _Proceedings of the 21th Annual Conference of the International Speech Communication Association (Interspeech 2020)_ , pp. 2472–2476, 2020.
  * [23] D. Baby and S. Verhulst, “SERGAN: Speech enhancement using relativistic generative adversarial networks with gradient penalty,” in _2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_. IEEE, 2019, pp. 106–110.
  * [24] R. Yamamoto, E. Song, and J.-M. Kim, “Parallel WaveGAN: A fast waveform generation model based on generative adversarial networks with multi-resolution spectrogram,” in _2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_. IEEE, 2020, pp. 6199–6203.
  * [25] C. Veaux, J. Yamagishi, and S. King, “The voice bank corpus: Design, collection and data analysis of a large regional accent speech database,” in _2013 international conference oriental COCOSDA held jointly with 2013 conference on Asian spoken language research and evaluation (O-COCOSDA/CASLRE)_. IEEE, 2013, pp. 1–4.
  * [26] L. Liu, H. Jiang, P. He, W. Chen, X. Liu, J. Gao, and J. Han, “On the variance of the adaptive learning rate and beyond,” in _the Eighth International Conference on Learning Representations (ICLR 2020)_ , Addis Ababa, Ethiopia, 2020\.
  * [27] A. Li, W. Liu, X. Luo, G. Yu, C. Zheng, and X. Li, “A simultaneous denoising and dereverberation framework with target decoupling,” _Proceedings of the 22th Annual Conference of the International Speech Communication Association (Interspeech 2021)_ , pp. 2801–2805, 2021.
  * [28] A. W. Rix, J. G. Beerends, M. P. Hollier, and A. P. Hekstra, “Perceptual evaluation of speech quality (PESQ)-a new method for speech quality assessment of telephone networks and codecs,” in _2001 IEEE International Conference on Acoustics, Speech, and Signal Processing (ICASSP)_ , vol. 2. IEEE, 2001, pp. 749–752.
  * [29] C. H. Taal, R. C. Hendriks, R. Heusdens, and J. Jensen, “A short-time objective intelligibility measure for time-frequency weighted noisy speech,” in _2010 IEEE international conference on acoustics, speech and signal processing (ICASSP)_. IEEE, 2010, pp. 4214–4217.
  * [30] X. Jiang, X. Peng, C. Zheng, H. Xue, Y. Zhang, and Y. Lu, “End-to-end neural audio coding for real-time communications,” _arXiv preprint arXiv:2201.09429_ , 2022.
  * [31] C. K. Reddy, V. Gopal, and R. Cutler, “Dnsmos: A non-intrusive perceptual objective speech quality metric to evaluate noise suppressors,” in _ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_. IEEE, 2021, pp. 6493–6497.


# Soft-Prompting with Graph-of-Thought for Multi-modal Representation Learning

###### Abstract

The chain-of-thought technique has been received well in multi-modal tasks. It
is a step-by-step linear reasoning process that adjusts the length of the
chain to improve the performance of generated prompts. However, human thought
processes are predominantly non-linear, as they encompass multiple aspects
simultaneously and employ dynamic adjustment and updating mechanisms.
Therefore, we propose a novel Aggregation-Graph-of-Thought (AGoT) mechanism
for soft-prompt tuning in multi-modal representation learning. The proposed
AGoT models the human thought process not only as a chain but also models each
step as a reasoning aggregation graph to cope with the overlooked multiple
aspects of thinking in single-step reasoning. This turns the entire reasoning
process into prompt aggregation and prompt flow operations. Experiments show
that our multi-modal model enhanced with AGoT soft-prompting achieves good
results in several tasks such as text-image retrieval, visual question
answering, and image recognition. In addition, we demonstrate that it has good
domain generalization performance due to better reasoning.

Keywords: chain-of-thought, aggregation-graph-of-thought, multi-modal,
representation learning

Soft-Prompting with Graph-of-Thought for Multi-modal Representation Learning

Juncheng Yang1,3,∗, Zuchao Li1,2,†, Shuai Xie4,∗, Wei Yu1, Shijun Li1,†, Bo
Du1,2
---
1School of Computer Science, Wuhan University, Wuhan, Hubei, China
2National Engineering Research Center for Multimedia Software, Wuhan, Hubei,
China
3School of Electronic Information Engineering, Henan Polytechnic Institute,
Nanyang, Henan, China
4JD Explore Academy, Beijing, China
<EMAIL_ADDRESS><EMAIL_ADDRESS>

Abstract content

††footnotetext: ∗Equal Contribution.††footnotetext: †Correspondence to: Zuchao
Li, Shijun Li.

## 1\. Introduction

Multi-modal representation learning includes the fusion of different modal
data such as visual, text, and sound which is essential for multi-modal tasks
like text-to-speech synthesis Ren et al. (2019), audio-visual-speech
recognition Afouras et al. (2018), text-image retrieval Gabeur et al. (2022);
Yu et al. (2022), vision-language recognition Zhou et al. (2022a), and visual
question answering Li et al. (2022); Zheng et al. (2021).

In vision-language models, prompt learning has been widely adopted as an
effective strategy Zhu et al. (2023); Ge et al. (2023).The prompt was first
applied in the field of NLP, and its role is to give a prompt to the pre-
trained language model to help it understand human problems better. The
earliest prompts were directly spliced with the original text into discrete
characters that remained unchanged during the training process, known as hard
prompts Schick and Schütze (2021); Han et al. (2022). This was later developed
into soft prompts Gu et al. (2022); Liu et al. (2022), which can be tuned
based on contextual semantics during training. While in multi-modal tasks,
CoOp Zhou et al. (2022b) first turns a fixed token prompt in CLIP Radford et
al. (2021) into a learnable token embedding. CoCoOp Zhou et al. (2022a) adds
an instance image feature to CoOp, and KgCoOp Yao et al. (2023) uses a fixed
text prompt and a learnable prompt as supervised signals to constrain the
training of the CLIP model. However, these approaches rely on a single
language cue to guide visual understanding but fail to be aware of a basic but
very important feature of language, reasoning, which allows multi-modal models
to quickly adapt to new scenarios.

Figure 1: Comparison between (a) ordinary step-by-step Chain-of-Thought and
(b) Aggregation-Graph-of-Thought.

The Chain-of-Thought (CoT) technique simulates the reasoning process of the
brain and imitates its behavior to improve the generalization of the model.
For complex problems, instead of providing the answer directly, a string of
reasoning sentences is used to generate the answer. As shown in Figure 1 (a),
a picture containing some books and a clock is given by CoT, and then reason
about the content of the picture step by step. From the initial description of
“There are some things on the table” to “There are some books on the
bookshelf”, we can see that the statement’s content becomes fuller. Finally,
the phrase “There are some books and a clock on the bookshelf” completes a
relatively complete description of the image. This step-by-step reasoning
improves the model’s visual understanding performance. CoT-PT Ge et al. (2023)
was the first to adapt CoT to the field of vision-language tasks, combining
visual and CoT cues to tune the model and achieve good results in downstream
tasks.

Although the CoT technique considers the progressively deeper understanding
process during reasoning, it fails to fully leverage the benefits of reasoning
for multi-modal comprehension due to its disregard of the fact that
comprehension can occur from multiple perspectives. In response to this
characteristic, we propose a new CoT reasoning method called Aggregation-
Graph-of-Thought (AGoT). As shown in Figure 1 (b), when inputting a picture,
each step in AGoT, from “There are two kinds of things on the table” to “There
are some books and a clock on the bookshelf” to “There are some different
types of books and a black clock on the bookshelf”, each step of the reasoning
process is more reasonable than Figure 1 (a) due to understanding aggregation
of multiple aspects. It’s apparent that at each reasoning node, AGoT
aggregates multiple levels of aspects, whereas CoT prompts consider only one
level. From the simulated results on the far right, it can be observed that
AGoT prompts, after aggregating multi-level aspects, provide a more
comprehensive understanding of the problem. Specifically, the AGoT models each
step of reasoning in the CoT as a reasoning aggregation graph to cope with the
multiple-view thinking problem that is overlooked in single-step reasoning.
Each aggregation node in the graph also serves as a node within the CoT.
Multiple nodes are interconnected to form a chain, with each node’s
representation aggregated from the previous node, multiple subnodes, and
vision features. The resultant node representation is then passed to the next
one. We design a flow controller to control the degree of information flow
between the previous node and the current one.

We used the strong baseline model CLIP as our base model and conducted
experiments in vision-language understanding tasks, text-image retrieval
(Flickr30k Plummer et al. (2015), MSCOCO Chen et al. (2015)), and visual
question answering (VQAv2 Shen et al. (2022)). Our AGoT method achieved better
results compared to CLIP +5.70%, +5.40%, +19.91% in each of these tasks, and
also showed improvement compared to CLIP + CoT, +1.70%, +0.80%, +0.88%
respectively. Additionally, since our AGoT can be used as a soft-prompting
method, it can improve the generalization ability of visual classification
tasks. We evaluated it on various visual classification tasks and found that
AGoT achieved better domain generalization ability than the previous state-of-
the-art method, demonstrating the effectiveness of the AGoT method.

In summary, our contributions are three-fold:

  * •

we propose a novel Aggregation-Graph-of-Thought mechanism for soft-prompt
tuning in multi-modal representation learning.

  * •

To cope with the overlooked multiple aspects of thinking in single-step
reasoning, we model the reasoning step as an aggregation graph, and turn the
whole process into a prompt aggregation and prompt flow operation.

  * •

AGoT exhibits strong multi-modal representation learning ability and achieves
good results on 18 datasets such as text-image retrieval, VQA, and image
classification and has good domain generalization ability.

Figure 2: Illustration of the proposed AGoT. It learns a high-quality soft
prompt with prompt aggregation (blue) and prompt flow operation (gray) for
multi-view thinking and adaptation to complex multi-modal tasks.

## 2\. Related Work

#### Multi-modal Representation Learning

Multi-modal representation learning has a significant number of models based
on the structure of vision-language encoding. CLIP is a text encoder that uses
Transformer Vaswani et al. (2017), while the image encoder is either ResNet He
et al. (2016) or Vision Transformer Dosovitskiy et al. (2021). The model is
trained on 400 million image-text pairs and performs well on 30 datasets.
ALIGN Jia et al. (2021) is another typical vision-language model that uses
contrast learning to train more than one billion noisy image-text pairs. It
achieves state-of-the-art performance on the MSCOCO and Flickr30k datasets.

These multi-modal foundational models are widely used for downstream tasks,
such as image classification, video-text retrieval, tracking, image
captioning, and point cloud understanding. CoOp enhances the generalization
performance of the model by converting the fixed text description of “a photo
of a” in CLIP into a learnable token embedding. To address the challenge of
poor performance on unseen data, CoCoOp incorporates instance features into
the text description, resulting in further improvements in the model’s
performance. KgCoOp enhances the model’s generalization by minimizing the
disparity between learnable prompts and hand-crafted prompts. DualCoOp Sun et
al. (2022) proposes using dual context optimization as a unified framework to
address the multi-label recognition tasks. ProGrad Zhu et al. (2023) updates
the prompt in a way that aligns its gradient with general knowledge, while
ProDA Lu et al. (2022) focuses on learning prompt prior distributions. These
methods utilize language cues to assist in visual tasks. However, they fall
short in replicating the innate reasoning process employed by humans, thereby
missing out on the advantages of deep visual understanding acquired through
vision-language reasoning.

#### Chain-of-Thought

CoT technology excels in Large Language Model (LLM) few-shot reasoning and
fine-tuning. Wei et al. Wei et al. (2022) was the first to use CoT as a
discrete cue learning in LLM, with good results. Self-Consistency-CoT Wang et
al. (2023) improves the performance of the CoT approach significantly by using
majority voting on answers. STaR Zelikman et al. (2022) proposes a boosting
approach that allows even small and medium-sized models to use CoT. AutoCoT
Zhang et al. (2023a) samples diverse problems and automatically generates
reasoning chains. Shi et al. Shi et al. (2023) evaluate the strong
multilingual reasoning power of CoT in solving multilingual grade school math
problems. ToT Yao et al. (2024) allows LLM to choose a course of action by
considering many different paths of reasoning and self-assessment. CoT also
has a wide range of applications in multi-modal tasks. Multimodal-CoT Zhang et
al. (2023b) integrates language and visual modalities into a two-stage
framework, using multi-modal information to infer answers better. He et al. He
et al. (2023) propose generating effective image features consistent with
language thoughts through the potential space diffusion process. CoT-PT Ge et
al. (2023) uses CoT prompt tuning for vision-language models. VCoT Rose et al.
(2023) incorporates visual augmentation into reasoning tasks, using CoT cues
and visuals to bridge logical gaps in sequential data recursively. However,
this method is geared towards generative tasks, while our method addresses
downstream discriminative tasks.

## 3\. Methodology

### 3.1. Contrastive Learning

To perform multi-modal representation learning, this work follows contrastive
learning in paired text and image. Consistent with CLIP, let $x$ be the input
image, and $t$ be the text sequence formed by the prompt and corresponding
class. The model computes the probability of each class $c_{k}$:

$p(c_{k}|x)=\frac{e^{cos(\mathcal{G}(t_{k}),\mathcal{H}(x))/\tau}}{\sum_{i=1}^{K}e^{cos(\mathcal{G}(t_{i}),\mathcal{H}(x))/\tau}},$
(1)

where $cos(\cdot,\cdot)$ denotes the cosine similarity, $\mathcal{H}(x)$
stands for the image representation from the image encoder,
$\mathcal{G}(t_{k})$ stands for the text representation from the text encoder
of concatenated prompt and class description sentences $t_{k}$, $K$ is the
total number of classes, and $\tau$ is the temperature parameter. The goal of
contrastive learning is to learn a generalized image encoder
$\mathcal{H}(\cdot)$ and text encoder $\mathcal{G}(\cdot)$ using limited pairs
of data, so that the model has a cross-modal semantic space and can better
transfer to downstream tasks such as text-image retrieval, VQA, and fine-
grained classification.

### 3.2. Prompt Learning and CoT

CoOp is the first to turn the fixed prompt, such as “a photo of a”, in the
CLIP model into a learnable embedding $h^{p}$, thus soft-prompting. This
overcomes the need for skill and engineering experience in manually creating
prompts and surpasses CLIP in overall performance across 14 datasets for image
recognition, fine-grained classification, and more. The probability for each
class can be expressed as:

$p(c_{k}|x)=\frac{e^{cos(\mathcal{G}([h^{p}_{k},t_{k}]),\mathcal{H}(x))/\tau}}{\sum_{i=1}^{K}e^{cos(\mathcal{G}([h^{p}_{i},t_{i}]),\mathcal{H}(x))/\tau}},$
(2)

where $[h^{p}_{k},t_{k}]$ denotes the concatenation of learnable prompt
embedding and text embedding.

To solve the problem that CoOp performs poorly on unseen classes, CoCoOp
further uses a simple MetaNet network to generate conditional tokens for each
image and then combines this image-side tokens with the text prompt, which
dynamically adapts to the changes of each image instance, overcoming the
problem of class shift. The prediction probability is computed as follows:

$p(c_{k}|x)=\frac{e^{cos(\mathcal{G}([h^{p}_{k},M(x),t_{k}]),\mathcal{H}(x))/\tau}}{\sum_{i=1}^{K}e^{cos(\mathcal{G}([h^{p}_{i},M(x),t_{i}]),\mathcal{H}(x))/\tau}},$
(3)

where $M(x)$ is the representation of conditional tokens from MetaNet.

CoT-PT introduces CoT for prompt learning in vision-language tasks. It
utilizes image features for soft prompt learning, and the prediction
probability is computed as follows:

$p(c_{k}|x)=\frac{e^{cos(\mathcal{G}([cot(h^{p}_{k}),M(x),t_{k}]),\mathcal{H}(x))/\tau}}{\sum_{i=1}^{K}e^{cos(\mathcal{G}([cot(h^{p}_{i}),M(x),t_{i}]),\mathcal{H}(x))/\tau}},$
(4)

where $cot$ is a CoT module for prompt learning.

### 3.3. Aggregation-Graph-of-Thought

The CoT achieves good results in various tasks due to its superior reasoning
abilities. However, it cannot gather information from different instances in
multiple views under a single reasoning step. This limitation restricts
further assistance of reasoning in multi-modal understanding. To address this
problem, we leverage the aggregation mechanism in the graph neural network and
employ a chain of subgraphs $G=\\{G_{1},\dots,G_{Z}\\}$ to construct the AGoT.
Each subgraph $G_{i}=\left(N_{i},E_{i}\right)$ is a directed weighted graph
built on a group of subnodes $n_{i}^{s}=\\{n_{i1}^{s},...,n_{iR}^{s}\\}$,
which learns meta-prompts from $R$ different views and aggregates to the
central node $N_{i}$ with the learned weights $E_{i}$ as edges, as shown in
Figure 2. Further, the aggregated prompt representation is combined with the
visual information extracted by the MetaNet and then flows to the next
subgraph with the flow controller. After $Z$ reasoning steps, the final AGoT
prompt representation $\mathcal{E}(G_{Z})$ is generated and combined with the
class label [CLASS] to build the final input prompt for the text encoder.
Finally, a CLIP-like multi-modal contrastive learning network is adopted to
make the classification or description of the image with the highest
similarity.

#### Prompt Aggregation

To better explore the prompt information from different perspectives, we
propose a prompt aggregation approach for each reasoning step of the AGoT to
generate high-quality prompt representation. Detailly, in each subgraph
$G_{i}$, $R$ Gaussian-initialized meta-prompt learners act as subnodes
$n_{i}^{s}$, and WeightNet learns the aggregation edge weights $E_{i}$ from
the image representation $\mathcal{H}(x)$. The central node prompt
representation $\mathcal{E}(N_{i})$ can be expressed as:

$\displaystyle\mathcal{E}(N_{i})$
$\displaystyle=\sum_{r=1}^{R}\mathcal{E}(n_{ir}^{s})*E_{ir}$ (5)
$\displaystyle where\ E_{ir}$
$\displaystyle=\textit{{WeightNet}}_{ir}(\mathcal{H}(x))$ $\displaystyle
subject$ $\displaystyle\leavevmode\nobreak\ to\leavevmode\nobreak\
\sum_{r=1}^{R}E_{ir}=1,$

where $\textit{{WeightNet}}_{ir}(\cdot)$ is implemented as 3-layer MLP network
with ReLU as activation function for subgraph $G_{i}$.

As discussed in CoCoOp, it adds a MetaNet network to encode the visual
features generated by the image encoder and adds to the text prompt as a
deviation term. We follow this design and introduce a MetaNet network to
aggregate the visual information into the central node representation. This
aggregation process can be expressed as:

$\mathcal{E}(G_{i})=\mathcal{E}(N_{i})+\textit{{MetaNet}}_{i}(\mathcal{H}(x)).$
(6)

#### Prompt Flow

In AGoT, we decompose problem-solving into a step-by-step reasoning process,
where each step relays the information from the previous steps to gradually
deepen understanding. Intuitively, as the complexity of the task increases, a
more intricate chain of reasoning graph becomes necessary. To cope with this
issue, we introduce a dynamic prompt information flow controller.
Specifically, given an input image, the encoded image features are fed to the
flow controller, which is also a 3-layer MLP network with ReLU activation that
generates an $\alpha$ to control the prompt fusion ratio between steps. The
prompt flow control operation can be represented as:

$\displaystyle\alpha_{i}$
$\displaystyle=\textit{{FlowControl}}_{i}(\mathcal{H}(x)),$ (7)
$\displaystyle\mathcal{E}(G_{i})$
$\displaystyle=(1-\alpha_{i})*\mathcal{E}(G_{i-1})+\alpha_{i}*\mathcal{E}(G_{i}).$

After the $Z$ reasoning steps, the final AGoT prompt representation is defined
as the last subgraph representation $\mathcal{E}(G_{Z})$, which is combined
with the image class label “[CLASS]” to build the input prompt for the text
encoder. Finally, we conduct a standard CLIP-like multi-modal contrastive
learning to train the AGoT network. To provide a more detailed description of
the entire model’s process, we use pseudocode to describe AGoT as shown in
Algorithm 1.

0: $G=\\{G_{1},\dots,G_{Z}\\}$: a chain of $Z$ subgraphs with initialized
parameter $\Theta$

$n_{i}^{s}=\\{n_{i1}^{s},...,n_{iR}^{s}\\}$: a group of $R$ meta-prompt
learners as subnodes for each $G_{i}$

WeightNet: the network that controls the aggregation weights of meta-prompts
in $G_{i}$

FlowControl: the network that controls the fusion ratio between steps in flow
operation

MetaNet: the network that encodes the visual information as a deviation term

$D$: the training dataset

$T$: the training iterations

0: Optimized $\Theta$

1: for $t=1,2,...,T$ do

2: Sample a batch of images $x$ from $D$

3: Extract the image features $\mathcal{H}(x)$ with CLIP image encoder

4: Compute the soft-prompt representation with AGoT

5: for each subgraph $G_{i}$ do

6: Compute the central node representation $\mathcal{E}(N_{i})$ with Eq. 5

7: Combine the visual information with Eq. 6

8: Combine the prompts from the last step with Eq. 7

9: end for

10: Get the final AGoT representation $\mathcal{E}(G_{Z})$

11: Build the final input prompt to CLIP text encoder with
$\mathcal{E}(G_{Z})$ and [CLASS]

12: Update $\Theta$ with CLIP-like multi-modal contrastive learning

13: end for

14: return Optimized $\Theta$

Algorithm 1 Aggregation-Graph-of-Thought Algorithm

## 4\. Experiments

We followed the exact experimental setup of previous works Zhou et al. (2022b,
a) and validated the multi-modal reasoning performance of AGoT against recent
state-of-the-art prompt learning models on widely-used benchmarks under
various settings. Our approach was mainly evaluated in the following five
primary tasks including text-image retrieval, VQA, cross label generalization,
cross dataset generalization and domain generalization. Extensive experiments
on a total of 18 datasets validate the effectiveness of our architecture.
Meanwhile, we devised four ablation experiments and investigated the best
practice of our method on various datasets.

### 4.1. Setup

#### Datasets

For the text-image retrieval task, we use Flickr30k, which contains 29K/1K/1K
for train/validation/test, and MSCOCO, which contains 113K/5K/5K for
train/validation/test. For the VQA task, we use the VQAv2 dataset, which
contains 1080K/210K/104K images paired with questions and answers for
training/validation/test.

For the cross-label generalization ability of our proposed AGoT, we adopt 11
classification datasets, encompassing 2 object classification datasets
(ImageNet Deng et al. (2009), Caltech101 Fei-Fei et al. (2004)), 5 fine-
grained image recognition datasets (OxfordPets Parkhi et al. (2012),
StanfordCars Krause et al. (2013), Flowers102 Nilsback and Zisserman (2008),
Food101 Bossard et al. (2014), FGVCAircraft Maji et al. (2013)), 1 scene
recognition datasets (SUN397 Xiao et al. (2010)), 1 texture classification
datasets (DTD Cimpoi et al. (2014)), 1 satellite image classification datasets
(EuroSAT Helber et al. (2019)), and an action classification (UCF101 Soomro et
al. (2012)). For the cross-dataset generalization setting, all models are
trained on ImageNet(Ima) in the source domain, and then the trained models are
tested on Caltech101(Cal), OxfordPets(Oxf), StanfordCars(Sta),
Flowers102(Flo), Food101(Foo), FGVCAircraft(FGV), SUN397(SUN), DTD,
EuroSAT(Eur), and UCF101(UCF) in the target domain.

For the domain generalization setting, we use ImageNet as the source domain
dataset and evaluate performance on ImageNetV2 Recht et al. (2019), ImageNet-
Sketch Wang et al. (2019), ImageNet-A Hendrycks et al. (2021b), and ImageNet-R
Hendrycks et al. (2021a). Text and image classification were used for cross-
dataset transfer evaluation and domain generalization evaluation.

#### Evaluation Metrics

For all tasks, we report the average result over three different random seeds.
For the text-image retrieval and VQA tasks, we evaluate using recall at 1
(R@1) as a metric††Note that we follow Ge et al. (2023) and consider VQA as a
classification problem where the model selects the answer from a set of
candidate answers based on the given question.. In the base-to-new text-and-
image classification setting, we report the harmonic mean
$H=2\times(\text{base}\times\text{new})/(\text{base}+\text{new})$, which
measures the generalization trade-off between the base and new sets. Finally,
for the cross-dataset transfer and domain generalization setting, we use
accuracy as the evaluation metric.

#### Training Details

Consistent with CoCoOp, we chose CLIP with the vision backbone ViT-B/16 and
loaded the pre-trained weights of CLIP, keeping them frozen during training.
In the image-text retrieval and VQA tasks, for training the model, we employed
image captions as class labels, set the learning rate to 0.02. The code can be
obtained from https://github.com/shishicode/AGoT. For the three experiments on
image classification, We utilized a chain length of 5 and initialized each of
them with the phrase “a photo of a” {class} using a pre-trained CLIP text
encoder. We used the class label as the class description and a learning rate
of 0.002 to train the model and sample 16 shots from the base classes.

#### Baselines

We compare our proposed approach with the following state-of-the-art models:
zero-shot CLIP with the fixed handcrafted prompt “a photo of a” {class}, CoOp,
CoCoOp, ProGrad, KgCoOp, and CoT-PT.

### 4.2. Results and Analysis

Training data | Method | Flickr30k | MSCOCO
---|---|---|---
0% | CLIP | 83.00 | 53.30
0.5% | CoCoOp | 82.80 | 53.50
CoT-PT | 83.50 | 55.80
AGoT | 84.30 | 57.10
1.0% | CoCoOp | 84.50 | 56.40
CoT-PT | 85.10 | 56.70
AGoT | 87.60 | 57.60
2.0% | CoCoOp | 85.00 | 57.00
CoT-PT | 86.00 | 57.90
AGoT | 88.70 | 58.70
Table 1: Comparison between CLIP, CoCoOp, CoT-PT, and our AGoT methods on the Flickr30k and MSCOCO datasets. Training data |  Method |  VQAv2
---|---|---
0% |  CLIP |  11.83
0.25% |  CoCoOp |  27.23
CoT-PT |  29.13
AGoT |  29.64
0.5% |  CoCoOp |  29.51
CoT-PT |  30.72
AGoT |  31.31
0.75% |  CoCoOp |  30.76
CoT-PT |  30.86
AGoT |  31.74
Table 2: Comparison between CLIP, CoCoOp, CoT-PT, and our AGoT on the VQAv2
dataset.

(a) Average over 11 datasets | Base | New | H
---|---|---|---
CLIP | 69.34 | 74.22 | 71.70
CoOp | 82.63 | 67.99 | 74.60
CoCoOp | 80.47 | 71.69 | 75.83
ProGrad | 82.48 | 70.75 | 76.16
KgCoOp | 80.73 | 73.60 | 77.00
CoT-PT | 80.23 | 74.20 | 77.10
AGoT | 80.50 | 75.05 | 77.68

(b) ImageNet | Base | New | H
---|---|---|---
CLIP | 72.43 | 68.14 | 70.22
CoOp | 76.46 | 66.31 | 71.02
CoCoOp | 75.98 | 70.43 | 73.10
ProGrad | 77.02 | 66.66 | 71.46
KgCoOp | 75.83 | 69.96 | 72.78
CoT-PT | 76.00 | 70.68 | 73.24
AGoT | 76.02 | 70.80 | 73.32

(c) Caltech101 | Base | New | H
---|---|---|---
CLIP | 96.84 | 94.00 | 95.40
CoOp | 98.11 | 93.52 | 95.76
CoCoOp | 97.96 | 93.81 | 95.84
ProGrad | 98.02 | 93.89 | 95.91
KgCoOp | 97.72 | 94.39 | 96.03
CoT-PT | 97.91 | 94.03 | 95.93
AGoT | 98.06 | 94.92 | 96.46

(d) OxfordPets | Base | New | H
---|---|---|---
CLIP | 91.17 | 97.26 | 94.12
CoOp | 94.24 | 96.66 | 95.43
CoCoOp | 95.20 | 97.69 | 96.43
ProGrad | 95.07 | 97.63 | 96.33
KgCoOp | 94.65 | 97.76 | 96.18
CoT-PT | 95.43 | 97.78 | 96.59
AGoT | 94.94 | 98.19 | 96.54

(e) StanfordCars | Base | New | H
---|---|---|---
CLIP | 63.37 | 74.89 | 68.65
CoOp | 76.20 | 69.14 | 72.49
CoCoOp | 70.49 | 73.59 | 72.01
ProGrad | 77.68 | 68.63 | 72.88
KgCoOp | 71.76 | 75.04 | 73.36
CoT-PT | 70.59 | 73.82 | 72.17
AGoT | 72.23 | 74.77 | 73.47

(f) Flowers102 | Base | New | H
---|---|---|---
CLIP | 72.08 | 77.80 | 74.83
CoOp | 97.63 | 69.55 | 81.23
CoCoOp | 94.87 | 71.75 | 81.71
ProGrad | 95.54 | 71.87 | 82.03
KgCoOp | 95.00 | 74.73 | 83.65
CoT-PT | 94.46 | 72.46 | 82.01
AGoT | 94.01 | 74.18 | 82.93

(g) Food101 | Base | New | H
---|---|---|---
CLIP | 90.10 | 91.22 | 90.66
CoOp | 89.44 | 87.50 | 88.46
CoCoOp | 90.70 | 91.29 | 90.99
ProGrad | 90.37 | 89.59 | 89.98
KgCoOp | 90.50 | 91.70 | 91.09
CoT-PT | 90.74 | 91.77 | 91.25
AGoT | 90.88 | 92.81 | 91.83

(h) FGVCAircraft | Base | New | H
---|---|---|---
CLIP | 27.19 | 36.29 | 31.09
CoOp | 39.24 | 30.49 | 34.30
CoCoOp | 33.41 | 23.71 | 27.74
ProGrad | 40.54 | 27.57 | 32.82
KgCoOp | 36.21 | 33.55 | 34.83
CoT-PT | 35.13 | 32.21 | 33.61
AGoT | 35.23 | 34.14 | 34.68

(i) SUN397 | Base | New | H
---|---|---|---
CLIP | 69.36 | 75.35 | 72.23
CoOp | 80.85 | 68.34 | 74.07
CoCoOp | 79.74 | 76.86 | 78.27
ProGrad | 81.26 | 74.17 | 77.55
KgCoOp | 80.29 | 76.53 | 78.36
CoT-PT | 79.44 | 77.20 | 78.30
AGoT | 79.59 | 77.60 | 78.58

(j) DTD | Base | New | H
---|---|---|---
CLIP | 53.24 | 59.90 | 56.37
CoOp | 80.17 | 47.54 | 59.68
CoCoOp | 77.01 | 56.00 | 64.85
ProGrad | 77.35 | 52.35 | 62.45
KgCoOp | 77.55 | 54.99 | 64.35
CoT-PT | 76.27 | 58.34 | 66.11
AGoT | 76.73 | 57.88 | 65.99

(k) EuroSAT | Base | New | H
---|---|---|---
CLIP | 56.48 | 64.05 | 60.03
CoOp | 91.54 | 54.44 | 68.27
CoCoOp | 87.49 | 60.04 | 71.21
ProGrad | 90.11 | 60.89 | 72.67
KgCoOp | 85.64 | 64.34 | 73.48
CoT-PT | 84.11 | 72.81 | 78.06
AGoT | 86.11 | 71.67 | 78.23

(l) UCF101 | Base | New | H
---|---|---|---
CLIP | 70.53 | 77.50 | 73.85
CoOp | 85.14 | 64.47 | 73.37
CoCoOp | 82.33 | 73.45 | 77.64
ProGrad | 84.33 | 74.94 | 79.35
KgCoOp | 82.89 | 76.67 | 79.65
CoT-PT | 82.47 | 75.09 | 78.61
AGoT | 81.90 | 78.68 | 80.25

Table 3: Comparison with existing methods (CLIP, CoOp, CoCoOp, ProGrad,
KgCoOp, CoT-PT) in the cross-label generalization setting with ViT-B/16 as the
backbone.

#### Text-Image Retrieval

| Source | Target
---|---|---
| Ima | Cal | Oxf | Sta | Flo | Foo | FGV | SUN | DTD | Eur | UCF | Avg
CoOp | 71.51 | 93.70 | 89.14 | 64.51 | 68.71 | 85.30 | 18.47 | 64.15 | 41.92 | 46.39 | 66.55 | 63.88
CoCoOp | 71.02 | 94.43 | 90.14 | 65.32 | 71.88 | 86.06 | 22.94 | 67.36 | 45.73 | 45.37 | 68.21 | 65.74
ProGrad | 70.21 | 94.43 | 93.21 | 71.75 | 89.98 | 85.77 | 32.93 | 71.17 | 57.72 | 70.84 | 77.82 | 74.21
KgCoOp | 70.19 | 94.65 | 93.20 | 71.98 | 90.69 | 86.59 | 32.47 | 71.79 | 58.31 | 71.06 | 78.40 | 74.48
AGoT | 72.12 | 95.92 | 98.25 | 75.11 | 75.73 | 91.20 | 32.01 | 76.53 | 58.75 | 71.46 | 79.48 | 75.44
Table 4: Comparison of CoOp, CoCoOp, ProGrad, KgCoOp and our AGoT methods in the cross dataset generalization setting. | Prompts | Source | Target
---|---|---|---
| ImageNet | ImageNetV2 | ImageNet-Sketch | ImageNet-A | ImageNet-R | Avg.
CLIP | hand-crafted | 66.73 | 60.83 | 46.15 | 47.77 | 73.96 | 57.17
UPT | vp+tp | 72.63 | 64.35 | 48.66 | 50.66 | 76.24 | 59.98
CoCoOp | vp+tp | 71.02 | 64.07 | 48.75 | 50.63 | 76.18 | 59.90
CoOp | tp | 71.51 | 64.20 | 47.99 | 49.71 | 75.21 | 59.28
ProGrad | tp | 72.24 | 64.73 | 47.61 | 49.39 | 74.58 | 59.07
KgCoOp | tp | 71.20 | 64.10 | 48.97 | 50.69 | 76.70 | 60.11
AGoT | vp+tp | 72.12 | 64.54 | 49.76 | 50.26 | 76.92 | 60.37
Table 5: Comparison of CLIP, UPT, CoCoOp, CoOp, ProGrad, KgCoOp and AGoT in
the domain generalization setting. Where “vp” and “tp” denote the visual
prompting and textual prompting, respectively.

Sample ratio | 0.2% | 0.5% | 1.0% | 2.0% | 5.0% | 8.0%
---|---|---|---|---|---|---
steps=3 | 83.30 | 84.30 | 87.60 | 88.70 | 89.50 | 90.10
steps=5 | 83.20 | 84.50 | 87.10 | 88.40 | 89.60 | 90.30

Table 6: Relationship between sample ratio and reasoning steps on the Flickr30k dataset. Number | 2 | 3 | 4 | 5 | 6
---|---|---|---|---|---
Recall | 86.90 | 88.20 | 87.90 | 88.70 | 88.20
Table 7: Comparison the effect of the number of aggregation subnodes on the Flickr30k dataset. $\alpha$ | 0.1 | 0.3 | 0.5 | 0.7 | Dynamic
---|---|---|---|---|---
Recall | 57.80 | 58.00 | 58.10 | 57.90 | 58.70
Table 8: Comparison of the effect of flow controller with the fixed and
dynamic prompt fusion ratios on the MSCOCO dataset.

In this experiment, we evaluated our model on two image captioning datasets:
Flickr30k and MSCOCO. We trained and tested our model using settings with
0.5%, 1.0%, and 2.0% of the training data and evaluated it on the test set. As
can be observed from Table 1, our method outperformed the CLIP, CoCoOp, and
CoT-PT in all tests. On both the Flickr30k and MSCOCO datasets, when the
training data consisted of only 0.5% of the total data, AGoT performed 0.80%
and 1.30% better than the current SOTA method CoT-PT; when the training data
was 1.0% of the total data, AGoT performed 2.5% and 0.9% better than CoT-PT;
and when the training data was 2.0% of the total data, AGoT performed 1.7% and
0.8% better than CoT-PT, respectively. These improvements demonstrate that our
proposed AGoT is more effective for multi-modal representation learning than
CoT, which was the previous state-of-the-art. Additionally, our proposed novel
soft-prompting is helpful for cross-modal representation learning without
tuning the pre-training parameters of CLIP.

Subnodes | 1 | 2 | $\Delta$ | 3 | $\Delta$ | 4 | $\Delta$ | 5 | $\Delta$ | 6 | $\Delta$
---|---|---|---|---|---|---|---|---|---|---|---
Recall | 86.00 | 86.90 | +0.90 | 88.20 | +2.20 | 87.90 | +1.90 | 88.70 | +2.70 | 88.20 | +2.20
Parameters(M) | 149.74 | 149.76 | +0.02 | 149.78 | +0.04 | 149.81 | +0.07 | 149.83 | +0.09 | 149.85 | +0.11
Reasoning Time(s) | 180 | 182 | +2 | 187 | +7 | 190 | +10 | 191 | +11 | 193 | +13
Training Time(s) | 208.11 | 216.42 | +8.31 | 222.57 | +14.46 | 227.34 | +19.23 | 231.24 | +23.13 | 235.63 | +27.52

Table 9: Evaluation of the impact on Recall, Parameters, Reasoning time, and
Training time with increasing aggregation subnodes on the Flickr30k dataset.

#### VQA

In VQA, we also sample 0.25%, 0.5%, and 0.75% from the training set following
Ge et al. (2023). The results are shown in Table 2. Similar to image-to-text
retrieval, our method outperforms CLIP, CoCoOp, and CoT-PT in all settings.
AGoT obtains 0.51%, 0.59%, and 0.88% gains over CoT-PT, respectively.
Experiments show that our AGoT approach can also benefit relatively more
complex VQA tasks due to better reasoning and multi-modal representation
learning.

#### Cross Label Generalization

In addition to being evaluated on multi-modal understanding tasks, based on
the conclusion of previous work Zhou et al. (2022a), better prompting methods
can lead to better domain generalization abilities. Therefore, we first
conducted experiments on the cross-label generalization setting. Specifically,
we conducted experiments using AGoT on 11 datasets, where each dataset was
divided into base and new classes to represent the source and target domains.
To train the model, we considered 16 samples per class and set the reasoning
step $Z=5$. The results, summarized in Table 3, revealed that AGoT
outperformed other models such as CLIP, CoOp, CoCoOp, ProGrad, KgCoOp, and
CoT-PT. Specifically, AGoT achieved accuracy improvements of 5.98%, 3.08%,
1.85%, 1.52%, 0.68%, and 0.58% over these strong baselines, respectively.
These experiments demonstrate that AGoT exhibits strong generalization
performance when encountering new data, which results from multiple aspects of
thought reasoning.

#### Cross Dataset Generalization

In addition to cross-label generalization, we also explore the cross-dataset
generalization setting. As shown in Table 4, AGoT outperforms the CoOp by
0.61% on ImageNet (source). Furthermore, when applying the model to the target
domain of 10 datasets, our results demonstrate a 0.96% improvement over
previous state-of-the-art KgCoOp. These experimental findings also provide
solid evidence of AGoT’s robust transferability compared to ordinary CoT.

#### Domain Generalization

Table 5 shows the results of domain generalization. Our model trained on
ImageNet was tested on four other ImageNet datasets of different types, and
our method outperforms CLIP, UPT, CoCoOp, CoOp, ProGrad, KgCoOp. This confirms
that CoT does have a good effect on domain generalization, and AGoT is an
effective improvement of CoT, bringing stronger reasoning generalization
ability.

### 4.3. Ablation Study

#### Reasoning Steps

Figure 3: Comparison of different reasoning steps on the Flickr30k dataset.

To demonstrate the impact of reasoning steps in AGoT, we conducted experiments
on the image-text retrieval task using the Flickr30K dataset with a training
ratio of 2.0%. Specifically, we varied the reasoning steps from 2 to 9. The
results, as shown in Figure 3, indicate that AGoT achieves its best
performance when the epoch is set to 80 and the reasoning steps are set to 3.
This can be attributed to the fact that when the data size is small, longer
reasoning steps in AGoT may be unstable and susceptible to overfitting. Thus,
increasing the reasoning steps does not always lead to improved performance.

As shown in Table 6, continuing to increase the proportion of training
samples, when the training sample proportion reaches 5.0%, the performance of
inference step 5 surpasses that of inference step 3, achieving optimal
performance. Similarly, the same conclusion is drawn when the sample
proportion increases to 8.0%. This indicates that as the dataset grows
relatively larger and tasks become more complex, AGoT with longer inference
steps can explore and embed more valuable features.

#### Number of Aggregation Subnodes

To demonstrate the impact of the number of subnodes (i.e., prompts) aggregated
by AGoT on reasoning performance, we conducted experiments on the Flickr30K
dataset using a train data ratio of 2.0%. Specifically, we varied the number
of subnodes from 2 to 6. As shown in Table 7, we found that the best result
was achieved when the number of aggregations was set to 5. This suggests that
having more aspect prompts tends to improve performance, but optimal
performance depends on task complexity.

#### Role of Flow Controller

In the AGoT model, we use a flow controller to dynamically regulate the degree
of message passing between central nodes. To evaluate the effectiveness of
this flow controller, we contrasted it with the fixed prompt fusion ratios
$\alpha$ of 0.1, 0.3, 0.5, and 0.7 on the MSCOCO dataset using 2.0% sampled
data. And the reasoing steps and subnodes are fixed to 3 and 4, respectively.
The results, presented in Table 8, demonstrate that the dynamic flow
controller outperforms the fixed ratios. This improvement can be attributed to
the varying features of different images and their requirement for distinct
reasoning processes. By employing input-based dynamic control, the model
achieves better reasoning and enhances its overall performance.

#### Role of Multi-View Thinking

To provide a more comprehensive analysis of the impact of the number of
aggregation subnodes on the computational and model complexity of AGoT, we
conduct experiments on the Flickr30k dataset with a sample proportion of 2.0%.
In these experiments, the number of subnodes is systematically varied from 1
to 6, we find that the best Recall is achieved when the number of aggregations
is set to 5. Notably, our method degenerates to COT-PT when the number of
subnodes is 1. The outcomes of these experiments are detailed in Table 9.

When increasing aggregation subnodes from 1 to 2, Recall improved by +0.90%,
parameters increased by +0.02M, reasoning time increased by +2s, and training
time increased by +8.31s. Similarly, when increasing aggregation subnodes from
1 to 5, Recall improved by +2.70%, parameters increased by +0.09M, reasoning
time increased by +11s, and training time increased by +23.13s. Hence, it is
evident that with the increasing number of aggregation subnodes, the model’s
parameters, reasoning time and training time exhibit minimal changes to
achieve a favorable trade-off. Please refer to Supplementary Material Section
A for the impact of different key values, different backbone architectures,
and flow controllers on the model across other datasets.

## 5\. Conclusion

In this paper, we introduce a novel AGoT mechanism for soft-prompt tuning in
multi-modal representation learning. The mechanism consists of a novel prompt
aggregation and prompt flow process designed to improve multi-modal reasoning.
We demonstrate the effectiveness of this method through experiments conducted
on 2 text-image retrieval tasks, 1 VQA task, and 15 image and text
classification tasks. In future work, we plan to introduce a reasoning early
exiting mechanism to address the varying complexity of inputs, which may
require different reasoning strategies.

## Acknowledgements

This research was supported by the National Natural Science Foundation of
China (No. 62306216, 62376200), the Natural Science Foundation of Hubei
Province of China (No. 2023AFB816), the Fundamental Research Funds for the
Central Universities (No. 2042023kf0133).

## Limitations

Applying AGoT in small models under fine-tuning scenarios. Typically, limited
by capacity, a small model requires fine-tuning to acquire thinking abilities.
A joint training schema may intensify the problem’s complexity. The stability
of CoT-based prompt technology currently lacks consistency and is sensitive to
hyper-parameters. Designing the prompt structure, determining its steps, and
establishing proper initialization methods all require extensive
experimentation and exploration. Furthermore, as the reasoning steps increase,
the overall speed of the system decreases significantly, and there is a higher
consumption of GPU memory.

## Ethics Statement

The mechanism presented in this paper is a robust tool with the potential to
enhance multi-modal representation learning and improve system generalization.
However, society needs to maintain a steadfast focus on data privacy and
remain cautious about the potential for technology to be maliciously misused.

## References

## References

  * Afouras et al. (2018) Triantafyllos Afouras, Joon Son Chung, Andrew Senior, Oriol Vinyals, and Andrew Zisserman. 2018.  Deep audio-visual speech recognition.  _IEEE transactions on pattern analysis and machine intelligence_ , 44(12):8717–8727.
  * Bossard et al. (2014) Lukas Bossard, Matthieu Guillaumin, and Luc Van Gool. 2014.  Food-101–mining discriminative components with random forests.  In _Computer Vision–ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part VI 13_ , pages 446–461. Springer.
  * Chen et al. (2015) Xinlei Chen, Hao Fang, Tsung-Yi Lin, Ramakrishna Vedantam, Saurabh Gupta, Piotr Dollár, and C Lawrence Zitnick. 2015.  Microsoft coco captions: Data collection and evaluation server.  _arXiv preprint arXiv:1504.00325_.
  * Cimpoi et al. (2014) Mircea Cimpoi, Subhransu Maji, Iasonas Kokkinos, Sammy Mohamed, and Andrea Vedaldi. 2014.  Describing textures in the wild.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 3606–3613.
  * Deng et al. (2009) Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. 2009.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , pages 248–255. Ieee.
  * Dosovitskiy et al. (2021) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. 2021.  An image is worth 16x16 words: Transformers for image recognition at scale.  In _9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021_. OpenReview.net.
  * Fei-Fei et al. (2004) Li Fei-Fei, Rob Fergus, and Pietro Perona. 2004.  Learning generative visual models from few training examples: An incremental bayesian approach tested on 101 object categories.  In _2004 conference on computer vision and pattern recognition workshop_ , pages 178–178. IEEE.
  * Gabeur et al. (2022) Valentin Gabeur, Arsha Nagrani, Chen Sun, Karteek Alahari, and Cordelia Schmid. 2022.  Masking modalities for cross-modal video retrieval.  In _Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision_ , pages 1766–1775.
  * Ge et al. (2023) Jiaxin Ge, Hongyin Luo, Siyuan Qian, Yulu Gan, Jie Fu, and Shanghang Zhan. 2023.  Chain of thought prompt tuning in vision language models.  _arXiv preprint arXiv:2304.07919_.
  * Gu et al. (2022) Yuxian Gu, Xu Han, Zhiyuan Liu, and Minlie Huang. 2022.  PPT: pre-trained prompt tuning for few-shot learning.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022_ , pages 8410–8423. Association for Computational Linguistics.
  * Han et al. (2022) Xu Han, Weilin Zhao, Ning Ding, Zhiyuan Liu, and Maosong Sun. 2022.  PTR: prompt tuning with rules for text classification.  _AI Open_ , 3:182–192.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 770–778.
  * He et al. (2023) Liqi He, Zuchao Li, Xiantao Cai, and Ping Wang. 2023.  Multi-modal latent space learning for chain-of-thought reasoning in language models.  _arXiv preprint arXiv:2312.08762_.
  * Helber et al. (2019) Patrick Helber, Benjamin Bischke, Andreas Dengel, and Damian Borth. 2019.  Eurosat: A novel dataset and deep learning benchmark for land use and land cover classification.  _IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing_ , 12(7):2217–2226.
  * Hendrycks et al. (2021a) Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, et al. 2021a.  The many faces of robustness: A critical analysis of out-of-distribution generalization.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 8340–8349.
  * Hendrycks et al. (2021b) Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. 2021b.  Natural adversarial examples.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 15262–15271.
  * 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. 2021.  Scaling up visual and vision-language representation learning with noisy text supervision.  In _International Conference on Machine Learning_ , pages 4904–4916. PMLR.
  * Krause et al. (2013) Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei. 2013.  3d object representations for fine-grained categorization.  In _Proceedings of the IEEE international conference on computer vision workshops_ , pages 554–561.
  * Li et al. (2022) Dongxu Li, Junnan Li, Hongdong Li, Juan Carlos Niebles, and Steven CH Hoi. 2022.  Align and prompt: Video-and-language pre-training with entity prompts.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 4953–4963.
  * Liu et al. (2022) Xiao Liu, Kaixuan Ji, Yicheng Fu, Weng Tam, Zhengxiao Du, Zhilin Yang, and Jie Tang. 2022.  P-tuning: Prompt tuning can be comparable to fine-tuning across scales and tasks.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers)_ , pages 61–68.
  * Lu et al. (2022) Yuning Lu, Jianzhuang Liu, Yonggang Zhang, Yajing Liu, and Xinmei Tian. 2022.  Prompt distribution learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 5206–5215.
  * Maji et al. (2013) Subhransu Maji, Esa Rahtu, Juho Kannala, Matthew Blaschko, and Andrea Vedaldi. 2013.  Fine-grained visual classification of aircraft.  _arXiv preprint arXiv:1306.5151_.
  * Nilsback and Zisserman (2008) Maria-Elena Nilsback and Andrew Zisserman. 2008.  Automated flower classification over a large number of classes.  In _2008 Sixth Indian Conference on Computer Vision, Graphics & Image Processing_, pages 722–729. IEEE.
  * Parkhi et al. (2012) Omkar M Parkhi, Andrea Vedaldi, Andrew Zisserman, and CV Jawahar. 2012.  Cats and dogs.  In _2012 IEEE conference on computer vision and pattern recognition_ , pages 3498–3505. IEEE.
  * Plummer et al. (2015) Bryan A Plummer, Liwei Wang, Chris M Cervantes, Juan C Caicedo, Julia Hockenmaier, and Svetlana Lazebnik. 2015.  Flickr30k entities: Collecting region-to-phrase correspondences for richer image-to-sentence models.  In _Proceedings of the IEEE international conference on computer vision_ , pages 2641–2649.
  * 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. 2021.  Learning transferable visual models from natural language supervision.  In _ICML_.
  * Recht et al. (2019) Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. 2019.  Do imagenet classifiers generalize to imagenet?  In _International Conference on Machine Learning_ , pages 5389–5400. PMLR.
  * Ren et al. (2019) Yi Ren, Yangjun Ruan, Xu Tan, Tao Qin, Sheng Zhao, Zhou Zhao, and Tie-Yan Liu. 2019.  Fastspeech: Fast, robust and controllable text to speech.  _Advances in neural information processing systems_ , 32.
  * Rose et al. (2023) Daniel Rose, Vaishnavi Himakunthala, Andy Ouyang, Ryan He, Alex Mei, Yujie Lu, Michael Saxon, Chinmay Sonar, Diba Mirza, and William Yang Wang. 2023.  Visual chain of thought: Bridging logical gaps with multimodal infillings.  _arXiv preprint arXiv:2305.02317_.
  * Schick and Schütze (2021) Timo Schick and Hinrich Schütze. 2021.  Exploiting cloze-questions for few-shot text classification and natural language inference.  In _Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume, EACL 2021, Online, April 19 - 23, 2021_ , pages 255–269. Association for Computational Linguistics.
  * Shen et al. (2022) Sheng Shen, Liunian Harold Li, Hao Tan, Mohit Bansal, Anna Rohrbach, Kai-Wei Chang, Zhewei Yao, and Kurt Keutzer. 2022.  How much can CLIP benefit vision-and-language tasks? In _The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022_. OpenReview.net.
  * Shi et al. (2023) Freda Shi, Mirac Suzgun, Markus Freitag, Xuezhi Wang, Suraj Srivats, Soroush Vosoughi, Hyung Won Chung, Yi Tay, Sebastian Ruder, Denny Zhou, Dipanjan Das, and Jason Wei. 2023.  Language models are multilingual chain-of-thought reasoners.  In _The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023_. OpenReview.net.
  * Soomro et al. (2012) Khurram Soomro, Amir Roshan Zamir, and Mubarak Shah. 2012.  Ucf101: A dataset of 101 human actions classes from videos in the wild.  _arXiv preprint arXiv:1212.0402_.
  * Sun et al. (2022) Ximeng Sun, Ping Hu, and Kate Saenko. 2022.  Dualcoop: Fast adaptation to multi-label recognition with limited annotations.  _Advances in Neural Information Processing Systems_ , 35:30569–30582.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  _Advances in neural information processing systems_ , 30.
  * Wang et al. (2019) Haohan Wang, Songwei Ge, Zachary Lipton, and Eric P Xing. 2019.  Learning robust global representations by penalizing local predictive power.  _Advances in Neural Information Processing Systems_ , 32.
  * Wang et al. (2023) Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc V. Le, Ed H. Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. 2023.  Self-consistency improves chain of thought reasoning in language models.  In _The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023_. OpenReview.net.
  * Wei et al. (2022) Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. 2022.  Chain-of-thought prompting elicits reasoning in large language models.  _Advances in neural information processing systems_ , 35:24824–24837.
  * Xiao et al. (2010) Jianxiong Xiao, James Hays, Krista A Ehinger, Aude Oliva, and Antonio Torralba. 2010.  Sun database: Large-scale scene recognition from abbey to zoo.  In _2010 IEEE computer society conference on computer vision and pattern recognition_ , pages 3485–3492. IEEE.
  * Yao et al. (2023) Hantao Yao, Rui Zhang, and Changsheng Xu. 2023.  Visual-language prompt tuning with knowledge-guided context optimization.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 6757–6767.
  * Yao et al. (2024) Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. 2024.  Tree of thoughts: Deliberate problem solving with large language models.  _Advances in Neural Information Processing Systems_ , 36.
  * Yu et al. (2022) Tan Yu, Hongliang Fei, and Ping Li. 2022.  U-bert for fast and scalable text-image retrieval.  In _Proceedings of the 2022 ACM SIGIR International Conference on Theory of Information Retrieval_ , pages 193–203.
  * Zelikman et al. (2022) Eric Zelikman, Jesse Mu, Noah D Goodman, and Yuhuai Tony Wu. 2022.  Star: Self-taught reasoner bootstrapping reasoning with reasoning.
  * Zhang et al. (2023a) Zhuosheng Zhang, Aston Zhang, Mu Li, and Alex Smola. 2023a.  Automatic chain of thought prompting in large language models.  In _The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023_. OpenReview.net.
  * Zhang et al. (2023b) Zhuosheng Zhang, Aston Zhang, Mu Li, Hai Zhao, George Karypis, and Alex Smola. 2023b.  Multimodal chain-of-thought reasoning in language models.  _arXiv preprint arXiv:2302.00923_.
  * Zheng et al. (2021) Wenfeng Zheng, Lirong Yin, Xiaobing Chen, Zhiyang Ma, Shan Liu, and Bo Yang. 2021.  Knowledge base graph embedding module design for visual question answering model.  _Pattern recognition_ , 120:108153.
  * Zhou et al. (2022a) Kaiyang Zhou, Jingkang Yang, Chen Change Loy, and Ziwei Liu. 2022a.  Conditional prompt learning for vision-language models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 16816–16825.
  * Zhou et al. (2022b) Kaiyang Zhou, Jingkang Yang, Chen Change Loy, and Ziwei Liu. 2022b.  Learning to prompt for vision-language models.  _International Journal of Computer Vision_ , 130(9):2337–2348.
  * Zhu et al. (2023) Beier Zhu, Yulei Niu, Yucheng Han, Yue Wu, and Hanwang Zhang. 2023.  Prompt-aligned gradient for prompt tuning.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 15659–15669.

## Appendix

##  Appendix A Ablation Study

#### Different Key Value

To investigate the influence of different sample quantities on model
performance, we selected three values for single category quantity: 4, 8, and
16, across 11 classification task datasets. We compared the model’s training
on the Base domain and testing on the New domain. The results were contrasted
with those of AGoT model against CoOp, CoCoOp, ProGrad, and KgCoOp, as
detailed in Table 10. When K is 4, AGoT achieved a value of 76.87%, surpassing
CoOp by +4.43% and KgCoOp by +0.97%. When K is 8, AGoT reached a value of
76.74%, outperforming CoOp by +3.24% and KgCoOp by +0.68%. When K is 16, AGoT
achieved a value of 77.68%, surpassing CoOp by +3.08% and KgCoOp by +0.68%.

Backbone | Methods | K=4 | K=8 | K=16
---|---|---|---|---
Base | New | H | Base | New | H | Base | New | H
ViT-B/16 | CoOp | 78.43 | 68.03 | 72.44 | 80.73 | 68.39 | 73.50 | 82.63 | 67.99 | 74.60
CoCoOp | 76.72 | 73.34 | 74.85 | 78.56 | 72.00 | 74.90 | 80.47 | 71.69 | 75.83
ProGrad | 79.18 | 71.14 | 74.62 | 80.62 | 71.02 | 75.20 | 82.48 | 70.75 | 76.16
KgCoOp | 79.92 | 73.11 | 75.90 | 78.36 | 73.89 | 76.06 | 80.73 | 73.60 | 77.00
AGoT | 79.34 | 74.56 | 76.87 | 78.19 | 75.34 | 76.74 | 80.50 | 75.05 | 77.68
Table 10: Comparing the effects of different K values on model performance

#### Different Backbone

To assess the impact of different backbone networks on model performance, we
conducted tests on four networks: ViT-B/16, ViT-B/32, ResNet-50, and
ResNet-101. The experimental results are shown in Table 11. From Table 11, it
can be observed that the model performs best when ViT-B/16 is used as the
backbone. Therefore, the choice of different backbones has a significant
impact on the model’s performance.

Backbones | Cal | Oxf | Stan | Flo | Foo | FGV | DTD | Eur | UCF | Ave
---|---|---|---|---|---|---|---|---|---|---
ViT-B/16 | 96.46 | 96.54 | 73.47 | 82.93 | 91.83 | 34.68 | 65.99 | 78.23 | 80.25 | 77.82
ViT-B/32 | 95.32 | 94.93 | 67.42 | 74.37 | 87.46 | 26.93 | 61.43 | 71.46 | 76.58 | 72.88
ResNet-50 | 93.61 | 94.37 | 64.39 | 76.18 | 85.19 | 24.64 | 58.36 | 49.28 | 72.39 | 68.71
ResNet-101 | 94.57 | 94.16 | 68.54 | 76.34 | 87.42 | 8.67 | 61.87 | 54.54 | 76.86 | 69.22
Table 11: Comparative analysis of the impact of different backbones on model
performance.

#### Flow controller on other datasets

To further validate the impact of dynamic and fixed values of the flow
controller on the nine datasets, As shown in Table 12, it is evident that
dynamic values are beneficial for improving the model’s performance.

Backbones | Cal | Oxf | Stan | Flo | Foo | FGV | DTD | Eur | UCF | Ave
---|---|---|---|---|---|---|---|---|---|---
ViT-B/16(dynamic) | 96.46 | 96.54 | 73.47 | 82.93 | 91.83 | 34.68 | 65.99 | 78.23 | 80.25 | 77.82
ViT-B/16(fix=0.5) | 95.72 | 95.81 | 72.36 | 81.37 | 90.64 | 33.93 | 64.61 | 77.45 | 79.48 | 76.79
Table 12: Impact of fixed and dynamic flow controller on model performance