Alessio Borgi

Alessio Borgi

PhD in Graph Neural Networks and Generative AI 👨🏻‍💻🤖

D-JEPA: Combining Feature Prediction with Generative Denoising

3 minute read

Published:

TL;DR: D-JEPA combines a feature-prediction objective with a conditional denoising objective. Predicted representations guide the generation of missing continuous tokens. It is a generative extension of JEPA, so its output and evaluation differ from a representation-only encoder.
First page of Denoising with a Joint-Embedding Predictive Architecture
Paper preview: Denoising with a Joint-Embedding Predictive Architecture.
Dengsheng Chen, Jie Hu, Xiaoming Wei, and Enhua Wu
Read the paper (PDF)

A useful feature does not specify every detail

Intuition first. Visible image regions may suggest that a hidden region contains fabric without determining its exact texture. A predicted feature can describe the context for a distribution of plausible completions. A generative model then supplies a way to sample one.

Chen et al. build D-JEPA around this separation. A context encoder and feature predictor estimate representations at masked positions. A conditional denoising network models continuous tokens using those predictions. Feature matching and generative denoising are trained together; the paper also explores a flow-matching alternative.

Keep the token and its predicted representation separate

Let \(x_i\) denote an input token and \(z_i\) the context-derived feature used to condition its generation. In the ImageNet implementation, input tokens come from a pretrained VAE, so both quantities live in latent spaces with different jobs. Official implementation.

For a noise-prediction formulation,

\[ \begin{aligned} x_i^t&=\sqrt{\bar\alpha_t}\,x_i+ \sqrt{1-\bar\alpha_t}\,\epsilon,\\ \mathcal L_d&=\mathbb E\left[\|\epsilon- \epsilon_\theta(x_i^t,t,z_i)\|_2^2\right],\\ \mathcal L&=\mathcal L_d+\mathcal L_p. \end{aligned} \]

\(\epsilon\) is Gaussian noise, \(t\) is diffusion time, and \(\mathcal L_p\) matches predicted features to target-encoder features. The paper’s formulation specifies the feature loss and the two objectives. Diffusion time here is a noise level, not a video’s physical timestamp.

Features condition a distribution over missing tokensFeature prediction supplies conditioning for the denoiser. Sampled token values then extend the visible context. Visible contextPredicted featuresConditional denoising
Feature prediction supplies conditioning for the denoiser. Sampled token values then extend the visible context.

Worked example: reconstruct one noisy coordinate

Take a scalar clean token \(x=2\), sampled noise \(\epsilon=-1\), and \(\bar\alpha_t=0.64\). Its corrupted value is

\[ x^t=0.8(2)+0.6(-1)=1.0. \]

If the denoiser predicts \(\hat\epsilon=-0.8\), its squared noise error is \((-1+0.8)^2=0.04\). An estimate of the clean coordinate is

\[ \hat x=\frac{x^t-0.6\hat\epsilon}{0.8} =\frac{1+0.48}{0.8}=1.85. \]

An exact noise prediction would recover 2 in this idealised calculation. This is one training example at one noise level; it is not the complete reverse sampling procedure.

The conditioning feature \(z_i\) helps distinguish which clean values are plausible in context. It is not itself the sampled noise or the clean token being recovered.

Sampling has two kinds of progress

The model gradually fills sets of masked token positions. Within a selected position, a conditional generative procedure produces the continuous token value. Filling additional positions supplies new context for subsequent predictions.

Consequently, the number of outer token-filling iterations and the number of denoising evaluations are different quantities. Compare both when interpreting generation cost. A smaller prediction network does not, by itself, establish lower end-to-end latency.

Compare the output required by the task

For an image generator, evaluate visual quality, diversity, conditioning, and computational cost. For a representation learner, evaluate what downstream predictors can extract from its features. Those measurements are related but not interchangeable.

A generated sample can look plausible while differing from an original hidden region, since multiple completions may be valid. Conversely, a useful representation may omit detail needed for faithful reconstruction. This makes the choice of evaluation especially important when comparing D-JEPA with I-JEPA.

Key Insight: D-JEPA assigns representation prediction and conditional generation separate roles. The predicted feature supplies context; the generative objective models variation in the token values consistent with that context.

Compare LLM-JEPA for another combination of representation prediction and a generative training objective.

References

  1. Chen, D., Hu, J., Wei, X., & Wu, E. Denoising with a Joint-Embedding Predictive Architecture. ICLR 2025.
  2. Authors’ D-JEPA implementation.

You May Also Enjoy

The Standard Library Modules Worth Knowing Before You pip install

8 minute read

Published:

Python ships with a counter, a queue, a memoiser, a CLI parser, a test framework’s smaller cousin, and a JSON-safe date library already installed, most of what people reach for a package to solve is one import away.

Files and Context Managers: Why with Is Not Optional

9 minute read

Published:

A file object that falls out of scope without with does get closed eventually, but ‘eventually’ means whenever the garbage collector gets around to it, which is not a promise any program handling more than a handful of files can live with.

Errors and Exceptions: EAFP, the Hierarchy, and Reading a Traceback

9 minute read

Published:

Exceptions in Python are not exceptional. They are a normal control-flow mechanism that the language leans on so heavily that the idiomatic style is to try the operation and handle the failure, rather than check first, which is faster, shorter, and free of race conditions.