JEPAhead: From Inpainting to Causal Forecasting for Urban Driving

Introduction

Motivation

World models have become fundamental for autonomous vehicles to comprehend the dynamics of urban environments and plan safe driving actions.

To overcome this limitation, this paper introduces an adaptation of the V-JEPA (Visual Joint Embedding Predictive Architecture) [1], framework tailored for the automotive context. Distinct from expensive conventional generative models, our approach bypasses pixel decoding entirely (although we tested it and showed results), performing future-horizon prediction directly via tokens within the latent space. By discretizing and predicting the temporal evolution of the scene at the embedding level, the model focuses strictly on the structural semantics of the environment (driver scenes).

Our results demonstrate that this formulation not only mitigates computational overhead but also optimizes the representation of spatio-temporal uncertainty in traffic scenarios. This work lays the foundation for highly efficient tokenized world models viable for real-time embedded systems.

The remainder of this section reviews related work and presents our contributions. We then describe the methodology, report the results, and discuss their implications. The paper concludes with the main findings and directions for future research.

Related Work

Joint-embedding predictive architectures

Joint-embedding predictive architectures learn by predicting target representations from context, with the prediction loss computed in embedding space. I-JEPA [2] applies this principle to masked image regions, while V-JEPA [1] extends it to spatiotemporal regions in video without pixel reconstruction or contrastive negatives. V-JEPA 2 [3] scales video pre-training and introduces an action-conditioned world model through additional training on robot interaction data, enabling planning in representation space.

Our work examines whether a V-JEPA predictor pre-trained for bidirectional latent inpainting can be adapted to causal future forecasting. V-JEPA 1 provides a starting point for studying this transfer, whereas V-JEPA 2-AC already incorporates action-conditioned predictive training.

*[email protected] †[email protected] ‡[email protected] §[email protected] ¶[email protected] ‖[email protected] ∗∗[email protected]

World models and latent dynamics

World Models [4] combines a variational autoencoder with recurrent latent dynamics, supporting controller learning within imagined environments. Dreamer [5] develops this approach by learning policies through latent rollouts, with representations shaped partly by observation reconstruction. For autonomous driving, MILE [6] jointly learns latent dynamics and a driving policy using geometric structure and bird’s-eye-view semantic prediction, while GAIA1 [7] combines autoregressive token prediction with video diffusion decoding to generate driving scenarios.

These approaches use observation reconstruction or generation within their learning pipelines. Our work instead studies the transfer of a predictor pre-trained solely through feature prediction to causal forecasting with a frozen encoder. A separately trained decoder supports visual evaluation without contributing to predictor training.

Latent prediction and representation analysis for driving

Feature-space prediction has been explored for driving and visual control. LAW [8] predicts future scene features conditioned on ego actions, while AD-L-JEPA [9] learns LiDAR representations through masked spatial prediction. Zhu and Choromanska [10] extend JEPA-based learning to spatiotemporal LiDAR prediction, evaluated through occupancy completion and forecasting. Drive-JEPA [11] combines video JEPA pre-training with trajectory planning. Complementing these approaches, Jeon et al. [12] use linear probes and targeted interventions to examine predictive information in driving policies. Beyond driving, DINO-WM [13] learns action-conditioned dynamics over frozen DINOv2 patch features for visual planning.

These studies motivate prediction and analysis in learned feature spaces. Our work specifically examines adapting a pre-trained bidirectional V-JEPA predictor to action-free causal forecasting.

Contributions

We study causal forecasting in the representation space of a frozen V-JEPA1 encoder on simulated urban driving, and report the following.

1. A causal adaptation of the pre-trained V-JEPA1 predictor. We initialize the predictor from pre-trained weights and fine-tune it using contiguous past-token context, future-position queries, and a group-level causal attention mask, while keeping the encoder frozen. It predicts eight future groups in one forward pass. A compact predictor trained from scratch provides a comparison with approximately six times fewer parameters.

2. Evidence of sequence-specific forecasting in representation space. Predicted tokens retrieve the correct future among 591 candidates in 69.0% of cases, compared with 55.3% for persistence, while retaining 98.4% of the target population’s spread. The advantage over persistence grows with the forecast horizon and is larger on tokens exhibiting greater change.

3. Evidence of driving-relevant information in predicted features. A fixed linear probe trained on real encoder tokens achieves 0.16 mIoU on patches undergoing semantic-class changes, compared with 0.13 for persistence. Complementary class-conditioned retrieval shows improvements over persistence across all six evaluated classes at the longest horizon.

4. A characterization of limited gains in pixel space. Improvements in latent forecasting yield modest L1 gains but little improvement under LPIPS with the evaluated decoder. We report reconstructions from true future tokens as an oracle reference and normalize pixel skill by the oracle-referenced achievable reduction to contextualize these results.

Evaluation covers CARLA Town12 and 16 future frames (eight temporal groups); generalization and action-conditioned counterfactual prediction remain untested.

Methodology

Problem formulation

We study causal future prediction in the representation space of a frozen video encoder. Following the joint-embedding predictive architecture (JEPA) formulation, a predictor maps a context representation $s_x = Enc(x)$ to an estimate of a target representation $s_y = Enc(y)$, conditioned on a variable $z$ that specifies what is to be predicted:

In V-JEPA pre-training, x is a sparse set of visible patches, y is a set of masked patches drawn from the same temporal window, and z encodes the position of the masked region. The task is therefore spatio-temporal inpainting: y may lie before, during, or after x in time.

We retain this formulation but change the temporal relation between x and y. In our setting x is a contiguous block of past observations, y is strictly posterior to x, and z addresses future instants that lie outside the observed window, as can be seen in Figure 1. The task becomes causal forecasting. Crucially, z remains available at inference time (it is a positional address, not content), whereas y is never provided to the model.

Frozen encoder and tokenization

We use a pre-trained V-JEPA1 ViT-L/16 encoder (24 blocks, embedding dimension 1024, 16 heads) and keep it frozen throughout. Input clips are 32 RGB frames at 224x224. A 3D patch embedding with kernel and stride (2, 16, 16) maps each tubelet of two frames and 16x16 pixels to one token, yielding 16 temporal positions x 14x14 = 196 spatial positions = 3136 tokens per clip. A 3D sine-cosine positional embedding is added, and 24 transformer blocks are applied with dense, unmasked attention.

Because the flattening order is temporal-major, the token sequence can be read as 16 consecutive groups of 196 tokens, where group $k$ corresponds to frames $(2k, 2k + 1)$. We use ”group” throughout to denote this slice.

To construct ground-truth future representations, we encode one 32-frame window per horizon, each ending at its target group. For horizon $h \in\{1, \ldots, H\}$, the window spans clip frames $[2h, 2h+32)$, and we retain only its last temporal group. This ensures that each target is computed without access to frames beyond its target group.

Dataset and sample construction

Our experiments use CARLA driving sequences. We used the PDM-Lite rule-based expert [14] to drive Carla Town12 entirely, its routes are partitioned at the route level with seed 42 into 70% train, 10% validation and 20% test, giving 1218 / 174 / 348 routes. Route-level splitting prevents frames of the same route from appearing in more than one split.

Figure 1 contains frames flows. The sampling window advances 16 frames, producing 6368 / 931 / 1773 samples for train / validation / test. In group units, each sample provides $K = 16$ context groups and $H = 8$ future groups. Global retrieval uses every third test sample, yielding 591 samples from 301 routes, with all 591 corresponding futures forming the candidate pool.

Causal predictor

We adapt the pre-trained V-JEPA1 predictor (12 blocks, embedding dimension 384, 16 heads; 23.9M parameters). The $K \times196 = 3136$ context tokens are projected to the predictor dimension and summed with positional embeddings. $H \times196 = 1568$ query tokens are formed from a single pre-trained mask token replicated and summed with the positional embeddings of the future instants; they carry no content. The two are concatenated into one sequence of 4704 tokens.

Causality is imposed by an attention mask built at group granularity and expanded by the group size, divided by context and future.

Within any permitted pair of groups, attention is unrestricted; the constraint is purely temporal. Since the reference implementation ignores the mask argument of its attention module, we reimplement the block forward pass to supply an explicit additive mask, and select a memory-efficient attention backend to avoid materializing the dense 4704 × 4704 matrix.

The positional embedding is regenerated for $K + H = 24$ groups rather than loaded from the checkpoint, since the pre-trained predictor was trained for 16 groups and has no representation for instants beyond the observed window.

The output head predicts a residual over the last observed group, $\hat{s}_y = s_x[K - 1] + \Delta$, so that the network models change rather than absolute state.

Compact predictor. As an efficiency baseline we train a compact predictor (3.97M parameters, 6.0x fewer) that factorizes space and time: features are projected to 256 dimensions, mixed spatially by depthwise 3 × 3 convolutions on the 14 × 14 grid, then attended temporally as 196 independent sequences of length $K + H$ under the same causal mask, and mixed spatially again before a linear head. It also predicts a residual over the last observed group.

Decoder

The GAN decoder used in the experiments receives a grid of tokens (14 × 14 × 1024) and generates two RGB frames of (224 × 224). Its generator has 6.97 million parameters and includes six residual blocks, a spatial self-attention module, four PixelShuffle magnification stages, and two output heads (21 convolutions in total). During training, it uses a PatchGAN discriminator with 2.76 million parameters, which is dispensed in inference.

Training

Both predictors are trained on cached tokens from a frozen encoder, so the encoder is never run during predictor training. The objective combines cosine and Smooth-L1 losses on predicted tokens, Smooth-L1 and cosine losses on consecutive-group temporal differences initialized with the last observed group, and a Smooth-L1 loss on second-order temporal differences:

For training parameters, after exploring parameters we converged to AdamW, batch size 32 without gradient accumulation, weight decay 1e-4, gradient clipping 1.0, fp16 mixed precision, 3 epochs of linear warm-up followed by cosine annealing to 1e-6, early stopping with patience 10. Peak learning rate 6e-4 for the adapted predictor (30 epochs) and 1.5e-3 for the compact predictor (50 epochs). Both runs stopped early, with the best epoch reached at 53% and 35% of peak learning rate respectively.

Baselines

Three analytic baselines are computed from the context tokens with no training. All are element-wise and treat each of the 196 spatial positions independently. With $g_k$ the $k$th context group and $h$ the horizon:

persistence : $g_{K-1}$

linear : $g_{K-1} + h(g_{K-1} - g_{K-2})$

EMA : $g_{K-1} + h \cdot\sum_i w_i(g_{i+1} - g_i)$, $w_i \propto decay^{(K-2-i)}$

Linear and EMA are the same constant-velocity extrapolation and differ only in how velocity is estimated. Persistence, the last group frozen through horizon, is the reference against which forecasting skill is measured.

Evaluation protocol

Latent metrics. To explore the performance on latent space we used metrics as Per-token cosine similarity between predicted and true future tokens, reported globally and per horizon. Because a driving scene is largely static, we additionally report cosine restricted to the top 10% and top 25% of tokens ranked by $\lVert y - s_x[K - 1]\rVert$, i.e. those that actually moved.

Distributional metrics. Over the evaluation population we report (i) batch retrieval, ranking every sample’s true future against all candidates by cosine and measuring top-1 and top-5 accuracy; (ii) a collapse ratio, the standard deviation of predictions across samples divided by that of

the true futures; and (iii) a Frechet distance between predicted and true token distributions, computed in a shared PCA basis because the sample count is small relative to the token dimension.

Pixel and perceptual metrics. Predicted tokens are decoded to frames by a GAN decoder fine-tuned on the same cached tokens, and compared to ground truth by L1, PSNR and LPIPS [15]. We also decode the TRUE future tokens, giving an oracle that bounds what any predictor can achieve through this decoder. LPIPS uses an AlexNet backbone: the decoder is trained with a VGG-16 perceptual loss, so a VGG-based LPIPS would evaluate it against a close relative of its own training objective.

Skill scores. For an error metric $e$, forecasting skill relative to persistence is $1 - e_{model}/e_{persistence}$. For pixel metrics we additionally report the fraction of the achievable reduction that is captured,

which separates predictor limitations from decoder limitations. This normalization is meaningful only in pixel space; in latent space the oracle error is zero by construction and the expression reduces to the ordinary skill score.

Experimental controls. Evaluation uses a stride that makes consecutive evaluated samples temporally disjoint, so that near-duplicate windows do not enter the retrieval pool.

Results

Forecasting in representation space

Both learned predictors outperform the analytic baselines in cosine similarity and retrieval Table 1. The adapted predictor achieves the highest global cosine similarity (0.8478 versus 0.8421 for the compact model and 0.7814 for persistence), corresponding to 30.4% skill relative to persistence. Its advantage increases on dynamic tokens, reaching 39.3% skill on the top 10%. Retrieval more clearly separates the learned models: the adapted predictor achieves 69.0% top-1 and 84.8% top-5 accuracy, compared with 58.8% and 77.3% for the compact model and 55.3% and 74.2% for persistence. EMA and linear extrapolation underperform persistence across all these measures.

Method	cos	top25 (%)	top10 (%)	retr. @1	retr. @5	std ratio	Frechet
original	0.8478	0.7113	0.6508	0.6904	0.8479	0.984	5.95
compact	0.8421	0.7021	0.6414	0.5882	0.7728	0.934	5.28
persistence	0.7814	0.5309	0.4249	0.5529	0.7418	1.018	2.52
EMA	0.6609	0.3573	0.2399	0.3780	0.5909	1.603	90.12
linear	0.4920	0.2161	0.1202	0.2481	0.4473	3.047	632.27
chance	—	—	—	0.00169	—	—	—

We additionally evaluated within-route temporal discrimination using two-alternative forced choice over 11,348 pairs separated by at least 48 frames. Accuracy was 91.9%, compared with 87.4% for persistence (chance: 50%). Because routes in CARLA Town12 share similar scenery, route restriction does not necessarily provide harder distractors than global retrieval.

Per-horizon behaviour

Retrieval accuracy decays with the horizon for every method, but not at the same rate, as seen in Table 2.

$h$	1	2	3	4	5	6	7	8
original	0.914	0.829	0.739	0.709	0.638	0.602	0.579	0.513
compact	0.848	0.753	0.645	0.604	0.548	0.462	0.443	0.403
persistence	0.858	0.714	0.623	0.563	0.492	0.415	0.399	0.359
ratio o/p	1.07	1.16	1.19	1.26	1.30	1.45	1.45	1.43

At $h = 1$ the adapted predictor barely exceeds trivial continuation; by $h = 6..8$ it exceeds it by roughly 45%. The advantage of prediction grows precisely where continuation fails.

Distributional properties

Also looking to Table 1, the adapted and compact predictors retain population variability close to that of the targets, with standard-deviation ratios of 0.984 and 0.934, respectively, providing evidence against population-level collapse. Their Fréchet distances are also substantially lower than those of EMA and linear extrapolation (5.95 and 5.28 versus 90.12 and 632.27), whose inflated variability indicates unstable extrapolation in feature space. Persistence nevertheless achieves the lowest Fréchet distance (2.52), consistent with reusing genuine encoder features. Thus, the learned predictors improve future-specific accuracy without surpassing persistence in distributional agreement.

Pixel and perceptual evaluation

Decoding true future tokens yields L1 0.0443 and PSNR 23.51 dB, revealing reconstruction error even without forecasting error (Table 3). The adapted predictor reduces L1 relative to persistence by 3.8% globally and 13.5% over dynamic pixels, capturing 25.6% of the oracle-referenced L1 headroom on dynamic pixels. However, LPIPS slightly worsens globally and improves by only 1.7% over dynamic pixels. Thus, latent forecasting gains translate into modest pixel-space improvements, while the relative contributions of forecasting and reconstruction limitations to perceptual quality remain unresolved.

method	L1	top10%	PSNR	LPIPS	LPIPS top10%
original	0.0612	0.1594	21.01	0.2926	0.3784
compact	0.0621	0.1656	20.89	0.2916	0.3802
persistence	0.0636	0.1843	20.63	0.2907	0.3851
EMA	0.0801	0.2116	18.92	0.3386	0.4245
linear	0.1297	0.2303	15.26	0.5157	0.5519
oracle	0.0443	0.0869	23.51	0.2269	0.2711

Capacity

The adapted V-JEPA predictor leads the compact predictor (6.0x fewer parameters) on every paired metric, most clearly on retrieval (0.690 vs 0.588). The gap in plain cosine is small (0.848 vs 0.842), so the discriminative measures separate the two models where the averaged similarity does not. The compact predictor remains attractive when capacity is constrained, like embedded systems.

Discussion

What the representation supports. The evidence that a frozen V-JEPA1 encoder carries usable predictive structure does not rest on beating a baseline. It rests on identifiability: predicted tokens retrieve the correct future among 591 candidates 69.0% of the time, 408x chance. Predictions neither collapse (98.4% of the true spread) nor leave the manifold (Frechet 5.95 against 632 for linear extrapolation). This is a statement about the representation, and it is independent of how strong any particular baseline happens to be.

Where prediction helps. Three independent measures show the same structure: the advantage of a learned predictor over trivial continuation grows with the difficulty of the discrimination. It grows with the horizon (retrieval ratio 1.07 at $h = 1$ to 1.43 at $h = 8$), it is larger on the tokens that actually moved (cosine 0.651 vs 0.425, against 0.848 vs 0.781 overall).

Averaged metrics understate this because driving scenes are dominated by static content, for which repeating the last observation is very nearly correct. Reporting only global cosine or global L1 in this domain will systematically understate the value of prediction.

Persistence is a strong baseline. Persistence reaches 0.553 on batch retrieval. A large fraction of what looks like moment identification is already available in the last observed token, with no prediction at all. The contribution of the predictor is the increment, and we report both figures throughout rather than the model’s value alone.

The gain does not reach pixels. As mentioned in Section 3.4, encoder–decoder pipeline does not fully preserve image detail. Reconstructions from predicted tokens yield additional smoothing, consistent with L1 improvements that do not translate into comparable perceptual gains. Increasing motion-weighted reconstruction loss also failed to improve decoder performance. These findings motivate alternative decoding approaches, but do not establish whether reconstruction limitations or forecasting errors primarily constrain perceptual quality.

Driving-relevant information in predicted features. Predicted features retain semantic information despite limited pixel reconstruction. A nine-class linear probe is trained on real, L2-normalized validation tokens; its weights and validation-fitted standardization remain fixed for route-disjoint testing. mIoU averages six classes with ≥1% validation frequency. On patches undergoing semantic-class changes, predictions achieve 0.16 mIoU, versus 0.13 for persistence and 0.36 for real future tokens.

For class-conditioned cosine retrieval, predicted and real future tokens are mean-pooled over each sample’s own patches, selected by majority CARLA labels in the last context frame. Galleries include samples with ≥8 class patches; 24 ego-hood patches are excluded throughout. Gains over persistence are positive across all six classes at h=8, with 95% route-bootstrap intervals excluding zero. Vehicle-associated gains are the largest observed: +0.22 at h=4 and +0.15 at h=8 (Figure 2), although unequal gallery sizes (176–572 samples) limit comparisons across classes. These findings motivate further investigation of latent forecasting for driving perception.

Limitations. Evaluation is restricted to CARLA Town12 and a 16-frame horizon. Generalization to real-world driving and longer horizons remains untested. The predictor is conditioned on past observations and target positions, but not on future ego actions, limiting its use for evaluating alternative action sequences in planning. Although semantic probing and class-conditioned retrieval indicate driving-relevant information, object-level forecasting and downstream utility remain unverified. Pixel-space conclusions depend on decoder quality, while LPIPS and the decoder’s perceptual training loss share ImageNet-based features, potentially affecting evaluation.

Conclusion

We show that an inpainting-pre-trained V-JEPA1 predictor can be fine-tuned for causal forecasting while keeping the encoder frozen. It retrieves the correct future among 591 candidates in 69.0% of cases, versus 55.3% for persistence, and outperforms the compact baseline in latent cosine similarity and retrieval. Gains over persistence are greater at longer horizons and on tokens exhibiting greater change.

Semantic analysis indicates driving-relevant information in predicted features. A fixed linear probe trained on real encoder tokens achieves 0.16 mIoU on patches undergoing semantic-class changes, versus 0.13 for persistence. Class-conditioned retrieval also improves over persistence across all six evaluated classes at the longest horizon.

These latent gains yield modest L1 improvements but little improvement under LPIPS. Reconstruction error remains even with true future tokens, motivating evaluation through latent and semantic measures alongside pixel metrics and an oracle reconstruction reference.

Results are restricted to simulated driving and a prediction horizon of 16 frames, represented by eight temporal groups. The predictor is not conditioned on future ego actions, while object-level forecasting and downstream utility remain unverified.

These findings motivate further analysis of predicted representations for driving-relevant tasks, including future semantic segmentation and object-motion anticipation, to determine whether latent forecasting gains translate into useful perception beyond persistence.