Thinkjepa Dual Temporal World Model

What This Skill Does

Enhance JEPA-based world models for egocentric trajectory prediction by swapping a single-pathway dense sampler with a dual-temporal architecture: a dense JEPA branch handling frame-by-frame dynamics, and a sparse VLM-guided branch providing semantic grounding across longer windows.

The Component Swap

Old component: Single JEPA pathway with dense frame sampling from short observation windows.

# Traditional JEPA: all dynamics from dense consecutive frames
class TraditionalJEPA(nn.Module):
    def forward(self, frames):
        # frames shape: [batch, T_dense, C, H, W]
        # All temporal reasoning happens in latent space
        z = self.encoder(frames)
        pred_z = self.predictor(z)
        return self.decoder(pred_z)

New component: Dual-temporal pathway with complementary sampling strategies and hierarchical VLM feature aggregation.

# ThinkJEPA: dual pathways with multi-layer VLM features
class ThinkJEPA(nn.Module):
    def __init__(self, encoder, predictor, decoder, vlm_extractor):
        super().__init__()
        self.encoder = encoder
        self.predictor = predictor
        self.decoder = decoder
        self.vlm_extractor = vlm_extractor  # Extracts multi-depth features

    def forward(self, dense_frames, sparse_frames, sparse_images):
        # Dense branch: all frames for fine-grained dynamics
        z_dense = self.encoder(dense_frames)

        # VLM branch: semantic guidance from intermediate layers
        vlm_features = []
        for layer_idx in [12, 18, 24]:  # Multi-depth extraction
            feat = self.vlm_extractor(sparse_images, layer=layer_idx)
            vlm_features.append(feat)

        vlm_context = torch.cat(vlm_features, dim=-1)

        # Fuse: predictor conditions on both dense dynamics and sparse semantics
        pred_z = self.predictor(z_dense, vlm_context)
        return self.decoder(pred_z)

The key architectural shift is replacing single-branch processing with a dual-stream design where the predictor receives both fine-grained pixel-space dynamics (from dense JEPA) and abstracted semantic guidance (from hierarchical VLM features), preventing loss of low-level control while anchoring predictions in semantic coherence.

Performance Impact

EgoDex egocentric trajectory prediction:

• ADE (Average Displacement Error): 0.071 → 0.061 (-14% error)
• FDE (Final Displacement Error): 0.066 → 0.056 (-15% error)
• Accuracy (at threshold): 0.471 → 0.596 (+27% absolute improvement)

Long-horizon behavior: Improvement increases with prediction horizon, indicating that the dual-pathway design mitigates error accumulation better than single-branch JEPA.

When to Use

• Egocentric hand-object trajectory prediction tasks
• When cached VLM features are available (Qwen3-VL or similar)
• When long-horizon predictions are critical and semantic context helps reduce drift
• Tasks requiring both fine-grained dynamics and high-level task understanding

When NOT to Use

• If only short-horizon predictions are needed (single JEPA may be sufficient)
• When VLM embeddings are not cached (overhead not justified for real-time systems)
• On non-egocentric datasets where semantic guidance may not transfer
• When computational budget precludes multi-pathway processing

Implementation Checklist

1. Data preparation:

• Ensure dense frame sequences are available for the JEPA encoder path
• Pre-compute and cache VLM features (Qwen3-VL) for sparse keyframes at layers [12, 18, 24]
• Verify frame alignment: dense frames should cover the same time window as sparse samples

2. Architecture integration:

# Minimal swap: replace your existing JEPA forward with:
model = ThinkJEPA(
    encoder=your_jepa_encoder,
    predictor=MultiPathPredictor(hidden_dim=256),  # Handles both inputs
    decoder=your_jepa_decoder,
    vlm_extractor=VLMFeatureExtractor(model_name='qwen3-vl')
)

3. Verification:

• Measure ADE/FDE on your validation set
• Check that VLM features improve long-horizon accuracy more than short-horizon
• Ablate: remove VLM branch to confirm contribution (~14-27% delta)

4. Hyperparameter tuning if needed:

• VLM layer selection: [12, 18, 24] works for Qwen3-VL; adjust for other models
• Feature fusion: try concatenation, attention-weighted combination, or gating
• Sampling ratio: sparse/dense ratio of 1:4 worked for EgoDex; adjust for your dataset

5. Known issues:

• VLM features may contain language-specific biases from training on vision-language tasks; validate on your domain
• Very short sequences (<4 frames dense): dual pathway overhead not justified
• GPU memory: caching multi-layer VLM features requires ~2-3× normal JEPA memory

Related Work

This builds on JEPA (Assran et al.) for masked future prediction and extends hierarchical vision-language fusion patterns from multi-modal transformers. The dual-pathway design parallels System 1/System 2 reasoning in cognitive science, where fast dynamics (JEPA) and slow semantic processing (VLM) jointly produce robust predictions.