Disentangled Representation Learning for Parametric Partial Differential Equations

TL;DR

DisentangO addresses the lack of physical interpretability in neural operators for parametric PDEs by learning disentangled latent factors from the operator parameters themselves. It combines a variational HVAE with a multi-task, meta-learned neural operator backbone to perform forward PDE solving and inverse physics discovery, with theoretical identifiability guarantees for the latent factors. Empirically, it demonstrates (i) competitive forward performance and significantly improved inverse interpretability in supervised settings, (ii) effective latent disentanglement and digit-aware representation in semi-supervised Mechanical MNIST, and (iii) interpretable latent factors governing microstructure in unsupervised heterogeneous materials, with latent traversals revealing meaningful physical controls. The framework promisingly bridges predictive accuracy and physical understanding, enabling robust generalization across diverse PDE systems and supervision regimes, while highlighting the role of encoding parameters in lifting layers for identifiability and scalability.

Abstract

Neural operators (NOs) excel at learning mappings between function spaces, serving as efficient forward solution approximators for PDE-governed systems. However, as black-box solvers, they offer limited insight into the underlying physical mechanism, due to the lack of interpretable representations of the physical parameters that drive the system. To tackle this challenge, we propose a new paradigm for learning disentangled representations from NO parameters, thereby effectively solving an inverse problem. Specifically, we introduce DisentangO, a novel hyper-neural operator architecture designed to unveil and disentangle latent physical factors of variation embedded within the black-box neural operator parameters. At the core of DisentangO is a multi-task NO architecture that distills the varying parameters of the governing PDE through a task-wise adaptive layer, alongside a variational autoencoder that disentangles these variations into identifiable latent factors. By learning these disentangled representations, DisentangO not only enhances physical interpretability but also enables more robust generalization across diverse systems. Empirical evaluations across supervised, semi-supervised, and unsupervised learning contexts show that DisentangO effectively extracts meaningful and interpretable latent features, bridging the gap between predictive performance and physical understanding in neural operator frameworks.

Figures (12)

• Figure 1: Overview of the DisentangO architecture. Each task correspond to a different (hidden) PDE parameter $\bm{b}$. For illustration, the same input function $f^\eta_i$ is shown for multiple tasks to highlight that different parameter fields $b^{\eta}$ can produce different output functions $u_i^{\eta}$ under identical input $f$; in practice, $f$ may vary across tasks. The task-specific lifting parameters $\theta_P^{\eta}$ are encoded and reconstructed through a VAE, and the reconstructed parameters $\hat{\theta}_P^{\eta}$ are fed into the iterative Fourier layers to form the task-specific neural operator $G^{\eta}$. Loss components are overlaid to indicate where each term in the objective $L_{loss}$ is computed.
• Figure 2: MMNIST scatterplot with DisentangO-2 and $\beta_{d}=1$: left: $(\beta_{kl}=1, \beta_{cls}=100, \text{data error } 18.81\%)$, middle: $(\beta_{kl}=10, \beta_{cls}=10, \text{data error } 16.94\%)$, right: fully unsupervised DisentangO-2 without classification loss $(\beta_{kl}=100, \text{data error } 12.65\%)$.
• Figure 3: MMNIST unsupervised scores against $\beta_{d}$ with DisentangO-2 (left) and DisentangO-15 (right). By comparing $\beta_{d}=1$ (solid lines) with $\beta_{d}=100$ (dashed lines), increasing $\beta_{d}$ forces the latent factors to maximize the contained information and in turn decreases MI, thus encouraging disentanglement.
• Figure 4: Unsupervised MI score against $\beta_{kl}$ with DisentangO-2 in heterogeneous material learning.
• Figure 5: Latent traversal of DisentangO-3 in unsupervised heterogeneous material learning, where the three latent dimensions control the border rotation between the two segments (top), the relative fiber orientation between the two segments (middle), and the fiber orientation of the top segment (bottom), respectively. Legend indicates fiber orientation ranging from 0 to $\pi$.

Theorems & Definitions (2)

• Theorem 1
• Theorem 2