Tadpole: Autoencoders as Foundation Models for 3D PDEs with Online Learning

Abstract: We introduce Tadpole, a novel foundation model for three-dimensional partial differential equations (PDEs) that addresses key challenges in transferability, scalability to high dimensionality, and multi-functionality. Tadpole is pre-trained as an autoencoder on synthetic 3D PDE data generated by an efficient online data-generation framework. This enables large-scale, diverse training without storage or I/O overhead, demonstrated by scaling to an equivalent of hundreds of terabytes of training data. By autoencoding single-channel spatial crops, Tadpole learns rich and transferable representations across heterogeneous physical systems with varying numbers of state variables and spatial resolutions. Although pre-trained solely as an autoencoder, Tadpole can be efficiently applied for multiple downstream tasks beyond reconstruction, including dynamics learning and generative modeling. For dynamics learning, we propose a novel parameter-efficient fine-tuning strategy that integrates low-rank adaptation, latent-space transformations, and reintroduced skip connections, achieving accurate temporal modeling with a minimal number of trainable parameters. Tadpole demonstrates strong fine-tuning performance across various downstream tasks, highlighting its versatility and effectiveness as a foundation model for 3D PDE learning. Source code and pre-trained weights of Tadpole are available at https://github.com/tum-pbs/tadpole

• The paper introduces an autoencoder-based method that leverages online data generation to overcome I/O bottlenecks in 3D PDE simulations.
• It employs a VAE with a patch discriminator and a translation-equivariant backbone to capture the latent solution manifold across diverse PDEs.
• Parameter-efficient fine-tuning via the Tadpole-DFT scheme markedly reduces trainable parameters while enhancing dynamics prediction and generative modeling.

Tadpole: Autoencoders as Foundation Models for 3D PDEs with Online Learning

Foundation Model Design for 3D PDEs

Tadpole introduces a new approach to scientific foundation models in the context of three-dimensional PDEs, targeting limitations in data scalability, transferability, and multi-functionality that persist in existing approaches. The core methodological innovation is the use of a variational autoencoder (VAE)-style architecture, trained on diverse, synthetic 3D PDE data generated via an efficient online pipeline using GPU-based pseudo-spectral solvers. This circumvents the prohibitive I/O/storage bottlenecks typically encountered with large 3D simulation datasets and enables procedural exposure to orders of magnitude more data diversity.

The Tadpole training regime focuses solely on reconstructing single-channel crops, rather than learning the dynamics directly. This explicitly structures the latent space to capture the full solution manifold, providing representations agnostic to the underlying equation, number of state variables, or spatial resolution (Figure 1).

Figure 1: Tadpole autoencoder pre-training on online-generated single-channel spatial crops, and its application to diverse downstream tasks via the learned latent space.

Online Learning Framework and Dataset Generation

To generate the training data, Tadpole leverages a procedural, on-the-fly data generation strategy: spectral solvers run on GPUs simulate short spatiotemporal trajectories for a broad set of semi-linear 3D PDEs, initialized from a maximally diverse range of initial conditions. A novel multi-stage buffering and asynchronous loading mechanism eliminates simulation–training throughput mismatches, allowing continuous streaming of simulation states to the model. The entire setup bypasses local storage and the simulation–I/O bottleneck, permitting scaling to the equivalent of hundreds of TB of data.

To address variable numbers of state variables across systems, single-channel crops are extracted by folding the channel dimension into the batch axis. This also facilitates batching across systems with arbitrary spatial dimensions and ensures maximal sample diversity. The approach is fundamentally distinct from data-scarcity-limited, trajectory-centric pre-training paradigms used in existing PDE foundation models, which have been restricted to mostly 1D/2D settings due to prohibitive storage and I/O when moving to large 3D datasets.

Pretraining Objective and Architectural Considerations

Unlike prior foundation models, Tadpole is not pre-trained on trajectory-based dynamics prediction tasks, but as a variational autoencoder with a patch discriminator to encourage sharp reconstructions in the output space. The underpinning hypothesis is that learning the solution manifold itself is substantially more tractable than learning the induced tangent and nearby flow—i.e., the dynamics—especially in the context of the highly variable, nonlinear nature of general PDEs.

The choice of a hybrid translation-equivariant backbone (based on P3D) enables flexible inference on arbitrary crop sizes and spatial resolutions, and allows efficient tile-wise embedding in high-resolution downstream settings. Notably, all skip connections are removed during pre-training to guarantee that the latent code contains all reconstructive information, facilitating robust transfer. The design is not tightly coupled to the P3D backbone and could be integrated with other transformer or convolution-based 3D latent variable models.

Parameter-Efficient Fine-tuning and Multi-task Adaptation

For downstream task adaptation, the authors introduce the Tadpole-DFT parameter-efficient fine-tuning scheme. Here, the pre-trained encoder–decoder is augmented with:

• A latent-space transformer sub-network aggregating variable interdependencies,
• Re-introduced and learnable skip connections (initialized to zero) between encoder and decoder,
• Low-Rank Adaptation (LoRA)-based PEFT applied only to a small subset of backbone parameters.

This combined mechanism enables adaptation to temporal dynamics tasks in the latent space with orders of magnitude fewer trainable parameters, preserves pretrained knowledge, and stabilizes learning relative to full-parameter fine-tuning.

Due to the generality of the latent code, Tadpole is readily adaptable for direct autoencoding across novel PDE domains as a zero-shot model, and more generally supports downstream generative modeling by learning a flow-matching model in latent space, then decoding synthetic samples in the high-dimensional output domain.

Empirical Evaluation

Autoencoding Tasks

Tadpole is benchmarked in both zero-shot and fine-tuned settings on four 3D systems: isotropic turbulence (Iso, $1024^3$), turbulent channel flow (TCF), magnetohydrodynamics (MHD), and transitional boundary layer (TBL), all at substantially higher spatial resolutions than seen during pretraining. The model demonstrates a robust scaling law: larger Tadpole variants consistently achieve lower normalized RMSEs, and higher diversity in pretraining (both in PDEs and initial conditions) directly translates to superior zero-shot accuracy.

Figure 2: Tadpole zero-shot and fine-tuned autoencoding performance across model size and pretraining diversity.

Figure 3: Zero-shot reconstructions by Tadpole-B on various high-resolution PDE datasets (velocity channel slices shown).

Fine-tuning, whether full or LoRA-based, further reduces reconstruction errors, with LoRA-32 adapters already yielding over 60% reduction in RMSE relative to from-scratch training, but with $<$10% of trainable parameters required.

Dynamics Learning

The adaptation of Tadpole using DFT (latent-space temporal modeling plus skip connections and LoRA) is evaluated on turbulent 3D flow dynamics—specifically, cropped TBL and Iso test cases, comparing against state-of-the-art 3D foundation models (Walrus, DPOT, MORPH) and non-foundational baselines (e.g., Swin3D, AViT).

Tadpole outperforms all comparably sized foundation models, and, notably, in TBL achieves better long-horizon enstrophy spectrum error than Walrus, despite a $>100\times$ reduction in parameter count (Figure 4). Ablations show that LoRA rank strongly determines adaptation quality, far more than sub-network capacity.

Figure 4: Downstream dynamics learning performance (NRMSE$^{ES}$) and trainable parameter count; Tadpole-DFT achieves state-of-the-art results with minimal parameter updates.

Figure 5: Relative error and trainable parameter fraction for various finetuning strategies on downstream dynamics tasks.

Further, the Tadpole-DFT adaptation is statistically more stable during optimization and less prone to local minima or loss spikes than full-parameter fine-tuning, likely due to better preservation of the pretrained autoencoding representation (Figure 6).

Generative Modeling

Tadpole is also tested as a foundation for latent-space generative modeling (via flow-matching) on 3D channel flow. Latent models built upon FPFT or LoRA-finetuned Tadpole convincingly outperform generative baselines trained directly in the physical space, not only in standard MMD and Wasserstein metrics but also in scientific scores such as $\chi^2_\mathrm{PQM}$, and achieve orders of magnitude faster sample generation.

Theoretical Perspective and Architectural Insights

A critical theoretical argument in the paper is the distinction between learning the solution manifold $\mathcal{M}$ (structural geometry) and dynamics on $\mathcal{M}$ (vector fields tangent to the manifold). The autoencoding objective approximates the manifold itself, which is typically low-dimensional and smooth due to the structure of PDE solutions. In contrast, learning temporal dynamics directly requires resolving both the tangent dynamics and maintaining manifold invariance under perturbations (Figure 7).

Figure 7: Illustration of the distinction between learning the solution manifold (geometry) and learning the induced dynamics (tangent and neighboring flows).

Empirically, this underlies Tadpole’s consistent superiority in out-of-distribution reconstruction and its parameter-efficient adaptation for spatio-temporal tasks.

Practical Implications and Future Directions

Tadpole demonstrates that autoencoder-based pretraining with diverse, unlimited procedural 3D data is an effective route to scalable, transferable scientific foundation models. The approach is hardware-friendly—suited to large-scale, distributed HPC settings—and decouples model capacity from costly data collection and storage.

Tadpole’s representation-centric paradigm suggests broader applicability beyond regular-grid 3D PDEs, including potential directions such as:

• Generalization to unstructured meshes or adaptive grids,
• Coupling with differentiable and physics-informed solvers,
• Long-range and global temporal dynamics, currently a noted limitation,
• Integration with active learning or curriculum-based online data pipelines,
• Expansion to multi-modal and symbolic-conditional scientific tasks.

Conclusion

Tadpole establishes that foundation models for scientific machine learning do not require trajectory-based, dynamics-centric pretraining; instead, reconstruction-based autoencoding in a maximally diversified online setting is sufficient to yield powerful, transferable latent representations. Through the Tadpole-DFT fine-tuning procedure, these representations are readily adaptable to diverse scientific tasks, including accurate 3D dynamics and generative modeling, using minimal parameter updates. This marks a significant methodological shift in large-scale PDE-based scientific ML and sets a new direction for high-dimensional, general-purpose modeling of physical phenomena (2605.15284).