Predictive and Generative Objectives › Contrastive and Masked Prediction

Michael BrenndoerferJuly 1, 20269 min read
World Models HandbookAll chapters
Part 6 of 12Learning World Models at Scale
Chapter 31Predictive and Generative Objectives › Contrastive and Masked Prediction
Reading section 4 of 12 · Contrastive and Masked Prediction · 9 minSection 4 of 12 · 9 min

Book contents

Section 4 of 12 · 9 min read

Predictive and Generative Objectives › Contrastive and Masked PredictionOne section · Sections in this chapter4 / 12 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Contrastive and masked objectives shape a representation through recognition or prediction tasks rather than requiring an unconditional full-image reconstruction. Contrastive and embedding-target methods need not reward pixel fidelity at all. Raw-target masked autoencoders do reconstruct masked pixels, so their pressure toward texture depends on the target, mask, and decoder design.

This is a deliberate trade. Instead of always defining the task as "produce the complete observation," these objectives ask the model to discriminate candidates or infer hidden targets. An embedding target can omit texture, whereas a raw-pixel target can retain it. The representation is therefore aligned with whatever information the chosen pretext task rewards. The price is that many such objectives lack a direct likelihood interpretation and are sensitive to how positives, negatives, masks, and targets are constructed.

Contrastive Objectives

Contrastive learning builds a representation by pulling together embeddings of similar examples and pushing apart different ones. The InfoNCE loss, introduced in the context of contrastive predictive coding, has the form

LInfoNCE=−log⁡exp⁡(sim(fθ(x),fθ(x+))/τ)∑i=1Nexp⁡(sim(fθ(x),fθ(xi))/τ).\mathcal{L}_{\mathrm{InfoNCE}} = -\log \frac{\exp(\mathrm{sim}(f_\theta(x), f_\theta(x^+)) / \tau)}{\sum_{i=1}^{N} \exp(\mathrm{sim}(f_\theta(x), f_\theta(x_i)) / \tau)}.

The numerator measures how similar the anchor is to its positive; the denominator sums the similarities to every candidate, including the positive. The ratio is therefore a softmax probability that the positive is the correct match, and the loss is the negative log of that probability. The temperature τ\tau sharpens or flattens the distribution. The variables are:

  • LInfoNCE\mathcal{L}_{\mathrm{InfoNCE}}: the InfoNCE loss
  • xx: the anchor
  • x+x^+: its positive
  • {xi}i=1N\{x_i\}_{i=1}^{N}: candidate set including that positive and N−1N-1 negatives
  • fθf_\theta: the encoder
  • sim\mathrm{sim}: a similarity such as cosine similarity
  • τ>0\tau>0: the temperature [@oord2018representation]
  • NN: number of candidates

The loss is an NN-way classification objective: identify the positive among all candidates.

Read as classification, the model must identify the positive future or view among candidates. A low temperature τ\tau sharpens differences in their logits; a high temperature flattens them. Temperature therefore changes both the probability distribution and gradient scale, and should be selected together with similarity normalization and the negative set.

The temperature reshapes the InfoNCE softmax without changing the underlying similarities, and fixing one similarity profile makes that effect concrete.

In[3]:
Code
## Fixed similarity profile for one anchor against six candidates.
## Candidate 0 is the positive; candidates 1 through 5 are negatives.
sims = np.array([0.90, 0.80, 0.55, 0.40, 0.20, 0.10])
temperatures = np.array([0.05, 0.20, 1.00])
infonce_probs = np.stack(
    [np.exp(sims / tau) / np.exp(sims / tau).sum() for tau in temperatures]
)
positive_prob = infonce_probs[:, 0]
Out[4]:
Visualization
Three bars show positive-candidate probabilities of about 0.88, 0.52, and 0.24 at temperatures 0.05, 0.20, and 1.00.
Probability assigned to the positive candidate for one fixed similarity vector. Lowering the temperature from 1.00 to 0.05 raises the positive probability from about 0.24 to 0.88 because the positive has the largest similarity. A lower temperature sharpens this ranking but would amplify the wrong candidate just as strongly if a negative ranked first.

In the temporal setting, positives may be future samples paired with their preceding context, while other samples serve as negatives. Contrastive predictive coding (CPC) trains a context representation to identify future latents under this loss [@oord2018representation]. The objective biases the representation toward features that help distinguish predictable future content; it does not directly require pixel-level reconstruction.

Why can this favor predictable structure? A feature of the future that can be inferred from context helps identify the correct candidate. Unpredictable variation is less useful for that discrimination, although finite data and correlations can still let the encoder represent it. This is often a useful bias for dynamics modeling, but "predictable from the past" is not the same as "relevant to the task." A contrastive representation can preserve predictable nuisance variables or omit variables that matter for control but not for its pretext task.

Contrastive approaches have two practical drawbacks. First, they need negatives, and the choice of negatives controls what the representation learns. Too-easy negatives give a weak discrimination task; false or ambiguous negatives can give a misleading gradient. Second, small batches provide fewer in-batch negatives and can weaken the signal. MoCo's queue and SimCLR's large batches address the supply of negatives in different ways [@he2020momentum; @chen2020simple]. Non-contrastive alternatives such as BYOL instead train an online network to predict the output of a slowly updated target network [@grill2020bootstrap], a pattern related to later joint-embedding predictive methods.

Negative construction is therefore part of the modeling problem rather than a minor implementation detail. A semantically valid match treated as a negative creates a false-negative error; a trivially different negative supplies little pressure to represent fine structure. Queues, large batches, and sampling rules trade off diversity against false-negative risk. Non-contrastive methods remove explicit negative sampling, but introduce their own requirements for target updates, asymmetry, or regularization.

Non-Contrastive Self-Distillation

Non-contrastive methods train an online or student branch to match a target or teacher branch on another view or masked region. Stop-gradient blocks direct optimization of the target branch. Avoiding a constant representation depends on the full asymmetric design, which may include a predictor, moving-average target, centering, or variance regularization. Stop-gradient alone is not sufficient.

The naive matching objective admits a constant solution: if both branches emit the same vector for every input, their discrepancy is zero. Successful methods combine stop-gradient with architectural or update asymmetry, and some add explicit anti-collapse statistics. The exact mechanism differs across BYOL, SimSiam, and JEPA-family implementations, so the safeguards should be evaluated as a system rather than attributed to one operator.

Masked Prediction

Masked prediction asks the model to reconstruct a hidden portion of the input from the visible portion. BERT popularized the approach for text [@devlin2018bert], and MAE adapted it to images, achieving strong representation quality by masking a large fraction of patches and reconstructing only those patches [@he2021masked]. In the world-model context, masked prediction has become a natural fit because the inputs and targets can live in different spaces. You can mask out a future observation and ask the model to predict its embedding, without ever producing a pixel.

The design space is captured by three questions:

  • What is masked? Tokens of text, patches of an image, frames of a video, or timesteps of a trajectory. In a world model, masking time steps forces the model to integrate information from the past and future to fill in what it has not seen.
  • What is the target? The raw masked content (as in MAE) or an embedding of it produced by a target encoder (as in JEPA). The latter avoids spending capacity on reconstructing unpredictable detail.
  • What is the loss? Typically MSE in embedding space, or cross-entropy over token identities for discrete data.

Masked prediction prevents direct access to a selected target region. In temporal data, a suitable mask can require the model to infer future content from earlier observations and actions rather than copy the target. The strength and relevance of that bottleneck depend on what is masked and what remains visible.

A reconstruction objective gives the model credit for every reproduced pixel, including pixels copied from visible context. A masked objective withholds its target region, so reducing its loss requires inferring that region from the remaining context. In a temporal world model, the analogous task is to infer a hidden future from the past and actions. Whether this matches deployment depends on how the mask, context, and prediction target are constructed.

JEPA-Style Predictive Embeddings

Joint-Embedding Predictive Architectures (JEPAs) combine masked prediction with a learned target embedding [@lecun2022path]. A JEPA trains a context encoder and a target encoder, often updated by exponential moving average, then predicts the target embedding of a masked observation from the context. No pixel decoder or explicit negatives are required. In the expression below, gθ(fϕ(xctx))g_\theta(f_\phi(x_{\mathrm{ctx}})) is the prediction, fϕˉ(xtgt)f_{\bar\phi}(x_{\mathrm{tgt}}) is the target, and sg\mathrm{sg} prevents gradients from flowing into the target encoder:

LJEPA=∥gθ(fϕ(xctx))−sg(fϕˉ(xtgt))∥22,\mathcal{L}_{\mathrm{JEPA}} = \lVert g_\theta(f_\phi(x_{\mathrm{ctx}})) - \mathrm{sg}(f_{\bar\phi}(x_{\mathrm{tgt}})) \rVert_2^2,

where:

  • LJEPA\mathcal{L}_{\mathrm{JEPA}}: the JEPA predictive loss
  • xctxx_{\mathrm{ctx}}: visible context
  • xtgtx_{\mathrm{tgt}}: masked target region
  • fϕf_\phi: context encoder
  • fϕˉf_{\bar\phi}: target encoder
  • gθg_\theta: predictor of the target representation
  • sg\mathrm{sg}: stop-gradient operator, blocking gradients through the target branch
  • ϕˉ\bar\phi: target encoder parameters

The target parameters ϕˉ\bar\phi are typically updated by an exponential moving average of ϕ\phi. As we discussed in Predictive Embeddings, JEPA, and Energy-Based Models, bypassing reconstruction removes the pixel-level capacity trap, but direct pixel samples are unavailable and additional asymmetry or regularization is needed to avoid collapse.

The pieces serve different roles. Predicting an embedding rather than pixels lets the target representation omit observation detail that its encoder does not preserve. A moving-average target changes gradually as the online encoder learns. Stop-gradient makes that target a fixed value for each online update. None of these choices alone rules out a constant representation; collapse avoidance comes from their interaction with masking, predictor asymmetry, normalization, and any explicit regularization in the particular method.

What Contrastive and Masked Objectives Buy You

Compared to reconstruction, these objectives offer three concrete benefits:

  • A configurable inductive bias. Embedding targets can make it cheap to ignore texture and lighting; raw masked-pixel targets retain more appearance pressure.
  • Prediction-sufficiency pressure. When the target is a future embedding, the encoder is pushed to preserve information needed by the chosen predictive task. This becomes control-sufficiency pressure only when that task is aligned with the downstream decision, as discussed in Sufficient State and State Abstraction.
  • Interface flexibility. Masked objectives can be adapted to any modality that has a natural masking operation, which is useful when the agent sees video, proprioception, and language.

The second benefit connects most directly to the control-theoretic notion of a sufficient state. A state is sufficient for a task if knowing it is as good as knowing the full history for the required prediction or decision. Contrastive and masked objectives reward information useful to their pretext tasks. If those tasks expose every distinction needed downstream, the learned representation may approach task sufficiency; if not, it can be sufficient for the wrong prediction.

The cost is a different optimization and evaluation burden. Temperature, mask ratio, target updates, and anti-collapse mechanisms can materially change the result, and many objectives do not come with the direct likelihood interpretation of reconstruction. Numerical comparisons across papers are also harder because there is no universal "bits per dimension" scale.