Adaptive Mesh Quantization for Neural PDE Solvers
Authors: Winfried van den Dool
Maksim Zhdanov
Yuki M Asano
Max Welling
Venue: Accepted by TMLR ()

Abstract: Physical systems commonly exhibit spatially varying complexity, presenting a significant challenge for neural PDE solvers. While Graph Neural Networks can handle the irregular meshes required for complex geometries and boundary conditions, they still apply uniform computational effort across all nodes regardless of the underlying physics complexity. This leads to inefficient resource allocation where computationally simple regions receive the same treatment as complex phenomena. We address this challenge by introducing Adaptive Mesh Quantization: spatially adaptive quantization across mesh node, edge and cluster features, dynamically adjusting the bit-width used by a quantized model.
We propose an adaptive bit-width allocation strategy driven by a lightweight auxiliary model that identifies high-loss regions in the input mesh. This enables dynamic resource distribution in the main model, where regions of higher difficulty are allocated increased bit-width, optimizing computational resource utilization. We demonstrate our framework's effectiveness by integrating it with two state-of-the-art models, MP-PDE and GraphViT, 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.

Published in Transactions on Machine Learning Research (11/2025)
Adaptive Mesh-Quantization for Neural PDE Solvers
Winfried van den dool
QUVA Lab
University of Amstendam
Maksim Zhdanov
m. m.zhdonov@@uu...l
AMLAB
University of Ammsterdam
Yuki M. Asano
muki.asamoouunnnco
PORAI Lab
University of Technology Nuremberg
Max Welling
m. www.wlinn@uv..ff ml
AMLAB
Universily of Amsterdam
Reviewed on OpenReview: https: // openreview. net/
forum? id= NN17y897WG
Abstract
Physical systems commonly exhibit spatially varying complexity, presenting a significant challenge for neural PDE solvers. While Graph Neural Networks can handle the irregular meshes required for complex geometries and boundary conditions, they still apply uniform computational effort across all nodes regardless of the undorlying physics complexity. This leads to inefficient resource allocation where computationally simple regions receive the same treatment complex phenomena. We address this challenge by introducing Adaptive Mesh Quantization: spatially adaptive quantization across mesh node, edge and cluster features, dynamically adjusting the bit-width used by a quantized model. We propose an adaptive bit-width allocation strategy driven by a wight auxiliary model that identifies high- loss regions in the input mosh. This enables dynamic resource distribution in the main model, where regions of higher difficulty are allocated increased bit-width, optimizing com- putational resource utilization. We demonstrate our framework's effectiveness by integrat- ing it with two state-of-thh-art models, MP-PDE and GraphViT, 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.
1 Introduction
In recent years, there has been surge in the application of machine learning algorithms to build neural surrogates for Partial Differential Equations (PDEs). Using neural networks, PDEs can be modeled as a next- frame prediction problem, where the dynamics of the system (such as fluid movements) are implicitly learned. In real-world applications, traditional numerical methods such the finite element method predominantly operate on meshes irregular grids (Mavriplis 1997). these representations naturally adapt to complex geometries. Graph Neural Networks (GNNs) (Brandstetter et al. 2022 Pfaff et al. 2020 Gilmer et al. 2017bb have emerged as a natural solution, since they can operate directly on unstructured meshes, thereby avoiding the mesh-to-grid interpolation otherwise needed to represent irregular domains on regular grids. However, since computation scales directly with input graph size, efficient inference of GNNs remains challenging (Zhang al. 2022). Furthermore, reducing the operational cost is crucial for the practical deployment of GNN--basd neural surrogates, as lower costs enable higher data resolution key factor in improving modeling accuracy (Tles et al. 2020 Bauer et al. 2015: Randall et al. 2007)

---

Published in Transactions on Machine Learning Research (11/2025)
areas of:
lower complexity higher complexity
main model
4-bit.
6 bit
8 bit
message passing
bit assigner
Figure 1: Overview of the proposed framework. Given point cloud graph, bit assigner returns node-wise quantization scheme for larger model. The goal is to assign higher precision to more difficult and complex regions.
A common way to lower the computational budget required for inference deep learning is (integer) quan- tization (Jacob et al. 2017). Converting model weights and activations from floating- peating-ppontt to integer repre- sentations can drastically reduce computing and memory requirements. Although quantization was initially associated with lesmteee-limited devices, is now gaining attention in broader applications including Large Language Models (Wang et al. 2023) and large-scale physical simulations (Lang et al. 2021). In the latter, reduced-precision arithmetic has shown particular promise in mathematical (climate) modeling (Kimpson) et al. 2023) and neural surrogates (Dool et al. 2023). enabling higher data resolutions within fixed compu- tational budgets.
While quantization typically applies uniform precision reduction across entire model, many physical systems exhibit complexity that varies both location and time. Phenomena such as turbulence, shock waves, and steep gradients create regions where accurate prediction requires substantially higher numerical precision than in other, smoother, more predictable arcas. (Strogatt 2000) Current neural PDE solvers treat all mesh nodes uniformly, allocating the same computational resources regardless of the underlying modeling difficulty at each location. This uniform approach leads to inefficient resource utilization, where computationally straightforward regions consume the same resources as challenging phenomena.
We address this inefficiency by introducing Adaptive Mesh Quantization (AMQ)- In mesh quantization, we apply quantization all the mesh-basod features including node features, odge fcatures, and possibly cluster represenlations processed by graph neural networks operating on meshes. Our adaptive approach dynamically varies this quantization allocating higher bit-width to regions of higher complexity.
More concretely, relatively casy nodes as to be selected by small auxiliary model will be processed using low bit-width, freeing resources for more complex regions of the input requiring high precision (Figure D. The wighh- auxiliary model trained to predict the loss the main model at each node, which serves as a proxy for local complexity. This proxy thus based on both alcatoric and epistemic uncertainty, signaling not only where prediction hard in general but also where prediction hard for the quantized main model specifically.
2

---

Published in Transactions on Machine Learning Research (11/2025)
A
imput
main model
prediction
B
C
MP layer
graph diffusion
Tij zjj
8-bit
W.
6-bit
4-bit
argsort binarize
auxiliary model
message MLP
bit passigner Caux
8-bit
complexaty weights
mannn 00
Agerogatc
bit assigner
predicted loss
OΦ
node-wise loss
optimize
update MLP
bit assignment
Figure 2: Overview of the training phase (A). ). A lightweight auxiliary model assigns a complexity weight to every mesh node to guide the resource allocation of the large main message-passien model. Gradients do not flow through the dotted black arrows, i.e., backpropagation happens through the solid black arrows. (B). The bit assignment process is schematically shown (C).
Overall, we make the following contributions:
We roduce Adaptive Mesh Quantization, a Tramework that learns to allocate resources locally on a mesh depending on the input data. We additionally propose an efficient implementation for the mixed-precision quantization of MLPs.
We investigate the particular setup where an auxiliary model learns which input regions are more complex based on the main model's loss landscape, training both models simultaneously.
We validate the framework on multiple PDE datasets, ranging in complexity and scale, and demon- strate consistent compute--erformance Pareto improvements.
2 Related work
2.1
Neural PDE solvers
The intersection of PDE solving and deep learning has emerged as active research area (Merreivy & Hakim
2024h. The usual task is to learn mapping from initial state to solution of PDE with neural network, in order to accelerate complex simulations in fields like fluid dynamics and climate modeling (Wang et al.) 2024). Originally focusing on regular grids (Li et al. 2020). multiple solutions were suggested to deal with irregular meshes, with message-passingg passing-based approaches (Brandstetter et al. 2022 most relevant our work. For these networks the input comes in the form of a graph and model learns evolution for every node. Lately, mesh attention (Janny et al. 2023) also introduced to capture long-range dependencies further enhancing the performance of mesh- based models.
We do not introduce new pde-solving model, but rather new method that is applicable to families of Graph Neural Networks in particular. To showcase our method we use both the original MPNN (Brandstetter let al. 2022), as well as the more elaborate Graph Attention architecture (Janny et al.] 2023). Apart from being state-of-tho-art on several datasots, the latter model has the advantage of covering many different layer types (c.g., graph pooling, rmn, attention), ensuring our method is broadly applicable across different GNN model architectures.
3

---

Published in Transactions on Machine Learning Research (11/2025)
2.2 Quantization
By reducing the precision of storage and arithmetic operations in neural networks, quantization offers signifi- cant reductions in memory overhead and computational costs while typically maintaining model performance (Jacob et al. 2017 Nagel et al. 2021). Conventional approaches implement uniform quantization across all model layers, either through post-training quanlization (PTQ) (Sung el al., 2015) or during the training phase via Quantization-Awars Training (QAT) or Quantized Training (QT) (Jacob et al. 2017 Nagel et al., 2021). Mixed-precision quantization (Wang et al., 2018 Pandey et al., 2023 has emerged as a more efficient approach, allowing different layers or operations works for bit-width selection (Wu al. 2018) have shown that optimized quantization strategies achieve higher compression rates while preserving model accuracy,
use different numerical precisions. Automated frame-
Of particular relevance to our work are the quantization of PDE solvers (van den Dool et al. 2023 ) and input-dependent quantization. van den Dool et al. (2023) show that quantization can be effectively applied to neural PDE solvers, and importantly, that spatial resolution modification and quantization represent complementary approaches that can be combined for optimal results.
In input-dependent quantization, Liu al. Liu et al. (2022) and [aaena & Roy (2023) employ auxiliary models to dynamically assign bit-widths to neural network layers based on input complexity, thereby achiev- ing superior model performance while maintaining the benefits of quantization. In contrast to these works, our method applies dynamic mixed quantization not per input instance, but on the complexity of particular areas within a single input.
2.3 pessage-paasing quantization
Quantization methods have recently gained popularity for MPNNs in particular, as inefficient inference is a known problem when scaling MPNNs up to real-world and large-scale graph applications (Ma et al., 2024). Tailor et al. (2020) propose a QAT method which selectively quantized nodes with low in-degree as less prone to the quantization noise. Zhu et al. (2023) further advanced the idea by learning the quantization bandwidth for each node based their in-degree. However, these methods based the graph architecture alone, which may not always be related to, or take into account, the complexity of the data that it represents.
3
Preliminaries
3.1
Quantization
To move from floating-pott point to efficient fixed-point operations, two quantization paramcters are needed: the scale factor 8, and the bit-width b. Together these parameters form a quantizer Q(-) = Q+: 8, which acts on a floating point tensor x as
Q(x) = clamp (Is x]; 2b + 1,,2-1 -1).
(1)
where 1-1 is the round-to-nearest operator and clamp(, 1, h) clamps an input between the values 1 and h. The approximate real value of x is then retrieved by
Q(x) x.
&
(2)
Note that do not apply shift, which helps us preserve the 0 value. Here can represent both weights matrix and an activation vector. The bit-width is chosen before training, thereby fixing the target range to [-20-1 + 1, 20 - 1]. The scale factor ensures efficient use of this target range, using statistics of x like max]x x in the most basic case. If x are activations, exponential moving averages of such statistics can be tracked during training, to be fixed after an initial calibration training phase has been completed. The scale must be chosen to balance outlier clipping and rounding errors. To reduce such errors, is often defined on per-channel basis, which can be seen as replacing with diagonal matrices S and S-1 in Equations 1 and respectively.
4

---

Published in Transactions on Machine Learning Research (11/2025)
The efficiency gains from quantization in multiplication like W á for weights matrix W and activations follows when taking the out of the matrix multiplication:
W à (Q(W)Q(a)) .a.
(3)
SW Sa
The cost of computing (Q(W)Q(a quantized using bw, be bits is then in the order of bw ba Multiply Accumulate operations (MACs) As the bias vector's 's impact on computational cost is commonly considered negligible, it is kept in floating point precision.
3.2
Message-passing neural networks
The neural network imput is graph G = (V,E), which comes from discretization of the D-dimensional computational domain Ω E R2, It might, for example, represent mesh or point cloud with induced connectivity (e.g. by k-nearest neighbours) and contains N nodes. Here V = [. is the set of nodes and & C V x V are the edges. Additionally, every node is endowed with its location Pt E Ω and, optionally, features x; E Idd. Message-passing neural networks are architectures developed specifically to operate on graph-structured data (Satorras et al., 2021 Braandsttter et al. 2022). Following the notation from Gilmer al. (0177)), we define the L-th layer a sequence of operations:
mij MLP.(X., -D))
message
m, Σ mij
aggregate
(4)
JEN(()
[+ MLP,, (x,, m;),
update
where (i) represents the set of neighbours of node e, MLP,,, MLP,, are message (edge) and update (node) MLPs respectively, and x Rd are latent d-dimensional feature vectors.
MP-PDE Brandstetter et al. 2022 follows the 20200) with the core element being M steps of learned message passing in the processor. More precisely, the processor is an MPNN with additional input (solution differenco) in the message network:
framework of Sanchez- Gonzalez GozzalzzzGoozzzeez et al.
my MLP, (x/, uy, uy, P,),
message
where u is vector of solution values at previous time steps. The encoder and decoder are applied node-wise and implemented using an MLP and shallow 1D convolutional network, respectively.
Janny et al. (023) leverages recent developments in computer vision to produce a scalable mesh transformer. The main idea is to capture long-term interactions in the domain, for which the mesh is coarsened and multi- head attention (Vaswani al. 2017) computed on the clusters. The final model (GraphViT) consists of 5 steps: 1) MPNN encoder (Eq. [), 2) clustering, 3) graph pooling the clusters, 4) attention, and 5) MPNN decoder. In experiments, we found that steps 1 and 5 are by far the most computationally demanding in the model (see Appendix C for the details), which makes the model suitable for our framework to test. its scalability.
It common MACs represeut compute cost, rather than wall-time. This partly bocause wall-time hardware- dependent, but also because often quantization simulation packages are used, optimal on-device implementation is not possible reserved for real-world applications (Siddegowda al. 2022).
5

---

Published in Transactions on Machine Learning Research (11/2025)
4 Method
4.1 Adaptive Mesh Quantization with fixed budget
We assume that for an ordered set V of N nodes weight w E RN is given that relates to the complexity of (predicting the output at) each node. We further assume a fixed total computational budget which we want to distribute across different quantization levels efficiently at inference time based on those complexity weights w.
Let Q be the used quantizers for different bit-widths, as defined in Equation 1 so that Qx has quantization parameters given by bit-width by and scale 8k. and let & be increasing in compute cost. For example, for setup with Int4, Int8 and Int12 quantization, Q = -Q Q2, Qa) with b = (4, 8, 12). Let associated budget allocation ratios a 2 0 be given such that Σx ar = 1.
To map the complexity weights to quantization levels we first sort the weights vector w, and then proceed by assigning the most expensive (i.e., having the highest bit-width) quantizer Qu to the or fraction of nodes with the highest weight, the hechne-highest Qx-1 to the next Qk-1 fraction, and so forth, ensuring that nodes with higher weights receive more precision, as shown in Algorithm 1
Algorithm 1 Weights-based Quantization Assignment
1: Input: Complexity weights w E RW Available quantization functions Q = [Q1,..., Qkl Allocation ratios [1 ".; og] with Σ. = 1, 2 0
I argsort(w) {Sort by incrcasing weight}
- [0] x K (Initialize empty buckets}
start t 0
5: for = 1 to K do
6: end < start + [N oc.]
7: B: + [start: end] {Assign indices to bucket)
8: start + end
end for
10: Output: /
The actual bit-width of node i is effectively used quantize activations the update MLP of Equation and the embedding and projection layers at the start and end of the GNN architecture. In Section 4.4 we describe how MLPs can be efficiently implemented to handle general multi-resolution bit assignment for all nodes.
We also use the node weights w to assign a similar complexity proxy weight to the edges and clusters (if available) by using the target node weight. for the edges and the mean weight of all cluster nodes, respectively. For both edges and clusters, we also follow the Alg, [ to get a bit assignment that follows the same budget allocation ratios a. The bit-width of edge j is used to quantize the activations of the message MLP of Equation and possibly edge embedding and update layers if included, whereas the cluster bits assignment is similarly used in the Gated Recurrent Unit of the RNN pooling layer and procossing MLPs. For further architecture details, we refer to the supplementary material.
For each MLP, only the activations are adaptively and non non- non-uniformly quantized. The weights are quantized to fixed bit-width, e.g., Int8 in our case, independent of the presumed complexity of the input elements. Assuning a fixed number of nodes, edges and clusters, our use a fixed budget allocation ratio a not only directly translates to fixed computational budget, but also enables the design of a predefined model architecture hardware.
4.2 Predicting spatial complexity
Let G (V,E) be the model input graph with node features x E Nxd, where N the number of nodes and d is the feature dimension, and let y e Nxd' be the target outputs. We aim assign a weight w; > to each node i, representing its complexity, using small auxiliary GNN model denoted by As (G,x) - w.
6

---

Published in Transactions on Machine Learning Research (11/2025)
To enforce the interpretation of representing the node complexity, propose to train the auxiliary model Aa to predict the loss of larger main model Mo (G,x) x) + g. Note that we ultimately seek good model Ma, but we first explain how Ag is trained. We define LM E 18 as the loss of the main model which has not been spatially reduced, i.c., if Ca (M)(x), is the actual loss of this model for input x, then
(M) Σ xxll
(6)
where [i] means selecting the i-th node, omitting the dependence on G for clarity. Tf a conventional loss function is used (e.g, MSE) LM[i] measures the error between the predicted value and the ground truth value of node i. We use this error for our proxy of node complexity: After applying smoothing function [0, 1 to the spatial loss L, use the smoothed result target for the auxiliary model Ag.
S S:R
More precisely, the parameters & are updated using Stochastic Gradient Descent (SGD) with learning rate as:
LM (x)
01t1 = Φ: 7AV (1) S(L
(7)
The smoothing function ensures that the target for the auxiliary model training is sufficiently well-behaved. It involves the following steps:
Normalization. We scale the loss Las by setting Las L max; Lil.). This transformation normalizes LM to the [0, 1] range, which both stabilizes the loss distribution and mitigates the influence of extreme values auxiliary model gradicnts. It also allows the auxiliary model A to use sigmoid activation in its final layer, simplifying training.
2. Graph Diffusion. We apply multiple rounds of graph diffusion (see Appendix, Equation 15 that loss information propagates across neighboring nodes. The underlying assumption is that the loss observed at any given node reflects not only its own operations but also the influence of nearby nodes, as a result of the message-passing framework. By performing diffusion steps, we effectively smoothen the loss values over the graph, reflecting this interdependence. The diffusion process smooths the loss landscape and further reduces the effect of local outliers.
4.3 Training
Instead of updating the auxiliary model on fixed (e.g., pre-trained) model M, suggested in equation we propose a single synchronous training procedure, using a framework that couples the auxiliary model Ag and the main model Mg. This coupling is based on:
Applying Alg. to the proxy loss outputs of As to obtain bit-width assignment B3 where B denotes combined representation for the bit-width assignment of the nodes, edges and clusters:
w
A(x; dr),
(8)
g
AssignQuant(w, Q,
(9)
with Q and a the available quantization functions and allocation ratios, defined in Section 4.1
Quantizing the main model as described in Section 4.1 We now define the main model M as
Ma.3 (G,x) + &.
(10)
Using the spatial loss of this quantized model as Las in equation
La Las (No,
(11)

---

Published in Transactions on Machine Learning Research (11/2025)
Algorithm 2 Basic mixed-precision linear layer
Input: activations
R Rxxxd weight matrix W Rdxd', available quantization functions Q [Q1.... Qxl: index buckets B .
"at....,Bkk] Bg] {actie < Distribute(a, 13) {Get a& E R Rlb, xdj
Initialize + [0]
for to K do
y& < Qx (at) (WT)
Y& sxew
Y + (ya) fappend Yε to Y}
end for
Y + Concatenate(Y)) fumpack buckets into array)
Y + Reorder(Y) {place in original order)
Output: Y
Before the actual training, we may start with an initial wo, potentially based on data uncertainty, reflecting the prior estimate of complexity on the mesh 2 parameters and as follows:
We then use SGD with learning rate 1) to update the
W++1 = A(x; et)
Bc+1 = Ausignnuant(We111..
(12)
0.+1 = 0t TM VoLN (M B141,
@c+1 = t+ - 7I V A(x;; Lu (x; B (,,1,0,,,y)
By basing our node complexity proxy A(x; Φi) on the loss of an existing model, instead of on the underlying data alone, we hypothesize that it better reflects epistemic (model) uncertainty, instead of just aleatoric (dataset) uncertainty. Training with an already quantized model, M (x; 0c), in the loop then allows the modeling complications caused by quantization be reflected in the complexity proxy that we use, providing direct feedback to the main model's precision requirements.
4.4 Hardware-efficient multiplications with non-uniform quantization
In this section, we describe how a non-uniform assignment of activation bit-widths can be used in a single quantized layer. Let B = [B.,.., Bx] be the buckets of indices of vectors quantized using Q [Q1..., Qk] The quantizers are defined in Equation and implicitly include their corresponding quantization param- eters, ie., the bit-width 6y and scale factor Sj.
Let the weight matrix of linear layer be given by WT E Rdxd' and assume it gets quantized using scale parameter Sw and bit-width bw- Let a E Blldd be the activations corresponding to bucket Bg, noting that these can refer to nodes, edges or clusters, depending on the particular layer the overall architecture. We can initially implement linear layer by separately computing the malrix products for each calegory, as described in Alg.
However, each time matrix multiplication task is passed to the GPU, there is overhead associated with transferring data and setting up the computation environment. If multiple smaller tasks are sent individually, this overhead is incurred repeatedly, which can reduce overall performance. Ideally, all tasks are combined into a single call, taking advantage of the GPU's architccture more effectively. This is why we propose a Assume the bit-widths b are given as
different approach for the mixexl-precision matrix multiplications. multiples of bitt-b bu, e.g., by 4,6 (4,8, 12) for Int4, Int8, Int12 quantization. Given
quantization level k, bucket index and channel we then split the bitstring of Qa (a into be/bo bitstrings
? Although the auxiliary network's predictions may be random during the initial training phase, thus unsuitable for asscssing regional complexity, we have observed that starting with a warm-up phase in which uniform weights are used until A is properly trained is unnecessary, leading better results.
8

---

Published in Transactions on Machine Learning Research (11/2025)
Algorithm 3 Optimized Mixed-Precision Linear Layer
Input: Activations R Weight matrix W Rdxd", Available quantization functions
QQ.....: Qk), Bit-widths b = [br,...,bxi .bl, Base bit-width bo, Index buckets B -B,, [ Bx] 2: {au]k=i < Distribute(a, B3) {Get ag E 1l 3: Initialize a + [0] (Encoded activations}
Initialize Scales < [0]
for : = to K do
Qa (ax) (Quantize activations in bucket k: ak Encode(Quus : (Encode as Int-bu values}
8: Scales + Scales U { th 2m m 0,..., by }} (Scale factors for each bit segment]
10: end for
11: a + Concatenate(a)
12: k diag(Scales)
13: Y S. Q" (WT) {Scaled matrix multiplication)
14: Y BeatterAdd()) (Combine original indices)
15: Output: Y
of size n. Omitting bucket and channel indices in Qx(ax) for case of notation, we then redefine those by /bo smaller bitstrings as Int-by values as follows:
Qa(ax) Σ
2
% {0.1)
(13)
1=0
(m+1)60
a. (a:) ==
2 mbo Σ
mho
(14)
by - 1). We can thus encode Qx(ax) E (o....,. 2b. 1] xd which consists of By vectors where m E (O,.... by
of d channels with Int-by precision, as 0g (ag)) E {0...., 2h - BB xd This flattened representation corresponds to B₁ by vectors of d channels with Int-ly precision. Since this transformation applies
all k, we can combine all buckets into single array a {O..... .,2. 1] Nxd where Ñ Σx BB /bx/ /b. We can then perform single matrix multiplication aQ"(WT). Finally, scale the output to reflect the effective contribution of each bit segment based its position within the original encoded bit-string, while simultaneously applying existing scaling factors. The steps are summarized in Alg. 3
5 Experiments and Results
Across all experiments, the auxiliary model is implemented channels. For Darcy and ShapeNet-Car, the main model is an MP-PDE with 6 hidden layers and 128 channels, containing about 550k parameters. For the EAGLE experiment, we employ GraphViT with a hidden dimension of 64 and 3 attention layers, totaling about 1M parameters. We implement our framework in flax (Hock ct al. 2024) and use the aqt Rescarchh (2023) library for quantization. Further details are given in Appendix and We run all experiments for all quantization regimes from Int4 until Int8, gradually increasing the fraction of Int8 nodes used for the mixed--reeisson runs. The main results are summarized in Figures 3 and 5 We also ran experiments at higher bit-width, but for brevity only show results until the quantization level at which increasing precision further has no significant impact on accuracy. Note that the uniform curves show exponential decay for lower bit- bit-width, with the convexity of the curves suggesting that it is nontrivial for combinations of categorics to result in lower losses (i.c., plain averaging of points on a convex curve leads to higher values by definition, see also the random assignment table 2
MPNN with 3 hidden layers and 32 hidden

---

Published in Transactions on Machine Learning Research (11/2025)
Darcy Flow
ShapeNettCor Car
EAGLE
Elasticity
0.105
4:0F
Uniform Adaptive
Uniffrmm Adaptive
Unifarm Adaptive
0.026
4-5it
Uniform Adaptive
0.100
0,0040
0.024
Validation Loss
0.095
e-bit
0.030
S-blt
0.0035
0.022
5-bit
0.090
5-bit
0.025
0.029
6-bit
0.085
0.020
7.07:
0,0025
6-bit
7-bit
u-bri
0.015
0.080
8-b1
250
350
MACs
MACS
Figure 3: Validation loss vs cost (estimated in 10 MACs) for uniform adaptive (ours) quantization. The latter is obtained through incrcasing the fraction of Int8 nodes used, and includes the cost of the auxiliary model. Our framework demonstrates consistent. Pareto improvements across all benchmarks. For Darcy, ShapeNet-Car and Elasticity, main models are MP-PDEs, for EAGLE- - GraphViT.
While Figurcs
D Bnd 5 present results across a broad range of compute budgets, we note that practical interest is often concentrated on the more elficient (low-bit-width) low-cost models corresponding to the left side of the curves, where MACs are lowest. The results indicate that our adaptive quantization method generally outperforms uniform quantization on the entire spectrum between Int4 and Int8. However, in practice, uniform quantization between those endpoints is not implemented on most hardware at all. Therefore, a further benefit. of our hod is that it makes this full regime accessible in the first place. We refer to Appendix D for more details on hardware implementation.
5.1 ShapeNet-Car
Dataset The dataset is generated by Umetani & Bickel (2018) and consists of 889 car shapes from ShapeNet. Every is represented by N 3600 mesh points its surface {PI,..,,PNy E g and its acrodynamics is simulated by solving the Navicr-Stokes equation. The task is thus to predict the pressure value y E RV for every mesh node. We follow the preprocessing from Alkin et al. (2024) and use k nearest neighbours to induce the connectivity. The train/test split. contains 700/189 samples. Each (main) model is trained by optimizing the MSE loss between predicted and ground truth pressure.
Results As shown in Fig 3 our adaptive quantization achieves Pareto-optimal performance compared to uniform quantization. The approach works especially well in the low bit-width regime, matching the performance of 8-bit models while using only half the computational resources.
Alternative auxiliary model As an alternative to training the auxiliary model directly on the main model's loss, we experimented with a setup where the auxiliary model is independently pretrained on the data itself. We let the auxiliary model predict both a mean and a variance for each node, with the variance representing aleatoric (data-inherent) uncertainty the sense of heteroskedastic regression (Kendall & Gal, 2017). We then use this variance to guide bit allocation, assigning higher prccision to nodes with greater predicted uncertainty.
This approach relies the assumption of Gaussian likelihood model, which provides tractable way to learn uncertainty but may not perfectly reflect the true noisc distribution in the datasct. Nevertheless, it offers a reasonable proxy for node complexity without requiring joint training with the main model. When comparing this approach to our surrogate-loss method, we do not see a big difference when a large ratio of Int8 is usod. However, we expect this because our assignment procedure (assigning Int8 to the top-k% tigh-complexity nodes) hides differences in ordering for larger Indoed, we slightly favor the surrogate-loss method, as it yields better results for the more challenging lower resource budgets (Table D.
10

---

Published in Transactions on Machine Learning Research (11/2025)
Table 1: Validation loss for surrogate loss-based (default) and uncertainty-based bit-width assignment on the ShapeNet Car dataset. For the latter the auxiliary model pre-trained on the data, and its varianco is used input the bit assigner. While the difference is negligible when there sufficient high-precision locations to distribute, it becomes clearer when only small fraction is available that the default method assigns them better.
Int4 / Int8 ratio
Uncertainty
Default
95%
5%
0.0330 + 0.0042
0.0285
0.0008
90%
10%
0.0232 0.0004
0.0224
0.0002
80%
20%
0.0205 1 0.0002
0.0204
0.0004
70% / 30%
0.0199 0.0002
0.0197 = 0.0002
Bitmap assignment
Input
Output
lint4 int8 (30%)
Figure 4: The proposed framework on an arbitrary example from the EAGLE dataset. The model predicts the field based on the PDE solution the previous timestep. The bitmap assignment process manages to successfully capture the airflow and allocate resources to the relevant region. In this example 30 % of the nodes is assigned Int8, leaving the in low (Int4) precision.
5.2 Darcy flow
Dataset We use the dataset from Li et al. (2020) which contains the solution of the steady-state of the 2D Darcy FLow equation the unit box. The data is defined regular grid 421 x 421, which sparsify and treat as directed graph with N = 2770 nodes. The task is to learn map from the dilfusion coefficient X; E R4 to the solution u, E R2 for node i. The (main) model is trained by optimizing the MSE loss between the prodicted and ground truth solution.
Results
The comparison between uniform and adaptive quantization is shown in Fig 3 method improves quantization across all MACs tested, but less significantly so in the high MACs regime. We focus on the low bit-width regime, however, which is often of most interest in NN quantization research as this is where compute savings are highest and accuracy hardest to retain.
The adaptive
5.3 EAGLE
Dataset EAGLE Janny et al. 2023 a large-scale dataset (1.1 1 106 meshes) resulting from non-steady fluid dynamics simulations. consists of 600 unique scenarios of a moving flow source interacting with non-linear domains (see Fig. for an example). Meshes are dynamic and vary in the number of nodes, with the average N 3388. The dataset is especially challenging as the dynamics are often highly turbulent and multi-resolutionall due to non-trivial environments. In particular, obstacle layouts vary between scenarios, so both training and validation include unseen configurations, providing inherent out-of-distribution test for both models. The task is defined as follows: given the simulation state at time t in the form of a graph G*, predict the future pressure and velocity (v/+1 E R4) for every node after a time-step dt. The model is trained by minimizing the MSE on pressure and velocity (see Appendix A for details).
Results We observe a noticeable effect of adaptive quantization (see Fig 3). Again, adding only 10% more MAC's (by converting small fraction of nodes to Int8) recovers more than half of the loss incurred
11

---

Published in Transactions on Machine Learning Research (11/2025)
when quantizing from Int8 to Int4. We also provide the example of the forward pass of our framework in 4i Notice that the bit assigner manages to successfully capture the area of rapid changes along the borders. The main model can then allocate fewer resources to center regions with stable pressure and velocity while focusing on highly turbulent arcas.
5.4 Elasticity
Dataset The dataset is introduced by Li et al. (2023) and contains solutions to a 2D plane-stress hyper- elasticity problem with a central void. The domain is a unit square with a randomly-shaped hole, clamped at the bottom and subjected to vertical tension at the top. The material follows the incompressible Rivlin Saandersss hyper-elastic model. The data consists of unstructured point clouds with N 1000 nodes per sample. Each sample consists of an input graph based the spatial coordinates and output graph based on the stress field at each node. The task is to learn map from geometry to stress response, with models trained using MSE loss between predicted and true stress tensors.
Results Here too, as shown in 3 our adaptive quantization achieves Pareto-optimal performance compared to uniform quantization in the low bit-width regime. When only a small portion of int4 nodes used in the mixed-precision ratio, the loss increases relatively much: the addition of only few int4 hurts the training more than changing to int7 precision. However, in practice, we expect interesting use cases to be on the opposite end of this setup: relatively few nodes getting assigned a high precision, with the bulk remaining in low precision.
5.5 Quantization assignment ablation
Dataset
Targeted Assignment (default)
Random Assignment
ShapeNet Car
+3.4%
+14.6%
EAGLE
+2.6%
+37.2%
Darcy
+ +29.1 199.%
+73.1%
Elasticity
+ +16.1% 1%
+72.2%
Table 2: Increase in loss result of quantizing half of the nodes to Int4, leaving the rest Int8. The default is to assign nodes with high (predicted) complexity to Int8 quantization. For the random assignment, we shuffle the node complexities. The numbers are normalized such that +0 % corresponds to the "base" loss of the all-Int8 uniform setup, and +100 % corresponds to the "fully quantized" loss of the all-Int4 uniform setup. This way, random assignment half the nodes to Int4 resulting in < 50% loss increase indicates relative robustness, whereas > 50% increase in loss indicates disproportionately large sensitivity quantizing nodes at Int4. Note that mitigating the increase in loss is all the more impressive for the latter category, where quantizing nodes to Int4 so drastically hurts performance.
One might attribute the improved performance to dynamic quantization acting as a form regularization. To test this effect, we conduct control experiment (see Table 2 where we randomly shuffle the weights coming to the bit assigned during training, which keeps the stochastic behavior but removes focused resource allocation. Our results indicate that the focused allocation is key to the framework's performance.
While the improvements on Darcy and Elasticity may at first appear less striking than those for ShapeNet Car and EAGLE, it is in fact notable that the largeted assignment strategy substantially miligates quantization loss in these cases. As shown in Table 2 random assignment leads to disproportionaaely large increases in loss, indicating that these datasets are highly sensitive to nodes being quantized to Int4. Against this challenging baseline, the much smaller increases under targeted assignment demonstrate that the method effectively preserves performance even where quantization most damaging.
12

---

Published in Transactions on Machine Learning Research (11/2025)
6 Discussion and Conclusion
We introduced Adaptive Mesh Quantization (AMQ), a general framework for improving the efficiency of quantized messag-passing neural networks (MPNNs) by enabling spatially adaptive bit-width allocation. Bit allocation is guided using a lightweight auxiliary GNN that estimates the local prediction loss of the main model, thereby serving as a proxy for spatial complexity. This mechanism allows the model to concentrate computational resources where they are most needed, yiclding more efficient and targeted use of precision. AMQ consistently achieves favorable compute-performant trade-ofls, outperforming uniform quantization baselines across diverse PDE surrogate tasks.
While larger datasets, i.c., meshes with more nodes, can be more practically relevant than the sizes often used in academic studies, we expect the benefits of AMQ to become even more pronounced in large-scale, high- resolution scenarios. Additional experiments on larger meshes (Appendix E confirm that AMQ continues to perform beyond the size of small-scale research datasets. Nevertheless, we expect the benefits to be especially pronounced on real-world, high-resolution data, where spatial complexity preserved in greater detail and the bit assigner exploit local variation more effectively. Coarser meshes, in contrast, tend to average complexity across regions, which can limit the effectiveness of adaptive schemes. This motivates exploration of AMQ in real-world, high-resolution simulation settings.
Our main results are reported in terms of MACs. To minimize potential discrepancies with real runtime, we implemented custom CUDA kernels and measured execution time directly (Appendix D These experiments indicate that the relative trends are consistent, suggesting that MAC-based analysis provides reasonable reflection of practical performance.
Finally, note that further gains may be possible through task-specific hyperparameter tuning. Also, alternative strategies for complexity estimation such as uncertainty-based metrics or learned priors cofors-oould complement or enhance the current loss-proxy approach.
7 Acknowledgements
This work is financially supported by Qualcomm Technologics Inc., the University of Amsterdam and the allowance Top consortia for Knowledge and Innovation (TKIs) from the Netherlands Ministry of Economic Affairs and Climate Policy.
13

---

Published in Transactions on Machine Learning Research (11/2025)
References
Benedikt Alkin, Andreas Fürst, Simon Schmid, Lukas Gruber, Markus Holzleitner, and Johannes Brandstet- ter. Universal physics transformers. ArXiu, 00ss/240..12600 2024. URL httpl/ttsp//ppi..maaaiiiiiio seaantcschoooa....
Peter Bauer, Alan Thorpe, and Gilbert Brunet. The quiet revolution of numerical weather prediction. Nature, 525(7567):47-55, 2015.
Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural pde solvers. arviv preprinl arXiv:2202.0333 2022.
Winfried van den Dool, Tijmen Blankevoort, Max Welling, and Yuki M Asano. Efficient neural pde-solvers using quantization aware training, arXiv preprint arXi2:2308. 07350, 2023.
Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry. In International Conference on Machine Learning, 201 7a. URL https:
org.Corpgo/D: Saapi.smmattisccoole.. D 9665943
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp. 1263-1272 PMT.R, 2017b.
Jonathan Heck, Ansolm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andrcas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2024. URL http://github. com/ com/ google/fll
Carley E Tles, Robert Vautard, Jane Strachan, Sylvie Joussaume, Bernd R Eggen, and Chris D Hewitt. The benefits of increasing resolution in global and regional climate simulations for european climate extremes. Geoscientific Model Development, 13(11):5583-5600 2020.
Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew G. Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient integer- arithmetic-only inference. 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2704 2713, 2017. URL httppp//pp...smmttiieebbaaa..rrp org/CorpusI 338676699
Steeven Janny, Aur'elien B'eneteau, Nicolas Thome, Madiha Nadri Wolf, Julie Digne, and Christian Wolf. Eagle: Large-scale learning of turbulent fluid dynamics with mesh transformers. ArXiv, abs/ hbs/2302.10003
2023. URL Wtttttts:://///// https://aptl wemanticscholar wwweaarrrrrr//eeee : 257050214
Alex Kendall and Yarin Gal. What uncertainties do we need in bayesian deep learning for com- puter vision? In I. Guyon, U. Von Luxburg, S. Bengio, Н. Wallach, R. Fergus, S. Vish- wanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Cur- ran Associates, Inc., 2017. URL. https://proceeding pdf
caturips.ccpaappp. httpri/ssscpppppe..ll///////eppp /t1les//appr//21///Eiee/
Tom Kimpson, E Adam Paxton, Matthew Chantry, and Tim Palmer. Climate-change modelling at reduced floating-point precision with stochastic rounding. Quarterly Journal of the Royal Meteorological Sociely, 149(752)::4333383 855, 2023.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun ), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URt.htttp://arxxw.orrg http://wwwwor//////1....... /ets//412..9999
Simon Lang, Andrew Dawson, Michail Diamantakis, Peter Düben, Sam Hatfield Martin Leutbecher, Tim Palmer, Fernando Prates, Christopher Roberts, Irina Sandu, and Nils Wedi. More accuracy with less precision. Quarterly Journal of the Royal Meleorological Society, 147, 10 2021. doi: 10.1002/qj.4181.
14

---

Published in Transactions on Machine Learning Research (11/2025)
Zong- Yi Li, Nikola B. Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew M. Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. ArXiu, abs/20 abs/2010.08885, 2020. URL
wttp:///cc..oooooooooooorraaaaaa 0 224705257
Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with
learned deformations for pdes on general geometries. Journal of Machine Learning Research, 24(388); 1-26, 2023.
Zhenhua Liu, Yunhe Wang, Kai Han, Siwei Ma, and Wen Gao. Instance-aware dynamic neural network quantization. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 12424-12433, 2022.
URL https:://apppp www.ssseeaanccooolaaa.m htt.ppppp//ppp....matc..cooooor. wwwpscchooaarrrg/orrrooee :rgg/CorpuSID:2440 249916643
Lu Ma, Zeang Sheng, Xunkai Li, Xin Gao, Zhezheng Hao, Ling Yang. Wentao Zhang, and Bin Cui. Acceler- ation algorithms in gnns: A survey. ArXiv, abs/ abs/2405.04414, 2024. URLhhttss: htttpp::/app...mmmaiiccc.c
DJ Mavriplis. Unstructured grid techniques. Annual Review of Fluid Mechanies, 29(1):473-51) 514,, 1997.
Nick McGreivy and Ammar Hakim. Weak baselines and reporting biases lead to overoptimism in machine lcarning for fluid-related partial differential equations. ArXiv, abs/2407.07218, 2024. URL htttps::///ap.
semanticscholar. wwcpllarr.rr/orrrrrreeeeee hhtp:/CorpsI:::::00 271088446
Markus Nagel, Marios Fournarakis, Rana Ali Amjad, Yelysei Bondarenko, Mart van Baalen, and Tijmen
Blankevoort. white paper on neural network quantization. ArXiu, Abs/2106.082200 2021. URL https:
235435934
Nilesh Prasad Pandey, Markus Nagel, Mart van Baalen, Yin Huang, Chirag Patel, and Tijmen Blankevoort. A practical mixed precision algorithm for post-training quantization. arXiv preprint arXiv:2302.05337, 2023.
Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzale and Peter W Battaglia. Learning mesh-based simulation with graph networks. arXiv preprinl ar Xxxv:20000003409 2020.
David A Randall, Richard A Wood, Sandrine Bony, Robert Colman, Thierry Fichefet, John Fyfe, Vladimir Kattsov, Andrew Pitman, Jagadish Shukla, Jayaraman Srinivasan, et al. Climate models and their evalu- ation. In Climate change 2007: The physical science basis. Contribution of Working Group I to the Fourth Assessment Report of the IPCC (FAR), pp. 589-662. Cambridge University Press, 2007.
Google Research. Aqt: Accurate quantized training, bttras///ggtaab..omm/gggl//////p 2023. URL https:
aqt
Alvaro Sanchez- Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter W. Battaglia. Learning to simulate complex physics with graph networks. ArXiv, abs/2002.09405,, 2020. URL https
/aapi.semanticscholaa. wwtp://larrrroooorrrreoe ht:::/CopuII::1111:500
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks.
ArXio, hass/1122.09844, 2021. URL
Utkarsh Saxena and Kaushik Roy. Mcqueen: Mixed precision quantization of early exit networks. In British Machine Vision Conference, 2023. URL httpp://app...enatt.ccbcb//////p wew.aaannieccoolaa...ooarraaaaaa
Sangeetha Siddegowda, Marios Fournarakis, Markus Nagel, Tijmen Blankevoort, Chirag Patel, and Ab- hijit Khobare. Neural network quantization with ai model efficiency toolkit (aimct). arXiv preprint urXiv:2201.08//4 2022.
Steven H. Strogatz. Nonlinear Dynamics and Chaos: With Applications to Physics, Biology, Chemistry and Engineering. Westview Press, 2000.
Wonyong Sung, Sungho Shin, and Kyuyeon Hwang. Resiliency of deep neural networks under quantization. Arxiv, abs/1511
164.0151 06488, 2015. URL pestps://apts semanticscholar org/CorpusID::422555)
15

---

Published in Transactions on Machine Learning Research (11/2025)
Shyam A. Tailor, Javier Fernández-Marquee and Nicholas D. Lane. Degroc-quant: Quantization-aware training for graph neural networks. ArXiu, 0oss/2008.00000 2020. URL schaa
http/CorpuIIID2211388555
Nobuyuki Umetani and B. Bickel. Learning three-dimensional flow for interactive aerodynamic design. ACM Transactions Graphics (TOG), 37:1 10, 2018. URL UtLl:ttt.s/ppp...ssaattcccooorro wttp:aarrrroroooorssi..S 21750671
W.V.S.O. van den Dool, Tijmen Blankevoort, Max Welling, and Yuki M. Asano. Efficient neural pde-solvers using quantization aware training. 2023 IEEE/CVF Internalional Conference on Compuler Vision Work- shops (ICCVW), pp. 1415-1424, 2023. URL ott.p//ttt//pppp.eeaaaaootooooo. http:/aataattttaoaoooooooooooo/o org//oorpusDD::209011.
Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Neural Information Processing Systems, 2017.
URL https://api. www.nticschollattt......ocoaaaaa Irg.aaauaarrrorrCaaaaa.. 13756448
Hongyu Wang, Shumiing Ma, Li Dong, Shaohan Huang, Huaijie Wang, Lingxiao Ma, Fan Yang, Ruiping Wang, Yi Wu, and Furu Wei. Bitnet: Scaling 1-bit transformers for large language models. ArXiv, abs/ 215/2230.11453, 2023. URL httpss//api...... http:/ppp...omanicccchhhollr http.//aaaa..........ooooooaaaaa
Kuan Wang, Zhijian Liu, Yujun Lin, Ji Lin, and Song Han. Haq: Hardware-aware automated quantization with mixed precision. 2019 IEEE/CVF Conference on Computer Vision and Pattern Recogmition (CVPR), URL https://apls www.seicscaninnooos........ http:::/ppp...eeaattttccooooooot
pp. 8604-8612, 2018.
wrw/CorpusID: 102350477
Sifan Wang, Jacob H. Seidman, Shyam Sankaran, Hanwen Wang, George J. Pappas, and Paris Perdikaris. Cvit: Continuous vision transformer for operator learning. 2024. URL https: semanticscoolar
org/CorpusID: 2699832222
Bichen Wu, Yanghan Wang, Peizhao Zhang, Yuandong Tian, Péter Vajda, and Kurt Keutzer. Mixed precision quantization of convnets via differentiable neural architecture search. ArXiv, abs/1812.00e 2018. URL
https://api.senaatt semanticscholart org/CorpusID 54441287
Wentao Zhang, Yu Shen, Zheyu Lin, Yang Li, Xiaosen Li, Wen Ouyang, Yangyu Tao, Zhi Yang, and Bin Cui. Pasca: A graph neural architecture search system under the scalable paradigm. In Proceedings of the ACM Web Conference 2022, pp. 1817-1825 2022.
Zeyu Zhu, Fanrong Li, Zitao Mo, Qinghao Hu, Gang Li, Zejian Liu, Xiaoyao Liang, and Jian Cheng, A2q: Aggregation-aware quantization for graph neural networks. ArXin, abs/2302.00199 2023. URL https://app semanticscoolar http://ww...............oop/oaaa 0rr/010991111-222555556
16

---

Published in Transactions on Machine Learning Research (11/2025)
A Training and architecture details
MPNN Apart from the default size of 6 hidden layers and 128 channels, we also implemented the MPNN used in the Darey and ShapeNet- Carpe experiments in a small variant (4 layers, 64 channels) and a large variant (256 channels). Results on these other variants are in 5 In each case, we use Adam (Kingma & Ba 2015 with a weight decay of 1 x 10-6, 500 epochs, and normed gradient clipping of 1.0. For Darcy 10 is used, respecitvely, and batch size of 16. For ShapeNet-Car, the learning rate is 1 10-3 and the batch size is 8. The (train and validation) loss in Darcy is the relative norm, as in Li ct al. 2020). ; and for Elasticity and ShapeNet-Car,
and Elasticity, learning rates of 1 x 10 and
MSE. For all meshes, we use K-nearest neighbors to get edges, with K = 5, self-loops included, and K = 10 for Elasticity. Mean aggregation is used to aggregate message features. Furthermore, a cosine learning rate decay with linear warmup of 5 epochs is used.
GraphViT For the GraphViT, we implemented quantized variants of each of the original modules used in Janny et al. (2023). We use 64 hidden channels for all node and edge-related operations and 128 hidden channels for the cluster state space. We apply message passing layers before the cluster processing layers. These update both the nodes and edge features, with the original edge features being defined as the norm and relative distance between nodes. After the cluster processing, which is done using gated recurrent unit to obtain cluster features followed by self-attention (sce Janny et al. (2023) for details), final messagge-asssing passing layer is applied on the nodes that have the cluster features appended to them. We use a cosine learning rate decay schedule with 10 linear warmup epochs and a K-means clustering algorithm with K = 400. We train for 1000 epochs, using Adam (Kingma & Ba 2015) with weight decay of 1 10-6, normed gradient clipping of 1.0, learming rate 5 10-4 and batch size of 8. The edges are provided with the EAGLE dataset. We also train large variant that uses 96 channels for nodes and edges and 192 for the clusters, as well as a small variant using only 2 message-passe pasggeggpasnnggggg layers before the cluster processing, 32 channels for the mossago-passing layers and 64 channels for the clusters. Results on these other variants are in For each network, the (train and valid) loss is MSE, with factor of a 0.1 applied to the pressure channels, as in (Janny et al. 2023
Auxiliary model The auxiliary model is a simple message-passing GNN with 3 hidden layers of 32 chan- nels. Its inputs are the same as the main model, and it has one output channel representing the proxy measure of prediction complexity in cach node. Training settings for the auxiliary model are identical to those for the main model that it operates with. The target is the normalized, detached loss of the main model, smoothened using 10 graph diffusion steps. Fach graph diffusion step is implemented as
4 -
2
1 4,.
L 2
(15)
jEN(i)
B Quantization
We implement quantization using the default settings of the aqt Research (2023) library. This means that the scale factor is determined per channel, based on maximum absolute value calibration, there are no learned quantization parameters, and activations are quantized symmetrically. We do note that this is not optimal because our frequent use of the GELU activation function means most activations are positive or relatively small when negative. However, the effect of finding a better quantization strategy would apply to both the uniform and the adaptive experiments. To keep things simple, we decided to stick with the default scttings, leaving general improvements open for future work.
We did not optimize quantization for speed or accuracy because the latest and developing techniques in that field are mostly orthogonal to our approach. Also, the quantization library aqt we used only simulates the reported bil-width while executing in Int8 on hardware. Consequently, we report MACs rather than latency measurements, as the latter would not accurately reflect the potential gains from optimized hardware implementation. For the uniformly quantized runs, we do note that intermodiate bit-widths like Int5 do not exist on hardware in practice, making our combined approach of Int4/In Int8 already more hardware- friendly. Our primary goal is to demonstrate theoretical feasibility, with hardware-ppecific optimizations
17

---

Published in Transactions on Machine Learning Research (11/2025)
being beyond our current scope soppe-particularly many such optimizations are complementary to core contribution. Nonetheless, for completeness, an indicatory analysis runtimes is given in Appendix Section D
Table 3: Validation loss for different auxiliary model sizes across various Int4/Ints ratios for MPNN on ShapeNet-Car.
Int4/ /It88 ratio
Default
Small
Tiny
90%
0.0239
0.0232
0.0230
80%
0.0207
0.0206
0.0207
70%
0.0201
0.0197
0.0196
50%
0.0192
0.0193
0.0190
Table 4: Validation loss for different auxiliary model sizes across various Int4/In computation ratios for GraphVit on EAGLE.
Int4/ Int8 Int4/Int8 ratio
Default
Tiny
80%
0.0028 0.0024
0.0031
50%
0.0023
C Results with different model sizes
To test the generalizability of our framework across different main model sizes, we reproduce the experiments from the main body and reduce/inctee inncrese the number of MACs considerably. The results shown in Fig. 5 Architecture details of small and big variants are in the previous Appendix Section A Overall, reproduce previous conclusions and demonstrate that the performance benefits are preserved across scales.
ShapeNet-Car
EAGLE
0.069
4-bit
Uniforrm
0.040
Uniform Adaptive
0.0070 Sit
Uniform
8.036
4-bt
Uniform Adaptive
0.055
Adaptive
6.035
Adaplive
0.034
0.032
0.0000
0.665
0.049
5/t
0.0253
Saz
0.026
0.025
G-r..
EbX
0.0045
0.0224
0.025
0.020
0.0000
0.022
0.029
0.0035
0.020
bit
MACs
800
1000
1200
MACs
1400
1500
MACs
100
120
500
700
MACs
800
900
1000
Figure 5: Performance of smaller (left) vs larger (right) main model on ShapeNet-Car and EAGLE lasks.
Additionally, we test how the size of auxiliary models affects the quantization result (sec Table 3) We introduce a small and tiny variant: The small variant uses Int4 quantization the auxiliary model, instead of Int8. The tiny variant also uses 24 features and hidden layers instead 36, 3 respectively. We found the framework to be robust to the design choice. Overall, even tiny model is sufficient for bit-width allocation, which is beneficial its inference provides little to no overhead. For completeness, Fig, 0 shows the overhead of the Auxiliary model compared to other layers in the GraphViT and the MPNN model that we used, indicating its negligible overhead.
18

---

Published in Transactions on Machine Learning Research (11/2025)
GraphVit MACs distribution
Attention
6. 6.7%
GraphPooling
6.8%
layer_o 0
6.7%
15.9%
16.8%
2.5%
Embedding
layer 1 15.9%
0:20
44%
Aux
15.9%
24.8%
layer fiaal
layer_2
72.6%
GNN Jlyers
MPNN MACs distribution
layer
16.5%
aaer 2
layer
16.5%
16.5%
GNN JNN_layers 98.9%
1.1%
Aux
leyer_3 3
16.5%
16.5%
layer
16.5%
layer_4 4
Figure 6: Distribution of MACss for default size GraphViT model on EAGLE datasct (top) and MPNN Darcy dataset (bottom). The method demonstrates robustness to the size of the auxiliary model, with even minimal overhead sufficient for guiding effective bit allocation.
19

---

Published in Transactions on Machine Learning Research (11/2025)
D
Runtime analysis
Because training costs of Neural Operators can be amortized (train once, deploy on many different PDE initial conditions) we conduct runtime analysis by focusing on inference, assuming weights are already quantizod. Our analysis targels malrix multiplications ions, which dominate the cost, and because other identical pipeline components cancel when comparing mixed precision with uniform quantization. While compute generally scales quadratically with matrix size and memory traffic scales lincarly, the latter can still become the bottleneck compute capability of modern GPUs keeps increasing. To investigate to what extend MACs is reasonable measure of efficiency, we profiled custom CUDA kernels (based on CUTLASS) for our mixed- precision GEMMs using NVIDIA Nsight Compute.
Quantized GEMM can be decomposed into: (i) quantization of activations (weights pre-quantized), (ii) GFMM, and (iii) dequantization of outputs. Kernel runtime depends on read, write, and compute operations. Since dequantization outputs are typically Int32, writes are costly and in practice GEMM and doquantization are often fused. However, because such customization options are limited for low-bitwidth GEMMs in CUTLASS implemented quantization, GEMM, and dequantization separately for both uniform and mixed-precision variants. We approximate fused runtime by omitting redundant writes and reads via mull pointers and dummy indices.
For runtime evaluation, we evaluate a range of settings designed to represent realistic deployment scenarios. In particular we assume inputs and outputs are float16, weights are already given and quantized to Int4. Low-prccision nodes use int4 activations, high-precision nodes Int8. While mixed scttings rely only on Int4x Int4 kernels and custom quantization/ quantzzaiion/eequanizziiin /dqquanizzztion,, uniform runtimes cannot be fully measured since Int5-7 kernels exist in hardware. This makes ability to operate between Int4 and Int8 inherent advantage. For comparison we linearly interpolate uniform runtimes between Int4 and Int8. We use 512 input/output channels and vary node counts from 4096 to 65536.
Ie10
4096 nodes
lell
16384 nodes
lel1
65536 nodes
150
3.5
MACs
MACS
5.5
MACs
Runting
Runtime
140
Runtime
5.0
3.0
Runtime (ji)
130
Runtime (18)
MACs
4.5
450
Runtime
120
4.0
400
2.0
110
3.5
350
100
1.5
300
0.0
0.2
1.0
0.6
0.8
Int8 fraction
Int8 fraction
IntB fraction
Figure 7: Runtime results vs. MACs. At 0% or 100% int8 the setup is uniform; intermediate points mixed- precision. Red squares represent the run-time results, and the corresponding blue circles their respective MAC. The different y-axes are scaled so that coincident (red and blue) lines would imply a linear relat conshio There are no uniform implementations between the two extremes, but if for comparison purposes one assumes linear trend from the Int4 to the Int8 mcasurement, this would correspond to the blue (MACs) line. We find that MACs appear to be reasonable, in some cases even slightly conservative, estimate of runtime savings.
E Results on larger meshes
As ablation to test our method on larger-scale problems, we experiments using MPNN on two extra datasets:
ShapeNet Car Large Umetani & Bickel (2018). containing 31,186 nodes. It is different from our default variant, as it also includes the air around the car. The model is trained predict both the velocity and pressure in this significantly larger larger 3D domain.
20

---

Published in Transactions on Machine Learning Research (11/2025)
Darey Large. This is higher resolution variant of the Darey dataset used in the other experiments, containing 11236 nodes. As we did not proportionally scale the model, in particular using the same number of neighbors per node, the task is significantly harder.
Table presents results for three key quantization configurations that represent the most practically relevant scenarios. These configurations correspond to hardware-impleennttabee quantization levels, as intermediate bit-widths (Int5-Int7) are typically unavailable.
ShapeNet Car Large
Darcy Large
Configuration
MACs
Loss
MACs
Loss
Int4 (uniform)
1619
0.0226 + 0.0002
645
0.132 = 0.0002
90% Int4, 10% Int8
1781
0.0167 + 0.0007
710
0.125 = 0.0003
Int8 (uniform)
3239
0.0103 = 0.0002
1290
0.120 = 0.0017
Table 5: Performance comparison large-scale datasets. The adaptive mixeddpreeissie setup recovers roughly half of the accuracy loss from Int8 to Int4 by converting just 10% of nodes to Int8.
The results demonstrate the effectiveness of our adaptive approach on larger problem sizes. The mixed- precision configuration shows that strategic allocation of higher precision to small fraction of nodes can bridge much of the performance gap between uniform Int4 and Int8 quantization, while maintaining com- putational efficiency close to the Int4 baseline.
21


Reviews:
[{"review_id": "oxFSSj5zY4", "reviewer": "Reviewer_Wnvc", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors introduce AMQ, a jointly trained loss-guided bit-width scheduler for GNN-based PDE solvers that adaptively ups precision in hard regions under a fixed compute budget, with an efficient mixed-precision implementation. The method is evaluated on several high-dimensional data sets, with generally encouraging results and improvements over uniform quantization at matched MACs.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The empirical results show Pareto gains over uniform quantization at matched MACs across four tasks and two architectures, and include ablations (random/low-weight assignment; surrogate-loss vs uncertainty proxy). That said, the presentation can be improved and some additional experiments can provide a clearer picture of the results (see requested changes). Additionally, limitations of the method need to be discussed.", "audience_yn": "Yes", "audience_explanation": "The paper is very well written and strikes a good balance between being self-contained and assuming background expertise. The topic is timely; adaptive precision for neural PDE solvers/graph operators is of broad interest to the ML-for-Science, systems/quantization, and GNN communities.", "requested_changes": "- The notation can profit from a general proofreading to ensure consistency. For instance, bold-font should be use consistently to denote vectors, including complexity weight vectors (**Algorithm 1**). Allocation ratios should be constrained to be positive. Matrices should be denoted either with an upper or a lower case, but not interchangeably so (consider also using bold font for matrices in keeping with standard notation). Some typos are also present, especially in the **Experiments and Results** section.\n- I feel that the role of the smoothing function in Eq. 7 needs to be explained in more detail.\n- $\\mathcal{Q}$ in Eq.9 was never defined previously, but was simply referred to as $Q$ (see point on consistency).\n- $B$ from **Algorithm 1** was redefined to $\\beta$ in Eq. 9.\n- Instead of fixing the network hyperparameters in all experiments, it would have been helpful to define model categories (e.g., \"tiny\", \"base\", \"large\"), as is typical in many DL fields and further explore performance as a function of model size.\n- Why do some of the loss curves in Figure 3 display a small uncertainty region whereas other do not?\n- Table 1 should definitely subdivide the results per data set and use NRMSE for comparison in case the outcomes have different scales. It should also show the loss for the adaptive method (and quantization ratio used) and specify that this is indeed validation loss. What is the error computed over currently? Data sets? Different model runs? Bootstrap? \n- The results displayed in Figure 4 (EAGLE) data set are impossible to interpret for non-experts. How was the input selected? What does the 30% refer to?\n- The surrogate uncertainty loss comparison needs a bit more motivation and explanation. It is important to note that this is aleatoric uncertainty coming from a heteroskedastic model, as described in Kendall & Gall (2017), and its quality is strongly dependent on the Gaussian assumption being legitimate.\n- There are no limitations of the method mentioned. Consider discussing these. Here is what I think and I am happy to discuss with the authors:\n1. The need to train an auxiliary model which may be difficult to train for complex inputs / low data settings).\n2. If spatial complexity shifts into an out-of-distribution (OOD) regime (e.g., new obstacle layouts), does the bit assigner still focus on the right regions? Consider adding a held-out scenario split for EAGLE and demonstrating a case where the method breaks.\n3. Did I understand correctly that all weights are fixed at Int8 and only activations are quantized? If that's the case, I would discuss differences between adaptive activations vs adaptive weights+activations and possibly provide an ablation to justify the choice.\n4. AQT simulates bit-widths, but in real hardware, MACs does not equal latency, as the authors point out. Perhaps a reiteration of this point in the discussion would be helpful.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "xvW3wIXp6c", "reviewer": "Reviewer_uJSi", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes an approach that adaptively allocates bit precisions to different grid resolution and improves computation efficiency.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The proposed algorithm is tested on several large-scale problems, and the methodology seems sound to me.", "audience_yn": "Yes", "audience_explanation": "Reducing computation cost is an important problem in Neural PDE methods.", "requested_changes": "+ In this work, the computation cost is measured by the total number of MACs executed during the training. While I agree with the authors view that wall-time is hardware-dependent, it feels that this metric is a bit obsolete. Actually, there are other aspects which may consume a significant amount of time, such as data reading, new model training and evaluating and GPU/CPU switching, that are rarely concerned in traditional methods. Therefore, it would be much more convincing to show that the proposed method achieves better wall-time efficiency, not just MAC efficiency.\n\n+ A minor concern: as shown in Figure 3, the validation loss for uniform and adaptive methods is not significant as they are on the same magnitude. I am fine with the results because the core contribution of this work is not the performance improvement, but a bit curious why the curve of Adaptive stops exactly at the intersection with the Uniform curve in the third and fourth plot. It seems that the Uniform curve will go under the Adaptive curve if MACs keep increasing.\n\n+ I recommend the authors to release the code upon the acceptance of this work.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "vTvVULtpgO", "reviewer": "Reviewer_n9Fj", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The high-level idea of this paper is straightforward: it aims to develop an efficient neural PDE solver by adaptively adjusting the quantization to accommodate the possibly different computational needs across various regions in a fixed mesh topology. The perspective of this paper largely focuses on advocating the power of quantization models, which is different from (and complements) many existing neural PDEs that advertise new network designs, algorithms, or data. The paper evaluates its method on typical PDE test problems in scientific machine learning and reports its performance against uniformly quantized baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The technical method looks sound to me at a high level. I am not an expert in quantization models, so I will leave it to other expert reviewers.\n\nThe experiments are convincing enough to showcase the proposed adaptive quantization model’s benefits over a uniformly quantized one. The experiments also cover a reasonably diverse set of typical PDEs in fluid and solid mechanics.", "audience_yn": "Yes", "audience_explanation": "The paper tells an interesting and, as far as I know, novel story. Incorporating quantization methods into neural PDEs is an underexplored topic, and the paper makes a timely contribution in this regard. I believe this combination is not the end of an idea but the start of several interesting follow-ups.", "requested_changes": "I will start with two high-level comments on the paper’s baselines and scalability, respectively.\n\nBaselines: The storytelling in the introduction emphasizes the shortcomings of adaptive mesh refinements (AMR) and uses them to motivate the proposed AMQ. Surprisingly, the paper does not plan to compare them (“While our method is inspired by and conceptually related to AMR, we do not aim for a direct comparison”). I am not entirely sure whether excusing the paper from such a comparison is fair to AMR, and I would like to be convinced by a simple motivating example that can present the pros and cons between AMR and AMQ, highlighting the paper’s claimed motivation from AMR to AMQ. Alternatively, I’d also be okay with an overhaul of rewriting the abstract/introduction to completely remove AMR from the story.\n\nScalability: Despite the positive results across various PDE problems, their problem sizes remain relatively small: The test problems contain at most ~3k node numbers, most of which are in 2D only. For a 3D problem, 3k nodes would be insufficient even for a coarse grid resolution of 20 x 20 x 20. Therefore, it is unclear to me whether the proposed model can be scaled up to a practically useful level. I appreciate that the paper has commented on this topic in the “Discussion and Conclusion” section, but I still would like to bring this up and welcome any additional experiments.\n\nThe following changes will have critical effects on my evaluation of this paper. The goal of these changes is to help us better understand the proposed method’s properties:\n1. Comparing AMR with AMQ in a simple, motivating example and discuss their pros and cons. See “Baselines” above.\n2. Select one experiment, scale its problem size up until the proposed method struggles to learn an effective quantization model, and discuss the results. See “Scalability” above.\n\nI would also like to request a few text changes: \n1. The title “adaptive mesh quantization” looks confusing to me because meshes are actually fixed in the proposed method. Please consider renaming or clarifying it.\n\n2. “Compared to models working on regular grids, they [graph neural networks] do not require interpolation, thereby reducing both cost and error.” I cannot see how the difference between regular grids and meshes (graphs) can lead to the conclusion that GNNs do not require interpolation. Please justify.\n\n3. “In certain domains, such as molecular data or graph-based representations of real-world systems, nodes represent specific, meaningful points (e.g., atoms in a molecule or locations in a sensor network) that cannot trivially be moved.” The molecular data argument here is a bit contrived because it actually does not apply to the paper’s problem setup (PDEs in fluid/solid mechanics). I suggest replacing it with a better motivation for quantization models.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)