UNITE-AUDIO: Joint Learning of Continuous Tokenization and Latent Flow Matching for Text-to-Audio Generation

1. INTRODUCTION

• Most TTA systems compress audio into latents for conditional diffusion or flow generation [1] [2] [3], while tokenizers are optimized for reconstruction independently of generation.The resulting latent space may be less suitable for generative modeling.Source: §1 · p. 1
• Speech-generation work adds pretrained semantic supervision [4] [5], while visual-generation studies shape representations through masked modeling, diffusion priors, or joint latent generation [7] [8].Other visual methods use pretrained representations or jointly model semantic features [10].Source: §1 · pp. 1–2
• Unite-Audio jointly optimizes continuous audio tokenization and latent generation, learning representations from scratch without external pretrained alignment models.An online encoder processes partial audio for flow prediction against a full-audio target from an EMA encoder, while reconstruction supports faithful decoding.Source: §1 · p. 2
• Experiments and ablations report competitive TTA performance with a compact model and support the effectiveness of the designs for learning a generation-friendly latent space.Figure 1 contrasts conventional separate tokenizer–generator training with joint single-stage training.Source: Figure 1 · p. 1, §1 · p. 2

2.2. Partial-Context Flow Matching

• Generation-path training divides waveforms into 2-second buckets, samples 0–30% retained context, and uses an EMA full-view representation as the target.The online context is noised to the same flow time as the EMA trajectory, so the flow model receives a complete noisy latent sequence while contextual positions retain a gradient path to the encoder.Source: §2.2 · p. 2
• The x-prediction flow-matching loss targets clean full-view EMA representations over positions complementary to online context.A shifted flow-time schedule emphasizes high-noise states, and gradients propagate to the online encoder through contextual representations.Source: §2.2 · p. 2

2.3. Training Objectives and Schedule

• The joint objective weights reconstruction loss by p_rec and flow-matching loss by 1 − p_rec, combining acoustic-fidelity and generative supervision.Reconstruction uses global and local spectral losses, adversarial learning, and feature matching.Source: §2.3 · p. 3

2.1. Overview

• The online encoder, decoder, and latent flow model are jointly optimized from scratch, with reconstruction and generation paths both shaping the latent space.The EMA encoder supplies a full-waveform target without gradients; later, the encoder and decoder are frozen while the flow model specializes in generation from Gaussian noise.Source: §2.1 · p. 2, §2.3 · p. 3

2.3. Training Objectives and Schedule

• Training proceeds from joint latent learning to flow-only full-generation training, then Flow-GRPO post-training with group-based rewards.Flow-GRPO uses CLAP similarity and PaSST KL divergence in its composite reward and a clipped policy objective; encoder and decoder remain fixed during post-training.Source: §2.3 · p. 3

3.1. Dataset

• Stage I uses 1,729,107 AudioSet and 246,324 WavCaps examples; Stages II and III use AudioCaps training data, with AudioCaps test entries removed from AudioSet-derived data.WavCaps samples shorter than 20 seconds are retained, and all audio is resampled to 16 kHz with a maximum duration of 20 seconds.Source: §3.1 · p. 3

3.2. Model Architecture

• The latent flow model has 117.2M trainable parameters, with 18 Transformer layers and 11 fused audio–text MMDiT blocks.The continuous tokenizer produces 64-dimensional representations at 25 Hz, and text conditioning uses a frozen FLAN-T5-Large encoder.Source: §3.2 · p. 3

3.3. Implementation Details

• Stage I trains for 240k iterations with reconstruction sampling probability 0.5; Stage II and Stage III use learning rates of 5 × 10^-5 and 1 × 10^-6, respectively.Stage III uses group size 4, and experiments use two NVIDIA H800 GPUs.Source: §3.3 · p. 3

3.4. Evaluation

• Evaluation follows the TangoFlux protocol on the 886-sample AudioCaps test split, reporting FD, KL, CLAP, IS, FD16k, and RTF.RTF is averaged across ten runs of 10-second generation on one NVIDIA H800 GPU; Table 1 compares systems and counts only latent generative-model parameters.Source: §3.4 · pp. 3–4, Table 1 · p. 4

4.1. Performance

• With a 117M-parameter latent flow model, Unite-Audio achieves the highest CLAP score (0.534) and lowest KLPaSST (1.057) among compared systems, with competitive IS.Flow-GRPO improves semantic alignment and event consistency, although FD degrades; Stage II offers a better distributional-fidelity trade-off, and 8 NFEs cause only modest degradation.Source: §4.1 · p. 4

4.2. Necessity of Noisy Partial-Context Flow Matching

• No context removes the online representation from flow matching, clean context adds clean latents to noisy flow states, and noisy context matches the online representation to the flow time.The three Stage-I checkpoints are evaluated at 200k iterations using EMA weights on the same 886-sample AudioCaps test set.Source: §4.2 · p. 4
• Noisy context performs best on most generation and representation metrics, achieves the lowest Mel distance, and maintains comparable SI-SDR.The table evaluates latent representations learned with different context-learning strategies at 200k iterations.Source: Table 2 · p. 4, §4.2 · p. 4

• Noise-matched partial-context flow matching lets the generative objective shape the encoder while avoiding clean-context information unavailable at inference.Full-generation training further aligns learning with generation from noise.Source: §5 · p. 4
• Flow-GRPO post-training further improves text-conditioned generation.Source: §5 · p. 4
• Unite-Audio achieves strong semantic fidelity with a compact 117M-parameter latent flow model and retains competitive performance under few-step sampling.Source: §5 · p. 4