# Enhancing Adversarial Training with Prior Knowledge Distillation for Robust
Image Compression

###### Abstract

Deep neural network-based image compression (NIC) has achieved excellent
performance, but NIC method models have been shown to be susceptible to
backdoor attacks. Adversarial training has been validated in image compression
models as a common method to enhance model robustness. However, the
improvement effect of adversarial training on model robustness is limited. In
this paper, we propose a prior knowledge-guided adversarial training framework
for image compression models. Specifically, first, we propose a gradient
regularization constraint for training robust teacher models. Subsequently, we
design a knowledge distillation-based strategy to generate a priori knowledge
from the teacher model to the student model for guiding adversarial training.
Experimental results show that our method improves the reconstruction quality
by about 9dB when the Kodak dataset is elected as the backdoor attack object
for psnr attack. Compared with Ma2023[1], our method has a 5dB higher PSNR
output at high bitrate points.

Index Terms—  Image Compression, Robustness, Adversarial Training, Gradient
Regularization, Prior Knowledge

##  1 Introduction

Deep learning has undergone extensive research in the field of image
compression[2, 3, 4, 5, 6, 7]. In terms of rate-distortion performance, the
performance of deep learning models has surpassed that of traditional methods
such as JPEG[8], JPEG2000[9], BPG[10], and even the recent Versatile Video
Coding (VVC)[11]. However, as deep learning-based image compression models are
applied in practice, considering the sensitivity of deep neural networks to
minor input perturbations, the robustness of these models against adversarial
perturbations has become an essential and unavoidable concern.

Fig. 1: Performance of the original model, Ma2023[1], and our model under
attack. Our model outperforms both Hyper model and Ma2023[1] model in terms of
both bpp and psnr performance.

Since the image compression models used in practice are mostly publicly
accessible, in this paper, we assume that attackers can access all parameters
of the neural network compression model. Therefore, the attacks mentioned in
this paper are all considered white-box attacks.

The concept of adversarial samples was first introduced by Szegedy[12], which
unveiled the sensitivity of deep neural networks to input perturbations.
Following this, Goodfellow introduced an efficient attack method called the
Fast Gradient Sign Method (FGSM) by utilizing gradient signs[13]. Kurakin
enhances the attack capability by breaking down the larger step in the FGSM
iteration into many smaller steps[14]. Madry proposed the Projected Gradient
Descent (PGD) method, which increases attack success rates through multiple
iterations[15]. In the context of image compression tasks, Ma introduced a
trainable noise as an input perturbation[1], using this noise during training
to update the model’s input and generate highly potent adversarial samples.
Liu leveraged the PGD method to efficiently attack the model’s bpp (bits per
pixel) parameter[16]. We have enhanced the techniques proposed in [1],
introducing two attack methods tailored for image compression: psnr attack and
bpp attack.

Fig. 2:  Prior knowledge-guided adversarial training framework for image
compression models. In the first stage, gradient regularization constraints
are used to train robust teacher models. Subsequently, in the second stage, a
prior knowledge distillation transfer strategy (PKDT) is used to transfer
prior knowledge generated by the teacher model to the finetune model to guide
adversarial training.

Adversarial Training[17, 18], serves as the most effective method to enhance a
model’s resilience against adversarial attacks. The core idea of this
technique is to introduce adversarial samples and use them as supplementary
data alongside the original dataset for model training. The training process
can be framed as a min-max problem, where the maximization problem involves
generating adversarial samples with strong attack capabilities and the
minimization problem focuses on training the model to minimize the output loss
for these adversarial samples. However, applying adversarial training directly
to image compression tasks resulted in relatively minor improvements in model
robustness. To further enhance the model’s robustness, we introduce prior
knowledge into the adversarial training process of the model through a
distillation framework[19]. Contributions of this paper are summarized as
follows:

$\bullet$ We propose a $smooth\\_loss$ to train the teacher model and utilize
this teacher model to generate strong robustness-oriented prior knowledge.

$\bullet$ We propose a distillation framework to transfer prior knowledge and
use this knowledge to guide the model’s adversarial training.

$\bullet$ Extensive experiments demonstrate that our proposed method
significantly enhances the model’s robustness. In the Kodak dataset, our
method can achieve a maximum robustness improvement of up to 23% against psnr
attack.

Algorithm 1 Adversarial Training with Prior Knowledge

1:Pre-trained model $M\\_pre$, Dataset $D$, Training epoch $N$, psnr attack
iterations $N_{p}$, bpp attack  iterations $N_{b}$

2:Finetune model $M\\_ft$

3:Stage I: Training a robustness teacher model $M\\_tc$

4:Let $M\\_tc=M\\_pre$

5:for $i<N$ do

6: $\vec{x}=Random\\_Sample(D,batchsize=16)$

7: $smooth\\_loss=\frac{{\partial{\rm{bpp}}\\_main}}{{\partial
x}}+\frac{{\partial y}}{{\partial x}}$

8: $loss=RD\\_loss+\alpha\cdot smooth\\_loss$

9: updating the parameters of model $M\\_tc$ by $loss$

10:end for

11:return teacher model $M\\_tc$

12:Stage II: Finetune network, adversarial training with prior knowledge

13:Let $M\\_ft=M\\_pre$

14:for $i<N$ do

15: $\vec{x}_{ori}=[x_{1},x_{2},x_{3}]$

16:$\quad\quad\quad=Random\\_Sample(D,batchsize=3)$

17: Generate adversarial sample $\acute{x_{1}}$ through psnr attack

18: on $x_{1}$ with $N_{p}=200$, generate adversarial sample

19: $\acute{x_{2}}$ through bpp attack on $x_{2}$ with $N_{b}=100$

20: $\vec{x}_{new}=[{x}^{\prime}_{1},{x}^{\prime}_{2},x_{3}]$

21: Input $\vec{x}_{ori}$ into model $M\\_tc$ to obtain $bpp\\_tc$

22: Input $\vec{x}_{new}$ into model $M\\_ft$ to obtain $bpp\\_ft$

23:
$D\\_loss={\left\|{{{\vec{x}}_{new}}-M\\_ft({{\vec{x}}_{new}})}\right\|_{1}}$

24: $R\\_loss={\left\|{bpp\\_tc-bpp\\_ft}\right\|_{1}}$

25: $smooth\\_loss=\frac{{\partial{\rm{bpp}}\\_main}}{{\partial
x}}+\frac{{\partial y}}{{\partial x}}$

26: $loss=D\\_loss+\alpha\cdot R\\_loss+\beta\cdot smooth\\_loss$

27: updating the parameters of model $M\\_ft$ by $loss$

28:end for

##  2 Method

Our method effectively improves the model’s robustness by introducing a
distillation framework for transferring prior knowledge. The framework of our
model is illustrated in Fig.2 and consists of two components: a teacher
network and a finetune network. The teacher network provides the prior
knowledge necessary for training the finetune network. A gradient
regularization term be used to train the teacher network, enabling it to
generate prior knowledge with robustness-related information. The finetune
network involves using prior knowledge generated by the teacher model to guide
adversarial training. Since the $psnr$ error caused by attacks is relatively
large, using $psnr$ as prior knowledge would make the adversarial training
process difficult to converge. At the same time, the error of $bpp$ is very
small, so we choose $bpp$ as the prior knowledge.

###  2.1 Stage I: Training a Robustness Teacher

Gradient regularization has been demonstrated as an effective method for
improving model robustness[20, 21]. However, due to the presence of two
performance metrics for the image compression task, we have designed two
gradient constraint terms. Regarding the $bpp$ term, we observed that the
changes in $bpp$ after the sample is attacked are mainly manifested in the
$bpp$ during the analysis transformation process, which we refer to as
$bpp\\_main$. Therefore, the constraint term is formulated as
$\partial{\rm{bpp}}\\_main/\partial x$. For the $psnr$ term, we also utilize
the gradient of $y$ with respect to $x$ after the analysis transformation as
another constraint. Therefore, the constraint term for gradient smoothness is
formulated as:

$smooth\\_loss=\underbrace{\frac{{\partial{\rm{bpp}}\\_main}}{{\partial
x}}}_{bpp\ constrain}+\underbrace{\frac{{\partial y}}{{\partial x}}}_{psnr\
constrain}$ (1)

Then the overall training loss of the model is:

$loss=RD\\_loss+smooth\\_loss$ (2)

where $RD\\_loss$ is the rate-distortion loss function of the pre-trained
model, and $RD\\_loss=R+\lambda D$, Distortion term $D$ is measured by peak
signal-to-noise ratio (PSNR) between original images and reconstructed images,
$R$ is the bit rate, $\lambda$ is hyper parameter to control the trade-off
between rate and distortion. Through training, we obtain the teacher model,
denoted as gradient model.

###  2.2 Stage II: Transferring Prior Knowledge

We transfer prior knowledge from the teacher network to the finetune network
through a distillation framework. The prior knowledge specifically refers to
the $bpp$ generated by input samples into the teacher model, denoted as
$bpp\\_tc$, We assume the input sample as $x$, and the reconstructed output
from the sample input into the finetuned model as $\hat{x}$, with $bpp$ as
$bpp\\_ft$. The loss for training the finetune network consists of three
components. First is the reconstruction error part:

$D\\_loss=\left\|x-\hat{x}\right\|_{1}$ (3)

Secondly, there is the introduction of prior knowledge:

$R\\_loss=\left\|bpp\\_tc-bpp\\_ft\right\|_{1}$ (4)

Lastly, the $smooth\\_loss$ in Equation 2. Then the overall training loss of
the model is:

$loss=D\\_loss+\alpha\cdot R\\_loss+\beta\cdot smooth\\_loss$ (5)

The entire training process is depicted as Algorithm 1.

###  2.3 Attack Method

The image compression task involves two evaluation metrics, $psnr$ (Peak
Signal-to-Noise Ratio) and $bpp$ (Bits Per Pixel), which directly impact the
model’s compression performance. Therefore, we conduct attacks on these two
metrics separately. When attacking $psnr$, we ensure that $bpp$ remains almost
unchanged to simulate a real-world scenario where image transmission
efficiency remains constant, but image decompression quality significantly
deteriorates. Conversely, when attacking $bpp$, we ensure that $psnr$ remains
almost unchanged to simulate a situation in the real world where image
decompression quality remains constant, but transmission efficiency
experiences a significant drop.

Assuming the input image is $x$, noise is $n$, adversarial samples
$x^{\ast}=x+n$. The input image and the adversarial sample are both
reconstructed to $\hat{x}$ and $\hat{x}^{\ast}$, simultaneously, the bit rates
are $bpp\\_ori$ and $bpp$, respectively. We improved the approach outlined in
[1] by modifying the loss function to:

$arg\min_{n}L_{d}=\left\\{\begin{matrix}\left\|n\right\|_{2}+bpp\\_error\quad\quad\left\|n\right\|_{2}^{2}\geq\epsilon\\\
1-\left\|x-\hat{x}^{\ast}\right\|_{2}^{2}\quad\quad\ \ \ \ \ \
\left\|n\right\|_{2}^{2}<\epsilon\end{matrix}\right.$ (6)
$arg\min_{n}L_{d}=\left\\{\begin{matrix}\left\|n\right\|_{2}+psnr\\_error\quad\left\|n\right\|_{2}^{2}\geq\epsilon\\\
-bpp\quad\quad\quad\quad\quad\quad\quad\left\|n\right\|_{2}^{2}<\epsilon\end{matrix}\right.$
(7)

where $\epsilon$ is the $l_{2}$ threshold of the input noise,
$bpp\\_error=\left\|bpp-bpp\\_ori\right\|_{1}$,
$psnr\\_error=\left\|\left\|x-\hat{x}^{\ast}\right\|_{2}^{2}-\left\|x-\hat{x}\right\|_{2}^{2}\right\|_{1}$.
Then implementing psnr attack through training Equation 6, and implementing
bpp attack through training Equation 7.

##  3 Experiment

Fig. 3: R-D performance of the model under different attack scenarios. The two
graphs on the left illustrate the robustness performance of our method applied
to the Hyper and Cheng models using the Kodak test dataset. The first graph on
the left pertains to psnr attack, while the second one is related to bpp
attack, with the attack iterations $N_{p}=N_{b}=1000$. The two graphs on the
right are based on the DIV2K valid dataset, with the attack iterations
$N_{p}=N_{b}=500$. It’s worth noting that in the second graph on the left, the
last two data points on curve Ours_Cheng(bpp attack) have bpp exceeding 9.
Fig. 4: R-D performance of the model under other $\epsilon$

$A.$ Experiment Setup

Datasets. The training dataset includes CLIC and DIV2K. Two datasets are used
for testing, including Kodak dataset with 24 images at a size of 768×512,
DIV2K valid dataset with 100 high-quality images.

Training details. We train each model for 100 epochs with Adam optimizer. For
the pre-train model, the learning rate is initially set to $10^{-4}$ and then
halved every 25 epochs after the initial 50 epochs. For the gradient model and
ours finetune model the learning rate is initially set to $5\times 10^{-5}$,
and then halved every 25 epochs after the initial 25 epochs. The pre-trained
model is trained using the CLIC dataset. The gradient model and our fine-tuned
model are trained using the DIVK dataset, more importantly, both of them are
fine-tuned based on the pre-trained model. In the adversarial training process
we proposed, the number of iterations for psnr attack $N_{p}$ is set to 200,
while for bpp attack, iteration counts $N_{b}$ is set to 100, and
$\epsilon=0.0005$.
$B.$ Results and Comparison

We tested two original models, Hyper[3] and Cheng[6], as well as the finetuned
models obtained using our method. We evaluated these models on various test
datasets and under different attack methods to obtain results. Fig.3 shows the
performance comparison of our model in different scenarios. It can be observed
that when our method is applied to different models, there is a noticeable
enhancement in robustness. Especially for models with large parameter sizes
like Cheng the improvement in robustness is more pronounced compared to the
Hyper model. Additionally, we can see that our method effectively enhances
model robustness in both low-resolution and high-resolution test datasets.
Fig.4 illustrates the robustness performance of our method concerning
different noise thresholds for $\epsilon=0.001$. It can be observed that the
model’s robustness is significantly enhanced when compared to smaller noise
thresholds $\epsilon=0.0005$. Particularly, in the case of bpp attack, the
robustness improvement can reach up to $23\%$. Table.1 contains the data
points for the low bitrate scenario shown in Fig.3, leftmost graph. It can be
observed that the robustness improvement brought by our method is similarly
higher than that of the gradient model, which is training with stage I.

Table 1: R-D performance of model

model | clean | psnr attack | bpp attack
---|---|---|---
bpp | psnr | bpp | psnr | bpp | psnr
Hyper | 0.4978 | 33 | 0.498 | 22.3502 | 0.9161 | 33.7702
Ours w/o Stage II | 0.4322 | 32.2592 | 0.4309 | 25.9039 | 0.7225 | 33.1432
Ours | 0.4323 | 32.0802 | 0.4286 | 28.4039 | 0.6822 | 32.9933

Adversarial training can potentially lead to a decrease in the model’s
generalization ability, resulting in less improvement in robustness against
other attack methods. However, we compared our approach with the attack method
mentioned in Ma2023. The results of the comparison can be seen in Fig.5. It is
evident that for adversarial samples where we did not undergo adversarial
training, our method still exhibits better robustness.

Fig. 5: R-D performance compared with Ma2023

$C.$ Ablation Experiments

We demonstrate the optimality of our method under the same training parameter
conditions in Table.2. The second row represents the most effective method
used in Algorithm 1. The third row represents the results obtained by training
without the $smooth\\_loss$ term in the stage II. Analogously, the results in
the last two rows correspond to the removal of $\partial y/\partial x$ and
$\partial{\rm{bpp}}\\_main/\partial x$ in stage II, respectively. Test dataset
is Kodak with iteration counts $N_{p}=400$, $N_{b}=200$, and noise threshold
$\epsilon=0.0005$. From the results, we can observe that removing any of the
constraint terms deteriorates the model’s robustness. Additionally, we have
noticed that the impact on robustness is relatively small. This further
validates that in our proposed method, the most significant improvement in
model robustness is achieved through the introduction of prior knowledge.

Table 2: Results of the ablation experiments

model | clean | psnr attack | bpp attack
---|---|---|---
bpp | psnr | bpp | psnr | bpp | psnr
best | 0.4323 | 32.0802 | 0.4273 | 28.6556 | 0.6953 | 33.0194
best w/o |  |  |  |  |  |
$smooth\\_loss$ | 0.4358 | 32.1936 | 0.4328 | 28.4497 | 0.7214 | 33.0731
best w/o |  |  |  |  |  |
$\partial y/\partial x$ | 0.4355 | 32.1549 | 0.4301 | 28.5293 | 0.7075 | 33.031
best w/o |  |  |  |  |  |
$\partial{\rm{bpp}}\\_main/\partial x$ | 0.4332 | 32.1463 | 0.4279 | 28.6224 | 0.7083 | 33.093

##  4 conclusion

In this paper, we propose a method of training the teacher model using
gradient regularization and use the prior knowledge generated by this teacher
model to guide the entire process of adversarial training. Experimental
results demonstrate that our approach can effectively improve the robustness
of the model. This is especially effective when dealing with large-parameter
models and large noise thresholds.

## References

  * [1] T. Chen and Z. Ma, “Towards robust neural image compression: Adversarial attack and model finetuning,” IEEE Transactions on Circuits and Systems for Video Technology, 2023.
  * [2] J. Ballé, V. Laparra, and E. P. Simoncelli, “End-to-end optimized image compression,” arXiv preprint arXiv:1611.01704, 2016.
  * [3] J. Ballé, D. Minnen, S. Singh, S. J. Hwang, and N. Johnston, “Variational image compression with a scale hyperprior,” arXiv preprint arXiv:1802.01436, 2018.
  * [4] D. Minnen, J. Ballé, and G. D. Toderici, “Joint autoregressive and hierarchical priors for learned image compression,” Advances in neural information processing systems, vol. 31, 2018.
  * [5] D. Minnen and S. Singh, “Channel-wise autoregressive entropy models for learned image compression,” in 2020 IEEE International Conference on Image Processing (ICIP), pp. 3339–3343, IEEE, 2020.
  * [6] Z. Cheng, H. Sun, M. Takeuchi, and J. Katto, “Learned image compression with discretized gaussian mixture likelihoods and attention modules,” in Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 7939–7948, 2020.
  * [7] Y. Bao, F. Meng, C. Li, S. Ma, Y. Tian, and Y. Liang, “Nonlinear transforms in learned image compression from a communication perspective,” IEEE Transactions on Circuits and Systems for Video Technology, vol. 33, no. 4, pp. 1922–1936, 2023.
  * [8] G. K. Wallace, “The jpeg still picture compression standard,” Communications of the ACM, vol. 34, no. 4, pp. 30–44, 1991.
  * [9] M. Rabbani and R. Joshi, “An overview of the jpeg 2000 still image compression standard,” Signal processing: Image communication, vol. 17, no. 1, pp. 3–48, 2002.
  * [10] D. Yee, S. Soltaninejad, D. Hazarika, G. Mbuyi, R. Barnwal, and A. Basu, “Medical image compression based on region of interest using better portable graphics (bpg),” in 2017 IEEE international conference on systems, man, and cybernetics (SMC), pp. 216–221, IEEE, 2017.
  * [11] B. Bross, J. Chen, J.-R. Ohm, G. J. Sullivan, and Y.-K. Wang, “Developments in international video coding standardization after avc, with an overview of versatile video coding (vvc),” Proceedings of the IEEE, vol. 109, no. 9, pp. 1463–1493, 2021.
  * [12] C. Szegedy, W. Zaremba, I. Sutskever, J. Bruna, D. Erhan, I. Goodfellow, and R. Fergus, “Intriguing properties of neural networks,” arXiv preprint arXiv:1312.6199, 2013.
  * [13] I. J. Goodfellow, J. Shlens, and C. Szegedy, “Explaining and harnessing adversarial examples,” arXiv preprint arXiv:1412.6572, 2014.
  * [14] A. Kurakin, I. J. Goodfellow, and S. Bengio, “Adversarial examples in the physical world,” in Artificial intelligence safety and security, pp. 99–112, Chapman and Hall/CRC, 2018.
  * [15] A. Madry, A. Makelov, L. Schmidt, D. Tsipras, and A. Vladu, “Towards deep learning models resistant to adversarial attacks,” arXiv preprint arXiv:1706.06083, 2017.
  * [16] K. Liu, D. Wu, Y. Wang, D. Feng, B. Tan, and S. Garg, “Denial-of-service attacks on learned image compression,” arXiv preprint arXiv:2205.13253, 2022.
  * [17] I. J. Goodfellow, J. Shlens, and C. Szegedy, “Explaining and harnessing adversarial examples,” arXiv preprint arXiv:1412.6572, 2014.
  * [18] A. Madry, A. Makelov, L. Schmidt, D. Tsipras, and A. Vladu, “Towards deep learning models resistant to adversarial attacks,” arXiv preprint arXiv:1706.06083, 2017.
  * [19] N. Papernot, P. McDaniel, X. Wu, S. Jha, and A. Swami, “Distillation as a defense to adversarial perturbations against deep neural networks,” in 2016 IEEE symposium on security and privacy (SP), pp. 582–597, IEEE, 2016.
  * [20] H. Drucker and Y. Le Cun, “Improving generalization performance using double backpropagation,” IEEE transactions on neural networks, vol. 3, no. 6, pp. 991–997, 1992.
  * [21] C. Finlay and A. M. Oberman, “Scaleable input gradient regularization for adversarial robustness,” Machine Learning with Applications, vol. 3, p. 100017, 2021.


# Scaling Up Music Information Retrieval Training with
Semi-Supervised Learning

###### Abstract

In the era of data-driven Music Information Retrieval (MIR), the scarcity of
labeled data has been one of the major concerns to the success of an MIR task.
In this work, we leverage the semi-supervised teacher-student training
approach to improve MIR tasks. For training, we scale up the unlabeled music
data to 240k hours, which is much larger than any public MIR datasets. We
iteratively create and refine the pseudo-labels in the noisy teacher-student
training process. Knowledge expansion is also explored to iteratively scale up
the model sizes from as small as less than 3M to almost 100M parameters. We
study the performance correlation between data size and model size in the
experiments. By scaling up both model size and training data, our models
achieve state-of-the-art results on several MIR tasks compared to models that
are either trained in a supervised manner or based on a self-supervised
pretrained model. To our knowledge, this is the first attempt to study the
effects of scaling up both model and training data for a variety of MIR tasks.

Index Terms—  Semi-supervised learning, large-scale training, music
information retrieval (MIR) benchmark

##  1 Introduction

Large-scale datasets together with increasing model sizes have been recognized
as one of the critical factors to build strong machine learning systems. This
direction has led to rapid success in many domains, such as natural language
processing [1], computer vision [2], and automatic speech recognition (ASR)
[3]. Plenty of resources from Internets and publicly available datasets enable
the models to scale up to billions of parameters and achieve state-of-the-art
results on several downstream tasks.

In the domain of music information retrieval (MIR), the lack of labeled
datasets has hindered the advances of several downstream tasks [4, 5].
Although large-scale datasets such as MSD [6] and Jamendo [7] have contributed
to significant progress in music auto-tagging, they are still considered
smaller-scale datasets in the aforementioned domains. Other MIR tasks, such as
chord recognition and structure segmentation, have an even smaller amount of
training data due to the labeling difficulty. For example, labeling time-
varying musical events at a millisecond scale such as beats requires
specialized training and is very labor-intensive.

To tackle this problem, several prior works [8, 9] have investigated
leveraging self-supervised approaches to train on large-scale data and
transfer the knowledge to downstream tasks. Castellon et al. [10] has shown
that probing the pretrained Jukebox model can achieve better results in
emotion recognition and comparable results in music auto-tagging compared to
supervised approaches. Li et al. [9] demonstrated that using large-scale data
to train a large-scale self-supervised model, the pretrained MERT can
outperform the supervised setting in several downstream tasks. Due to the
characteristics of different MIR downstream tasks, pretrained models have not
been observed to outperform some supervised approaches for tasks such as key
detection.

Data Sources |  Hours
---|---
RWC [11], Ballroom [12], Beatles [13], USPOP [14], Billboard [15], Hainsworth [16], SMC [12], Isophonics [17], HookTheory [18], HJDB [19], SALAMI [15], Simac [20], HarmonixSet[21], Internal-Structure, Internal-Beat/Key |  1000
Audio-92K |  92,000
Music-151K |  151,000
Table 1: Statistics of each data source.

In this work, we attempt to scale up the training data and model sizes from a
different angle. We leverage the semi-supervised learning to train on both
labeled and unlabeled datasets. To scale up the model size, we use the
training framework of teacher-student knowledge expansion to iteratively
expand the size of student models. Teacher-student training has demonstrated
its usefulness in several MIR tasks [4, 22, 5]. We push this idea to a new
limit by expanding the models from as small as less than 3 million to almost
100 million parameters. To scale up the data size, we collect labeled datasets
for various MIR tasks and mix them together, as shown in Table 1. To further
expand the training data, two internal datasets (with 1.3 million and 2.4
million full songs respectively) are included, yielding 300 times more than
any other publicly available datasets (e.g. Jamendo [7] only has 3760 hours)
in the MIR domain.

In the experiments, we investigate the performance correlation between data
size and model size. We test the proposed method on four MIR tasks and show
that by scaling up both model and training data, the proposed approaches can
outperform not only existing self-supervised approaches but also achieve
state-of-the-art results on these MIR tasks.

| Round | Params | $R$ | Pool | Hop | Layer | Spec D. | Temp D. | Heads | L. Arr | Augment
---|---|---|---|---|---|---|---|---|---|---|---
Beat | 1st | 11M | 1 | (2, 2) | 160 | 5 | 32 | 128 | 4 | 4 | A1, A2
2nd | 43.3M | 5 | 64 | 256 | 4 | 4 | A1, A2, A4
3rd | 91.3M | 6 | 256 | 256 | 4 | 4 | A1, A2, A4
Chord | 1st | 2.9M | 2 | (4, 4) | 500 | 5 | 64 | 64 | 4 | 8 | A1, A2
2nd | 11.4M | 5 | 128 | 128 | 4 | 8 | A1, A2, A3
3rd | 91.9M | 6 | 256 | 256 | 8 | 8 | A1, A2, A3
Key | 1st | 1.5M | 3 | (4, 100) | 320 | 5 | 128 | 32 | 4 | 4 | A1, A2
2nd | 5.9M | 5 | 256 | 64 | 4 | 4 | A1, A2, A3
3rd | 42.8M | 8 | 512 | 128 | 8 | 8 | A1, A2, A3
Structure | 1st | 2.9M | 1 | (4, 6) | 512 | 5 | 64 | 64 | 4 | 8 | A1, A2
2nd | 11.4M | 5 | 128 | 128 | 4 | 8 | A1, A2, A3
3rd | 91.9M | 6 | 256 | 256 | 8 | 8 | A1, A2, A3
Table 2: Parameters of each task for three rounds of training. “Pool”: the
maxpooling sizes (F, T) for frequency and time axes. “Hop”: hop size. “Spec
D.” and “Temp D.”: the latent dimensions for spectral and temporal
Transformers, respectively [23]. “Heads”: number of attention heads. “L. Arr”:
number of latent arrays. “Augment”: augmentation methods (see Section 2.3).

##  2 Methodology

###  2.1 Training Scheme

We use the noisy student training following previous works [24]. The learning
process includes four steps:

  1. 1.

Train a supervised teacher $M$ using small training set.

  2. 2.

Use the trained teacher model to assign pseudo-labels to unlabeled samples.

  3. 3.

Apply data augmentation to the audio of pseudo-labeled samples, then mix them
with the supervised samples to create a new training set.

  4. 4.

Train a student model with the same or larger size using the new training set.

After the student model is well-trained and outperforms the teacher, we
replace the old teacher with the student and repeat steps 2 – 4 to train a new
student model.

Fig. 1: Model architecture overview.

###  2.2 Model architecture

Fig. 1 depicts the proposed model architecture. The input audio is first
transformed into spectrogram. We use six trainable harmonic filters [25, 26]
to enhance the spectrogram. Then, $R$ residual units [27] followed by a stack
of spectro-temporal max-pooling layers are used to produce a time-frequency
representation. Each residual unit has a kernel size of 3 and $F$ feature
maps, where $F$ is the spectral dimension in Table 2. We use PerceiverTF [23],
one from the Transformer family, as the model architecture. PerceiverTF has
shown state-of-the-art performance in multi-instrument music transcription
[23], demonstrating its ability to characterize useful information for
modeling both pitch- and time-related events. More importantly, PerceiverTF
incorporates the Perceiver structure [28] to improve training scalability,
making it a good fit for this work over SpecTNT [29]. We use exactly the same
Perceiver TF blocks as in [23]. Lastly, a dense layer is added to predict the
frame-level labels from the PerceiverTF output. In each training round, we
gradually scale up the parameters of PerceiverTF until it is close to 100M, as
shown in Table 2.

###  2.3 Data Augmentation

We use the following augmentations to train the noisy student:

  1. 1.

Pitch shifting (A1): we randomly shift the pitches in a range of $\pm 3$
semitones in the experiments.

  2. 2.

Torchaudio_augmentations111https://github.com/Spijkervet/torchaudio-
augmentations (A2): with a probability $p=0.8$, we add audio transformations
including white noise (SNR in the range between 0.3 and 0.5 dB), random gains,
high/low pass filters, and polarity inversion.

  3. 3.

Frequency masking (A3): we apply two frequency channels masking blocks; each
randomly masks out 0–25% of the input frequency bins.

  4. 4.

Time masking (A4): we apply two time steps masking blocks; each randomly masks
out 0–25% of the frames.

We use torchaudio library222https://pytorch.org/audio/stable/transforms.html
for the frequency and time masks.

##  3 Data Preparation

###  3.1 Data Sources

As summarized in Table 1, we used the datasets in the first row for supervised
training. Except for publicly available datasets, we include two in-house
datasets for supervised training for two reasons. First, we want to maximize
the size of labeled datasets, as Transformer-based architecture is data-
hungry. Second, a strong teacher model can build a better foundation to boost
the semi-supervised learning process. If the teacher can generate high-quality
pseudo-labels, requirements for post-processing and post-filtering can be
reduced. Our in-house dataset for structure segmentation (called _Internal-
Structure_) contains around 160 hours of music across a variety of styles. Our
in-house dataset for downbeat and key detections (called _Internal-Beat/Key_)
contains 50k clips and a total size of 580 hours. This dataset covers a wide
variety of genres.

For unlabeled datasets, we categorize them into three types. _Datasets of
other MIR tasks (Other-MIR)_ : for example, songs in downbeat datasets do not
have chord annotations, so they can be part of the unlabeled dataset for chord
recognition. _Audio-92K_ : audio recordings collected internally, covering a
wide range of languages, genres, and styles. _Music-151K_ : music recordings
collected internally from a different source than Audio-92K. It contains
mostly English songs. All the datasets used in this work are for research
purposes only.

###  3.2 Pseudo-labeling

Data filtering during the teacher-student learning process is critical to the
success of training a student model. In this work, we use the “soft labels” as
the pseudo-labels. Specifically, the teacher’s output probabilities are used
directly as supervision signal for the student, using Cross-Entropy loss. We
observe that when the teacher is more confident in a binary prediction, it
tends to output a probability close to one or zero. If the teacher is less
confident, it tends to output a probability around 0.5. Using soft labels
enables the student model to learn this “degree of uncertainty” information
from the teacher as well, so a rule-based filtering scheme is not needed.

##  4 Experiment

###  4.1 Benchmark Tasks

Downbeat Tracking aims to predict the first beat of each bar, and is typically
conducted together with beat tracking. In this work, we focus on downbeat
tracking, because it still remains a challenge compared to beat tracking.
Following previous works [26], the model outputs frame-level probabilities of
beat, downbeat, and non-beat per 50 ms. A dynamic Bayesian network implemented
in madmom [30] is used as post-processing to decode the timestamps of
downbeats. Ballroom, Beatles, Hainsworth, HJDB, RWC, Simac, Harmonix, SMC, and
Internal-Beat/key are used as supervised datasets. The training process
randomly samples a chunk of 6 seconds from each training song. We select 74
songs from HarmonixSet as validation set. GTZAN [31] is used as test set. We
use F-measure implemented in mir_eval [32] for evaluation metric.

Chord Recognition is a challenging MIR task due to the lack of training data
as well as the complexity of modeling the harmonic relationship in music. The
model outputs frame-level probabilities of 25 classes (12 pitches of major and
minor chords plus one “none”) per 125 ms. We use Isophonics, Billboard,
HookTheory training set, RWC, and USPOP for training, and the 2000 songs of
HookTheory test set for validation. We use JayChou, one of the benchmark
datasets in MIREX,333https://www.music-ir.org/mirex/wiki/MIREX_HOME as our
testing set. JayChou is not used for training by the compared models in the
experiments. The evaluation metric is major/minor weighted accuracy in
mir_eval [32].

Key Detection aims to predict the tonal scale and pitch relation across the
entire song. It has been a task in MIREX and studied for years. The model
outputs the frame-level probabilities of 25 classes (12 major and 12 minor
keys plus one “none”) per 2 seconds. We use Isophonics, Billboard, HookTheory
training set, and Internal-Key/Beat for training, and the 2000 songs of
HookTheory test set for validation. Giantsteps [33] is used as test set. The
evaluation metric is a refined accuracy (weighted accuracy) in mir_eval [32],
with error tolerance that gives partial credits to reasonable errors.

Structure Segmentation aims to partition a music recording into non-
overlapping segments and predict the functional label (e.g. ’verse’ and
’chorus’) for each segment. Following previous work [34, 35], the model
outputs frame-level label probabilities of 7 classes and a frame-level
boundary probability, with a frame hop of 192 ms. We use Billboard, SALAMI,
RWC, HarmonixSet, and Internal-Structure for training. 200 pieces of songs are
randomly picked for validation. Isophonics is used for testing. We use the
F-measure of boundary hit rate at 0.5 seconds (_HR.5F_ in mir_eval [32]) for
evaluation.

(a) Downbeat Tracking

(b) Chord Recognition

(c) Key Detection

(d) Structure Segmentation

Fig. 2: Results of four MIR tasks with different sizes of training data and
model parameters. Yellow dashed line represents the results of supervised
teacher $M$. Red dotted line represents the SOTA results of non-semi-
supervised training.

###  4.2 Experiment Details

As summarized in Table 2, the parameters are chosen as suggested in previous
works or through hyperparameters tuning based on the first round of supervised
training. We use the same model architecture in the first round to train the
supervised teacher model $M$. Input audio is downsampled into 16k Hz. We use
Adam optimizer, 0.9 weight decay, and 20 epochs of patience. Each epoch runs
500 mini-batches, and early stopping is applied when validation scores do not
decrease for 80 epochs. Model with the best validation result is used for
testing. We train the models using Nvidia V100-32G GPUs. Total three rounds of
training are conducted with parameters shown in Table 2. We consider three
different data sizes for each round of training: (1) _D1_ is the Other-MIR
dataset depending on the task; (2) _D2_ is D1 plus Audio-92K; (3) _D3_ is D2
plus Music-151K.

###  4.3 Baseline System

We compare the proposed work with two types of baseline systems: one is the
supervised teacher $M$ trained from scratch with the labeled data only; the
other is based on the pre-trained model _MERT-330M_ [9], a state-of-the-art
large-scale self-supervised model for music audio.

For the baseline systems based on MERT-330M, two training schemes are
considered, namely _MERT-P_ and _MERT-F_. In MERT-P, we freeze the MERT model
to extract the deep embeddings, and then add two GRU layers for frame-level
classification. We use GRU layers instead of Multilayer Perception (MLP)
because they show better performance. In MERT-F, we add a dense layer (i.e.,
the same as the proposed model) and then fine-tune the MERT model for a fair
comparison.

| Downbeat | Key | Chord | Structure
---|---|---|---|---
SpecTNT-TCN [26] | 0.756 | - | - | -
CNN [36] | - | 0.743 | - | -
JLCX1 [37] | - | - | 0.845 | -
MuSFA [35] | - | - | - | 0.598
MERT-P | 0.736 | 0.664 | 0.723 | 0.616
MERT-F | 0.769 | 0.717 | 0.797 | —
Supervised $M$ | 0.753 | 0.731 | 0.802 | 0.594
Proposed | 0.776 | 0.764 | 0.862 | 0.632
Table 3: Results for four MIR tasks. Fine-tuning MERT on structure could not
fit into a V100 GPU.

##  5 Results and Discussion

Figure 2 shows the results with different sizes of models and data on four MIR
tasks. For most of the tasks, we find the performance is improved as the
training round unfolds. Larger models are more likely to benefit from more
data (see D2 or D3). Models trained on D3 perform the best except for chord
recognition. This could be due to data mismatch: our D3 comprises mainly
English songs, while JayChou contains only Chinese-pop songs. Nevertheless, it
achieves the largest relative improvement in chord recognition as compared to
its supervised teacher $M$. Since we do not have larger internal datasets for
chord recognition, transformer could not perform well in the supervised
setting. Among the tasks, key detection benefits the least from the proposed
approach. As key is a relatively global musical event that does not vary with
finer temporal resolution, there is less temporal dependency that the temporal
Transformer of PerceiverTF could take advantage of. Due to this constraint, we
also couldn’t see any benefits when scaling up the model more.

Table 3 summarizes our best results compared to the baselines. It is clear
that the proposed models all outperform their supervised teachers $M$ as well
as their self-supervised counterparts, demonstrating the effectiveness of the
semi-supervised framework. The proposed models also achieve state-of-the-art
performance compared to existing SOTA systems, showing the proposed approach
is task-independent. MERT performs well on downbeat tracking and structure
segmentation that require more temporal information, but is less successful on
chord recognition and key detection that require capturing harmonic
information. Fine-tuning MERT seems to be a better choice than probing MERT.
The improvement is larger especially on chord recognition and key detection
since it allows the pre-trained MERT to learn task-dependent information.

##  6 Conclusion

In this paper, we investigate the effectiveness of large-scale semi-supervised
MIR training by scaling up models to almost 100M parameters and training data
to 240k hours. Addressing the challenge of scarce labeled MIR data, our
experiment demonstrate that our proposed approach offers a better solution for
advancing MIR performance beyond supervised or self-supervised learning
approaches. For future research, we aim to include more downstream tasks.
Similar to other domains, we also plan to explore the impact of further
scaling up the training data and model size.

## References

  * [1] OpenAI, “GPT-4 technical report,” _ArXiv_ , vol. abs/2303.08774, 2023.
  * [2] P. Goyal _et al._ , “Self-supervised pretraining of visual features in the wild,” _ArXiv_ , vol. abs/2103.01988, 2021.
  * [3] A. Xiao _et al._ , “Scaling asr improves zero and few shot learning,” in _Proc. Interspeech_ , 2022, p. 5135–5139.
  * [4] M. Won _et al._ , “Semi-supervised music tagging transformer,” in _Proc. ISMIR_ , 2021, pp. 769–776.
  * [5] I. Simon _et al._ , “Scaling polyphonic transcription with mixtures of monophonic transcriptions,” in _Proc. ISMIR_ , 2022, pp. 44–51.
  * [6] T. Bertin-Mahieux _et al._ , “The million song dataset,” in _Proc. ISMIR_ , 2011, pp. 591–596.
  * [7] D. Bogdanov _et al._ , “The mtg-jamendo dataset for automatic music tagging,” in _ICML Workshops_ , 2019.
  * [8] P. Dhariwal _et al._ , “Jukebox: A generative model for music,” _ArXiv_ , vol. abs/2005.00341, 2020.
  * [9] Y. Li _et al._ , “MERT: Acoustic music understanding model with large-scale self-supervised training,” _ArXiv_ , vol. abs/2306.00107, 2023\.
  * [10] R. Castellon _et al._ , “Codified audio language modeling learns useful representations for music information retrieval,” in _Proc. ISMIR_ , 2021, pp. 88–96.
  * [11] M. Goto _et al._ , “RWC music database: Popular, classical and jazz music databases.” in _Proc. ISMIR_ , vol. 2, 2002.
  * [12] F. Gouyon _et al._ , _A computational approach to rhythm description-Audio features for the computation of rhythm periodicity functions and their use in tempo induction and music content processing_. Ph.D. dissertation, Universitat Pompeu Fabra, 2005.
  * [13] M. E. Davies _et al._ , “Evaluation methods for musical audio beat tracking algorithms,” _Queen Mary University of London, Centre for Digital Music_ , 2009.
  * [14] A. Berenzweig _et al._ , “A large-scale evaluation of acoustic and subjective music-similarity measures,” _Computer Music Journal_ , vol. 28, pp. 63–76, 2004.
  * [15] J. B. L. Smith _et al._ , “Design and creation of a large-scale database of structural annotations,” in _Proc. ISMIR_ , 2011, pp. 555–560.
  * [16] F. Gouyon _et al._ , “An experimental comparison of audio tempo induction algorithms,” _IEEE Trans. Audio, Speech, and Language Process._ , vol. 14, no. 5, 2006.
  * [17] M. Mauch _et al._ , “OMRAS2 metadata project 2009,” in _ISMIR Late Breaking and Demo_ , 2009.
  * [18] C. Donahue _et al._ , “Melody transcription via generative pre-training,” in _Proc. ISMIR_ , 2022, pp. 485–492.
  * [19] J. Hockman _et al._ , “One in the jungle: Downbeat detection in hardcore, jungle, and drum and bass,” in _Proc. ISMIR_ , 2012, pp. 169–174.
  * [20] A. Holzapfel _et al._ , “Selective sampling for beat tracking evaluation,” _IEEE Trans. ASLP_ , vol. 20, no. 9, pp. 2539–2548, 2012.
  * [21] O. Nieto _et al._ , “The Harmonix Set: Beats, downbeats, and functional segment annotations of western popular music,” in _Proc. ISMIR_ , 2019, pp. 565–572.
  * [22] Y.-N. Hung _et al._ , “A large TV dataset for speech and music activity detection,” _EURASIP Journal on Audio, Speech, and Music Processing_ , vol. 2022, pp. 1–12, 2022.
  * [23] W.-T. Lu _et al._ , “Multitrack music transcription with a time-frequency perceiver,” in _Proc. ICASSP_ , 2023.
  * [24] Q. Xie _et al._ , “Self-training with noisy student improves imagenet classification,” in _CVPR_ , 2019.
  * [25] M. Won _et al._ , “Data-driven harmonic filters for audio representation learning,” in _ICASSP_ , 2020, pp. 536–540.
  * [26] Y.-N. Hung _et al._ , “Modeling beats and downbeats with a time-frequency transformer,” in _Proc. ICASSP_ , 2022, pp. 401–405.
  * [27] K. He _et al._ , “Identity mappings in deep residual networks,” in _Proc. ECCV_. Springer, 2016.
  * [28] A. Jaegle _et al._ , “Perceiver: General perception with iterative attention,” in _Proc. ICML_ , 2021, pp. 4651–4664.
  * [29] W.-T. Lu _et al._ , “SpecTNT: a time-frequency transformer for music audio,” in _Proc. ISMIR_ , 2021, pp. 396–403.
  * [30] S. Böck _et al._ , “Joint beat and downbeat tracking with recurrent neural networks,” in _ISMIR_ , 2016, pp. 255–261.
  * [31] U. Marchand _et al._ , “GTZAN-Rhythm: Extending the gtzan test-set with beat, downbeat and swing annotations,” in _ISMIR Late Breaking and Demo_ , 2015.
  * [32] C. Raffel _et al._ , “MIR_EVAL: A transparent implementation of common mir metrics,” in _ISMIR_ , 2014, pp. 367–372.
  * [33] P. Knees _et al._ , “Two data sets for tempo estimation and key detection in electronic dance music annotated from user corrections,” in _Proc. ISMIR_ , 2015, pp. 364–370.
  * [34] J.-C. Wang _et al._ , “To catch a chorus, verse, intro, or anything else: Analyzing a song with structural functions,” in _Proc. ICASSP_ , 2022, pp. 416–420.
  * [35] ——, “MuSFA: Improving music structural function analysis with partially labeled data,” in _ISMIR Late Breaking and Demo_ , 2022.
  * [36] F. Korzeniowski _et al._ , “End-to-end musical key estimation using a convolutional neural network,” in _Proc. EUSIPCO_ , 2017, pp. 966–970.
  * [37] J. Jiang _et al._ , “MIREX 2018 submission A structural chord representation for automatic large-vocabulary chord transcription,” in _MIREX, ISMIR_ , 2018.