JavisDiT: Joint Audio-Video Diffusion Transformer with Hierarchical Spatio-Temporal Prior Synchronization

In plain words

Making a video with sound from a text prompt is often done in steps: one model draws the picture, another adds a soundtrack, or two models run side by side. The hard part is timing and place. A door slam has to land on the frame where the door moves, at the moment the sound starts. A text prompt rarely says when or where anything happens, so the sound and the picture drift apart.

JavisDiT generates the video and the sound in one model, at the same time. Each part works on small chunks of its data, and the two parts exchange information as they go. A small extra network reads the prompt and produces a few hints: where things should appear in the picture, and when they should sound. Both parts are steered by those hints so they line up. That network is trained by showing it real, matching video-audio pairs and deliberately broken ones.

The authors also release a benchmark of 10,140 sounding videos, many of them with several sounds at once or in sequence, plus a metric that scores how well the sound and picture line up.

TL;DR

• JavisDiT is a diffusion transformer with one branch for video and one for audio, linked by bidirectional cross-attention §3.1.
• A learned ST-prior estimator maps the prompt to 32 spatial and 32 temporal tokens, sampled from a Gaussian. Cross-attention injects them into both branches §3.2 app. E.3.
• Training has three stages: audio pretraining, contrastive ST-prior training, and joint training of only the cross-attention layers §3.3 tab. A1.
• Author-reported: JavisScore 0.154 vs 0.151 for FoleyCrafter on JavisBench, but UniVerse-1 is better on FVD, KVD and TV-IB tab. 1.
• Official code (Apache-2.0) and checkpoints are public. Several config values differ from the paper text; see Paper vs code.

Problem

• Joint audio-video generation (JAVG) takes a text prompt s and produces video v and audio a together. A good result needs coarse semantic agreement with s and fine spatio-temporal agreement between v and a §2.
• Spatial alignment means an event in a region of the frame matches frequency content in the audio. Temporal alignment means onsets and offsets in the video match the start and stop of the sound §2.
• A text prompt gives only a coarse event description. It does not fix when or where the event happens, and the authors avoid prompting an LLM for explicit timings §3.2.
• The paper compares against cascaded pipelines (text-to-audio then audio-to-video, or text-to-video then video-to-audio) and joint models (MM-Diffusion, UniVerse-1) §5.1.
• Existing evaluation is weak. AIST++ and Landscape have only 20 to 100 test samples and narrow scenarios app. D.1. AV-Align scores near chance on the authors' 3,000-pair check (AUROC 0.5296) app. D.4.2 tab. A4.

Method

Overview

• Two DiT branches. The video branch is initialized from OpenSora. The audio branch starts from the video branch weights and is then trained app. C.1.
• The model has 28 layers, 16 attention heads, hidden width 1152 and FFN ratio 4. It has 3.14B parameters in total app. C.1.
• Each branch uses spatial and temporal self-attention in sequence, which the authors use to cut attention cost app. C.1. Coarse text enters through T5 cross-attention §3.1.
• The fine priors enter through a spatial and a temporal cross-attention per branch §3.2. The video and audio tokens then pass through bidirectional cross-attention §3.1.
• Sampling is a denoising step at time t, mapping the noisy video and audio latents to time t-1 given the prompt eq. 1. The scheduler is rectified flow app. C.1.

(\mathbf{v}_{t-1}, \mathbf{a}_{t-1}) = \mathcal{G}_\theta(\mathbf{v}_{t}, \mathbf{a}_{t}, s, t) eq. 1

• The video latent comes from the OpenSora VAE, which downsamples time, height and width by (4, 8, 8) app. C.2.2. The audio is a mel spectrogram encoded by the AudioLDM2 VAE app. C.1. Both VAEs are frozen during training app. C.1.

ST-prior estimator

• Coarse prior: the T5 embedding of the prompt is reused as is §3.2.
• Fine prior: the 77 hidden states of ImageBind's text encoder feed a 4-layer transformer. It is queried by 32 learned spatial tokens and 32 learned temporal tokens §3.2 app. C.2.1.
• The transformer outputs a mean and a variance for each token. The priors are sampled from that Gaussian, so one prompt can yield different priors §3.2 app. C.2.1.
• Default token count and width are 32 + 32 tokens of 128 channels, chosen after the ablations in app. E.3.

# ST-prior estimator: prompt -> spatial and temporal prior tokens  [§3.2]
# Input:  prompt (str)
# Output: spatial_prior (32, 128), temporal_prior (32, 128)

def estimate_st_prior(prompt):
    # 77 hidden states of ImageBind's text encoder, width 1024  [app. C.2.1]
    text_states = imagebind_text_states(prompt)       # (77, 1024)
    # Linear projection to the transformer width  [code: base_model.py#L221]
    memory = linear_in(text_states)                   # (77, 512)
    # 32 spatial + 32 temporal learned queries  [code: base_model.py#L222]
    queries = concat(spatial_queries, temporal_queries)   # (64, 512)
    # 4-layer transformer over memory and queries  [§3.2]
    hidden = prior_transformer(memory, queries)       # (64, 512)
    # Mean and variance per token, then sample  [§3.2] [app. C.2.1]
    gaussian = diagonal_gaussian(linear_out(hidden))  # 2 x 128 channels
    tokens = gaussian.sample()                        # (64, 128)
    # First 32 tokens are spatial, last 32 temporal  [code: base_model.py#L233]
    spatial_prior = tokens[:32]
    temporal_prior = tokens[32:]
    return spatial_prior, temporal_prior

Contrastive training of the estimator

• Anchor: the prompt and its estimated priors. Positive: the real, synchronized video-audio pair from the training data app. C.2.1.
• Negatives come from asynchronous pairs. Easy negatives pair the real video with audio that AudioLDM2 generated from the caption alone. Hard negatives corrupt either modality with the augmentations listed in app. C.2.4.
• A VA-Fuser encodes each pair. It maps the video and audio latents to matching token counts, applies bidirectional attention, and projects to spatial and temporal embeddings app. C.2.2.
• The total objective is given in eq. A1 and the per-prior contrastive term in eq. A2.

\mathcal{L}_{prior} = \mathcal{L}_{contrast}(\mathbf{p}_s,\mathbf{e}_s^+,\mathbf{e}_s^-) + \mathcal{L}_{kl}(\mathbf{p}_s) + \mathcal{L}_{contrast}(\mathbf{p}_t,\mathbf{e}_t^+,\mathbf{e}_t^-) + \mathcal{L}_{kl}(\mathbf{p}_t) eq. A1

\mathcal{L}_{contrast} = \mathcal{L}_{token} + \mathcal{L}_{disc} + \mathcal{L}_{vad} + \mathcal{L}_{reg} eq. A2

# Contrastive loss for one prior type, spatial or temporal  [app. C.2.3]
# Input:  prior (N, C); positive embedding (N, C); negatives (K, N, C)
# Output: scalar loss, the sum of the four terms of [eq. A2]

def contrastive_loss(prior, positive, negatives):
    # Token-level term on cosine similarity  [app. C.2.3]
    sim_pos = cosine(prior, positive)
    sim_neg = cosine(prior, negatives)
    loss_token = abs(1.0 - sim_pos) + abs(1.0 + sim_neg)
    # Auxiliary discriminator: 1-layer attention, [CLS] token  [app. C.2.3]
    # Written as in the paper; see Pitfalls  [app. C.2.3]
    logit = discriminator(prior, positive)
    loss_disc = bce(logit, 1.0) + bce(logit, 0.0)
    # Separate positive and negative video-audio embeddings  [app. C.2.3]
    loss_vad = abs(1.0 + cosine(positive, negatives))
    # Pull the prior toward the positive embedding, L2 norm  [app. C.2.3]
    loss_reg = norm(prior - positive, p=2)
    return loss_token + loss_disc + loss_vad + loss_reg

Joint layer

• The paper describes one DiT layer as a sequence of modules and says it is repeated 28 times app. C.1. The code runs the same modules in a different order; see Paper vs code.

# One joint audio-video layer; the model stacks 28  [app. C.1]
# Input:  video (B, T*S, C), audio (B, R*M, C), text (B, 300, C)
# Output: video and audio tokens with the same shapes

def joint_layer(video, audio, text, spatial_prior, temporal_prior):
    # Video: spatial then temporal self-attention with text  [§3.1]
    video = video_spatial_block(video, text)
    video = video_temporal_block(video, text)
    # Audio: the same pattern on the spectrogram tokens  [§3.1]
    audio = audio_spatial_block(audio, text)
    audio = audio_temporal_block(audio, text)
    # Spatial prior attends along space, temporal prior along time  [§3.2]
    # [code: javisdit/models/stdit/vastdit3.py#L111-L137]
    video = video_prior_block(video, spatial_prior, temporal_prior)
    audio = audio_prior_block(audio, spatial_prior, temporal_prior)
    # Bidirectional cross-attention between the two modalities  [§3.1]
    video, audio = bi_cross_attention(video, audio)
    return video, audio

Training stages

1. Audio pretraining: the audio branch learns from 788K audio-caption pairs drawn from public datasets. The video branch is frozen §3.3 app. C.1 tab. A1.
2. ST-prior training: 611K synchronized video-audio-text triplets (136K from MMTrail, 475K from TAVGBench). Contrastive loss with synthesized negatives app. C.1 tab. A1.
3. JAVG training: the video and audio branches and the prior estimator are frozen. Only the ST-cross-attention and bidirectional cross-attention modules train, on 611K samples §3.3 tab. A1.

JavisBench and JavisScore

• JavisBench has 10,140 video-audio pairs with captions and a 5-dimension, 19-category taxonomy. A 1,000-sample subset is called JavisBench-mini §4.1.
• Of the samples, 75% have several sounding events, 28% have sequential events and 57% have simultaneous events §4.1.
• Sources: test sets of Landscape, AIST++ and FAVDBench, plus YouTube videos uploaded June to December 2024. About 30K candidates were filtered down to 10,140 app. D.3.
• JavisScore cuts each clip into 2 s windows with 1.5 s overlap. It uses frozen ImageBind encoders and averages the window scores §4.2 app. D.4.1.

S_{Javis} = \frac{1}{W}\sum_{i=1}^{W}\sigma\left(E_v(V_i), E_a(A_i)\right) eq. 2

\sigma(V_i,A_i) = \frac{1}{k}\sum_{j=1}^{k}\operatorname{top\text{-}k}_{\min}\left\{\cos\left(E_v(V_{i,j}),E_a(A_i)\right)\right\} eq. A3

# JavisScore for one clip  [§4.2] [app. D.4.1]
# Input:  video frames, audio waveform; frozen ImageBind encoders
# Output: scalar (author scale, cosine based)

def javis_score(frames, audio):
    # 2 s windows, 1.5 s overlap  [app. D.4.1]
    windows = sliding_windows(frames, audio, size=2.0, overlap=1.5)
    window_scores = []
    for clip_frames, clip_audio in windows:
        # Cosine between each frame embedding and the window's audio  [eq. A3]
        sims = cosine(vision_embed(clip_frames), audio_embed(clip_audio))
        # Keep the 40% least similar frames and average them  [app. D.4.1]
        # [code: eval/javisbench/src/metrics.py#L427-L430]
        window_scores.append(mean(lowest_fraction(sims, 0.4)))
    # Average over windows  [eq. 2]
    return mean(window_scores)

Hyperparameters

Source: arXiv v2 LaTeX as fetched. "Not stated" means the paper text and appendix do not give the value. Code citations are from the repository at code_commit.

Results

Author-reported. Setup: JavisBench, 4 s clips at 240p and 24 fps, audio at 16 kHz tab. 1 app. C.1. The baselines were reproduced by the authors, except that AIST++ and Landscape numbers are copied from their papers §5.1 tab. 2.

Main table on JavisBench tab. 1:

Landscape and AIST++, with numbers copied from the original papers tab. 2:

Length generalization, JavisDiT only tab. 5:

Ablations, run on JavisBench-mini with a 60K-sample stage-3 subset §5.3. Quality (lower is better), consistency and synchrony (higher is better):

• Qualitative: spatial priors focus on the sounding subject, and temporal priors spread evenly over a continuous sound §5.3 fig. 7.
• Human study: 100 JavisBench prompts, three raters, blind win/tie/lose on video quality, audio quality and alignment. The authors report video quality slightly below UniVerse-1, and audio and alignment above it §5.4 fig. 8. No numbers are given in the text.
• Category breakdown: the authors report that scenes with several sounding objects or events score lower on JavisScore §5.2 fig. 5 app. D.5.

Paper vs code

Code: github.com/JavisVerse/JavisDiT, Apache-2.0, at code_commit. The repository's README names this paper's model JavisDiT-v0.1 and also describes JavisDiT++ (v1.0). This card cites only the v0.1 configs (configs/javisdit-v0-1/).

• Stage 1 epochs: 55 in the paper tab. A1, 50 in the config stage1_audio.py.
• Stage 2 epochs: 1 in the paper tab. A1 §5.1, 2 in the config stage2_prior.py.
• Stage 2 warm-up: 1000 steps in the paper tab. A1, 100 in the config stage2_prior.py.
• Stage 3, the same learning rate, warm-up, EMA and epoch count appear in both tab. A1 stage3_jav.py.
• Token-level loss: the paper writes an unweighted sum of two absolute terms app. C.2.3. The code computes 0.5 times the mean of two ReLU hinge terms base_model.py base_model.py.
• Discriminator loss: the paper writes two BCE terms on the same logit, labelled 1 and 0 app. C.2.3. The code uses one BCE over the positive and all negatives, with the positive labelled 1 base_model.py.
• Regularizer: the paper's term is the L2 norm of the difference app. C.2.3. The code uses the mean squared error over all elements base_model.py.
• KL term: the paper lists one KL term per prior with weight 1 eq. A1. The code adds a single KL over both priors, weighted by 1e-6 base_model.py.
• Prior transformer depth: the paper says 4 layers §3.2. The code builds 4 encoder and 4 decoder layers base_model.py base_model.py.
• VA embedding width: the appendix says the fused embedding is projected to width 512 app. C.2.2. The VA encoder is built at out_dim, which the v0.1 configs set to 128 base_model.py stage2_prior.py. The 512 in the config belongs to the prior transformer only.
• The appendix mentions a 2-layer MLP for the VA projection app. C.2.2. This card did not find it in base_model.py. The only MLP on that path is inside the cross-modality block.
• Layer order: the appendix lists spatial self-attention, cross-attention, spatial cross-attention, bidirectional attention, FFN, then the temporal sequence app. C.1. The code runs video spatial and temporal self-attention, audio spatial and temporal self-attention, the two prior blocks, then bidirectional attention vastdit3.py.
• Stage 3 freezing: the main text says the self-attention blocks are frozen §3.3. The appendix says the whole video branch is frozen app. C.1. The code freezes the full spatial and temporal block stacks and final layers of both branches, and trains only the prior-injection and cross-attention modules (the text embedder is frozen by freeze_y_embedder) vastdit3.py stage3_jav.py. This matches the appendix, not the main-text wording.
• Random masking: the paper describes a 6 x 6 grid app. C.2.4. The code uses grid cells of size min(H, W) // 6 pixels, placed at random positions, possibly overlapping augment.py.
• The quality formula in §5.3 uses S_KFD, but the tables report KVD. This card assumes KVD.
• Inference settings match: 30 steps, CFG 7.0, 240p, 4 s app. C.1 sample.py sample.py.

Pitfalls

• Prior sampling at inference: the estimator draws a fresh prior on every call base_model.py. Repeated calls with the same prompt give different priors unless the random state is reseeded. Seeding the whole run makes it reproducible, but the paper does not say whether it fixes the prior per prompt.
• Unconditional branch: the code uses a learned null embedding for the prior base_model.py. The paper gives only the CFG scale and does not describe the unconditional prior. This comes from code only.
• Config values not in the paper: the inference config sets an aesthetic score of 6.5 and leaves the motion score unset sample.py. The paper does not mention either. This card did not trace how the aesthetic value is used.
• Headline margin: JavisScore 0.154 vs 0.151 is a small gap on one benchmark. The main table reports single values with no error bars or repeated runs tab. 1.
• Mixed picture on quality: UniVerse-1 is better on FVD, KVD, TV-IB and CLIP tab. 1. On AIST++, AV-DiT has lower FVD (68.8 vs 86.7) tab. 2. Baseline numbers in tab. 2 are copied from other papers, so their protocols may differ.
• Ablation vs main table: ablations use a 60K-sample stage-3 subset and JavisBench-mini §5.3. The full-model sync score is 0.153 there and 0.154 in the main table tab. 3 tab. 1.
• JavisScore library defaults: the function in the evaluation code defaults to 0.5 s windows with no overlap metrics.py. The CLI passes 2.0 s and 1.5 s main.py. Calling the function directly silently changes the metric.
• Metric accuracy: on the authors' 3,000-pair check, JavisScore reaches AUROC 0.6533 and accuracy 75.14%. AV-Align reaches 0.5296 and 52.54% app. D.4.2 tab. A4. A quarter of pairs are still misordered.
• Audio alone is weaker than AudioLDM2 in-domain: FAD 5.19 vs 2.01 on AudioCaps at 10 s tab. A7 app. E.2. The authors attribute this partly to variable-length training and to AudioLDM2's exposure to AudioCaps app. E.2.
• Hidden cost: the three stages sum to 328 H100 GPU-days tab. A1. The benchmark captions use Qwen2-VL-72B, Qwen2-Audio-7B and Qwen2.5-72B-Instruct, and the hard negatives use an external Qwen API for sound-source extraction app. D.3.3 app. C.2.4.
• Data provenance: part of the training data is YouTube footage the authors collected and filtered app. C.1. The paper does not state the license or whether that footage will be released. Do not redistribute it without checking.
• Benchmark selection: the content-based post-filter removes clips with only background music or only speech, which the authors say is nearly half of the videos. Human checking then yields the 10,140 samples app. D.3.4. Results on the benchmark therefore reflect that filter.
• Repository scope: the repo also holds JavisDiT++ configs and a v1.0 model collection, which this card does not cover.

Lineage

• Builds on: OpenSora for the video branch, VAE and dynamic masking §3.3 app. C.1; AudioLDM2 for the audio VAE and easy negatives app. C.1 app. C.2.4; ImageBind for the fine-prior text encoder and the JavisScore encoders app. C.2.1 app. D.4.1; T5 for the coarse prior §3.1; rectified flow app. C.1.
• Bidirectional cross-attention follows an earlier audio-video grounding design, cited as §3.1.
• Compared against: TempoTokens and TPoS (text-to-audio then audio-to-video); ReWaS, See&Hear and FoleyCrafter (text-to-video then video-to-audio); MM-Diffusion and UniVerse-1 (joint); AV-DiT and MM-LDM on Landscape and AIST++ §5.1 tab. 2.
• Metric lineage: JavisScore's frame selection is motivated by a prior benchmark's approach §4.2. It replaces AV-Align, which the authors find unreliable on multi-event clips §4.2.