D-JEPA: Design-Recoverable JEPA Representation with Swappable Physics Decoders

Nitin Nagesh Kulkarni; Aashwin Anand Mishra; Yin Yu; Peter Lyu

cs.LGSep 27, 2026

D-JEPA: Design-Recoverable JEPA Representation with Swappable Physics Decoders

Authors: Nitin Nagesh Kulkarni, Aashwin Anand Mishra, Yin Yu, Peter Lyu

Organizations: Luminary San Mateo, CA, USA

Abstract

Joint-Embedding Predictive Architectures (JEPAs) provide a framework for learning compact representations without directly reconstructing high-dimensional observations. However, in parameterized physical systems, learned representations can entangle geometry with operating conditions and task-specific physical responses, limiting their reuse across prediction tasks. We introduce D-JEPA (Design-recoverable JEPA), a geometry-centric JEPA that computes a compact representation from geometry alone and reuses it across operating conditions and physical response spaces through lightweight physics-specific decoders. An explicit design-recoverability objective encourages the geometry latent to preserve information about the underlying design variables, enabling the representation to support design analysis and optimization. We further identify a case-level collapse failure mode in which target representations become nearly invariant across distinct geometries despite low reconstruction error, and mitigate it using case-level variation constraints and auxiliary target reconstruction. Across four 3D aerodynamic, hydrodynamic, and structural benchmarks, D-JEPA maintains or improves full-field prediction accuracy while achieving near-perfect linear recoverability of design parameters. The frozen geometry representation can be reused at held-out operating conditions and transferred to a structural response task with fewer trainable parameters. Finally, the representation supports differentiable design optimization, with designs validated using high-fidelity CFD, preserving the predicted ranking of candidate designs. These results demonstrate that separating a reusable geometry representation from physics-specific prediction provides a practical representation for scientific surrogate modeling and design.

Figures & tables

Figure 1: Overview of D-JEPA: (a) D-JEPA architecture and training pathways. (b) Design recoverability, condition reuse and cross-physics transfer.

Model	Dataset	Pressure	WSS x	WSS y	WSS z	$C_{D}$	$C_{L}$
DoMINO	SHIFT-Wing	0.9678 $\pm$ 0.001	0.8972 $\pm$ 0.001	0.9456 $\pm$ 0.004	0.9654 $\pm$ 0.002	0.940 $\pm$ 0.002	0.891 $\pm$ 0.001
	SHIFT-Submarine	0.9321 $\pm$ 0.000	0.9332 $\pm$ 0.000	0.9421 $\pm$ 0.003	0.9243 $\pm$ 0.002	0.973 $\pm$ 0.001	–
	SHIFT-SUV	0.9765 $\pm$ 0.000	0.9475 $\pm$ 0.000	0.9241 $\pm$ 0.002	0.9342 $\pm$ 0.002	0.980 $\pm$ 0.003	–
GeoTransolver	SHIFT-Wing	0.9950 $\pm$ 0.001	0.9350 $\pm$ 0.001	0.9645 $\pm$ 0.002	0.9899 $\pm$ 0.001	0.960 $\pm$ 0.002	0.956 $\pm$ 0.001
	SHIFT-Submarine	0.9921 $\pm$ 0.001	0.9432 $\pm$ 0.001	0.9621 $\pm$ 0.001	0.9743 $\pm$ 0.002	0.990 $\pm$ 0.002	–
	SHIFT-SUV	0.9931 $\pm$ 0.001	0.9698 $\pm$ 0.001	0.9641 $\pm$ 0.001	0.9542 $\pm$ 0.001	0.996 $\pm$ 0.002	–

Table 1: Surrogate accuracy across parameterized 3D aerodynamic benchmarks. Full-field

R^{2}

is reported for surface pressure and wall-shear-stress components, together with integrated force coefficients

C_{D}

and

C_{L}

. Values are mean

\pm

std over three seeds; bold indicates the best value for each dataset and quantity.

Dataset	None	Token SIGReg	Case SIGReg	+ Variance	+ Target Recon
SHIFT-Wing	0.0003	0.0003	0.6211	1.0346	1.0317
SHIFT-Submarine	0.0004	0.0005	0.7566	1.0633	1.0620
SHIFT-SUV	0.0002	0.0002	0.0003	0.9892	0.9889

Table 2: Case-level variation in the target representation under progressively stronger representation constraints. Values report

\mathrm{std}_{\mathrm{case}}(\mathrm{mean}_{m}Z_{t})

, measuring variation across geometry cases.

	Design R 2 $\uparrow$				Pressure R 2 $\uparrow$
Dataset	Full	$-\mathcal{L}_{\mathrm{lat}}$	Learned-only	$-\mathcal{L}_{\mathrm{dis}}$	Full	$-\mathcal{L}_{\mathrm{lat}}$	Learned-only	$-\mathcal{L}_{\mathrm{dis}}$
SHIFT-Wing	0.998	0.993	0.976	0.722	0.9971	0.9912	0.9885	0.9823
SHIFT-Submarine	0.998	0.982	0.757	0.564	0.9967	0.9896	0.9852	0.9872
SHIFT-SUV	0.989	0.984	0.976	0.622	0.9937	0.9851	0.9801	0.9821

Table 3: Component ablation of D-JEPA across the SHIFT datasets, showing the contributions of latent prediction and design supervision to design recoverability and pressure prediction. Learned-only removes the parameter free global and regional descriptors.

	Trained operating conditions										Held-out
Mach	0.5000	0.6000	0.7000	0.7500	0.8000	0.8250	0.8375	0.8500	0.8750	0.9000	0.7750	0.8125
Pressure $R^{2}$	0.9902	0.9935	0.9948	0.9944	0.9941	0.9956	0.9960	0.9967	0.9952	0.9939	0.9922	0.9944
Std	0.001	0.002	0.001	0.003	0.002	0.001	0.001	0.002	0.001	0.002	0.001	0.002

Table 4: SHIFT-Wing pressure prediction accuracy (

R^{2}

) across trained and held-out Mach numbers using a single frozen geometry representation

g

. The held-out Mach values are interpolation points within the trained Mach range. Bold values denote held-out conditions.

Training epoch budget	Frozen geometry encoder (0.56 M params)	From scratch (2.30 M params)
10	0.3316 $\pm$ 0.020	0.0396 $\pm$ 0.007
100	0.8901 $\pm$ 0.013	0.8740 $\pm$ 0.034
200	0.9327 $\pm$ 0.002	0.9297 $\pm$ 0.012
300	0.9385 $\pm$ 0.001	0.9313 $\pm$ 0.002

Table 5: Structural displacement prediction accuracy (

R^{2}

) on AutoHood 3D comparing frozen geometry encoder transfer versus training from scratch.

Benchmark	Design candidate	Predicted [N]	CFD [N]	Relative error	Predicted $\to$ CFD rank
SHIFT-Submarine (drag)	OPT-wide	645.8	637.69	1.3%	1 $\to$ 1
	OPT-tau	677.0	688.36	1.7%	2 $\to$ 2
	MID	718.6	723.08	0.6%	3 $\to$ 3
	ANTI	758.7	773.75	2.0%	4 $\to$ 4
SHIFT-Wing (lift)	OPT-Wide	3,426,080	3,193,866	7.27%	1 $\to$ 1
	OPT-Tau	3,108,536	3,086,237	0.72%	2 $\to$ 2

Table 6: High-fidelity CFD validation of surrogate-optimized designs. Integrated drag is reported for SHIFT-Submarine and integrated lift for SHIFT-Wing. Predicted and CFD rankings agree for all evaluated designs.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Cases	Design parameters	Operating conditions	Surface fields
SHIFT-Wing	1,966	7	angle of attack, Mach	pressure, WSS x , WSS y , WSS z
SHIFT-Submarine	4,000	7	fixed	pressure, WSS x , WSS y , WSS z
SHIFT-SUV	1,996	7	fixed	pressure, WSS x , WSS y , WSS z
AutoHood 3D	12,000	–	load case	pressure, displacement

Appendix

Table A.1: Benchmark dataset properties, parameter dimensions, operating conditions and surface targets. A dash indicates that the quantity is not applicable.

Loss term	Latent	Recon.	Design	Token	Case	Var.	Tgt. rec.
	( $\mathcal{L}_{\mathrm{lat}}$ )	( $\lambda_{r}\mathcal{L}_{\mathrm{rec}}$ )	( $\lambda_{d}\mathcal{L}_{\mathrm{dis}}$ )	( $\lambda_{s}\mathcal{L}_{\mathrm{sig}}$ )	( $\lambda_{c}\mathcal{L}_{\mathrm{sig\text{-}case}}$ )	( $\lambda_{v}\mathcal{L}_{\mathrm{var}}$ )	( $\lambda_{t}\mathcal{L}_{\mathrm{rec\text{-}t}}$ )
Weight	1.0	1.0	1.0	0.01	1.0	1.0	1.0

Appendix

Table A.2: Hyperparameter weights used for the composite D-JEPA objective.

Setting	Value
Optimizer	AdamW ( Loshchilov and Hutter, 2019 )
Base LR	$1\times 10^{-3}$
Weight decay	$1\times 10^{-3}$
LR schedule	Cosine decay
Gradient clipping	10.0
Epochs	400

Appendix

Table A.3: Optimization configuration and training hyperparameters.

Dataset	Configuration	$\sigma_{\mathrm{case}}$	$L_{\mathrm{lat}}$	$L_{\mathrm{rec}}$	Pressure $R^{2}$	Design $R^{2}$
SHIFT-Wing	None	0.0003	0.0000	0.0354	0.9906	0.992
	Token SIGReg	0.0003	0.0000	0.0352	0.9906	0.992
	Case SIGReg	0.6211	0.0403	0.0373	0.9893	0.987
	+ Variance	1.0346	0.0498	0.0381	0.9886	0.986
	+ Target Recon.	1.0317	0.0447	0.0379	0.9971	0.998
SHIFT-Submarine	None	0.0004	0.0000	0.0348	0.9896	0.982

Appendix

Table A.4: Complete representation-collapse diagnostics. The case-level standard deviation measures between-case variation, while

L_{\mathrm{lat}}

and

L_{\mathrm{rec}}

report latent prediction and physical reconstruction losses. Pressure

R^{2}

and design-recovery

R^{2}

quantify downstream performance.

Figure A.1: Qualitative surface pressure prediction for SHIFT-SUV. Predicted pressure and absolute point-wise error are shown for representative median and worst-case held-out geometries.

Figure A.2: Qualitative surface pressure prediction for SHIFT-Submarine. Predicted pressure and absolute point-wise error are shown for representative median and worst-case held-out geometries.

Figure A.3: Qualitative surface pressure prediction for SHIFT-Wing. Predicted pressure and absolute point-wise error are shown for representative median and worst-case held-out geometries.

SHIFT-SUV		SHIFT-Wing		SHIFT-Submarine
Parameter	$R^{2}$	Parameter	$R^{2}$	Parameter	$R^{2}$
hood_height	0.995	AR (aspect ratio)	0.999	cross_section_shape	1.000
hood_base_height	0.987	Lambda_c_4 (quarter-chord sweep)	0.998	fore_body_shape	0.996
windshield_angle	0.986	twist_20	0.997	fore_length	0.999
vehicle_height	0.991	twist_k	0.997	mid_length	0.999
backlight_angle	0.990	twist_60	0.999	aft_tail_half_angle	0.997

Appendix

Table A.5: Linear-probe recovery accuracy (

R^{2}

) for individual geometric design parameters from the frozen 128-dimensional geometry representation

g

Figure A.4: Design-variable probe recoverability across benchmark datasets. True design parameters are compared with values recovered from the frozen geometry representation using linear probes.

Figure A.5: Operating-condition generalization across Mach numbers on SHIFT-Wing. Surface pressure $R^{2}$ is evaluated across trained and held-out Mach conditions using the same geometry representation.

Execution pipeline stage	Compute hardware	Wall-clock time
D-JEPA/AeroJEPA training (400 epochs)	1 $\times$ NVIDIA T4 (16 GB)	4–8 hours / dataset
Decoder transfer fine-tuning (400 epochs)	1 $\times$ NVIDIA T4 (16 GB)	$\sim$ 2 hours
GeoTransolver baseline (500 epochs)	8 $\times$ NVIDIA H100 (80 GB)	$\sim$ 12 hours / dataset
Mach-condition fine-tuning (200 epochs)	1 $\times$ NVIDIA T4 (16 GB)	$\sim$ 3 hours

Appendix

Table A.6: Primary computational hardware configurations and approximate wall-clock runtimes across pipeline stages.

Predictive Representation Learning

Predictive representation learning trains models to encode useful structure by predicting unseen or future representations from context.

See all 176 papers

Sep 28, 2026

PDE-JEPA: Predictive Representation Learning of Latent Dynamics Modeling for Parametric PDEs

Physical trajectories contain more than snapshots of a system: they also reveal how its states evolve under governing conditions. However, representation learning for parametric partial differential equations (PDEs) has largely relied on reconstruction-based objectives that emphasize recovering observed physical fields. In this paper, we investigate predictive representation pretraining as an alternative to reconstruction-based learning. We find that predictive representations preserve rich physical information, yet this advantage alone does not ensure accurate field evolution. Based on these observations, we introduce PDE-JEPA for parametric PDE dynamics. Specifically, we first train an encoder using a masked-latent prediction to capture the underlying regularities of PDE dynamics. To explicitly adapt the pretrained representation toward a more dynamics-aligned state space, we then introduce a geometry projector that aligns latent trajectory geometry with the evolution geometry of physical fields. Finally, building on this geometry-aligned latent space, we further develop a physics-structured latent predictor that decomposes the dynamics into parameter-independent evolution and parameter-dependent response components. Extensive experiments on nine widely used PDE benchmarks demonstrate that our framework outperforms existing state-of-the-art methods by an average of 33.4% in-distribution, while achieving an average improvement of 51.4% when extrapolating to unseen governing parameters. The project page is available here.

Zhentao Tan, Jianrong Zhang, Ruijie Quan +1

Zhejiang University.

Latent Dynamics Modeling Joint-Embedding Predictive Architecture PDE Surrogate Modeling Predictive Representation Learning Physics-Informed ML

Ask PDF

May 7, 2026

AeroJEPA: Learning Semantic Latent Representations for Scalable 3D Aerodynamic Field Modeling

High-fidelity CFD is essential for aerodynamic design, but repeated simulations are computationally expensive, motivating surrogate models for rapid evaluation across geometries and operating conditions. Most existing surrogates are designed for direct field regression, requiring the evaluation of millions of field points even when only an aerodynamic quantity or a localized region is needed, while their internal representations are not intended for direct use in downstream tasks. We introduce AeroJEPA, a framework inspired by joint-embedding predictive architectures that represents the problem in two distinct latent spaces: context tokens encode geometry, while predicted tokens encode the aerodynamic state. Both representations remain directly accessible for downstream tasks, such as linear readouts of design variables and aerodynamic quantities without decoding and integrating the full field. When spatial detail is needed, a continuous implicit decoder evaluates the field only at the requested coordinates while reusing the encoded geometry. We evaluate AeroJEPA on HiLiftAeroML, with multi-million-point fields, and SuperWing, which spans a broad family of transonic wings. Compared with state-of-the-art direct-regression surrogates, AeroJEPA trades peak full-field accuracy for compact, reusable representations. In our selective-decoding experiment, however, AeroJEPA substantially outperforms the evaluated direct-regression surrogates while avoiding predictions over the remainder of the aircraft. The learned representations further support controlled interpolation, concept-vector arithmetic, and preliminary constrained latent-space optimization. These results show how predictive representations can support aerodynamic analysis with or without full-field reconstruction.

Francisco Giral, Abhijeet Vishwasrao, Andrea Arroyo Ramo +8

Universidad Politécnica de Madrid · University of Michigan · Universitat Politècnica de València +3

Joint-Embedding Predictive Architecture PDE Surrogate Modeling Predictive Representation Learning Computational Fluid Dynamics

Ask PDF

Sep 9, 2026

Semigroup-JEPA: Latent Dynamics Consistency for Zero-Shot Physics Generalization

Joint-Embedding Predictive Architecture (JEPA) world models learn a compact latent representation of the world that supports prediction and planning, but their capability to learn physics and generate physically realistic dynamics remains hitherto untested. In this work, we introduce SemiGroup-JEPA (SG-JEPA), which extends the LeWorldModel framework by supplying the parameter governing the physics to the temporal model via action-conditioning and jointly training an encoder and predictor through an autoregressive latent rollout. To evaluate the model's ability to generalize out of distribution, we design dynamical tasks under different gravitational fields that, despite obeying the same physical law, exhibit qualitatively different dynamics, ranging from floating motion in weak gravitational fields to rapid bouncing in strong ones. In contrast to DINO-WM, SG-JEPA reduces open-loop prediction error by up to 2 times on two-dimensional datasets, and increases control success rate up to 2.5 times for three-dimensional robotic datasets, for which we train independent diffusion policies. To explain this advantage, we develop a linear feature model that separates local law-conditioned error from its recursive amplification under rollout. Guided by this model, we find that back-propagating the multi-step rollout loss into the representation trains the encoder to keep the features that the predictor can carry forward, and that those are the features the dynamics depend on, so most of the gain comes from the encoder learning better features rather than from the predictor learning better dynamics. See project page at https://sg-jepa.github.io.

Andy Zeyi Liu, Haoran Sun, Lucas Baker +2

Yale University · Jump Trading · Brown University

World Models for Robotics Latent Dynamics Modeling Joint-Embedding Predictive Architecture

Ask PDF

Abstract

Predictive Representation Learning

Explore similar work

PDE-JEPA: Predictive Representation Learning of Latent Dynamics Modeling for Parametric PDEs

AeroJEPA: Learning Semantic Latent Representations for Scalable 3D Aerodynamic Field Modeling

Semigroup-JEPA: Latent Dynamics Consistency for Zero-Shot Physics Generalization