ImmerseDiffusion: A Generative Spatial Audio Latent Diffusion Model

The Ambisonic surround-sound system encodes and reproduces sound directions and amplitudes to deliver immersive 3D audio experiences. For simplicity and considering the fact that the first-order Ambisonics (FOA) format generally provides adequate spatial resolution [11], this paper focuses exclusively on generating FOA components. However, it can be generalized to higher-order Ambisonics for improved spatial fidelity, albeit with increased complexity. FOA achieves spatial rendering using four-channel components, denoted as W, X, Y, and Z as follows:

	$\displaystyle W$	$\displaystyle=\frac{1}{\sqrt{2}}p,$	$\displaystyle\quad X$	$\displaystyle=p\cos(\theta)\cos(\phi)$		(1)
	$\displaystyle Y$	$\displaystyle=p\sin(\theta)\cos(\phi),$	$\displaystyle\quad Z$	$\displaystyle=p\sin(\phi)$

where $p$ is sound pressure, $W$ is the omnidirectional component, $X$ is the front-back directional component, defined by the azimuth angle $\theta$$\;\in(-\pi,\pi]$, and elevation angle $\phi$$\;\in(-\pi/2,\pi/2]$, $Y$ is the left-right directional component, defined by the azimuth angle $\theta$ and elevation angle $\phi$, and $Z$ is the up-down directional component, defined by the elevation angle $\phi$. Both the training data and the generated audio will conform to this standard.

Large model sizes, high-dimensional audio data, iterative diffusion processes, and the need to capture long-term structures necessitate efficient computation in generative audio models. Autoencoder architectures are commonly employed to manage these demands by operating within compressed latent domains. For our model, which generates four FOA components simultaneously, compression is even more critical.

To address compression, various approaches have been explored in the literature. Some models leverage autoencoders with quantized latents such as Vector Quantized Variational Autoencoders VQ-VAE (e.g., [12]) or Residual Vector Quantization RVQ-VAE (e.g., [13]). Other models leverage continuous VAE latents to avoid quantization information loss [2, 3, 9]. Following [3], we utilize the 1D convolutional U-Net autoencoder from the Descript Audio Codec (DAC) [14] and replace the discrete RVQ bottleneck with a continuous VAE bottleneck. We also eliminate the decoder’s $tanh()$ activation to prevent harmonic distortion, as noted in [9]. Additionally, we modify the training loss functions to suit the FOA components. Unlike stereo codecs, which leverage losses like Mid-Side Short Time Fourier Transform (MS-STFT) for proper stereo reconstruction, FOA does not need such losses due to the inherent orthogonality of the $X$, $Y$, and $Z$ channels. The training loss function of our FOA codec is as follows:

$$\mathcal{L}_{\text{codec}}=\frac{\lambda_{\text{mrstft}}}{4}\mathcal{L}_{\text{mrstft\_W}}+\frac{\lambda_{\text{mrstft}}}{4}\mathcal{L}_{\text{mrstft\_X}}+\frac{\lambda_{\text{mrstft}}}{4}\mathcal{L}_{\text{mrstft\_Y}}+\frac{\lambda_{\text{mrstft}}}{4}\mathcal{L}_{\text{mrstft\_Z}}+\lambda_{\text{kl}}\mathcal{L}_{\text{kl}}+\lambda_{\text{adv}}\mathcal{L}_{\text{adv}}+\lambda_{\text{fm}}\mathcal{L}_{\text{fm}}$$

(2)

In this context, $\mathcal{L}_{\text{mrstft}}$ represents the Multi-Resolution STFT loss for each channel, collectively contributing to the generation loss. $\mathcal{L}_{\text{kl}}$ denotes KL divergence loss applied to the VAE bottleneck, $\mathcal{L}_{\text{adv}}$ signifies the adversarial (discrimination) loss, while $\mathcal{L}_{\text{fm}}$ refers to the feature matching loss. The $\lambda$ parameters are weights perceptually assigned to adjust each corresponding loss contribution in total loss [3]. Our ambisonic codec model transforms a 4-channel FOA signal with length $L$ into $64$ channel latent with length $L/2048$ with an overall compression factor of $128$.

Many generative audio methods use text prompts for user control [12, 13]. Recent approaches have transitioned from pure language models to CLAP-based models, such as LAION CLAP [5] for audio and MuLan [15] for music, due to their enhanced performance [3, 2, 16]. These models are effective in aligning language and audio embeddings, enabling cross-modality training and inference [2]. Additionally, various conditioning types are explored, including control over temporal order, pitch, energy [17], music genre [12], and rhythm and chords [18].

ImmerseDiffusion offers spatial control in audio generation via two models: descriptive and parametric. The descriptive model utilizes text-based prompts to define the sound source, its spatial attributes, and environmental acoustics. In contrast, the parametric model uses non-spatial text prompts only for sound source descriptions while providing numerical values for spatial and environmental parameters. Moreover, in line with previous stereo diffusion models  [3, 9], both models employ temporal conditioning, including start time and total time as well. In Figure 1, (I) and (II) illustrate how descriptive and parametric conditioning incorporate location and environmental parameters through text or a combination of text and numerical values, along with temporal conditioning.

To generate conditioning spatial text embeddings for the descriptive model, we use the Embeddings for Language and Spatial Audio (ELSA) [4] model’s text encoder. Trained on spatial audio-text pairs using synthetic spatial variants of AudioCaps [19], FreeSound [20], and Clotho [21], ELSA demonstrated superior performance in spatial audio retrieval tasks [4]. ELSA’s design for comprehending spatial qualities makes it ideal for text conditioning in our model. Thanks to ELSA text conditioner and the mentioned spatial training datasets, our proposed descriptive model considers spatial description details (e.g., right corner, far left, front top) and environment acoustics (e.g., medium-sized acoustically damped room, large highly reverberant hall) in spatial audio generation. In addition to that, two cross-attention conditioners are employed to include the mentioned temporal conditions (i.e. start time and total time).

For the parametric model, we replace ELSA text encoder with the LAION CLAP [5] text encoder, which is solely trained on non-spatial prompts to validate parametric conditioning and avoid potential spatial bias from ELSA embeddings. Alongside text embeddings, this model includes three numerical spatial conditioning parameters (azimuth, elevation, distance) and two environmental acoustic conditioning parameters (room size, 30dB decay time). It also comprises two temporal conditioning parameters, namely, start time and total time, similar to the descriptive model.

Several generative audio latent diffusion models utilize 1D convolutional U-Net structures for their diffusion blocks (e.g., [2, 3]). However, the Diffusion Transformer (DiT), initially proposed for image generation [22], has recently been adapted for music generation [9, 23]. In our work, we employ a DiT that features a cascade of blocks that include self-attention, cross-attention, and gated MLP components with layer normalization and skip connections, similar to that used in [9], with structural adoptions for FOA-domain. We adjust conditioning token dimensions to $512$ and introduce a projecting conditioning layer to meet dimensionality requirements. The model is trained using the v-objective approach [24] to minimize the Mean Squared Error (MSE) between the true $v$ and the predicted $v_{\theta}(z_{t},t,c)$ at timestep $t$, conditioned on $c$ and given the latent variable $z_{t}$ as follows:

$$\mathcal{L}_{\text{Diffusion}}=\left\|v_{\theta}(z_{t},t,c)-v\right\|_{2}^{2}$$

(3)

We trained the ambisonic codec and both descriptive and parametric ImmerseDiffusion models using synthetically spatialized versions of the FreeSound [20], AudioCaps [19], and Clotho [21] datasets, as provided in [4]. To improve generalization and performance on speech audio, we also included a spatialized version of LibriSpeech [25] and the corresponding FOA noise dataset, detailed in [26]. Each dataset contains spatialized audio, original captions, spatial and environmental parameters, and spatially derived captions, as described in [4]. Datasets information is provided in Table 2.

We train our models with a cluster of A100-80GB GPUs. For the ambisonic audio codec, we use a batch size of 96 and an excerpt length of 1.48 seconds for 600K steps. The spatial codec generator and discriminator were trained with base learning rates of $10^{-4}$ and $2.\times 10^{-4}$ respectively, using the AdamW optimizer, applying a $10^{-3}$ decay factor and momentum parameters of 0.8 and 0.99.

Both descriptive and parametric models share the same training setup. We trained the Diffusion transformers with frozen ambisonic codec and conditioner weights with a batch size of 576 and a context window of 5.95 seconds for 500K steps. The base learning rate was set to $10^{-4}$ with the AdamW optimizer, employing momentum values of 0.9 and 0.99. Spatial and non-spatial text conditioning for descriptive and parametric models utilized the official pre-trained ELSA [4] and LAION CLAP [5] models, respectively.

The quality of FOA codec models is evaluated using STFT and MEL distances between the original and reconstructed FOA audio from the test set using AuraLoss [27] with default settings similar to [3, 9].

We evaluate the plausibility of generated FOA audio clips by computing FAD which compares the statistical properties of the generated and reference audio embeddings. To compute the reference and generated FOA audio embeddings, we used ELSA audio encoder that supports FOA audio at 48 kHz. Additionally, we report KL divergence using a pre-trained ELSA.

The CLAP score is computed as the cosine similarity between the conditioning spatial text embeddings and the generated FOA audio embeddings, using the ELSA model. For the parametric model, KL divergence and CLAP score are computed using corresponding spatial captions from the test set, despite the model being trained on non-spatial captions and spatial parameters.

To measure the spatial accuracy of FOA audio generated by ImmerseDiffusion, we report metrics based on ground truth and estimated azimuth ($\theta$), elevation ($\phi$), and distance (d) using equation set (5). The intensity vectors $I_{x}$, $I_{y}$ and $I_{z}$ are derived by multiplying the omnidirectional channel with the corresponding directional channels, as outlined in (4).

$I_{x}=W\cdot X,\quad\quad I_{y}=W\cdot Y,\quad\quad I_{z}=W\cdot Z$

(4)

$\begin{aligned} \text{$\theta$}=\tan^{-1}\left(\frac{I_{y}}{I_{x}}\right),\quad\text{$\phi$}=\tan^{-1}\left(\frac{I_{z}}{\sqrt{I_{x}^{2}+I_{y}^{2}}}\right),\quad\text{d}=\sqrt{I_{x}^{2}+I_{y}^{2}+I_{z}^{2}}\end{aligned}$

(5)

For all mentioned three parameters, we report the L1 norm between the ground truth and estimated values for azimuth, elevation and distance, across the test set. To calculate the L1 norm for azimuth, we use the circular difference method (see Equation 6). This method accounts for the circular nature of azimuth angles, ensuring we measure the shortest angular distance between the ground truth $\theta$ and the estimated azimuth $\hat{\theta}$. It avoids issues with linear differences, where angles like -3.13 and +3.13 are in nearly the same direction but appear far apart in a linear measurement.

$$\text{L1}_{\text{$\theta$}}=\|\min\left(\left|\text{$\theta$}-\hat{\theta}\right|,2\pi-\left|\text{$\theta$}-\hat{\theta}\right|\right)||_{1}$$

(6)

To quantify the difference between the ground truth and estimated direction of arrival (DoA), we also include their spatial angle ${\Delta}_{\text{Angular}}$ [28] calculated using equations (7) and  (8) as follows:

$$a=\sin^{2}\left(\frac{{\Delta}_{\phi}}{2}\right)+\cos(\phi)\cdot\cos({\hat{\phi}})\cdot\sin^{2}\left(\frac{{\Delta}_{\theta}}{2}\right)$$

(7)

$${\Delta}_{\text{Spatial-Angle}}=2\cdot\arctan 2\left(\sqrt{a},\sqrt{1-a}\right)$$

(8)

where ${\Delta}_{\phi}$ and ${\Delta}_{\theta}$ denote the linear and circular differences for elevation and azimuth respectively, and ${\phi}$ and $\hat{\phi}$ represent the ground truth and estimated elevations.