Category: Hard

  • DL0133 VLM Fusion: Cross-Attention, Q-Former, and MLP

    Compare Cross-Attention fusion, Perceiver Resampler / Q-Former, and Linear/MLP Projection (e.g., LLaVA) for vision-language alignment.

    Answer

    All three designs answer one question: where does the frozen vision encoder’s patch grid meet the language model? Linear/MLP projection (LLaVA, LLaVA-1.5, Qwen2-VL, InternVL) maps every patch embedding into the LLM’s token space with a two-layer MLP and concatenates the result onto the text sequence, so vision and language share one self-attention stack. Perceiver Resampler / Q-Former (Flamingo, BLIP-2) inserts a small cross-attention module in which a fixed set of K learned queries (64 latents in Flamingo, 32 in BLIP-2) attends over the N patch tokens and emits exactly K vectors, decoupling the LLM’s sequence length from image resolution. Gated cross-attention fusion (Flamingo’s xattn-dense blocks, Llama 3’s vision adapter, NVLM-X) never puts image tokens in the LLM sequence at all: new cross-attention layers are interleaved into the language stack (every fourth layer in Llama 3) and read the patch tokens as an external key/value memory, with a tanh gate initialized to zero so the model starts out behaviorally identical to the text-only LLM. The trade-off is compute versus fidelity versus intrusiveness: MLP concatenation is the simplest and preserves the most visual detail but makes prefill grow as O((N+T)^2); resampling is cheapest and constant-cost but imposes a hard information bottleneck; cross-attention keeps text throughput almost untouched and protects a frozen LLM, at the price of new parameters and a more complex training recipe.

    (1) Where The Visual Tokens Live: MLP projection puts them inside the LLM sequence, a resampler puts a compressed K of them inside, and cross-attention keeps them outside as an external memory.
    (2) Prefill Cost Scaling: concatenation pays O((N+T)^2 d) per layer, a resampler pays O((K+T)^2 d) with K \ll N, and cross-attention pays O(T N d) only in the inserted layers.
    (3) Added Parameters: LLaVA-1.5’s connector is roughly 20M parameters of MLP, BLIP-2’s Q-Former is about 188M with a BERT-base initialization, and Llama 3’s cross-attention adapter adds tens of billions at the 405B scale.
    (4) Information Bottleneck: a fixed K caps how much of a high-resolution image can survive, which is why resampler-based models underperform on dense OCR, charts, and counting.
    (5) Frozen Versus Tuned Backbones: the Idefics2 ablation found cross-attention wins when the LLM is frozen, while the fully autoregressive concatenation design wins once the LLM is unfrozen.
    (6) Token Count Beats Connector Design: MM1’s ablations show the connector type matters far less than image resolution and visual token count, which is why the field converged on MLP plus cheap token compression.

    Three-panel architecture diagram. Panel A: ViT encoder produces 576 patch tokens, a two-layer MLP maps them into the LLM token space, and they are concatenated with T text tokens so self-attention runs over N plus T positions. Panel B: 32 learned queries cross-attend over the 576 patch tokens to produce 32 resampled tokens, which are concatenated with the text so self-attention runs over K plus T positions. Panel C: patch tokens stay outside the LLM as an external key value memory read by gated cross-attention layers inserted every fourth self-attention block, so self-attention still runs over T text tokens only.

    Figure 1: The three fusion families differ only in where the patch tokens are allowed to go. Concatenation admits all N of them into the shared sequence, a resampler admits a fixed K, and gated cross-attention admits none, reading them instead as an external memory through layers whose tanh gate starts at zero.

    The training recipe follows from the architecture. An MLP connector is so small that a two-stage schedule with about 558K caption pairs for alignment and 665K instruction samples for tuning is enough to reach state-of-the-art benchmark scores, which is what made LLaVA-1.5 reproducible on a single node. A Q-Former is a real transformer that must learn what to query, so BLIP-2 needs a dedicated representation-learning stage (contrastive, matching, and captioning objectives) before the generative stage, and remains the least data-efficient of the three per unit of final accuracy. Cross-attention sits in between: the gate makes optimization stable and lets you keep the language model frozen, so text benchmarks cannot regress, but you are training new layers that must learn to be useful without ever seeing visual tokens in the residual stream. NVLM’s controlled comparison is the cleanest evidence on the trade-off, finding the decoder-only concatenation variant stronger on OCR and multimodal reasoning while the cross-attention variant gave better throughput on high-resolution inputs.

    Mathematical Formulation:
    Z = E_v(I) \in \mathbb{R}^{N \times d_v}
    H_v = W_2\,\sigma(W_1 Z) \in \mathbb{R}^{N \times d}
    L = N + T
    R = \mathrm{Attn}(Q, Z, Z) \in \mathbb{R}^{K \times d}
    H_t \leftarrow H_t + \tanh(\alpha)\,\mathrm{Attn}(H_t, Z, Z)
    C_{\mathrm{self}} = O(L^2 d)
    C_{\mathrm{cross}} = O(T N d)

    Where:

    • I is the input image, E_v the vision encoder (typically a frozen CLIP or SigLIP ViT), and Z its patch token grid.
    • H_v is the projected visual sequence, with W_1, W_2 the connector weights and \sigma a GELU nonlinearity; a single W recovers the original linear projection of LLaVA-1.
    • N is the patch count (576 for ViT-L/14 at 336 px), T the text length, L the LLM sequence length, d the LLM width, and d_v the vision width.
    • Q \in \mathbb{R}^{K \times d} are the learned latent queries and R the resampled output; K is fixed at design time, so R has the same size for a thumbnail and for a 4K page.
    • H_t is the text hidden state inside the LLM and \alpha the scalar gate parameter, initialized so that \tanh(\alpha) = 0 and the pretrained language behavior is exactly preserved at step zero.
    • C_{\mathrm{self}} is the per-layer self-attention cost that concatenation inflates, and C_{\mathrm{cross}} the per-layer cost of a cross-attention block, which is linear in N rather than quadratic.
    Log-scale line chart of relative prefill attention cost versus number of visual tokens from zero to 2880, for a 512-token text prompt and a 32-layer language model. The MLP concatenation curve rises quadratically to about forty-four times the text-only baseline, the gated cross-attention curve rises linearly to about two point four times, and the 64-latent resampler curve stays nearly flat at about one point four times.

    Figure 2: Prefill attention cost relative to a text-only forward pass at T = 512. At 2880 visual tokens, concatenation costs about 44\times the text-only baseline because the whole stack pays O((N+T)^2), gated cross-attention costs about 2.4\times since only 8 of 32 layers see the image and they scale as O(TN), and a 64-latent resampler stays near 1.4\times because the LLM never sees more than K + T positions.

    PropertyLinear / MLP projectionPerceiver Resampler / Q-FormerGated cross-attention
    Tokens entering the LLM sequenceAll N patch tokens (576 at 336 px, thousands with tiling)Exactly K latents (32 in BLIP-2, 64 in Flamingo)None; patch tokens are external keys and values
    Prefill scaling in image sizeQuadratic, O((N+T)^2 d) in every layerConstant for the LLM, linear O(KN d) in the resamplerLinear, O(TN d) in the inserted layers only
    Added parametersSmallest; a 2-layer GELU MLP, roughly 20MMedium; about 188M for the BERT-base Q-FormerLargest; new attention plus FFN blocks scaled to the LLM width
    Training complexityLowest; align then instruction-tune on about 1.2M samplesHighest; needs a separate representation-learning stage to convergeModerate; the zero-init tanh gate makes the warm start stable
    Effect on a frozen LLMWeakest option when frozen; usually the LLM must be tunedWorks frozen, which is exactly why BLIP-2 used itBest; text-only behavior is provably unchanged at initialization
    Main weaknessContext blowup with high resolution, video, or many imagesFixed-K information bottleneck hurts OCR, charts, and countingExtra parameters and weaker reported OCR and reasoning transfer
    Representative modelsLLaVA, LLaVA-1.5, Qwen2-VL, InternVL, NVLM-DFlamingo resampler, BLIP-2, InstructBLIP, IdeficsFlamingo xattn-dense, Llama 3.2 Vision, NVLM-X

    Login to view more content
  • DL0130 MoE Router Collapse and Auxiliary Loss

    What is Router Collapse in Sparse Mixture-of-Experts (MoE) LLMs? Derive the auxiliary load-balancing loss and Router Z-loss used to stabilize MoE training.

    Answer

    Router collapse is the failure mode in which the learned router of a sparse MoE layer stops using most of its experts and concentrates almost all token assignments on a small subset, so a layer you paid N experts’ worth of memory for behaves like a much smaller model. The cause is a positive feedback loop: an expert that happens to win slightly more tokens early in training receives more gradient signal, becomes genuinely better on those tokens, so the router raises its logit further, while starved experts never see enough tokens to become useful. Under token-choice top-k the visible symptoms are a skewed load histogram, a large dropped-token fraction once the hot experts hit their capacity buffer, and a validation loss that tracks a dense model of the active-parameter size rather than the total-parameter size. Production training stacks suppress it with two extra loss terms rather than with a new routing algorithm: an auxiliary load-balancing loss that pushes router probability mass away from overloaded experts, and a router z-loss that penalizes the squared log-sum-exp of the router logits so the logits cannot grow without bound and destabilize the softmax in low precision. Both are differentiable surrogates for quantities that are not: the actual assignment counts are piecewise constant, and the actual numerical blow-up is a hardware property, so each loss attacks a proxy the gradient can reach.

    (1) Collapse Is Self-Reinforcing: a small early advantage in router logits compounds through the gradient the winning expert receives, which is why collapse typically happens in the first few thousand steps and is nearly irreversible afterwards.
    (2) The Assignment Is Non-Differentiable: the load fraction f_i comes from a top-k selection and has zero gradient almost everywhere, so the balance loss must pair it with the differentiable mean gate probability P_i.
    (3) Balance Loss Is A Normalized Dot Product: N \sum_i f_i P_i equals 1 under a uniform assignment and grows toward N under total collapse, giving a scale-free objective independent of expert count.
    (4) Its Gradient Is Load-Proportional: the derivative with respect to P_i is exactly \alpha N f_i, so probability mass is pushed down in proportion to how overloaded each expert already is.
    (5) Z-Loss Bounds The Logits: penalizing the squared log-partition function caps \max_j h_j, which matters because the relative round-off in e^{h} grows linearly with |h| in bfloat16 and can flip top-k selection.
    (6) Coefficients Are A Quality Trade: \alpha \approx 10^{-2} and \beta \approx 10^{-3} are the common settings, because a large \alpha buys perfect balance by actively fighting the language-modeling objective.

    Cycle diagram of router collapse: router logits for one expert edge above the others, top-k dispatches more tokens to it, that expert receives most of the expert gradient, the starved experts stay undertrained and score lower, which feeds back into the logit gap; two green intervention boxes on the right show the router z-loss acting on the logits node and the auxiliary balance loss acting on the dispatch node

    Figure 1: Router collapse is a closed loop, not a single bad step: the logit gap, the dispatch skew, and the gradient imbalance each amplify the next. The two stabilizers cut the loop at different points, with the z-loss constraining the logit magnitudes and the balance loss constraining the dispatch distribution.

    The derivation of the balance loss starts from what you actually want to penalize, namely the variance of the per-expert token counts, and then asks which part of that quantity carries a gradient. The counts themselves come from a top-k over the router logits, so they are a step function of the parameters and give nothing to backpropagation. The fix used by GShard and simplified by Switch Transformers is to pair the non-differentiable load vector f with the differentiable importance vector P, the mean softmax probability per expert, and minimize their inner product. Treating f as a constant, the objective is linear in P with coefficient N f_i, so each step lowers the router’s probability for exactly the experts that were overloaded in the current batch, and because f is recomputed every step this becomes a self-correcting controller whose fixed point is the uniform assignment. The factor N is a normalization choice: it makes the minimum value 1 regardless of how many experts you have, so the same \alpha transfers from an 8-expert layer to a 256-expert layer. The z-loss has an entirely different motivation: it is a numerical guard, derived from the observation that a logit stored with relative precision \epsilon produces an absolute error of about |h|\epsilon, and exponentiation converts that absolute error into a relative error of the same size, so large logits make the softmax and therefore the selected expert set unreliable.

    Mathematical Formulation:
    h_t = W_r x_t, \quad p_t = \mathrm{softmax}(h_t)
    f_i = \frac{1}{T}\sum_{t=1}^{T} \mathbb{1}[i \in \mathcal{T}_t]
    P_i = \frac{1}{T}\sum_{t=1}^{T} p_{t,i}
    \mathcal{L}_{\mathrm{bal}} = N \sum_{i=1}^{N} f_i P_i
    \frac{\partial \mathcal{L}_{\mathrm{bal}}}{\partial P_i} = N f_i
    1 \leq \mathcal{L}_{\mathrm{bal}} \leq N
    \mathcal{L}_{z} = \frac{1}{T}\sum_{t=1}^{T}\left(\log \sum_{j=1}^{N} e^{h_{t,j}}\right)^{2}
    \max_j h_{t,j} \leq \mathrm{lse}(h_t) \leq \max_j h_{t,j} + \log N
    \mathcal{L} = \mathcal{L}_{\mathrm{LM}} + \alpha \mathcal{L}_{\mathrm{bal}} + \beta \mathcal{L}_{z}

    Where:

    • x_t \in \mathbb{R}^{d} is the hidden state of token t, W_r \in \mathbb{R}^{N \times d} the router matrix, and h_t the router logits before any selection.
    • T is the number of tokens the statistics are aggregated over (micro-batch, device batch, or global batch), N the expert count, and \mathcal{T}_t the set of k experts selected for token t.
    • f_i is the load, the fraction of tokens dispatched to expert i; it is piecewise constant in the parameters and therefore contributes no gradient.
    • P_i is the importance, the mean router probability assigned to expert i; it is smooth, so all of the balance-loss gradient flows through it and into W_r.
    • \mathcal{L}_{\mathrm{bal}} = 1 exactly when f_i = P_i = 1/N for all i, and approaches N when one expert takes every token, so the value is directly readable as an imbalance factor.
    • \mathrm{lse}(h_t) = \log\sum_j e^{h_{t,j}} is the log-partition function; squaring it penalizes large logits in either direction and, by the sandwich bound, keeps \max_j h_{t,j} within \log N of a small target.
    • \alpha and \beta are the balance and z-loss coefficients, commonly \alpha = 10^{-2} and \beta = 10^{-3}; \mathcal{L}_{\mathrm{LM}} is the ordinary next-token cross-entropy.
    Two-panel chart: left panel plots the maximum expert load fraction against training step for an eight-expert layer with balance-loss coefficients zero, one thousandth, and one hundredth, showing full collapse toward one expert without the loss and a curve close to the uniform 0.125 line with alpha one hundredth; right panel plots relative round-off error in exp of a router logit against logit magnitude on a log scale for bfloat16 and float32 storage, with bfloat16 crossing one percent error at a logit magnitude near 2.6

    Figure 2: The two losses guard different quantities. Left: without a balance term the maximum load fraction runs from the uniform 1/N = 0.125 to near 1.0 within a few thousand steps, while \alpha = 10^{-2} holds it close to uniform. Right: the relative round-off in e^{h} grows linearly with |h|, which is why bfloat16 routers become unreliable at moderate logits and the z-loss keeps \mathrm{lse}(h) near O(1).

    PropertyAuxiliary balance lossRouter z-lossBias-based loss-free balancing
    Quantity penalizedThe dot product of load and importance, scaled by the expert countThe squared log-sum-exp of the router logits, averaged over tokensNothing; a per-expert bias is added to the selection logits only
    Failure it preventsExpert collapse, wasted parameters, and dropped tokens at the capacity bufferLogit blow-up, low-precision softmax round-off, and loss spikesExpert collapse, without perturbing the gate weights used in the output
    Typical settingCoefficient 1e-2, aggregated per device batch or per global batchCoefficient 1e-3, with the router itself computed in float32Bias update rate around 1e-3, driven by observed per-expert load error
    Effect on the LM objectiveAdds an interference gradient that trades some quality for balanceMild regularizer, usually neutral or slightly positive for qualityNo interference term at all, since the bias is not part of the output gate
    Where it is usedGShard, Switch, Mixtral, OLMoEST-MoE and most later open MoE training stacks, including OLMoEDeepSeek-V3 and its loss-free-balancing follow-ups

    Login to view more content
  • DL0127 DDPM Training Objective

    What is the DDPM training objective, and why do we train the network to predict the noise rather than the clean image directly?

    Answer

    A denoising diffusion probabilistic model is trained by maximizing a variational bound on the data log-likelihood, and because both the forward corruption and the reverse posterior are Gaussian, that bound collapses into a sum of KL divergences whose only learnable content is a mean-matching term at each noise level. Reparameterizing that mean through the closed-form forward marginal turns every term into a weighted squared error between the injected noise \epsilon and a network prediction, and Ho et al. then drop the per-step weight \lambda_t to get the simple objective L_{\mathrm{simple}} = \mathbb{E}\|\epsilon - \epsilon_\theta(x_t,t)\|^2. Predicting the noise is not a different model from predicting the clean image: \epsilon-prediction, x_0-prediction, and direct mean prediction are affine reparameterizations of one another given x_t and t, so they define the same optimum but different implicit loss weightings and different numerical conditioning. The practical reasons to regress noise are that the target is always unit-variance no matter how noisy the input is, that the residual structure lets the network pass the low-frequency content through instead of re-synthesizing it, and that discarding \lambda_t under this parameterization down-weights the easy low-noise steps by roughly an order of magnitude, which empirically buys much better FID. It also makes the network a rescaled score estimator, which is what connects DDPM to score matching and to every ODE and SDE sampler built on top of it.

    (1) The Objective Is One Scalar MSE: sample x_0, sample t \sim \mathcal{U}\{1,\ldots,T\} and \epsilon \sim \mathcal{N}(0,I), form x_t in closed form, and regress \epsilon. No per-step model, no adversarial term, no sequential rollout during training.
    (2) It Comes From The ELBO, Not From Heuristics: each L_{t-1} is a Gaussian KL whose learnable part is a mean difference, and substituting the forward marginal converts it into \lambda_t\|\epsilon-\epsilon_\theta\|^2.
    (3) Dropping \lambda_t Is A Reweighting Choice: \lambda_t is largest at small t, so setting all weights to 1 shifts capacity toward the high-noise steps that decide global structure.
    (4) The Target Is Scale-Stable: \epsilon is standard normal for every t, so a single output head with fixed normalization works across the whole schedule, whereas the difficulty of predicting x_0 varies enormously with t.
    (5) Residual Prediction Is Easier For A U-Net: at low noise, most of x_t is already the answer, so predicting the small perturbation avoids forcing the network to reconstruct an image it was handed.
    (6) Conditioning Flips At The Two Ends: converting \epsilon_\theta to \hat x_0 amplifies error by 1/\sqrt{\mathrm{SNR}_t}, which is tiny at low noise and roughly 150x at t=T, the reason v-prediction exists.

    Left-to-right flow diagram of one DDPM training step: sample a clean image, sample a timestep and Gaussian noise, form the corrupted sample in closed form, pass it through the U-Net to predict the noise, compute the squared error against the sampled noise, and a dashed feedback arrow carrying the gradient step back to the U-Net

    Figure 1: One training step. The closed-form forward marginal means an arbitrary timestep can be sampled directly, so training never simulates the chain; the loss is a single unit-variance regression shared by all T noise levels.

    The reweighting argument is the one that most candidates miss. Under the ELBO, the term at t=1 carries roughly 50x the weight of the term at t=T for the standard linear schedule, and those low-t terms correspond to removing almost imperceptible noise, a task that contributes little to perceptual quality but a lot to likelihood. Setting every weight to 1 in \epsilon-space is therefore a deliberate trade of likelihood for sample quality, and it is why DDPM reports strong FID with mediocre bits-per-dimension. The conditioning argument cuts the other way at high noise: since \hat x_0 is recovered by dividing by \sqrt{\bar\alpha_t}, a small noise error becomes a large image error near t=T, so pure \epsilon-prediction is a poor target for few-step or distilled samplers that must produce a usable \hat x_0 from the very first step. Finally, the identity \epsilon_\theta = -\sqrt{1-\bar\alpha_t}\,s_\theta shows the trained network is a scaled score function, which is exactly what a probability-flow ODE or an annealed Langevin sampler needs, so the same checkpoint serves DDPM, DDIM, and higher-order solvers.

    Mathematical Formulation:
    x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1-\bar\alpha_t}\,\epsilon
    L_{t-1} = \lambda_t\,\|\epsilon - \epsilon_\theta(x_t,t)\|^2
    \lambda_t = \frac{\beta_t^2}{2\sigma_t^2\alpha_t(1-\bar\alpha_t)}
    L_{\mathrm{simple}} = \mathbb{E}\left[\|\epsilon - \epsilon_\theta(x_t,t)\|^2\right]
    \hat x_0 = \frac{x_t - \sqrt{1-\bar\alpha_t}\,\epsilon_\theta}{\sqrt{\bar\alpha_t}}
    \epsilon_\theta(x_t,t) = -\sqrt{1-\bar\alpha_t}\,s_\theta(x_t,t)

    Where:

    • x_0 is a clean training sample, x_t its corrupted version at step t, and \epsilon \sim \mathcal{N}(0,I) the injected noise that serves as the regression target.
    • t \in \{1,\ldots,T\} is sampled uniformly, \beta_t is the schedule, \alpha_t = 1-\beta_t, and \bar\alpha_t = \prod_{s=1}^{t}\alpha_s gives the closed-form marginal that makes single-step training possible.
    • \epsilon_\theta is the U-Net with parameters \theta, conditioned on t through a sinusoidal timestep embedding.
    • \lambda_t is the ELBO weight and \sigma_t^2 the reverse-process variance (commonly \sigma_t^2 = \beta_t); L_{\mathrm{simple}} is L_{t-1} with \lambda_t replaced by 1.
    • \mathrm{SNR}_t = \bar\alpha_t/(1-\bar\alpha_t) is the signal-to-noise ratio, which decreases monotonically in t and sets the error amplification 1/\sqrt{\mathrm{SNR}_t} from noise space into image space.
    • s_\theta(x_t,t) \approx \nabla_{x_t}\log q(x_t) is the score, so a trained \epsilon_\theta doubles as a score model for ODE and SDE samplers.
    Two-panel chart: left panel plots the ELBO per-step weight relative to its value at the final timestep on a log axis against the diffusion step, showing roughly fifty times more weight at small t than at large t, next to the flat unit weight used by the simple objective; right panel plots on a log axis the factor by which a prediction error is amplified into clean-image space for noise prediction, clean-image prediction, and v-prediction, with noise prediction rising from about 0.01 to about 150

    Figure 2: The two quantitative reasons the parameterization matters. Left: dropping \lambda_t removes a roughly 50x preference for the easiest low-noise steps. Right: the same trained error costs almost nothing in image space at low noise but is amplified by 1/\sqrt{\mathrm{SNR}_t} near t=T, which is where v-prediction stays bounded.

    PropertyNoise prediction (DDPM)Clean-image predictionv-prediction
    TargetThe sampled noise, standard normal at every tThe data sample itself, with data-dependent statisticsA schedule-dependent mix of noise and data
    Implicit weighting under uniform MSEDown-weights low-noise steps, favors perceptual quality over likelihoodUp-weights low-noise steps by the SNR, closer to the raw ELBORoughly balanced across the schedule, close to SNR-plus-one weighting
    Numerical weak spotError into image space amplified by about 150x near t = TError into noise space amplified near t = 0, where the residual is tinyAmplification stays at or below 1 in both directions
    Typical usePixel and latent diffusion with many sampling steps, Stable Diffusion 1.x and 2.0-baseVery short schedules, consistency-style objectives, some latent decodersProgressive distillation, high-resolution and upscaling models, SD 2.1-v

    Login to view more content
  • DL0125 Action-Conditioned World Model

    Explain the architecture of an Action-Conditioned World Model for autonomous driving, of the kind Wayve describes for its GAIA models.

    Answer

    An action-conditioned world model is a generative model of future sensor data whose predictions are explicitly steered by the ego vehicle’s own control commands, so it can answer “what would the road look like in three seconds if I brake at 5 m/s² instead of holding speed?”. Architecturally it factors into four blocks that are trained in two stages: a video encoder that compresses multi-camera frames into a small latent grid, an action and context encoder that embeds the low-dimensional control signal, a latent dynamics model that rolls the compressed state forward in time, and a decoder plus auxiliary heads that render latents back to pixels or BEV occupancy. Almost every production-scale system follows this shape and differs only in the dynamics block: GAIA-1 uses a 6.5B-parameter causal transformer over discrete video tokens with a separate diffusion video decoder, GAIA-2 replaces it with a flow-matching diffusion transformer over continuous latents so five surround views stay geometrically consistent, and Dreamer-style recurrent state-space models keep a compact deterministic-plus-stochastic state for reinforcement learning in imagination. The reason compression comes first is the same as in latent diffusion: at 6.25 Hz a 26-frame context of 288 \times 512 frames is 3.8M pixels per frame-stack but only 14,976 latent tokens, and the dynamics model’s cost is quadratic in that number. What makes the model a simulator rather than a video generator is the rollout loop plus faithful action adherence, which is precisely where these models are weakest and where interviewers push.

    (1) Observation Encoder: a VQ or KL-regularized video autoencoder with spatial factor f = 16 to 32 maps each frame to an 18 \times 32 grid; it is trained first with perceptual and adversarial losses, then frozen.
    (2) Action And Context Conditioning: a 2-D action (curvature and acceleration) is embedded by a small MLP and injected per frame at every block, as prefix tokens or through AdaLN modulation, alongside text prompts, agent bounding boxes, and camera calibration.
    (3) Latent Dynamics Backbone: a causal spacetime transformer predicting the next token, or a diffusion/flow-matching transformer denoising a whole latent chunk conditioned on the past; this is where nearly all parameters live.
    (4) Decoder And Auxiliary Heads: a diffusion or convolutional decoder for video, plus cheap heads for BEV occupancy, collision cost, and value that planning actually consumes.
    (5) Two-Stage Training With Teacher Forcing: the dynamics model always sees ground-truth latents during training, which creates the exposure-bias gap that dominates long-horizon rollouts.
    (6) Three Consumers: closed-loop policy evaluation against counterfactual actions, planning or RL in imagination, and generation of rare corner-case scenarios that fleets rarely log.

    Left to right pipeline: five surround cameras feed a video tokenizer producing an 18 by 32 latent per frame, a dynamics transformer conditioned on action, text, agent boxes and camera pose predicts future latents with an autoregressive rollout loop, and a decoder renders predicted video and BEV occupancy

    Figure 1: The four blocks of an action-conditioned world model. Only the frozen codec ever touches pixels, the action modulates every block of the dynamics model, and the rollout loop is what turns a video generator into a driving simulator.

    How the action enters is the part worth rehearsing. A two-number control signal has to influence a token stream dominated by appearance, so injecting it once at the input is not enough: it is embedded per frame and re-applied at every layer, and it must be time-aligned to the interval it causes rather than the frame it was logged with. Two failure modes follow directly. First, shortcut learning: logged actions are almost perfectly predictable from the visible road geometry, so a model can minimize training loss while ignoring a_t entirely, and you only detect this by rolling out counterfactual actions that contradict the scene. Second, weak controllability at sampling time, which is why conditioning is dropped for 10-20% of training samples so that classifier-free guidance can later amplify action adherence. Evaluation therefore needs an action-following metric (does the rendered ego trajectory match the commanded one?) next to FID or FVD, because generative realism and control fidelity move independently.

    Mathematical Formulation:
    z_t = \mathcal{E}(o_t)
    h_t = f_\theta(h_{t-1}, z_{t-1}, a_{t-1})
    \hat{z}_t \sim p_\theta(z_t \mid h_t, a_{t-1})
    \hat{o}_t = \mathcal{D}(h_t, \hat{z}_t)
    \mathcal{L} = \mathcal{L}_{\text{obs}} + \beta \, \mathcal{L}_{\text{dyn}}
    \mathcal{L}_{\text{obs}} = -\log p_\theta(o_t \mid h_t, z_t)
    \mathcal{L}_{\text{dyn}} = D_{KL}\!\left(q_\phi(z_t \mid h_t, o_t) \,\|\, p_\theta(z_t \mid h_t)\right)
    p(z_{1:T} \mid a_{1:T}) = \prod_{t=1}^{T} p(z_t \mid z_{1:t-1}, a_{1:t-1})
    N = T \cdot (H/f) \cdot (W/f)
    N = 26 \cdot 18 \cdot 32 = 14{,}976

    Where:

    • o_t is the multi-camera observation at step t and \hat{o}_t its prediction; a_t is the ego action (curvature and acceleration, or steering and pedal).
    • z_t is the compressed observation latent and h_t the deterministic recurrent or attention-carried state that summarizes the past.
    • \mathcal{E} and \mathcal{D} are the frozen encoder and decoder, f_\theta the dynamics backbone, p_\theta the action-conditioned prior, and q_\phi the posterior that also sees the true observation.
    • \beta balances reconstruction against the dynamics (KL) term; in practice it is annealed and often free-bits clipped so the prior does not collapse onto the posterior.
    • T is the context length in frames, H \times W the frame resolution, f the spatial compression factor, and N the token count whose square drives attention cost.
    • Required initial conditions for a rollout: h_0 = 0 and a real context z_{1:k} encoded from logged frames, after which only actions are supplied and j > k steps are pure imagination.
    Left panel: bird's-eye view of three ego trajectories branching from the same start under hold-speed, brake and lane-change action sequences, with a braking lead vehicle and a predicted collision marker. Right panel: prediction error versus rollout horizon for one-step re-encoding, free-running rollout, and a variant with action tokens ablated

    Figure 2: Left: fixing the initial latent and varying only the action sequence produces counterfactual rollouts, the property that makes the model usable for closed-loop evaluation. Right: because training is teacher-forced, free-running error compounds superlinearly, and ablating the action tokens makes it worse still.

    Dynamics blockDiscrete-token autoregressiveLatent diffusion / flow matchingRecurrent state-space (RSSM)
    Latent representationVQ codebook indices, 576 tokens per frameContinuous latent grid, tens of channels, temporally compressedSmall vector state: deterministic GRU plus categorical stochastic units
    How the action entersPer-frame prefix tokens in the causal sequenceAdaLN modulation or cross-attention at every block, with conditioning dropoutConcatenated into the recurrent transition at each step
    Rollout costHundreds of sequential token decodes per frame; slowestTens of denoising steps per chunk, parallel across positionsOne cheap matrix step per frame; fast enough for RL in imagination
    StrengthExact likelihoods, easy long-context scaling, LLM tooling reuseBest photorealism and multi-view consistency; controllable via guidanceCompact enough to train a policy on millions of imagined steps
    Typical failure modeQuantization artifacts and drift after a few seconds of rolloutPlausible but unfaithful scenes; ignores the action unless guidedBlurry reconstructions; posterior collapse hides rare agents
    Representative systemGAIA-1GAIA-2, Vista, NVIDIA CosmosDreamerV3, MILE

    Login to view more content
  • DL0124 3D Occupancy Flow

    What is 3D Occupancy Flow? What is the tradeoff for predicting dense spatiotemporal occupancy grids instead of discrete 3D bounding-box trajectories?

    Answer

    3D Occupancy Flow is a joint perception-and-forecasting output format that replaces the list of tracked objects with a dense grid: for every voxel (or BEV cell) and every future waypoint the network predicts an occupancy probability plus a flow vector describing how the mass in that cell moves. Waymo’s Occupancy Flow Fields formulation predicts three quantities per waypoint on a 256 \times 256 BEV grid covering roughly 80\ \text{m} \times 80\ \text{m}: observed occupancy, occluded occupancy, and backward flow. Camera-only 3D variants such as Occ3D on nuScenes predict a 200 \times 200 \times 16 voxel grid at 0.4\ \text{m} resolution, and Tesla presented an occupancy network with an occupancy-flow head at its 2022 AI Day. The appeal is that free space and obstacle geometry become class-agnostic and non-parametric: a tipped-over mattress, an articulated trailer, a swinging crane boom, and an overhanging branch all get represented without appearing in a detector taxonomy, and probability mass can sit on both branches of a fork at once without a mode head or non-maximum suppression. The cost is that a grid has no notion of an object, so instance identity, track continuity, and per-agent attributes disappear, and the output tensor grows by roughly two orders of magnitude, which pushes cost onto compute, memory, label pipelines, and the loss function’s handling of a grid where the overwhelming majority of voxels are empty. In practice this is why most production stacks run occupancy flow alongside a box pipeline rather than as a drop-in replacement.

    (1) Dense Spatiotemporal Output: the head emits o_t(v) and f_t(v) for every cell and every future timestep, not a parametric box with a heading and a velocity per agent.
    (2) Class-Agnostic Geometry: anything that occupies space is representable, which removes the long-tail detection failure where an unlisted object class becomes invisible to the planner.
    (3) Backward Flow, Not Forward: predicting motion from t back to t-1 makes warping a gather with one source per cell, so mass never collides during the warp and a flow-grounded occupancy consistency check becomes well defined.
    (4) Non-Parametric Multimodality: a grid holds several futures simultaneously as spread probability mass, but that same property makes averaged modes look like blur or ghost occupancy rather than a ranked set of hypotheses.
    (5) Identity Is What You Give Up: without instance IDs, right-of-way logic, interaction-aware conditioning, and per-object intent signals such as turn indicators lose their handle on the scene.
    (6) Cost Scales With The Grid: output size grows as H W D T, dense labels require multi-sweep LiDAR accumulation and voxelization, and roughly 95% or more of voxels are empty, so class imbalance dominates the occupancy loss.

    Diagram contrasting two forecasting pipelines from the same sensor input: an object-centric branch running detector, tracker, and trajectory predictor to emit sparse box waypoints, and an occupancy-centric branch running a BEV or voxel encoder with occupancy and backward-flow heads to emit a dense spatiotemporal grid

    Figure 1: The object-centric branch (detect → track → forecast) produces a sparse, identity-carrying output limited by its taxonomy; the occupancy branch skips detection and data association entirely and produces a dense, class-agnostic grid with no instance IDs.

    The deeper tradeoff is not really compute, it is what the downstream planner can express. A box trajectory is a commitment: this vehicle, with this ID, will be here in 3 seconds with this probability, which lets a planner reason about yielding to a specific agent, replay a scenario in simulation, and produce an auditable explanation for a maneuver. Occupancy flow is a statement about space, which is exactly what collision checking and drivable-free-space queries want, but a grid that hedges between “the cyclist goes straight” and “the cyclist turns” paints both corridors at moderate probability, and a naive cost function that treats any occupancy above a threshold as blocked yields the freezing-robot behavior. Occupancy flow partially recovers correspondence without identity: warping the previous occupancy along the predicted backward flow and multiplying it against the current occupancy gives a differentiable consistency term, so the model is penalized for teleporting mass even though it never names an object.

    Mathematical Formulation:
    0 \leq o_t(v) \leq 1
    \hat{o}_t = o_t \odot \mathcal{W}(o_{t-1}, f_t)
    \mathcal{L} = \mathcal{L}_{occ} + \lambda \mathcal{L}_{flow}
    \mathcal{L}_{occ} = \sum_{t=1}^{T} \sum_{v} \mathrm{BCE}(o_t(v), y_t(v))
    \mathcal{L}_{flow} = \sum_{t=1}^{T} \sum_{v \in \Omega_t} \lVert f_t(v) - f_t^{*}(v) \rVert_1
    N_{occ} = 200 \cdot 200 \cdot 16 \cdot 8 = 5.12 \times 10^{6}
    N_{box} = 50 \cdot 6 \cdot 16 \cdot 5 = 2.40 \times 10^{4}

    Where:

    • o_t(v) is the predicted occupancy probability of cell v at future waypoint t, and f_t(v) is the backward flow vector pointing to where that mass sat at t-1.
    • \mathcal{W} is the warp operator that gathers o_{t-1} along f_t (bilinear or trilinear), and \hat{o}_t is the flow-grounded occupancy used both as a loss term and as an evaluation metric.
    • y_t(v) is the voxelized ground-truth occupancy label, f_t^{*} the ground-truth flow, and \Omega_t the set of genuinely occupied cells to which the flow loss is masked.
    • t \in \{1, \ldots, T\} indexes waypoints and v indexes the H \times W \times D grid; \lambda balances the two terms, and \mathrm{BCE} is usually replaced by a focal or class-balanced variant because empty voxels dominate.
    • N_{occ} counts predicted occupancy values for a 200 \times 200 \times 16 grid over 8 waypoints, and N_{box} counts a comparable box head with 50 agents, 6 modes, 16 waypoints, and 5 numbers per waypoint.
    • Required initial condition: o_0 and the whole grid must be expressed in the ego frame at t = 0, so ego motion is compensated before flow is interpreted as agent motion.
    Grid diagram showing occupancy at time t minus one in light dashed cells and predicted occupancy at time t in solid cells, with backward flow arrows drawn from each occupied cell at time t to the cell its mass came from, alongside the flow-grounded occupancy product equation

    Figure 2: Backward flow assigns each occupied cell at time t a single source cell at t-1, so the warp is a gather rather than a scatter and the consistency product o_t \odot \mathcal{W}(o_{t-1}, f_t) penalizes mass that appears without a plausible origin.

    Resolution and horizon are the two knobs that make or break the design. Halving the voxel size multiplies the tensor by 8 in 3D, and every additional waypoint is another full grid, so a 0.2\ \text{m} grid over a 5-second horizon at 2\ \text{Hz} is far beyond a real vehicle compute budget once flow channels and semantics are added. Going the other way is not free either: at 0.4\ \text{m} voxels a pedestrian walking at 1.4\ \text{m/s} moves less than one voxel per 0.2\ \text{s} frame, so the flow target is sub-voxel and the occupancy channel alone cannot express the motion, which is precisely why the flow head is kept as a continuous regression rather than a discrete cell-to-cell assignment.

    Log-scale bar chart of predicted output size per inference: 24000 numbers for multi-modal box trajectories, 524288 for a 256 by 256 BEV occupancy grid over 8 waypoints, 5.12 million for a 200 by 200 by 16 voxel grid over 8 waypoints, and 20.48 million once three flow channels are added

    Figure 3: Output size is where the tradeoff becomes concrete: the 3D voxel grid predicts about 213x more numbers than a multi-modal box head, and adding a 3-channel flow field multiplies that by another 4 before any semantic classes are included.

    Property3D occupancy flowBox trajectoriesHybrid stack
    Output per inferenceMillions of per-cell values, occupancy plus flow per waypointTens of thousands of numbers, a few modes per tracked agentBoth, sharing one BEV or voxel backbone
    Unlisted geometryRepresented, since occupancy is class-agnosticDropped if no detector class fits itCovered by the occupancy branch
    Instance identityNone; only flow-based correspondenceExplicit IDs, attributes, and track historyIDs from the box branch, geometry from the grid
    MultimodalityImplicit in the probability field, no mode count to tuneExplicit ranked modes with confidencesExplicit modes for interaction, field for collision checks
    SupervisionAccumulated multi-sweep LiDAR voxelized into dense labelsHuman box and track annotations onlyBoth label pipelines must be maintained
    Typical failure modeBlurred or ghost occupancy that makes the planner over-conservativeMissed detection or ID switch removes an obstacle entirelyDisagreement between branches needs an arbitration policy

    Login to view more content
  • DL0123 3D-Aware Visual Pretraining

    How does 3D-aware visual pre-training bridge spatial intelligence between 2D generative video models and 3D physical world simulators?

    Answer

    3D-aware visual pre-training trains a visual backbone or a video generator with objectives that can only be solved by recovering scene geometry: per-pixel pointmaps, metric depth, camera rays, and cross-view correspondence, instead of appearance statistics alone. The two ends it connects have complementary deficits. A 2D video diffusion model absorbs enormous visual and dynamic diversity from internet footage, but its geometry is implicit, so it has no queryable state: you cannot ask it for the distance to a mug in centimeters, and a long rollout quietly changes room layout because nothing anchors the scene. A physics simulator has the opposite profile, exact meshes, contacts, and metric scale, but its asset library and rendering diversity are tiny compared with the real world. Pre-training with geometric targets creates the shared interface: the same representation gives the generator a camera and depth channel it can be conditioned on, and gives the simulator a path to ingest real scenes as reconstructions, so the pipeline runs video → pointmaps → 3D Gaussians or mesh → simulator asset in one direction and simulator depth, pose, and contact labels back into the pre-training loss in the other.

    (1) Geometry-Grounded Pretext Tasks: pointmap, ray, and depth regression plus multi-view matching force the backbone to encode where surfaces are, not just what they look like; ordinary video supplies the multiple views for free.
    (2) Camera As A First-Class Input: Plücker ray embeddings or pose tokens turn a text-to-video model into a controllable renderer, which is exactly the interface a simulator viewport exposes.
    (3) Persistent State Beats Frame Memory: caching an explicit 3D representation and re-projecting it into each new frame removes the layout drift and object-identity loss that pure 2D autoregressive rollouts accumulate.
    (4) Two-Way Data Flow: real-to-sim reconstruction fills the simulator’s asset gap, while simulator renders provide perfect metric depth, pose, and contact supervision that no internet video can give.
    (5) 3D Awareness Is Measurable: linear probes on frozen features show that self-supervised 2D backbones encode single-view depth surprisingly well yet fail at multi-view consistency, and 3D-aware fine-tuning lifts both plus downstream segmentation.
    (6) Metric Scale Is The Robot Requirement: scale-invariant depth is enough for a nice-looking novel view and useless for a grasp, so metric anchoring is the part embodied policies actually consume.

    Diagram: internet video corpora and simulator renders both feed a shared backbone whose pointmap, camera-pose, and differentiable-render heads serve two consumers, a camera-controlled video generator and a 3D physics simulator

    Figure 1: One backbone, two supervision sources, two consumers. Unlabeled video supplies visual diversity and self-consistency signals, simulator renders supply exact depth and pose, and the geometry heads are what the video generator and the simulator both read from.

    The concrete form most systems settle on is a feed-forward multi-view transformer that predicts a pointmap per frame in a common coordinate frame, so camera intrinsics, extrinsics, depth, and correspondence all fall out of one prediction rather than a fragile SfM plus MVS chain. DUSt3R established the pointmap formulation for two views, MASt3R added a matching head, and VGGT scaled it to many frames with attention alternating between frame-local and global blocks. That output is directly liftable: initialize 3D Gaussians from the points for a renderable scene, or run meshing plus material estimation for a body a rigid-body engine can collide against. Two caveats decide whether the bridge is load-bearing. First, monocular training recovers geometry only up to an unknown similarity transform, so metric scale has to be injected from calibrated stereo, a known camera baseline, simulator ground truth, or object-size priors. Second, the reprojection and correspondence losses assume a static scene, so dynamic content needs motion masks or a 4D formulation or the model folds object motion into depth as flying pixels.

    Mathematical Formulation:
    u = \pi(K, T_c, X)
    \bar{X}_i = X_i / s
    s = \frac{1}{N}\sum_i \lVert X_i \rVert_2
    \mathcal{L}_{geo} = \sum_i \lVert \bar{X}_i - \bar{X}^{*}_i \rVert_1
    \mathcal{L}_{corr} = \sum_i \lVert u'_i - \pi(K', T_{c'}, X_i) \rVert_2
    \mathcal{L}_{rgb} = \lVert I_c - \mathcal{R}(G, K, T_c) \rVert_1
    \mathcal{L} = \mathcal{L}_{2D} + \lambda_g \mathcal{L}_{geo}
    +\ \lambda_c \mathcal{L}_{corr} + \lambda_r \mathcal{L}_{rgb}

    Where:

    • u is a pixel coordinate and X \in \mathbb{R}^3 the predicted scene point for that pixel, expressed in the reference camera’s frame.
    • \pi is the pinhole projection, K the intrinsics, and T_c the pose of camera c; primed symbols refer to a second view.
    • i \in \{1, \ldots, N\} indexes valid pixels, and X^{*} is the ground-truth pointmap from a simulator, LiDAR, or a reconstruction pipeline.
    • s is the mean point norm used for scale normalization, so \mathcal{L}_{geo} on \bar{X} is scale-invariant; dropping the division makes the loss metric and requires metric labels.
    • \mathcal{L}_{corr} is the reprojection error against a matched pixel u' in the other view, and it is valid only where the scene is static between the two frames.
    • \mathcal{R} is a differentiable renderer (splatting or volume rendering) over a lifted representation G, letting 2D images supervise 3D structure through \mathcal{L}_{rgb}.
    • \mathcal{L}_{2D} is the original appearance objective (diffusion denoising or masked prediction) and \lambda_g, \lambda_c, \lambda_r \geq 0 weight the geometric terms.
    Log-scale bar chart of global attention cost relative to a single view for 1, 2, 4, 8, 16, and 32 views at 1369 tokens per view, growing as the square of the view count

    Figure 2: Joint multi-view geometry is what makes a representation 3D-aware, but attention over all views at once costs O((VN)^2), so a 32-view window is roughly 1,000x a single view. This is why practical models alternate frame-local and global attention and slide a bounded view window.

    Property2D generative video model3D-aware pre-trained representation3D physics simulator
    Scene stateImplicit in activations and past frames; not queryableExplicit pointmaps, depth, poses, optionally 3D GaussiansExact meshes, transforms, joints, contact manifolds
    Camera controlText prompts only, unless pose conditioning is addedRay or pose embeddings make trajectory a direct inputArbitrary and exact by construction
    Metric scaleAbsentAvailable only if metric supervision or calibration is injectedExact, in SI units
    Visual and asset diversityInternet scale, effectively unboundedInherits video-scale diversity, adds geometry headsLimited to authored assets and shaders
    Typical failure modeLayout drift, object teleporting, impossible occlusionsScale ambiguity, flying pixels on moving objects, cost of joint viewsAppearance gap and unmodeled deformable or fluid behavior

    Login to view more content
  • DL0122 V-JEPA vs Video World Models

    Compare V-JEPA (Joint-Embedding Predictive Architecture) with video generation world models such as OpenAI’s Sora or NVIDIA Cosmos. What does each objective actually learn, and when would you deploy one over the other?

    Answer

    Both families are trained by predicting the unseen part of a video, and both are called world models, but they differ in where the prediction target lives. V-JEPA masks a large 3D region of a clip and trains a predictor to regress the embeddings of that region produced by an EMA target encoder, using an L1 loss in representation space; there is no decoder anywhere in the system, so the model is never asked to reproduce a pixel. Sora and Cosmos Predict instead train a generative model (a diffusion transformer over spacetime latent patches, or an autoregressive token model) to reconstruct the actual frames, which forces the network to spend capacity on every detail, including detail that is fundamentally unpredictable, such as the exact texture of foliage or the phase of a reflection. That single design choice cascades: V-JEPA gets a compact motion-aware representation and a rollout that costs one forward pass per step, which is what makes energy-based planning tractable, while generative world models get a renderable future that can serve as a simulator, a synthetic-data generator, or an input to an existing camera-based perception stack. Neither dominates; the question to ask in an interview is whether your downstream consumer is a policy or a pixel pipeline.

    (1) Target Space: V-JEPA minimizes distance between predicted and EMA-encoded latents; Sora and Cosmos minimize a reconstruction or denoising loss defined on tokenizer latents that decode back to RGB.
    (2) Discarded Information Is A Feature: because the target is learned, JEPA can drop aleatoric high-frequency detail instead of averaging over it, which is exactly what produced blurry futures in older pixel-space predictors.
    (3) Collapse Risk: a learned target can be gamed by a constant function, so JEPA needs stop-gradient plus an EMA teacher and aggressive 3D block masking; a generative model has a fixed data target and cannot collapse.
    (4) No Rendering: V-JEPA cannot show you its prediction, so debugging and human review happen through probes and downstream task metrics, not by watching a video.
    (5) Rollout Cost: one predictor call per latent step versus tens of denoiser calls plus a decoder pass, a difference of one to two orders of magnitude inside a model-predictive control loop.
    (6) Evaluation Protocol: JEPA is scored with frozen attentive probes on motion-heavy benchmarks and with planning success rate; generative world models are scored with FVD, human preference, physics-consistency suites, and sim2real transfer of policies trained on their output.

    Two stacked architecture rows: V-JEPA encodes a masked clip, runs a predictor, and compares predicted latents against stop-gradient EMA target latents with an L1 loss; a video world model encodes a clip with a causal tokenizer, iterates a DiT denoiser for tens of steps, and decodes back to frames

    Figure 1: The two families share a video encoder but differ in the prediction target: V-JEPA regresses EMA target embeddings and never instantiates a pixel, while a generative world model passes through a tokenizer, an iterative denoiser, and a decoder that must reconstruct every frame.

    The masking strategy is what makes the latent objective non-trivial. V-JEPA masks large 3D blocks that span the full temporal extent of the clip, removing on the order of 90 percent of the tubelet tokens, so the predictor cannot solve the task by interpolating from neighboring patches and is pushed toward object identity, motion, and rough physics. V-JEPA 2 scaled this recipe to a ViT-g encoder over more than a million hours of internet video and then post-trained an action-conditioned predictor on roughly 62 hours of unlabeled robot video, after which planning is just a search over action sequences whose predicted latent lands closest to a goal image embedding. Cosmos takes the opposite bet at similar scale, curating on the order of tens of millions of hours of video into diffusion and autoregressive world foundation models plus tokenizers, precisely because a physical-AI developer wants renderable, controllable rollouts they can feed into an existing autonomy stack. The two are complementary in practice: Cosmos ships a separate reasoning model alongside its generators, and several robotics stacks now use a generative model for data augmentation and a latent predictor for the control loop.

    Mathematical Formulation:
    s_y = E_{\bar{\theta}}(y)
    \hat{s}_y = P_{\phi}(E_{\theta}(x), m)
    \mathcal{L}_{\mathrm{JEPA}} = \lVert \hat{s}_y - \mathrm{sg}(s_y) \rVert_1
    \bar{\theta} \leftarrow \tau \bar{\theta} + (1 - \tau)\,\theta
    \mathcal{L}_{\mathrm{gen}} = \mathbb{E}\left[\lVert \epsilon - \epsilon_{\theta}(z_t, t, c) \rVert_2^2\right]
    a^{*} = \arg\min_{a} \lVert P_{\phi}(s_t, a) - s_g \rVert_1

    Where:

    • x is the visible (context) part of a clip and y the masked target region; m carries the positional mask tokens telling the predictor where to predict.
    • E_{\theta} is the online encoder, E_{\bar{\theta}} the EMA target encoder, and P_{\phi} the narrow predictor discarded after pretraining.
    • \mathrm{sg}(\cdot) is the stop-gradient that blocks the trivial solution, and \tau is the EMA momentum, typically ramped from about 0.998 toward 1.
    • z_t is the noised tokenizer latent at diffusion step t, \epsilon the sampled Gaussian noise, \epsilon_{\theta} the denoiser, and c the conditioning (text, past frames, or actions).
    • s_t is the current latent state, a a candidate action sequence, and s_g the goal embedding; the planner requires an initial condition s_t = E_{\theta}(o_t) from the current observation.
    Log-scale bar chart of forward passes per planning step for a 200-candidate 16-step CEM search: 3200 for a one-pass latent predictor, 12800 for a distilled four-step video diffusion model, and 96000 for a thirty-step video diffusion world model

    Figure 2: For the same 200-candidate, 16-step CEM search, a single-pass latent predictor needs 3,200 network calls while a 30-step generative rollout needs 96,000, before the decoder is even invoked; this 30x gap is why planning loops favor latent prediction.

    PropertyV-JEPA (latent predictive)Sora / Cosmos (generative world model)
    Training targetEmbeddings from an EMA copy of the encoder, L1 lossNoise or next token on VAE latents that decode to RGB
    Handles stochastic detail byDropping it from the representationSampling it, which costs capacity and steps
    Can render a futureNo decoder; futures exist only as vectorsYes, watchable and consumable by any vision stack
    Cost per rollout stepOne predictor forward passTens of denoiser passes plus a decode
    Main pathologyRepresentation collapse and shortcut solutions under weak maskingPhysically implausible but photoreal rollouts, plus error accumulation over long horizons
    Natural deploymentFrozen backbone for recognition and anticipation; goal-conditioned planningSynthetic data, neural simulators, scenario replay, content generation

    Login to view more content
  • DL0120 VLA Video Pretraining

    How do you leverage internet-scale human video datasets (e.g., Ego4D, Something-Something) for pre-training robotic VLA policies? How do you overcome the embodiment/kinematic gap?

    Answer

    Human video is enormous but action-free: Ego4D contributes roughly 3,670 hours of egocentric recording across 74 locations and Something-Something V2 adds about 220,000 clips over 174 manipulation-style classes, yet none of it carries joint commands, gripper states, or proprioception. The practical recipe is therefore to never try to learn a robot action space directly from video, but to choose an intermediate interface that video can supervise and that a small teleoperated dataset can later decode into commands. Three interfaces dominate: visual representation pretraining (R3M, VIP, VC-1 style encoders trained on Ego4D with time-contrastive or masked objectives), latent action pretraining (a VQ inverse-dynamics model turns each frame pair into a discrete code, the VLA is pretrained to predict those codes from image plus instruction, and only the final decoder head sees real actions), and embodiment-free intermediate predictions such as future 2D/3D point tracks, subgoal images, or hand-contact affordance maps. The embodiment gap is then attacked on four separate axes rather than as one problem: appearance (human arms and hands in frame), kinematics (a roughly 27-DoF hand versus a 1-DoF parallel gripper), viewpoint and calibration (head-mounted camera with unknown scale versus a fixed wrist or shoulder camera), and dynamics (contact forces and speeds a rigid gripper cannot reproduce). In every working system the last mile is still robot data, typically 50 to 200 hours of teleoperation, mixed into training with a weight that keeps human video as a prior rather than a target.

    (1) Pick The Interface, Not The Dataset: transferability and actionability trade off directly, so the design decision is which quantity human video supervises, not whether to use it.
    (2) Latent Actions As A Shared Vocabulary: a quantized inverse-dynamics model maps (o_t, o_{t+k}) to a discrete code that describes observed scene change, which both a hand and a gripper can produce.
    (3) Inverse Dynamics Pseudo-Labeling: the VPT approach trains an IDM on a small labeled set and pseudo-labels the rest, which works when the action space is shared and fails badly across embodiments.
    (4) Retargeting Where Geometry Allows: a MANO hand fit gives wrist pose plus a scalar aperture, enough to drive a parallel gripper but not a multi-finger in-hand manipulation.
    (5) Visual Alignment Beats Visual Luck: masking or inpainting human hands, matching camera placement with head-mounted rigs, and cropping to the same field of view remove most of the appearance gap cheaply.
    (6) Co-Training Ratio Is A Hyperparameter: human video enters with weight \lambda alongside robot batches; too high and the policy predicts plausible video instead of executable actions.

    Pipeline diagram: human video from Ego4D and Something-Something enters a VQ inverse-dynamics latent action model, the resulting discrete codes pretrain a VLA backbone, and a small robot demonstration set finetunes an action head that outputs joint-space commands

    Figure 1: The latent-action route: human video → discrete latent codes → VLA pretraining, with the embodiment gap crossed only in the final stage, where a few tens of hours of teleoperated demonstrations train the head that emits real commands.

    Treating the gap as four separate problems makes each one tractable. For appearance, hand segmentation plus inpainting, or observation-space alignment using head-mounted capture on the human side and a matching camera on the robot, is usually enough that the frozen backbone stops keying on skin texture. For kinematics, the honest move is to restrict what human video is asked to teach: it supplies where to go, what to touch, and in what order, while grasp synthesis and force control come from robot data or from a classical controller. For viewpoint and scale, monocular depth or 3D point tracking lifts predictions into a metric frame, and predicting relative motion rather than absolute pose removes the calibration dependency. For dynamics, the temporal stride k matters more than people expect: a stride of a few frames captures human motion an order of magnitude faster than the robot will execute, so codes must be interpreted as subgoals and re-planned, not replayed. Cross-embodiment action tokenization, as used in the Open X-Embodiment and RT-X work over 22 embodiments, is the complementary trick on the robot side: a shared tokenizer plus per-embodiment heads lets one backbone absorb heterogeneous action spaces.

    Mathematical Formulation:
    z_t = q\big(E(o_t, o_{t+k})\big)
    \hat{o}_{t+k} = D(o_t, z_t)
    \mathcal{L}_{\mathrm{LAM}} = \lVert o_{t+k} - \hat{o}_{t+k} \rVert^2 + \beta \mathcal{L}_{\mathrm{VQ}}
    \mathcal{L}_{\mathrm{pre}} = -\log \pi_\theta(z_t \mid o_t, \ell)
    \mathcal{L}_{\mathrm{ft}} = \lVert a_t - g_\phi(h_\theta(o_t, \ell)) \rVert^2
    \mathcal{L} = \mathcal{L}_{\mathrm{robot}} + \lambda \mathcal{L}_{\mathrm{human}}

    Where:

    • o_t is the observation (usually a single RGB frame) and \ell the language instruction; a_t is the true robot action, such as an end-effector delta or joint target.
    • z_t is the discrete latent action for the transition, produced by encoder E and quantizer q against a codebook of size C with m tokens per transition, giving C^m distinct codes.
    • D is the forward decoder that must reconstruct o_{t+k} from o_t and z_t alone, which is what forces z_t to carry the motion information.
    • k > 0 is the temporal stride in frames and \beta weights the VQ commitment loss; small k captures noise, large k makes the code ambiguous.
    • \pi_\theta is the VLA backbone during latent pretraining, h_\theta its features after pretraining, and g_\phi the small action head trained only on robot data.
    • \lambda is the co-training weight on the human-video term; a typical schedule anneals \lambda from near 1 toward 0 so that the deployed policy is dominated by real action supervision.
    Qualitative scatter plot of transfer interfaces: joint torques and end-effector poses are directly executable but robot-specific, latent action codes and subgoal images sit in a shaded sweet spot, and frozen visual features are fully embodiment-agnostic but need a full policy on top

    Figure 2: The interface ladder. Human video can only supervise quantities that are not tied to a specific body, so the useful designs sit in the middle band, where latent codes, point tracks, and subgoal images are both learnable from video and decodable into commands with modest robot data.

    PropertyFrozen visual representationsLatent action pretrainingTracks and subgoal images
    What human video supervisesEncoder weights and a temporal value or progress signalA discrete motion vocabulary shared by hands and grippersFuture 2D or 3D point trajectories, or the next subgoal frame
    Robot data still requiredA full policy trained on top of the frozen featuresA small action head mapping codes to commands, tens of hoursA track-following or goal-reaching low-level policy
    Handles hand versus gripper byIgnoring actions entirely, so the mismatch never appearsAbsorbing it in the quantizer: codes describe scene change, not jointsStaying in task space, where points and pixels carry no body
    Main failure modeTask-agnostic features help perception but little on long horizonsCodes drift toward camera motion and background dynamicsTracker failure under occlusion, and no notion of contact force
    Best fitSmall robot datasets and cluttered scenes needing better perceptionGeneralist VLAs that must follow unseen instructionsQuasi-static rearrangement and tool use with visible motion

    Login to view more content
  • DL0119 VLA Imitation vs RL

    Explain imitation learning vs RL for VLA training. When would you reach for each one on a real robot, for example when post-training a model like Physical Intelligence’s \pi_0?

    Answer

    Imitation learning trains a vision-language-action model by maximum likelihood on teleoperated demonstrations: every timestep carries a ground-truth action label, so the objective is the same supervised next-token or flow-matching loss used to pretrain the underlying VLM, and no reward, reset, or simulator is required. Reinforcement learning replaces that per-step label with a scalar return and optimizes the policy on its own state distribution, which is the only way to fix the failure mode imitation cannot fix and to exceed the demonstrator’s performance. The asymmetry is practical rather than philosophical: demonstrations are cheap to collect and trivially parallel across operators, while a single real-robot RL run needs a success detector, an automated reset, and 10^4 to 10^6 environment steps of physically risky exploration. That is why essentially every published generalist VLA (RT-2, OpenVLA, \pi_0) is imitation-pretrained on cross-embodiment data such as Open X-Embodiment, and RL appears only as a narrow post-training stage on a handful of tasks where reward is machine-checkable. The interview answer is therefore not “which is better” but “imitation for coverage and language grounding, RL for the last 10 to 20 points of success rate on tasks you can actually score”.

    (1) Different Objectives: imitation maximizes \log \pi_\theta(a \mid o) under a fixed data distribution, while RL maximizes expected discounted return under the policy’s own visitation distribution.
    (2) Compounding Error: behavior cloning suffers covariate shift, and its regret against the expert grows as O(\epsilon T^2) in the horizon, whereas interactive or on-policy training reduces this to O(\epsilon T).
    (3) Supervision Density: a demonstration gives a full action vector at every control step; a sparse success reward gives one bit per episode, so RL’s credit assignment problem is orders of magnitude harder.
    (4) Performance Ceiling: imitation is capped at the demonstrator and inherits their pauses, jerk, and inconsistent strategies; RL can discover non-human solutions but will hack a badly shaped reward.
    (5) Multimodality: human demos are multimodal, so a naive MSE head averages incompatible strategies, which is exactly why modern VLAs use discretized action tokens, diffusion, or flow-matching action experts plus action chunking.
    (6) Hybrid Is The Default: production recipes pretrain with imitation, then apply human-in-the-loop corrections (HG-DAgger), offline RL with a conservatism penalty, or advantage-weighted fine-tuning with a KL leash to the cloned policy.

    Flow diagram: teleoperated demonstrations feed imitation pretraining of a VLA, the resulting policy produces on-robot rollouts, a success detector or human takeover produces a reward signal, and an RL fine-tuning stage with advantage weighting and a KL penalty updates the policy in a loop

    Figure 1: The standard VLA recipe is imitation first, RL second: demonstrations supply dense per-step labels with no reward machinery at all, and the RL loop is only viable once the cloned policy already succeeds often enough for a scalar success signal to be informative.

    The sample-cost gap dominates every design decision. Collecting 100 demonstrations of a new task takes a single operator under an hour and immediately yields a policy that succeeds sometimes; reaching the same point with online RL from a random initialization means exploring a 7-DoF continuous action space with sparse reward, which is hopeless on hardware and merely expensive in simulation. Once an imitation-pretrained policy exists, RL becomes tractable because exploration starts near the solution manifold: human-in-the-loop RL methods that keep demonstrations in the replay buffer and let an operator intervene on impending failures report near-perfect success on precise, contact-rich assembly tasks after roughly one to three hours of real-world training. The corresponding risk is catastrophic forgetting: unconstrained RL on one task will happily destroy the language grounding and cross-task generalization that the imitation stage paid for, so practitioners freeze most of the VLM backbone, fine-tune the action expert, and add an explicit KL term back to the behavior-cloned reference policy.

    Mathematical Formulation:
    \theta_{BC} = \arg\max_{\theta} \mathbb{E}_{(o,a) \sim \mathcal{D}} \left[\log \pi_{\theta}(a \mid o)\right]
    J(\pi) = \mathbb{E}_{\pi}\left[\sum_{t=0}^{T} \gamma^{t} r(s_t, a_t)\right]
    J(\pi^{*}) - J(\pi_{BC}) = O(\epsilon T^{2})
    J(\pi^{*}) - J(\pi_{DAgger}) = O(\epsilon T)
    \max_{\theta} \mathbb{E}_{\mathcal{D}}\left[\log \pi_{\theta}(a \mid o) \exp\!\left(A(o,a)/\beta\right)\right]

    Where:

    • \pi_{\theta}(a \mid o) is the VLA policy mapping an observation o (images, proprioception, language instruction) to an action or action chunk a.
    • \mathcal{D} is the fixed demonstration dataset, and \pi^{*} is the human expert whose behavior it samples.
    • s_t and t \in \{0, \ldots, T\} are the environment state and control step, r the reward, and \gamma \in (0,1] the discount.
    • \epsilon is the per-step supervised action error under the expert’s distribution; the quadratic term comes from errors pushing the policy into states absent from \mathcal{D}.
    • A(o,a) is the estimated advantage and \beta > 0 the temperature; as \beta \to \infty the advantage-weighted objective degenerates to plain behavior cloning, which is why the hybrid is a continuum rather than a binary choice.
    • Required initial condition for the RL stage: \pi_{\theta} is initialized at \theta_{BC} and the task admits an automated reset plus a reward oracle, otherwise J is not estimable on hardware.
    Line chart of regret against the expert versus episode horizon for a per-step error of 0.02, comparing a quadratic behavior-cloning bound that reaches 200 at horizon 100 with a linear interactive-training bound that reaches only 2

    Figure 2: With a fixed per-step error of \epsilon = 0.02, the quadratic behavior-cloning bound separates from the linear on-policy bound by a factor of T, which is why long-horizon manipulation degrades far faster than single-step prediction accuracy suggests.

    PropertyImitation learning (BC)Offline RLOnline RL
    Supervision per episodeOne action label at every control stepLogged actions plus reward or success labelsScalar reward, often a single success bit
    What it needs to runTeleoperation rig only; no reset, reward, or simulatorA labeled buffer of mixed-quality dataAutomated reset, success detector, safety envelope
    Typical budget for one task50 to 2,000 demonstrations, hours of human timeReuses existing logs; compute-bound, not robot-bound10^4 to 10^6 real steps, or a sim-to-real gap to close
    Distribution optimized onExpert states only, so covariate shift is unaddressedBuffer states, with pessimism outside their supportThe policy’s own states, which is the point
    Performance ceilingThe demonstrator, minus compounding errorBest behavior recoverable from the logged dataCan exceed the human on speed and precision
    Main failure modeDrift into unseen states; mode averaging on multimodal demosValue overestimation on out-of-support actionsReward hacking, hardware damage, forgetting language grounding

    Login to view more content
  • DL0116 Video VLM Temporal Sampling

    How do video VLMs balance temporal sampling resolution against KV cache memory limits when processing multi-hour video streams, as in Google’s Gemini video understanding?

    Answer

    A video VLM has no separate temporal-resolution dial: it flattens the stream into one token sequence, so frame rate, tokens per frame, and clip length all collapse into a single token count N = f \cdot T \cdot t_f, and that count fixes both the KV cache size (linear in N) and the prefill cost (quadratic in N). Two hours sampled at 1 fps with 64 tokens per frame is 460,800 tokens; on a Llama-3-8B-class decoder with 32 layers, 8 KV heads of head dimension 128, and a bf16 cache, each token costs 128 KB of KV, so the cache alone is roughly 60 GB, against the roughly 62 GB an 80 GB H100 has free after weights and activations. Multiply the frame rate by four or keep the encoder’s full 256 tokens per frame and you are at 240 GB, which no single device holds. Production systems therefore invert the question: fix a token budget from the memory and latency SLO first, then decide how to spend it across time and space. The three levers are temporal subsampling (fewer frames), token compression (pooling, merging, or resampling each frame), and bounded state (a sliding window plus a compressed memory bank or an offline retrieve-then-read pass), and only the third one makes memory independent of stream length.

    (1) Tokens Are The Currency: memory scales as O(N) and prefill attention as O(N^2), so a 4x frame-rate increase is a 4x memory bill and a 16x prefill bill.
    (2) Per-Token KV Cost Is Architectural: m = 2 L n_{kv} d_h b bytes, so GQA with 8 KV heads instead of 32 already cuts it 4x, and INT4 KV quantization cuts it another 4x before any sampling change.
    (3) Equal Memory, Different Failures: halving f and quartering t_f can cost the same tokens, but the first causes temporal aliasing of short events while the second destroys small text and fine spatial detail.
    (4) Slow-Fast Splitting: strong systems decouple the two axes, keeping many frames at very few tokens for motion and a handful of keyframes at full spatial resolution for detail.
    (5) Adaptive Beats Uniform At Fixed Budget: dropping near-duplicate frames by feature similarity and selecting query-relevant segments spends the same tokens on the informative part of the timeline.
    (6) Time Grounding Must Survive Downsampling: position encodings need absolute timestamps, not token indices, or a variable frame rate makes the model’s answers about “when” drift.

    Log-scale bar chart of KV cache size for a two-hour video under six configurations: 242 GB at 1 fps with 256 tokens per frame, 60 GB at 1 fps with 64 tokens, 30 GB at 0.5 fps, 15 GB at 16 tokens per frame, 15 GB with INT4 cache, and under 1 GB for a streaming window plus memory bank, against a 62 GB device budget line

    Figure 1: KV cache for the same two-hour video under six budgets. Memory depends only on the product f \cdot T \cdot t_f and the per-token cost, so sampling changes and cache quantization are interchangeable for memory, but only a bounded streaming state stops growth with stream length.

    The reason the two sampling axes are not interchangeable for accuracy is that they alias different things. Uniform sampling at 1 fps is enough for plot-level questions over an hour of footage but silently deletes any event shorter than a second, which is why counting, ordering, and “who handed what to whom” questions collapse at low frame rates. Cutting tokens per frame is safe for gist and scene recognition and catastrophic for anything requiring OCR of on-screen text or small-object detail. Public systems make these trade-offs explicit: Gemini samples video at 1 fps with about 258 tokens per frame by default, which is roughly 0.93M tokens per hour and effectively consumes a 1M-token context, and offers a low-resolution mode at about 66 tokens per frame that stretches the same window to several hours. For genuinely unbounded streams the only stable design is bounded state: keep the last few tens of seconds as an exact KV window with attention sinks so the distribution stays in-domain, merge evicted frames into a fixed-size memory bank, and for offline archives replace dense prefill with a two-pass pipeline that indexes cheap summaries and then re-decodes only candidate segments at high frame rate.

    Mathematical Formulation:
    N = f \cdot T \cdot t_f
    m = 2 L n_{kv} d_h b
    M_{kv} = m \cdot N
    C_{prefill} = O(N^2 d)
    f_{max} = \frac{M_{budget}}{m \cdot T \cdot t_f}
    m = 2 \cdot 32 \cdot 1024 \cdot 2 = 131072
    N = 1 \cdot 7200 \cdot 64 = 460800

    Where:

    • N is the total visual token count fed to the decoder, f the sampling rate in frames per second, T the stream duration in seconds, and t_f the tokens retained per frame after pooling or resampling.
    • m is the KV bytes per token, with L layers, n_{kv} key-value heads, head dimension d_h, and b bytes per stored element; the factor 2 counts keys and values.
    • M_{kv} is cache size and M_{budget} the free HBM after weights, activations, and the text prompt; f_{max} is the highest frame rate a dense prefill can afford at a given t_f.
    • C_{prefill} is prefill attention cost with model width d; it usually binds before memory does, since doubling N quadruples time to first token.
    • The numeric example uses L = 32, n_{kv} d_h = 1024, and b = 2 bytes, giving 128 KB per token, and T = 7200 s at f = 1, t_f = 64, so M_{kv} \approx 60 GB.
    • A streaming design replaces N with N_w + N_{mem}, the exact window plus a fixed memory bank, which is O(1) in T.
    Pipeline diagram: a two-hour 30 fps stream is decoded at 1 fps, encoded per frame into 256 patch tokens, reduced to 64 tokens per frame, then fed into a sliding-window KV cache whose evicted frames are merged into a fixed-size compressed memory bank, with both feeding an LLM decoder that answers a timestamped user query

    Figure 2: A streaming video VLM: decode rate and per-frame token count shrink the sequence by two orders of magnitude, then eviction into a fixed memory bank makes resident state constant in stream length while the exact sliding window preserves fine detail for recent seconds.

    PropertyUniform dense samplingQuery-adaptive keyframe retrievalStreaming window plus memory bank
    Token growth with durationLinear and unboundedBounded at read time, but the index still scales with durationConstant resident state
    KV cache for 2 h60 GB at 1 fps and 64 tokens per frame; 240 GB at full 256 tokensA few GB, set by the number of retrieved segmentsUnder 1 GB for roughly 5,900 live tokens
    Where information is lostBetween sampled frames, uniformly across the timelineIn segments the retriever scored as irrelevantIn merged history: detail decays with age
    Best fitClips of a few minutes and dense temporal groundingOffline archives with a known question per requestLive monitoring, assistants, always-on capture
    Typical failureOut-of-memory, or time to first token in the tens of secondsRetriever misses the one relevant second, so the answer is confidently wrongQuestions about hour-old fine detail that the merge already discarded
    Latency profileOne huge quadratic prefill per requestCheap index build offline, small prefill per queryAmortized per frame, answers available at any instant

    Login to view more content