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

PDE-JEPA: Predictive Representation Learning of Latent Dynamics Modeling for Parametric PDEs
Sep 28, 2026Zhentao Tan, Jianrong Zhang, Ruijie Quan +1Latent Dynamics Modeling Joint-Embedding Predictive Architecture
Zhejiang University.
AeroJEPA: Learning Semantic Latent Representations for Scalable 3D Aerodynamic Field Modeling
May 7, 2026Francisco Giral, Abhijeet Vishwasrao, Andrea Arroyo Ramo +8Joint-Embedding Predictive Architecture PDE Surrogate Modeling
Universidad Politécnica de Madrid · University of Michigan · Universitat Politècnica de València +3
Semigroup-JEPA: Latent Dynamics Consistency for Zero-Shot Physics Generalization
Sep 9, 2026Andy Zeyi Liu, Haoran Sun, Lucas Baker +2World Models for Robotics Latent Dynamics Modeling
Yale University · Jump Trading · Brown University

Sep 28, 2026J/K move · Enter open · S save

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

Zhentao Tan, Jianrong Zhang, Ruijie Quan, Yi Yang

Zhejiang University.

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

Open paper AskPDF arXiv

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.

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