QA-MDT: Quality-aware Masked Diffusion Transformer for Enhanced Music Generation

1 Introduction

Text-to-music (TTM) generation aims to transform textual descriptions of emotions, style, instruments, rhythm, and other aspects into corresponding music segments, providing new expressive forms and innovative tools for multimedia creation. According to scaling law principles [Peebles and Xie, 2023; Li et al., 2024a], effective generative models require a large volume of training data. However, unlike image generation tasks [Chen et al., 2024a; Rombach et al., 2021], acquiring high-quality music data often presents greater challenges, primarily due to copyright issues and the need for professional hardware to capture high-quality music. These factors make building a high-performance TTM model particularly difficult.

In the TTM field, high-quality music signals is scarce. This prevalent issue of low-quality data, highlighted in Figure 1 (a), manifests in two primary challenges. Firstly, most available music signals often suffer from distortion due to noise, low recording quality, or outdated recordings, resulting in diminished generated quality, as measured by pseudo-MOS scores from quality assessment models [Ragano et al., 2023]. Secondly, there is a weak correlation between music signals and captions, characterized by missing, weak, or incorrect captions, leading to low text-audio similarity, which can be indicated by CLAP scores [Wu et al., 2023]. These challenges, especially the inherent quality in the music signal itself, significantly hinder the training of high-performance music generation models, resulting in poor rhythm, noise, and inconsistencies with textual control conditions in the generated audio. Additionally, as shown in Figure 1 (b), directly filtering low-quality music, which robustly reduces the dataset size, leads to a consistent decline in model performance. Therefore, finding an effective training strategy for largescale datasets with low-quality waveforms, mismatches, and missing labels has become an urgent challenge.

In this paper, we introduce a novel quality-aware masked diffusion transformer (QA-MDT) to enhance music generation, aiming to tackle the aforementioned problems while making further improvements through architectural studies. We made efforts on effectively leveraging extensive opensource music databases, which often contain data of varying quality and style, to produce high-quality, diverse and high text-audio consistency music. For music quality enhancement, we innovatively inject music quality into the denoising stage with multiple granularities to foster quality awareness during training, while high-quality music can be obtained by setting a quality threshold during inference. Regarding the modeling architecture, in preliminary experiments, we found that the Diffusion Transformer (DiT) framework, which has been successful in the image domain [Peebles and Xie, 2023], is not directly effective for modeling music spectrograms. However, injecting a masking strategy significantly enhances the spatial correlation of the music spectrum and further accelerates convergence. Additionally, we utilize large language models (LLMs) and the CLAP model to synchronize music signals with captions, thereby enhancing text-audio correlation in extensive music datasets. Our ablation studies on public datasets confirm the effectiveness of our methodology, with the final model surpassing previous works in both objective and subjective measures. In summary, we focus on developing better training strategies and network architectures to enhance the quality and aesthetic of music generation. At the same time, we address the longstanding issue of text-audio consistency in the field of TTM, which can be listed as:

• We propose a quality-aware training paradigm that enables the model to perceive the quality of the dataset during training, thereby achieving superior music generation in terms of both musicality and audio quality.
• We innovatively introduced the Masked Diffusion Transformer to music signals, demonstrating its unique efficacy in modeling music latent space and its capability in perceiving quality control, thereby further improving both the generated quality and musicality.
• We address the issue of low text-audio correlation in large-scale music datasets for TTM, effectively improving text alignment and generative diversity.

2 RelatedWork

Text to music generation. Text-to-music generation aims to create music clips that correspond to input descriptive or summary text. Previous efforts have utilized either language models (LMs) or diffusion models (DMs) to model quantized waveform representations or spectral features. Models like MusicLM [Agostinelli et al., 2023], MusicGen [Copet et al., 2024], MeLoDy [Lam et al., 2024], and Jen-1 [Li et al., 2024b] leverage LMs and DMs on residual codebooks obtained via quantization-based codecs [Zeghidour et al., 2021; D´efossez et al., 2022]. Moˆusai [Schneider et al., 2023], Noise2Music [Huang et al., 2023a], Riffusion [Forsgren and Martiros, 2022], AudioLDM 2 [Liu et al., 2023a], and Stable Audio [Evans et al., 2024a] use U-Net-related diffusion to model mel-spectrograms or latent representations obtained through compression networks. Although some approaches attempt to guide the model towards generating high-quality content by setting negative prompts like "low quality" [Liu et al., 2023a; Chen et al., 2024b], few explicitly inject quality information during training. This results in the model's inability to effectively perceive and control content quality.

Transformer based diffusion models. Traditional diffusion models typically use U-Net as the backbone, where the inductive biases of CNNs do not effectively model the spatial correlations of signals and are insensitive to scaling laws [Li et al., 2024a]. However, transformer-based diffusion models (DiT) [Peebles and Xie, 2023] have effectively addressed these issues. This advantage is particularly evident in fields such as video generation [Brooks et al., ], image generation [Peebles and Xie, 2023; Chen et al., 2024a; Bao et al., 2022], and speech generation [Liu et al., 2023c]. To expedite training and foster inter-domain learning of correlations, the masking strategy has proven effective, yielding SOTA class-conditioned performances on ImageNet [Gao et al., 2023]. Additionally, a simpler architecture [Zheng et al., 2023] incorporating reconstruction losses and unmasked finetuning further enhances model training speed. However, these models have not yet been verified for text-controlled music generation on large-scale music datasets, and their adaptability with additional control information remains an open question. Make-an-audio 2 [Huang et al., 2023b] and, more recently, Stable Audio 2 [Evans et al., 2024b], have explored the DiT architecture for audio and sound generation. However, their approach models latent tokens by segmenting only along the time dimension to control and extend generation duration. In contrast, our focus is on finer segmentation within the latent space across both time and frequency, aiming for more precise modeling of music signals.

Quality enhancement in audio domain. Previous research has made efforts to improve the quality of generated audio, particularly in two key areas: waveform fidelity and the consistency between input text and generated content. Waveform quality can be compromised by issues like aliasing from low sampling rates and limited expressiveness due to monophonic representations, while models like MusicGen [Copet et al., 2024] and Stable Audio [Evans et al., 2024a; Evans et al., 2024b], which directly model 32k and 44.1k stereo audio, have significantly enhanced perceptual quality. Despite higher sampling rates and channels, the quality of audio in training datasets remains inconsistent, often suffering from noise, dullness, and a lack of rhythm or structure. These problems, often reflected by the Mean Opinion Score (MOS), are rarely addressed. In terms of text-audio consistency, Makean-audio 2 [Huang et al., 2023b] and WavCaps [Mei et al., 2024] have employed ChatGPT-assisted data augmentation to improve temporal relationships and accuracy in audio effect generation. Although studies like Music-llama [Liu et al., 2024] and LP-musiccaps [Doh et al., 2023] have introduced captioning approaches for music, few have explored the augmentation and utilization of synthetic data in large-scale music generation tasks.

3 Preliminary

Latent diffusion model. Direct application of DMs to cope with distributions of raw signals incurs significant computational overhead [Ho et al., 2020; Song et al., 2020]. Conversely, studies [Liu et al., 2023b; Liu et al., 2023a] apply them in a latent space with fewer dimensions. The latent representation z0 is the ultimate prediction target for DMs, which involve two key processes: diffusion and reverse processes. In the diffusion process, Gaussian noise is incrementally added to the original representation at each time step t, described by zt+1 = √1-βtzt + √βtϵ, where ϵ is drawn from a standard normal distribution N(0, I), and βt is gradually adapted based on a preset schedule to progressively introduce noise into the state zt. The cost function [Ho et al., 2020; Liu et al., 2023b] is formalized as arg minθ E(z0,y),ϵ [∥∥ϵ - Dθ (√αtz0 + √1 - αtϵ, t, y)∥∥2 ]. where Dθ, the denoising model, strives to estimate the Gaussian noise ϵ, conditioned on the latent state zt, the time step t, the conditional embedding y, and where αt represents a predefined monotonically increasing function. In the reverse process, we obtain zt-1 via the recursive equation: zt-1 = 1 √1-βt (zt - βt √1-αt ϵθ ) + √ 1-αt-1 1-αt βtϵ, where ϵθ represents the estimated Gaussian noise.

Classifier-free guidance. Classifier-free guidance (CFG), introduced by [Ho et al., 2020], increases the versatility and flexible control ability of DMs by both considering conditional and unconditional generation. Typically, a diffusion model generates content based on specific control signals y within its denoising function Dθ(zt, t, y). CFG enhances this mechanism by incorporating an unconditional mode Dθ(zt, t, ∅), where ∅ symbolizes the absence of specific control signals. The CFG-enhanced denoising function is then expressed as DCFG θ (zt, t, y) = Dθ(zt, t, y) + w(Dθ(zt, t, y)-Dθ(zt, t, ∅)), where w ≥ 1 denotes the guidance scale. During training, the model substitutes y with ∅ at a constant probability puncond. In inference, ∅ might be replaced by a negative prompt like "low quality" to prevent the model from producing such attributes [Liu et al., 2023a].

4 Method

4.1 Quality Information Injection

At the heart of our work lies the implementation of a pseudoMOS scoring model [Ragano et al., 2023] to meticulously assign music quality to quality prefixes and quality tokens.

We define our training set as Do = {(Mi, T o i ) | i = 1, 2, . . . ,ND}, where each Mi represents a music signal and T o i is the corresponding original textual description. To optimize model learning from datasets with diverse audio quality and minimize the impact of low-quality audio, we initially assign p-MOS scores to each music track using a model finetuned with wav2vec 2.0 [Baevski et al., 2020] on a dataset of vinyl recordings for audio quality assessment, and achieve the corresponding p-MOS set S = {s1, s2, . . . , sND }. These scores facilitate dual-perspective quality control for enhanced granularity and precision.

First, We analyze this p-MOS set S to identify a negative skew normal distribution with mean µ and variance σ2 . We define text prefixes based on s as follows: prepend "low quality" if s < µ - 2σ, "medium quality" if µ - σ ≤ m ≤ µ + σ, and "high quality" if s > µ + 2σ. This information is prepended before processing through the text encoder with cross-attention, enabling the initial separation of qualityrelated information.

To achieve a more precise awareness and control of waveform quality, we synergize the role of text control with quality embedding. We observed that the distribution of p-MOS in the dataset is approximately normal, which can be shown in Figure 1, allowing us to use the Empirical Rule to segment the data accordingly. Specifically, we define the quantization function Q : [0, 5] → {1, 2, 3, 4, 5} to map the p-MOS scores to discrete levels based on the distance from the mean µ in terms of standard deviation σ:

where r = 2 for s > µ, otherwise, r = 1. Subsequently, Q(s) is mapped to a d-dimensional quality vector embedding using the embedding function E, such that

This process provides finer granularity of control within the following model and facilitates the ability of interpolative quality control during inference, enabling precise adjustments in Rd . In later stages, the quality embedding is treated as a token on par with every latent audio patch, participating in the attention computations to enable interaction.

4.2 Quality-aware Masked Diffusion Transformer

In a general patchify phase with patch size pf × pl and overlap size of × ol, patchified token sequence X = {x1, x2, . . . , xP } ⊂ Rpf×pl are obtained through spliting the music latent space Mspec ∈ RF×L, as described in Section 5.2. The total number of patches P is given by:

A 2D-Rope position embedding [Su et al., 2024] is added to each patch for better modeling of relative position relationship while a binary mask m ∈ {0, 1}P is applied during the training stage, with a variable mask ratio γ. This results in a subset of ⌊γP ⌋ patches being masked that ∑PN i=1 mi = ⌊γP⌋, leaving P - ⌊γP⌋ patches unmasked. The subset of masked tokens is invisible in the encoder stage and replaced with trainable mask tokens in the decoder stage following the same strategy utilized in AudioMAE [Huang et al., 2022] and MDT [Gao et al., 2023].

The transformer we use consists of N encoder blocks, M decoder blocks, and an intermediate layer to replace the masked part with trainable parameters. We treat the embedding of the quantized p-MOS score as a prefix token, concatenated with each stage's music tokens. Let Xk = [xk1 , xk2 , . . . , xkP ] ∈ RP×d represent the output of k-th encoder or decoder block, where the initial input of the encoder X0 = zt = αtz0 + √1 - αtϵ, and the final decoder block estimate XN+M = z0 = [x1, x2, . . . , xP ]. For k < N, indicating the encoder blocks, the sequence transformation focuses only on unmasked tokens:

where m ∈ {0, 1}P is the mask vector, with 1 indicating masked positions and 0 for visible tokens.

For N < k < N +M, indicating the decoder blocks, the full sequence including both unmasked tokens and learnable masked tokens is considered:

where the previously masked tokens are now subject to prediction and refinement. In the decoding phase, the portions that were masked are gradually predicted, and throughout this entire phase, the quality token qvq(s) is progressively infused and optimized. Subsequently, the split patches are unpatchified while the overlapped area is averaged to reconstruct the output noise and every token contributes to calculating the final loss:

In the inference stage, the model can be guided to generate high-quality music through modified CFG:

Here qhigh vq and qlow vq indicate quantified p-MOS for guiding the model in a balance between generation quality and diversity. After obtaining the sampled latent ˆ z0 with DDIM sampler, we can finally generate the music using the VAE decoder.

4.3 Music Caption Refinement

We divided the caption refinement stage into three steps including text information enriching with music caption model Fcap, caption adjustment with CLAP cosine similarity function S and caption diversity extension with LLMs which we denoted as Fllm.

Initially, pretrained music caption model [Doh et al., 2023] is employed to re-annotate each music signal Mi to T g i , shown as Dg = {(Mi, T g i ) | T g i = Fcap(Mi), i = 1, 2, . . . ,N}. CLAP text-audio similarity is applied to filter Dg with a threshold of ρ1, resulting in

In this context, we meticulously filter out generated captions that do not correspond with their respective audio files. This misalignment may be attributed to inaccuracies within the captioner's insufficient training. For the filtered data pairs, we opt to retain the use of the original captions.

To ensure that valuable information from the original captions is not overlooked when using only the generated captions, we adapt a fusing stage to combine the original caption and generated pseudo prompt. Firstly, we need to filter out original captions that is useless or inaccurate, formulated as:

The issue can stem from the original data being improperly labeled with terms such as 'speech, car' from datasets like AudioSet [Gemmeke et al., 2017] and also may be because of desperately missing of the original labels.

Finally, only the original caption that suffers low CLAP text similarity score should be merged with the generated ones, for redundant, repetitive parts result in long and verbose final captions. Thus, we set the threshold to ρ3 and merge them by LLMs to Tfusion = Fllm(T o , T g ):

5 Experimental Setup

5.1 Datasets

For training, we used the following databases for our training: AudioSet Music Subset (ASM) [Gemmeke et al., 2017], MagnaTagTune (MTT) [Law et al., 2009], Million Song Dataset (MSD) [Bertin-Mahieux et al., 2011], Free Music Archive (FMA) [Defferrard et al., 2016], and an additional dataset 1 . Each track in these databases was clipped to 10second segments and sampled at 16kHz to ensure uniformity across the dataset. The final training set was developed through a process of caption refinement, as detailed in Section 4.3. Finally, we got our training set totaling 12.5k hours of diverse music data. The specific composition of these datasets is further elaborated in the Appendix. For evaluation, we test our model on the widely used MusicCaps benchmark [Agostinelli et al., 2023] and the Song-DescriberDataset [Manco et al., 2023]. MusicCaps consists of 5.5K 10.24-second clips sourced from YouTube, each accompanied by high-quality music descriptions provided by ten musicians. The Song-Describer Dataset is made up of 706 licensed high quality music recordings.

5.2 Models and Hyperparameters

Audio compression. Each 10.24-second audio clip, sampled at 16 kHz, is initially transformed into a 64 × 1024 mel-spectrogram with mel-bins of 64, hop-length of 160 and window length of 1024. Subsequently, this spectrogram is compressed into a 16 × 128 latent representation Mspec using a Variational Autoencoder (VAE) pretrained with AudioLDM 2 [Liu et al., 2023a] with series of quantization loss and adversarial loss. We use pretrained Hifi-GAN [Kong et al., 2020] vocoder to reconstruct the waveform from the generated mel-spectrogram.

Caption processing and conditioning. We utilize the LPMusicCaps [Doh et al., 2023] caption model for ASM, FMA, and subsets of MTT and MSD that have weak or no captions. We use the official checkpoint from LAION-CLAP [Wu et al., 2023] 2 for text-to-text and text-to-audio similarity calculations. Based on small scale subjective experienment, thresholds are set at ρ1 = ρ2 = 0.1 to ensure any generated text or original caption not aligned well with the corresponding waveform is filtered out. Additionally, after filtering, generated text that fall below a threshold of ρ3 = 0.25 are merged with original tags with the prompt: Merge this music caption "generated caption" with the ground truth tags "original tags", and do not add any imaginary elements.. We use FLAN-T5-large [Peebles and Xie, 2023] as text encoder for all models.

1 We use 55k music tracks from https://pixabay.com, which is a large scale copyright free dataset.

2 music speech audioset epoch 15 esc 89.98.pt

Diffusion backbone. We train our diffusion model with three backbones for comparison: U-Net [Ronneberger et al., 2015] based at 1.0B parameters; our proposed Quality-aware Masked Diffusion Transformer (QA-MDT) with N = 20 encoder layers and M = 8 decoder layers at 675M. We study the impact of the patch size and overlap size in the Appendix, and apply a patch size of 1 × 4 without overlap for the training of our final model. We train on 10.24-second audio crops sampled at random from the full track, maintaining a total batch size of 64, learning rate of 8e-5, and a condition drop of 0.1 during training. The final model was trained for a total of 38.5k steps. During inference, we use Denoising Diffusion Implicit Models (DDIM) [Song et al., 2020] with 200 steps and a guidance scale of 3.5, consistent with AudioLDM [Liu et al., 2023b]. We begin by presenting our approach to refining captioning, which includes the capability for quality awareness, transitioning from text-level control to token-level control. Finally, we compare proposed model with previous works subjectively and objectively.

5.3 EvaluationMetrics

We evaluate the proposed method using objective metrics, including the Fr´echet Audio Distance (FAD) [Kilgour et al., 2018], Kullback-Leibler Divergence (KL), Inception Score (IS) 3 . We also utilize pseudo-MOS scoring model [Ragano et al., 2023] to estimate generation quality, with more accurate assessments derived from subjective metrics.

6 Results

6.1 Quality Awareness

This subsection explores the effects and interactions of model control over quality tokens and quality text prefixes during the training phase, as well as their comparative effects across different models. In our previous MTT dataset of 1,000 test pairs, we filtered out pairs labeled with low quality or quality is poor to avoid confusion when applying quality prefixes, resulting in a new subset of 519 entries, which we refer to as the MTT Filter Set (MTT-FS). Figure 4 illustrates the impact of different quality prefixes during inference when quality is used as a text prefix during training for U-Net and MDT-based backbones. It was observed that U-Net, when inferred with different quality prefixes, showed only minor changes in p-MOS scores and did not adhere to the threshold set during training. In contrast, MDT demonstrated better learning of quality information from prefixes, achieving p-MOS scores significantly higher than those of U-Net and the test set. Additionally, by decoupling quality information from the training set, we achieved better FAD (5.602 vs 5.757) and higher p-MOS (4.039 vs 3.796) compared to training and inference without quality text prefixes. Given that quality tokens are specifically designed for the Transformer architecture, Figure 3 (left) shows the controlled outcomes when different quality tokens are used after integrating quantified quality as a token during training. Remarkably, using quality tokens alone provided more precise and accurate pMOS score control. In our ablation study, we compared the effects of using only text prefixes against combining both approaches. As shown in Figure 3 (right), as the quantized control level gradually increased, the model steadily improved in p-MOS scores, which represent the quality of generation. Concurrently, FAD and KL also progressively optimized until a turning point at level 4, where a higher average p-MOS was achieved than when solely using prefixes. This turning point may be due to the scarcity of examples with quality level 5 in the dataset. Moreover, by combining two types of quality information injection, the refined decoupling and interaction allowed the model to more accurately perceive audio data quality features during training, leading to significant reductions in FAD and KL compared to using only one of them.

3 We strictly follow the comparison method and evalution code in AudioLDM2 and ensure that the indicators in the paper follow consistent sampling rates and durations for fair comparison. All above metrics are computed using the audioldm eval library [Liu et al., 2023b], ensuring standardized evaluation.

We also compare our approach with the traditional "negative prompt" strategy [Liu et al., 2023a] in Appendix, highlighting our approach's significant improvement in quality and reduction in FAD.

6.2 Impact of Music Caption Refinement

We conducted our ablation study on a subset of our training set, which includes ASM and FMA, totaling approximately 3,700 hours and 1.1 million clips. For evaluation, we utilized an out-of-domain set with 1,000 samples randomly selected from MTT [Law et al., 2009]. Table 2 compares the model's performance using different textual representations: sentences formed by merging original tags with commas (Do), generated captions (Dg), and generated captions refined through filtering and fusion (Dmerge). During the filtering and fusion stage, 8.9% of the generated captions were filtered out, and 15.1% were fused with original tags using ChatGPT. Each model underwent training for 60,000 steps with a batch size of 64.

From Table 2 we can also observe consistent trends: employing a captioner to transform audio annotations from sparse words into detailed sentences significantly improved the models' generalization and diversity. This indicates that detailed annotations are essential for learning the relationship between the models and spectral features. Moreover, the filter and fusion stages led to enhancements across all metrics, highlighting the significance of precise, comprehensive annotations for generalization ability and control ability. We also found that compared to U-Net, the MDT architecture shows stable improvements in basic modeling metrics, making it a better backbone for music spectral modeling.

6.3 Compared with Previous Methods

We compared our proposed method with the following representative previous methods: AudioLDM 2 [Liu et al., 2023a], Mousai [Schneider et al., 2023] and Jen-1 [Li et al., 2023] which model music using spectral latent spaces, MusicLM [Agostinelli et al., 2023], and MusicGen [Copet et al., 2024], which focus on modeling discrete representations.

We re-inferred AudioLDM2-Music and MusicGen-1.5B using their official checkpoints to compare additional metrics under the same environment. The results are presented in Table 2. For Ours (U-Net), we inferred all text with the prefix "high quality", while for Ours (QA-MDT), we used the same prefix along with a p-MOS quality token set to level 5. When calculating the CLAP score, we evaluated the generated music with original prompt, which did not include any quality prefix. The experimental results show significant advantages in both subjective and objective metrics for our models. Since KL divergence measures the distance between audio samples, higher quality audio often results in deviations from the original waveform of Musiccaps, which can lead to lower performance. Although Ours (U-Net) showed a slight FAD advantage on the Song-Describer-Dataset, this may be due to instabilities arising from the small scale test dataset, and we further demonstrated the superiority of QA-MDT in subsequent subjective experiments. Additionally, since MusicGen was trained on non-vocal tracks, it may underperform on captions that include vocals.

Based on subjective evaluation shown in Table 3, our proposed method significantly improves overall audio quality and text-audio consistency, thanks to the label optimization for large music datasets and the quality-aware training strategy. By analyzing the backgrounds of the evaluators and their corresponding results, we can also see that for beginners, the comparison between different systems is not sensitive, which is related to their lack of music background experience and knowledge. However, from the perspective of our method in product operators, video editors, and audio producers, our method offers considerable enhancements, underscoring its potential value to audio industry professionals.

7 Conclusion and Discussion

In this study, we address the key challenges in the music generation domain, including model architecture design, largescale uneven audio quality, and unaligned textual annotations, all of which impede the progress of TTM with quality, musicality, and text alignment. In the future, we aim to further enhance and expand our model to achieve long-duration, highsampling-rate, controllable, and highly interactive music generation.

Acknowledgements

This work was supported by the National Natural Science Foundation of China under Grant 62171427.

Contribution statement

Chang Li and Ruoyu Wang contributed equally. Jun Du supervised the entire project and serves as the corresponding author.