Adaptive Mesh Quantization for Neural PDE Solvers
Authors: Winfried van den Dool, Maksim Zhdanov, Yuki M Asano, Max Welling
TMLR 2025 | Venue PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | We demonstrate our framework s effectiveness by integrating it with two state-of-the-art models, MP-PDE and Graph Vi T, to evaluate performance across multiple tasks: 2D Darcy flow, large-scale unsteady fluid dynamics in 2D, steady-state Navier Stokes simulations in 3D, and a 2D hyper-elasticity problem. Our framework demonstrates consistent Pareto improvements over uniformly quantized baselines, yielding up to 50% improvements in performance at the same cost. |
| Researcher Affiliation | Academia | Winfried van den dool EMAIL QUVA Lab University of Amsterdam Maksim Zhdanov EMAIL AMLAB University of Amsterdam Yuki M. Asano EMAIL Fun AI Lab University of Technology Nuremberg Max Welling EMAIL AMLAB University of Amsterdam |
| Pseudocode | Yes | Algorithm 1 Weights-based Quantization Assignment Algorithm 2 Basic mixed-precision linear layer Algorithm 3 Optimized Mixed-Precision Linear Layer |
| Open Source Code | No | The paper mentions using 'flax' and the 'aqt Research (2023) library' which are third-party tools. There is no explicit statement or link provided by the authors for the release of their own code implementation of the proposed methodology. |
| Open Datasets | Yes | We validate the framework on multiple PDE datasets, ranging in complexity and scale, and demonstrate consistent compute-performance Pareto improvements. 5.1 Shape Net-Car Dataset The dataset is generated by Umetani & Bickel (2018) 5.2 Darcy flow Dataset We use the dataset from Li et al. (2020) Dataset EAGLE (Janny et al., 2023) is a large-scale dataset (1.1 106 meshes) 5.4 Elasticity Dataset The dataset is introduced by Li et al. (2023) |
| Dataset Splits | Yes | 5.1 Shape Net-Car Dataset ...The train/test split contains 700/189 samples. |
| Hardware Specification | No | The paper mentions profiling custom CUDA kernels (based on CUTLASS) for GEMMs using NVIDIA Nsight Compute and discusses 'modern GPUs', but does not specify any particular GPU model (e.g., NVIDIA A100, RTX 3090, etc.) or other detailed hardware specifications for running the experiments. |
| Software Dependencies | Yes | We implement our framework in flax (Heek et al., 2024) and use the aqt Research (2023) library for quantization. ... In each case, we use Adam (Kingma & Ba, 2015). |
| Experiment Setup | Yes | MPNN Apart from the default size of 6 hidden layers and 128 channels... we use Adam (Kingma & Ba, 2015) with a weight decay of 1 10 6, 500 epochs, and normed gradient clipping of 1.0. For Darcy and Elasticity, learning rates of 1 10 3 and 5 10 4 is used, respecitvely, and a batch size of 16. For Shape Net-Car, the learning rate is 1 10 3 and the batch size is 8. Graph Vi T For the Graph Vi T... We train for 1000 epochs, using Adam (Kingma & Ba, 2015) with weight decay of 1 10 6, normed gradient clipping of 1.0, learning rate 5 10 4 and a batch size of 8. |