Task-related Joint-embedding Pre-Training

• Task-related Joint-embedding Pre-Training (TJP) is a representation-learning method that pre-trains models to predict latent targets aligned with downstream task distinctions.
• It leverages strategies like semantically aligned masking and auxiliary regression heads to prevent collapse and preserve critical task-specific information in the embedding space.
• TJP has been applied in domains such as ECG analysis, grasp prediction, and medical imaging, resulting in state-of-the-art metrics by ensuring task-relevant structure in learned representations.

Task-related Joint-embedding Pre-Training (TJP) denotes a family of representation-learning strategies in which a model is pre-trained to predict latent targets in a shared embedding space, while the pretext is explicitly arranged so that the learned representation preserves distinctions needed by a downstream task. In recent JEPA-based work, task-relatedness is introduced by semantically aligned masking, auxiliary regression or classification heads, joint supervised probes during pre-training, or retained latent predictors whose rollouts are later consumed by downstream probes. The common objective is not raw-input reconstruction, but the learning of embeddings that are predictive of clinically relevant morphology, grasp-relevant geometry, control-relevant state distinctions, or other task-defining structure (Weimann et al., 2024, Yu et al., 12 Sep 2025, Waggener et al., 30 Jun 2026, Yang et al., 11 May 2026).

1. Conceptual scope and intellectual lineage

Within the JEPA literature, a joint-embedding predictive architecture is built around a context representation and a target representation that inhabit the same latent space. A predictor then maps context-side information to target-side embeddings. TJP is the task-specialized form of that pattern: the predictive target, the masking policy, the auxiliary signals, or the data source are chosen so that the latent space retains distinctions that matter for downstream inference rather than merely supporting generic invariance or reconstruction (Yu et al., 12 Sep 2025).

This orientation has precursors outside the contemporary JEPA vocabulary. In product recommendation, Content2Vec re-specialized image, text, and collaborative-filtering encoders with a product-pair objective and then fused them late to produce a joint embedding optimized for retrieval-based recommendation (Nedelec et al., 2017). In image–text grounding, IMAGINATOR constructed a word-grounded joint space by combining object–object co-location, word–object co-location, and word–object correlation, explicitly targeting captioning, retrieval, and social-media generation (Krishna et al., 2023). These earlier systems did not use the modern JEPA teacher–student formulation, but they already exhibited the central TJP principle: pre-training or specialization is most effective when the joint embedding is shaped by the semantics downstream tasks actually consume.

Recent JEPA work makes that principle more explicit. In ECG modeling, the pretext is task-related because predicting masked patch features in latent space forces the model to encode P–QRS–T morphology, rhythm structure, and inter-lead relationships that underpin diagnostic labels (Weimann et al., 2024). In model-based reinforcement learning, the formal role of task-relatedness is assigned to an auxiliary function $p$, which anchors the equivalence relation that the latent space must preserve (Yu et al., 12 Sep 2025). In multimodal medical imaging, task-relatedness is operationalized by training supervised auxiliary heads jointly with the JEPA objective during pre-training (Waggener et al., 30 Jun 2026).

2. Canonical objectives and architectural patterns

A common masked-prediction JEPA formulation uses a context encoder $E_\theta$, an EMA target encoder $\bar{E}_\theta$, and a predictor $P_\phi$ conditioned on mask information. In the ECG formulation, the loss is written as

$$\mathcal{L}_{\theta,\phi} = \left\| P_\phi(E_\theta(x), z) - \mathrm{sg}(\bar{E}_\theta(y)) \right\|_1,$$

where $\mathrm{sg}(\cdot)$ denotes stop-gradient and $z$ encodes mask positions via learnable tokens and positional encodings (Weimann et al., 2024). The architectural consequence is that the model predicts latent features for masked targets rather than reconstructing raw waveforms, thereby emphasizing semantically relevant patterns over point-level fidelity.

A second major formulation arises in latent-dynamics JEPA. In the P-JEPA variant studied theoretically for deterministic MDPs, an encoder $E_\phi$ maps observations to latent codes, a latent transition model $T_\psi$ predicts the next latent, and an auxiliary head $P_\theta$ regresses a task-relevant function $E_\theta$0. The total loss is

$E_\theta$1

with

$E_\theta$2

Here task-relatedness is not supplied by masking, but by the auxiliary target that determines which distinctions cannot collapse (Yu et al., 12 Sep 2025).

Other variants modify the same template rather than replacing it. T-JEPA for tabular data predicts the latent representation of one subset of features from the latent representation of a different subset within the same sample, using a Transformer context encoder, an EMA target encoder, and a predictor that consumes context latents together with target-mask tokens (Thimonier et al., 2024). Clin-JEPA co-trains a Qwen3-8B-based encoder and a retained latent trajectory predictor under an $E_\theta$3 teacher-forcing loss plus a short-horizon rollout loss, rather than discarding the predictor after pre-training (Yang et al., 11 May 2026). AEGIS for mammography uses an I-JEPA-style teacher–student setup with contiguous-block masking, a cross-attention predictor, a second predictor pass over CLS tokens, and a JEPA loss instantiated as an MSE between predictor outputs and selected teacher features (Waggener et al., 30 Jun 2026).

Across these systems, three ingredients recur. First, context and target are encoded into a common space. Second, prediction occurs in latent space rather than input space. Third, collapse prevention is handled either by EMA targets and stop-gradient, by contrastive alignment in a shared embedding space, or by the presence of auxiliary or generative terms that constrain the embedding geometry (Weimann et al., 2024, Chen et al., 11 Dec 2025, Huang et al., 11 Sep 2025).

3. How task-relatedness is imposed

Task-relatedness is introduced through the construction of the predictive problem itself. In ECG, contiguous temporal masking with a 75–85% mask ratio forces the model to infer longer-range dependencies such as rhythm and conduction patterns, while preserving amplitude- and morphology-dependent cues that augmentation-heavy invariance-based methods may suppress (Weimann et al., 2024). In RoboPEPP, masks are placed over robot joint locations rather than arbitrary image regions, so the encoder–predictor must infer embeddings for articulated parts from surrounding links and adjacent joints; this makes the representation explicitly relevant to joint-angle and pose estimation (Goswami et al., 2024). In Point-JEPA for grasp prediction, a sequencer imposes an order on point-cloud patches so spatially contiguous context and target blocks can be sampled, aligning the predictive task with object geometry rather than unordered point sets (Guzelkabaagac et al., 13 Sep 2025).

A second route is auxiliary anchoring. The P-JEPA analysis shows that the auxiliary function $E_\theta$4, together with transition dynamics, induces the largest bisimulation $E_\theta$5. Under perfect training, any pair of observations not equivalent under that bisimulation must map to distinct latent codes. The paper states this as

$E_\theta$6

This theorem establishes a precise mechanism by which auxiliary tasks determine which distinctions the representation must preserve (Yu et al., 12 Sep 2025). AEGIS adopts the same general idea empirically: pre-training combines the JEPA objective with supervised auxiliary probes for triage, detection, density, view classification, and implant detection, while SigREG regularization explicitly shapes CLS tokens for downstream probeability (Waggener et al., 30 Jun 2026).

A third route is task-aligned source data and retained inference structure. Point-JEPA is pre-trained on ShapeNet point clouds aligned with the categories used to construct DLR–Hand II, then the context encoder is fine-tuned as the backbone for multi-hypothesis grasp joint-angle prediction (Guzelkabaagac et al., 13 Sep 2025). Clin-JEPA makes task-relatedness even tighter by retaining the latent trajectory predictor at inference time; downstream probes consume both history embeddings and rollout-derived features, so the pre-training objective and the deployment interface are the same object (Yang et al., 11 May 2026). This suggests that in TJP, “task-related” may refer not only to auxiliary labels, but also to whether the pre-training predictor remains operational in the downstream system.

4. Domain-specific instantiations

The breadth of TJP is visible in the diversity of domains where the same design principle has been instantiated.

Domain	Task-related design	Selected reported outcome
ECG JEPA (Weimann et al., 2024)	Latent prediction on masked 12-lead temporal patches without hand-crafted augmentations	PTB-XL all statements macro AUC 0.945
Point-JEPA for grasping (Guzelkabaagac et al., 13 Sep 2025)	ShapeNet-pretrained backbone fine-tuned with a lightweight multi-hypothesis head	At 25% labels, top-logit RMSE 0.246 ± 0.012 vs 0.332 ± 0.002 from scratch
AEGIS mammography (Waggener et al., 30 Jun 2026)	JEPA pre-training plus supervised auxiliary probes for triage, detection, density, view classification, and implant detection	Triage AUC 0.949; density AUC 0.953
Clin-JEPA for EHR (Yang et al., 11 May 2026)	Five-phase co-training of encoder and retained latent trajectory predictor	Mean AUROC 0.883 on 8 binary risk tasks
T-JEPA for tabular data (Thimonier et al., 2024)	Latent prediction between disjoint feature subsets within each sample	DCNv2 + T-JEPA reaches 0.420 RMSE vs 0.502 on raw data

In ECG, JEPA pre-training on 1,011,849 records drawn from ten public ECG databases yields superior representations under both linear evaluation and fine-tuning. The strongest PTB-XL all statements result is macro AUC 0.945 for ViT-S with two-stage fine-tuning, while JEPA linear evaluation reaches 0.940 and the ST-MEM benchmark reaches 0.935 AUC. Notably, pre-training only on the PTB-XL training split remains beneficial, indicating that the latent predictive task itself contributes beyond simple scale effects (Weimann et al., 2024).

In grasp prediction, Point-JEPA provides the backbone for a label-efficient pipeline in which the pooled object embedding is concatenated with a 7D wrist/hand pose and fed to a lightweight MLP that emits $E_\theta$7 joint-angle hypotheses plus logits for top-logit selection. The gains are concentrated in low- and moderate-label regimes: at 1%, 10%, and 25% labels, JEPA pre-training improves top-logit RMSE and reduces the selection gap relative to scratch training, while reaching parity at 100% labels (Guzelkabaagac et al., 13 Sep 2025).

In mammography, AEGIS frames JEPA pre-training as a multi-task clinical representation learner rather than a purely self-supervised backbone. Pre-training uses 71,103 studies from 14 clinical sites, fine-tuning scales resolution progressively to $E_\theta$8, and the largest model attains AUC 0.949 for breast cancer triage with 93% sensitivity and 75% specificity at the optimal operating point. An ensemble with an FDA-cleared baseline improves discrimination to 0.952 AUC, and external zero-shot evaluation on VinDr-Mammo yields triage AUC 0.871 (Waggener et al., 30 Jun 2026).

In EHR modeling, Clin-JEPA extends the paradigm from masked prediction to latent patient-trajectory forecasting. Its five-phase curriculum addresses collapse and online/target drift, enabling a single backbone to support rollout and multi-task probing without per-task backbone fine-tuning. The framework reports mean AUROC 0.851 on ICareFM EEP and 0.883 on eight binary risk tasks, together with convergent 48-hour rollout drift and a clinically discriminative latent geometry (Yang et al., 11 May 2026).

TJP has also expanded into broader multimodal settings. VL-JEPA predicts continuous text embeddings rather than tokens and reports stronger performance than matched token-space training while using about 50% fewer trainable parameters, with the additional property that the learned embedding space directly supports open-vocabulary classification, retrieval, and discriminative VQA (Chen et al., 11 Dec 2025). LLM-JEPA augments standard autoregressive training with a cosine-energy term between embeddings of matched views such as text and code, improving finetuning and pretraining across several model families and datasets (Huang et al., 11 Sep 2025).

5. Representation geometry, collapse prevention, and optimization behavior

The strongest theoretical account of TJP is the “No Unhealthy Representation Collapse” theorem. In deterministic MDPs, if training drives both the latent-transition consistency loss and the auxiliary regression loss to zero, then any pair of observations that do not share the same transition dynamics or auxiliary label must map to distinct latent representations (Yu et al., 12 Sep 2025). The practical implication is that auxiliary tasks are not merely regularizers; they define the equivalence classes that the latent space is allowed to merge.

Empirical studies reinforce that point. In the controlled counting environment used for P-JEPA, dynamics-only training collapses into a single compact cluster, whereas joint training with a reward auxiliary produces nine clear clusters aligned with object counts. A random 256-D auxiliary prevents collapse but preserves many distinctions unrelated to count, showing that auxiliary signals can anchor the wrong geometry if they are misaligned (Yu et al., 12 Sep 2025). This is a central TJP result: collapse prevention and task alignment are coupled, not independent.

Other domains report analogous optimization phenomena. In ECG JEPA, lower pretext loss does not strictly correlate with better downstream linear evaluation over long pre-training; best validation checkpoints often occur early, around 10k–30k steps, and linear evaluation can decline with over-training even when fine-tuning remains robust (Weimann et al., 2024). In T-JEPA for tabular data, regularization tokens are described as critical: without the learned $E_\theta$9 token, optimization often remains in an early collapsed regime, whereas adding one or more regularization tokens enables convergence to non-trivial equilibria. The resulting representations show increased intra-feature variance, decreased inter-feature variance, increased KL divergence, and increased uniformity over training, and the embedding-variance score correlates with supervised feature importance with Kendall’s $\bar{E}_\theta$0 and $\bar{E}_\theta$1-value $\bar{E}_\theta$2 (Thimonier et al., 2024).

Clin-JEPA highlights a different failure mode: even when collapse is avoided, autoregressive rollout can diverge because of online/target mismatch. Its curriculum decomposes stabilization into predictor warmup, co-training, EMA alignment, hard sync, and predictor finalization. Removing warmup drives $\bar{E}_\theta$3 and produces +367% drift accumulation, while removing alignment and hard sync yields +4951% drift. The full system instead achieves mean drift $\bar{E}_\theta$4 over 48 hours and a cohort centroid displacement ratio of $\bar{E}_\theta$5 for deteriorating versus stable patients (Yang et al., 11 May 2026).

6. Limitations, misconceptions, and open problems

A common misconception is that TJP names a single architecture or a fixed loss. The literature instead supports a broader reading. Some systems use EMA teachers and stop-gradient (Weimann et al., 2024); some update encoder, dynamics, and auxiliary head jointly without stop-gradient (Yu et al., 12 Sep 2025); some retain the predictor for inference-time rollout (Yang et al., 11 May 2026); some combine JEPA with auxiliary supervised heads during pre-training (Waggener et al., 30 Jun 2026); and some hybridize embedding-space prediction with autoregressive language modeling rather than replacing token prediction outright (Huang et al., 11 Sep 2025). This suggests that TJP is better understood as a design paradigm than as a single canonical algorithm.

Another misconception is that any auxiliary task improves JEPA representations. The theoretical and experimental record is more restrictive. Trivial or nearly constant auxiliary functions weakly constrain the representation and risk collapse, while misaligned auxiliaries preserve the wrong distinctions (Yu et al., 12 Sep 2025). Similarly, ECG results indicate that high mask ratios and contiguous spans work well, but masking sensitivity remains an open issue, and best downstream checkpoints may occur early during pre-training (Weimann et al., 2024). In mammography, the combined loss weights for auxiliary heads are not reported, and comparative ablations against alternative JEPA variants are absent (Waggener et al., 30 Jun 2026).

Generalization also remains uneven. AEGIS shows zero-shot transfer to VinDr-Mammo, but external AUC is lower than internal AUC because of domain shift and a different reference standard (Waggener et al., 30 Jun 2026). Point-JEPA is evaluated on a synthetic, ShapeNet-aligned dataset and explicitly leaves broader cross-domain evaluation, including real-robot tests, for future work (Guzelkabaagac et al., 13 Sep 2025). Clin-JEPA does not evaluate cross-institution transferability or counterfactual action-conditioned rollouts (Yang et al., 11 May 2026). These limitations indicate that task-relatedness can improve in-domain alignment without eliminating the usual problems of data coverage, reference-standard shift, and deployment mismatch.

Model scale and objective choice are likewise not monotone. In ECG, ViT-B underperforms ViT-S at the current data scale, suggesting that larger capacity may require more data to be fully utilized (Weimann et al., 2024). In AEGIS, model scaling benefits are preserved internally, but only the largest ViT-L significantly outperforms the CNN baseline on triage (Waggener et al., 30 Jun 2026). In language and vision-language settings, embedding-space objectives appear to improve efficiency and multitask transfer, but they depend strongly on the quality of the target embedding space and do not directly optimize long-form generative reasoning (Chen et al., 11 Dec 2025, Huang et al., 11 Sep 2025).

Open problems follow directly from these observations. Theoretical guarantees beyond deterministic MDPs remain incomplete. Stabilizing collapse within bisimilar classes, rather than merely preventing collapse of non-equivalent observations, remains open (Yu et al., 12 Sep 2025). Automated selection of task-related masks, auxiliary targets, and curriculum schedules is still largely heuristic. More broadly, the literature suggests that the key question for TJP is not whether latent prediction should replace reconstruction or contrastive learning in the abstract, but how the predictive target, the conditioning signal, and the deployment interface can be made to share the same task-defining structure.