Inside NVIDIA Cosmos 3’s Two-Tower Architecture

NVIDIA Cosmos 3 is combination of autoregressive reasoning and continuous world generation inside one Mixture-of-Transformers architecture.

Every transformer layer holds two full sets of weights.

The Reasoner tower processes language tokens and visual-understanding tokens. It uses causal attention and next-token prediction, so on its own it behaves like a standard VLM.

The Generator tower processes continuous image, video, audio and action latents. It uses full attention and iterative denoising.
Both towers start from the same pretrained VLM weights. The Reasoner is trained first and stays frozen while the Generator is pretrained.

The attention mask defines how the two towers communicate
Reasoner queries attend causally to Reasoner keys and values only. Generator queries attend bidirectionally to both Reasoner and Generator keys and values.
Reasoner: Attention_causal(Q_AR, K_AR, V_AR) Generator: Attention_full(Q_DM, [K_AR; K_DM], [V_AR; V_DM])

The Reasoner never reads Generator tokens.

Understanding and generation use different visual encoders
A ViT encoder produces the visual-understanding tokens for the Reasoner. It is trained together with the transformer.

A separate video VAE produces the latents for the Generator. It is frozen. On the generation path, video is compressed 4x in time and 32x32 in space, which includes a 2x2 patch merge.

The Generator turns noise into video, audio and actions

It does not predict the next token. It learns how to move a noisy latent toward the clean one.

During training, each real example is encoded to a latent and mixed with noise. Some examples get a little noise, some are mostly noise. The model predicts how the latent should change at that noise level. Because the training code knows both the original latent and the noise it added, it can compute the exact target and score the prediction against it.

Video, audio and action latents in the same sequence do not have to be equally noisy.

Position encoding connects the modalities
Cosmos 3 uses 3D multimodal rotary position embeddings with separate time, height and width axes. Video tokens use all three. Audio and action tokens use only the time axis.

The time axis is scaled to real time. Video at 24 fps, audio at 25 tokens per second and actions at their own sample rate all land on the same physical time axis, so they line up when generated together.

Actions are inputs and outputs. The same model supports 3 action modes

Forward dynamics: given current observations and actions, generate future observations.
Inverse dynamics: given an observed transition, generate the actions that caused it.
Policy: given observations and an instruction, generate both the action sequence and the future observations it should produce.

Inference
A reasoning-only request runs the Reasoner tower alone. The Generator weights are not activated. A generation request runs both towers.

The data strategy is well thought too and I will post that next.

Discover more from AI & Data Notes for Leaders

Subscribe now to keep reading and get access to the full archive.

Continue reading