MM-Diffusion: Learning Multi-Modal Diffusion Models for Joint Audio and Video Generation

In plain words

Most diffusion models make one kind of output at a time, either a picture or a sound. MM-Diffusion makes a short video and its soundtrack in a single process. Both start as random noise and are cleaned up step by step. Two connected networks do this work: one handles the moving pictures, the other the sound. At each step the two networks look at each other through a limited attention window, which helps a visible event and its sound stay in step.

The same trained model can also be guided by one kind of signal. Given a video, it can produce matching sound, or the other way round, without training for that task. The authors report better scores than single-modality baselines, and in a human test most viewers judged its Landscape clips to be real footage. These results come from two datasets, nature scenes and street dance, and were measured by the authors.

TL;DR

• The model is a coupled pair of denoisers: a video branch with 2D and 1D convolutions and an audio branch with dilated 1D convolutions. Both are trained with the usual noise-prediction loss §3.2 §3.3 eq. 7.
• Video and audio talk through random-shift windowed cross-attention, which is much cheaper than full cross-attention §3.3.
• Conditional generation (audio from video, video from audio) is zero-shot, by replacing the conditioning signal or by gradient guidance. Only figures show it; no table measures it §3.4 §4.5 fig. 5.
• Main numbers are author-reported on Landscape and AIST++ and compare against single-modality models, not against a joint baseline §4.4. The joint baseline is a two-stage pipeline with a user study §4.6.
• The code is official and MIT-licensed. Several defaults, and the gradient guidance formula, differ from the paper (see Paper vs code).

Problem

Single-modality diffusion models cannot produce synchronized video and sound. The authors' baseline chains two single-modality models instead (sound first, then video conditioned on it) §1 §4.6. Video and audio also have different shapes: video is a 3D signal over space and time, audio is a 1D waveform, so one network must handle both §1.

Full cross-attention between all video and audio positions is too expensive. The authors estimate its cost as O((F×H×W)×T), with F frames, H×W pixels per frame and T audio samples, and they note that both signals are temporally redundant §3.3.

Fixed clip sizes are assumed: 16 video frames at 64×64 and 25,600 audio samples, which the authors describe as 10 fps video and 16 kHz audio tab. A.1 §4.1.

Method

Preliminaries §3.1

The forward process adds Gaussian noise over T steps with a linear beta schedule §3.1 eq. 1 eq. 2. The reverse process uses a Gaussian whose mean a network predicts. The authors predict the noise (epsilon-prediction) and drop learned variance because it gives only a small gain §3.1. The SR network does learn variance (see Pitfalls).

Multi-modal diffusion §3.2

1. Audio and video get independent forward processes that share one beta schedule §3.2 eq. 5.
2. The reverse step for each modality is conditioned on both current noisy inputs: the audio update uses (a_t, v_t) to produce a_{t-1}, and the video update does the same in the other direction §3.2 eq. 6.
3. The training loss is the noise-prediction MSE for each modality, with an optional time weighting λ(t) that the paper does not fix §3.2 eq. 7.

Coupled U-Net §3.3

• Video branch: spatial and temporal dimensions are separated. Stacked 1D temporal convolutions are followed by 2D spatial convolutions, and attention is likewise 2D plus 1D §3.3 app. A tab. A.1.
• Audio branch: dilated 1D convolutions, with dilation doubling from 1 to 2^N (N is a hyperparameter) §3.3. Temporal attention is removed from audio blocks because it is expensive and gave little benefit in preliminary runs §3.3.
• Four scales. Each scale has two multi-modal blocks and one down- or upsampling block §4.1 tab. A.1.
• Random-shift multi-modal attention (RS-MMA) runs only at scales 2, 3 and 4, with window sizes 1, 4 and 8 §4.1 tab. A.1 fig. 3.

RS-MMA in three steps:

• Step I: split the audio into F segments, one per video frame §3.3.
• Step II: pick a window of S frames (S much smaller than F) and a random shift R in [0, F−S]. Audio segment i attends to video frames from (i+R) mod F through (i+R+S) mod F §3.3.
• Step III: cross-attention from the audio segment to the video window, softmax(Q K^T / sqrt(d_k)) V, with K = linear(flatten(v_j)) §3.3 eq. 8 eq. 9. The paper omits the video-to-audio direction as symmetric §3.3.

# Random-shift cross-modal attention at one U-Net scale  [§3.3]
# Input:  video (B, F, C, H, W), audio (B, C, T), window S
# Output: video and audio, each updated by a residual attention term

def rs_mma(video, audio, window):
    # Step II: draw one shift per call, uniform in [0, F - window]  [§3.3]
    shift = sample_uniform_int(0, F - window)
    video_out = copy(video)
    audio_out = copy(audio)
    # Step I: split audio into one segment per video frame  [§3.3]
    audio_segments = split_by_frame(audio, num_frames=F)  # (B, F, C, T/F)
    for frame in range(F):
        # Step II: window of frames, wrapping modulo F  [§3.3]
        frames = [(frame + shift + k) % F for k in range(window)]
        keys_values = video[:, frames]                    # (B, S, C, H, W)
        # Step III: audio segment attends to the window  [eq. 8] [eq. 9]
        audio_out[:, frame] = attend(
            query=audio_segments[:, frame],
            keys=keys_values,
            values=keys_values,
        )
    # The video-to-audio direction is not written out in the paper;
    # the code implements it symmetrically
    # [code: mm_diffusion/multimodal_unet.py#L498-L564]
    return video_out, audio_out

Each block adds its attention output to its input as a residual, so the attention starts from an identity mapping; see Pitfalls for the zero-initialised projections in the code.

Training step

# One joint noise-prediction training step  [§3.2] [eq. 7]
# Input:  video (B, 16, 3, 64, 64) in [-1, 1]; audio (B, 1, 25600)
# Output: scalar loss

def training_step(video, audio):
    # One diffusion step per example, drawn uniformly  [app. A]
    t = sample_uniform_steps(batch_size=B, num_steps=1000)
    # Independent Gaussian noise per modality  [§3.2]
    noise_video = random_normal_like(video)
    noise_audio = random_normal_like(audio)
    # Forward noising with the shared schedule  [§3.2] [eq. 2] [§3.4]
    # signal_scale = sqrt(alpha_bar_t), noise_scale = sqrt(1 - alpha_bar_t)
    video_t = signal_scale[t] * video + noise_scale[t] * noise_video
    audio_t = signal_scale[t] * audio + noise_scale[t] * noise_audio
    # One network reads both noisy inputs and predicts both noises  [§3.3]
    pred_video, pred_audio = coupled_unet(video_t, audio_t, t)
    # Sum of the two MSE terms, each averaged over its own elements  [eq. 7]
    loss_video = mean((pred_video - noise_video) ** 2)
    loss_audio = mean((pred_audio - noise_audio) ** 2)
    return loss_video + loss_audio

Sampling

• Unconditional sampling starts from Gaussian noise for both modalities and runs the joint reverse process §3.2. The base sampler is DPM-Solver by default; complete DDPM sampling (1,000 steps) is used for the rows marked * §4.1 §4.4. The number of DPM-Solver steps is not stated in the paper.
• Super-resolution maps 64×64 frames to 256×256 with a separate model, sampled with DDIM in 25 steps §4.1 tab. A.1.

Zero-shot conditional generation §3.4

Two methods, both without extra training §3.4:

• Replacement: at each reverse step, the conditioning modality is replaced by a forward-noised copy of the real condition §3.4.
• Gradient-guided: the replacement is kept, and the target is also nudged by the gradient of how far the predicted condition is from the real condition §3.4 eq. 10 eq. 11 eq. 12. The paper calls λ a gradient weight and gives no value.

# Gradient-guided step for audio given video  [§3.4] [eq. 10] [eq. 11] [eq. 12]
# Input:  audio_next (noisy audio at step t+1), video_real (clean video)
# Output: audio_t (guided audio at step t)

def guided_audio_step(audio_next, video_real, t, lam):
    # Replace the video by its forward-noised version at step t+1  [§3.4]
    video_hat = forward_noise(video_real, step=t + 1)
    # Joint denoising step gives audio and video at step t  [eq. 11]
    audio_t, video_t = unet(audio_next, video_hat, step=t + 1)
    # Gradient of the video mismatch with respect to the audio  [eq. 12]
    mismatch = sum((video_t - video_hat) ** 2)
    grad = grad_of(mismatch, wrt=audio_t)
    # Move the audio against the gradient; lam is the guidance weight  [eq. 12]
    return audio_t - lam * sqrt(1 - alpha_bar[t]) * grad

The paper's formula is not the one the code runs. See Paper vs code before implementing it.

Hyperparameters

Coupled U-Net and super-resolution (SR) settings from the supplement, plus sampling and evaluation settings from the main text. "Not stated" means the paper does not say it.

Results

Author-reported. Setup: Landscape (1,000 clips of 10 s from 928 YouTube videos) and AIST++ (1,020 clips, center crop 1024×1024 for training) §4.2. Metrics are computed at 64×64 for all methods §4.1. FVD and KVD use an I3D classifier trained on Kinetics-400. FAD uses AudioCLIP features and is printed scaled by 1e4 §4.3. Baselines are DIGAN and TATS-base (video) and DiffWave (audio) §4.4. "Ours-v" and "Ours-a" are the video and audio sub-networks of the coupled U-Net §4.4. "Ours*" uses complete 1,000-step DDPM sampling tab. 1 tab. 2.

Landscape tab. 1 (lower is better)

AIST++ tab. 2 (lower is better)

Window-size ablation tab. 3 (80k steps, scales 2, 3, 4)

Random-shift ablation tab. 4 (FVD / FAD at 60K / 80K / 100K steps)

Mean opinion score tab. 5 (1 to 5, higher is better; VQ = video quality, AQ = audio quality, A-V = alignment)

Turing test tab. 6 (percent of clips judged real)

Authors' conclusions (from the text): Ours beats the single-modality models on video and audio §4.4. Ours-v beats DIGAN and TATS-base on most metrics §4.4. Joint training beats the separate sub-networks §4.4. Full sampling beats DPM-Solver on quality §4.4. The text does not quantify the gradient-guided method's advantage over replacement; it is claimed from figures only §4.5 fig. 5 fig. C.4.

Paper vs code

Official repository: github.com/researchmm/MM-Diffusion, linked from the abstract. MIT licence, commit 7250222114ee1eca36e32d90e8a41f12d92651e2. The README names the paper as CVPR 2023 README.md. Discrepancies:

1. Gradient guidance formula. The paper scales the update by sqrt(1−ᾱ_t) and compares the predicted condition at step t with the forward-noised condition at step t §3.4 eq. 12. The code scales by sqrt(ᾱ_t) [code: mm_diffusion/multimodal_gaussian_diffusion.py#L817]. It compares the model's prediction at step t−1 with the condition noised to step t−1 [code: mm_diffusion/multimodal_gaussian_diffusion.py#L788] [code: mm_diffusion/multimodal_gaussian_diffusion.py#L811]. It also skips the update at t = 0 multimodal_gaussian_diffusion.py.
2. Gradient weight. The paper gives no value for λ. The shipped scripts use classifier_scale 3.0 audio2video_sample_sr.sh [code: ssh_scripts/video2audio_sample.sh#L13].
3. Replacement method. The code re-noises the given condition at every step, with one fixed noise tensor for the whole sample [code: mm_diffusion/multimodal_gaussian_diffusion.py#L642-L700]. The paper does not say whether noise is reused.
4. Random shift. The code draws one shift per attention call, shared across the batch, with random.randint(0, f - local_window) [code: mm_diffusion/multimodal_unet.py#L614-L622]. It has no training-only check, so sampling is also randomised. The paper does not say whether the shift is used at inference.
5. Undocumented bottleneck attention. When the window list is [1, 4, 8], the code adds an extra cross-attention block in the middle of the U-Net with a window of all 16 frames and no shift [code: mm_diffusion/multimodal_unet.py#L875-L897]. The paper describes none.
6. Scale numbers. The code keys attention on downsampling factors: 2, 4, 8 for the U-Net multimodal_script_util.py and 8, 16, 32 for SR audio2video_sample_sr.sh. The paper's "2, 3, 4" and "4, 5, 6" are consistent with these only if they are 1-based level numbers. This is my reading; the paper does not say (see Pitfalls).
7. Sampling steps. The paper gives no DPM-Solver step count. The code hard-codes 20 steps with order 3 for the base model [code: py_scripts/multimodal_sample_sr.py#L125-L128]. Its SR default is 50 steps with order 2 multimodal_sample_sr.py. The shipped SR script uses DDIM in 25 steps instead multimodal_sample_sr.sh.
8. Motion term. A motion-variance loss (weight 0.05) is defined but has no call site in the training loss multimodal_gaussian_diffusion.py. The paper does not mention it.
9. Batch and hardware. The paper uses batch 128 on 32 V100s for the U-Net. The shipped script runs 8 processes with batch 4 multimodal_train.sh multimodal_train.sh. The SR script's 8 processes with batch 6 gives 48, matching the paper image_sr_train.sh image_sr_train.sh.
10. Augmentation. Table A.1 lists no augmentation for the U-Net. The dataset code flips video horizontally by default, and audio is not flipped [code: mm_diffusion/multimodal_datasets.py#L23] [code: mm_diffusion/multimodal_datasets.py#L168].
11. Learned variance. The SR model learns variance, as the shipped script shows image_sr_train.sh. The U-Net does not. This matches Table A.1, but §3.1 says variance is dropped in general.
12. Evaluation kernel. The paper does not specify the KVD kernel. The code uses scikit-learn's polynomial_kernel with no parameters set [code: evaluations/compute_fvd_kvd.py#L18-L34].
13. Evaluation split. The README says only the training set is used for training and evaluation README.md. The paper does not state the split.

Pitfalls

• Best numbers come from a different sampler. The best FVD on both datasets, and the best AIST++ FAD, come from the complete-DDPM row (Ours*). The best Landscape FAD comes from the DPM-Solver row tab. 1 tab. 2. The text says complete sampling is better overall §4.4, but Ours* is worse on Landscape FAD.
• Zero-shot is not measured. The abstract claims zero-shot conditional generation. The paper has no table for it, only figures and a claim that gradient guidance is better §4.5 fig. 5 fig. C.4.
• Replacement vs gradient. The paper gives no λ. The code uses 3.0 for both conditional scripts. Conditional sampling uses DDPM with 1,000 steps and batch size 1 or 2 in the shipped scripts audio2video_sample_sr.sh video2audio_sample.sh. Each step backpropagates through the U-Net, so this is slow.
• Equation 10 vs 12. Equation 10 has 1/sqrt(1−ᾱ_t) as its prefactor, but equation 12 has sqrt(1−ᾱ_t). v̂_t is used but never defined §3.4 eq. 10 eq. 12. Implement the code's version (Paper vs code, item 1), or ask.
• Scale numbering. "Scales 2, 3, 4" (U-Net) and "4, 5, 6" (SR) are not defined in the paper. The code uses downsampling factors. The SR numbers are out of range if read as 0-based indices into six channel multipliers tab. A.1.
• Random shift is per call. Shifts are drawn anew for each attention block and each denoising step, and shared across the batch. Results are not reproducible without seeding random, and the paper does not say how the shift is handled at inference.
• Random shift ablation is mixed on video. With RS, FVD is worse at 60K and 80K steps (440.78 vs 415.50; 306.42 vs 303.19) and better only at 100K tab. 4. FAD is worse at 60K (13.29 vs 12.55) and better at 80K and 100K. The text's claim that RS helps audio matches the later checkpoints §4.5.
• Window choice trades off metrics. The chosen [1, 4, 8] gives the best FVD and KVD but the worst FAD at 80K steps tab. 3.
• Sub-network baselines. The paper says the coupled U-Net is "decomposed" into audio and video sub-networks, but does not say whether each was retrained alone §4.4.
• Two-stage baseline is under-specified. The text says DiffWave generates audio and TATS "transfers audio to video" §4.6. Checkpoints and conditioning details are not stated. The gap to Ours in Table 5 may partly reflect the baseline setup.
• Metrics at 64×64. Scores are computed at 64×64, while the final output is 256×256 after SR §4.1. Do not compare them with 256×256 numbers from other papers.
• Scaled FAD. Table FAD values are already multiplied by 1e4 §4.3. Divide by 1e4 to compare with unscaled FAD.
• Ground-truth KVD is negative (−0.12 on Landscape) tab. 1. The paper does not explain the negative value.
• Commented-out row in the LaTeX source. The AIST++ table source has a commented-out alternative row (FVD 51.70, KVD 10.22, FAD 9.85). It is not in the typeset table and is not used here [raw/src/2212.09478/sections/Experiments.tex].
• Learned variance. §3.1 says variance is dropped, but the SR model learns it §3.1 tab. A.1. Follow Table A.1 for the SR model.
• Initialisation. Output projections of the cross-attention block are zero-initialised, so the block starts as an identity [code: mm_diffusion/multimodal_unet.py#L609-L610]. The paper does not say this.
• Optimizer defaults. The paper does not state the optimizer. The code's AdamW call passes only learning rate and weight decay; other settings stay at library defaults [code: mm_diffusion/multimodal_train_util.py#L107-L109]. Learning rate is constant, because lr_anneal_steps defaults to 0 and disables decay [code: mm_diffusion/multimodal_train_util.py#L336-L338].
• Data and checkpoints are external. The README points to Google Drive and Baidu links for datasets, checkpoints and the i3d and AudioCLIP weights [code: README.md#L65]. Scripts hard-code absolute paths and must be edited.
• Size. 115.13 M (U-Net) plus 311.03 M (SR) parameters, trained on 32 V100s for 100,000 steps, plus a separate SR model trained for 270,000 steps on 8 V100s §4.1 tab. A.1.
• Clip timing. 16 frames at 10 fps and 25,600 audio samples at 16 kHz both give 1.6 s (derived from the stated values) tab. A.1. Clips of other lengths need the shapes changed everywhere.

Lineage

• Builds on: DDPM (Ho et al.) for the diffusion core and DDIM for SR sampling; DPM-Solver and DPM-Solver++ for sampling; DiffWave for the 1D audio U-Net and dilations; ADM for the image U-Net and SR; Video Diffusion for the 2D+1D factorisation and the replacement and gradient conditioning §2 §3.3 §3.4.
• Compared against: DIGAN and TATS-base (video), DiffWave (audio), and a two-stage DiffWave plus TATS pipeline §4.4 §4.6.
• Related conditional work named in the paper: CDCD, CMT, Sound2Sight, and CMCGAN §2.
• Superseded by: not stated in the content.