siddik-lanl/spus-pde-unet-36m

SPUS: Small PDE U-Net Solver (36M)

A lightweight residual U-Net foundation model for solving partial differential equations, released by Los Alamos National Laboratory.

This is the 36M-parameter convolutional U-Net based foundation model pretrained on Compressible Euler equations from the PDEgym benchmark (https://arxiv.org/pdf/2405.19101)

• Paper: arXiv:2510.01370 (PDF)
• Code: https://github.com/lanl/SPUS-Small-PDE-U-net-Solver
• Model: https://huggingface.co/siddik-lanl/spus-pde-unet-36m
• Authors: Abu Bucker Siddik, Diane Oyen, Alexander Most, Michal Kucer, Nathan A. DeBardeleben, Earl Lawrence, Ayan Biswas (Los Alamos National Laboratory)

Overview

SPUS (Small PDE U-Net Solver) is a compact and efficient foundation model designed as a unified neural operator for solving a wide range of partial differential equations. Unlike existing state-of-the-art PDE foundation models — primarily based on large transformer architectures with high computational and parameter overhead — SPUS leverages a lightweight residual U-Net architecture that has been largely underexplored as a foundation model architecture in this domain.

The model is pretrained autoregressively, predicting the next timestep given the current timestep, which closely replicates the behavior of numerical solvers and helps the model learn the underlying physics.

Architecture

Residual U-Net (Unet2D from the model.py) with:

Hyperparameter	Value
final_insize	5
final_out_dim	5
hidden_channels	32
n_blocks	2
ch_mults	(1, 2, 2, 4)
act_fn	GELU
Total parameters	~36M

The architecture uses residual convolutional blocks, GELU activations, strided-conv downsampling, and transposed-conv upsampling with skip connections in the standard U-Net pattern.

Usage

The model.py file (containing the Unet2D class) is included in this repository. Download it alongside the weights, or import it from your local clone of the GitHub repo.

import torch
from huggingface_hub import PyTorchModelHubMixin
from model import Unet2D

class SPUSUnet2D(Unet2D, PyTorchModelHubMixin):
    pass

model = SPUSUnet2D.from_pretrained("siddik-lanl/spus-pde-unet-36m")
model.eval()

# Input shape: (batch, 1, channels=5, H, W)
# The leading singleton dim is squeezed inside forward().
x = torch.randn(1, 1, 5, 128, 128)
with torch.no_grad():
    pred = model(x)   # shape: (batch, 5, H, W)

The 5 input channels correspond to the Compressible Euler state variables used during pretraining. To use this model on a different PDE system, fine-tune on your target dataset (see the paper for details on downstream task evaluation).

License & Copyright

© 2026. Triad National Security, LLC. All rights reserved.

This program was produced under U.S. Government contract 89233218CNA000001 for Los Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC for the U.S. Department of Energy / National Nuclear Security Administration. All rights in the program are reserved by Triad National Security, LLC, and the U.S. Department of Energy / National Nuclear Security Administration. The Government is granted for itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license in this material to reproduce, prepare derivative works, distribute copies to the public, perform publicly and display publicly, and to permit others to do so.

Released under the BSD-3-Clause License. See the LICENSE file in the source repository for full terms.

O# (O5089)

Contact

For questions about the model or paper, contact: Abu Bucker Siddik — siddik@lanl.gov