PaperlayerView exact revision on arXiv

Source-linked AI summary

The importance of tensor handling and implementation of physics-informed DeepONets

Nahum Banks, Gustavo A. M. de Almeida, Sergio Maldonado

arXiv:2610.22157v1cs.LGphysics.comp-ph

TL;DR

Conventional PINNs typically require retraining when problem conditions change, while PI-DeepONets face computational costs that grow with dimension. This paper modifies Branch-network processing and analyzes the resulting scaling; ATenH reduces training time with comparable accuracy and enables a reported 3D case that Vanilla could not train within memory limits.

  • Problem

    PI-DeepONet training is computationally costly, particularly for high-dimensional systems and higher-order derivatives, while conventional PINNs typically need retraining when conditions change.

    Source: §1 · p. 2

  • Method

    ATenH evaluates the Branch network once per input function and reuses its encoding across collocation points, alongside a theoretical model of training-time scaling.

    Source: §4 · p. 16

  • Results

    The 3D Vanilla model failed due to memory limitations, while ATenH completed training with a mean predictive error of 5.44%.

    Source: §5 · p. 16

  • Takeaways & Limitations

    Implementation choices can materially affect PI-DeepONet practical performance, and ATenH demonstrates computational gains without sacrificing predictive accuracy.

    Source: §5 · pp. 16–17

    ATenH requires all output functions to use the same number of collocation points, which may restrict data-driven settings with varying labelled-data amounts.

    Source: §4 · p. 16

Abstract

from arXiv · show

Physics-informed neural networks (PINNs) reduce or eliminate the need for large labelled datasets by enforcing governing equations during training. Physics-informed Deep Operator Networks (PI-DeepONets) extend this approach to learning solution operators, enabling rapid prediction across families of problems. Despite these advantages, we show that their application to higher-dimensional problems remains limited by the computational cost of training. To address this issue, we propose a modified PI-DeepONet architecture that reduces this cost by removing redundancy in how the Branch network processes input functions across collocation points. The proposed formulation is mathematically equivalent to the original architecture but reduces input tensor sizes and the number of Branch network evaluations. We analyse the computational complexity of both architectures and derive theoretical scaling laws for training time with respect to the number of collocation points and input functions. The proposed architecture is evaluated on 1D, 2D, and 3D linear advection problems. In all cases, it achieves comparable predictive accuracy to the original PI-DeepONet while consistently reducing training time. In 2D experiments, a maximum speed-up of $3.75\times$ is observed, with empirical training-time scaling agreeing with the predicted linear dependence on collocation points and input functions. Furthermore, training a 3D model on an NVIDIA A100 GPU is only feasible using the proposed architecture, as the original formulation exceeds memory limits. These results demonstrate that careful consideration of how PI-DeepONets are implemented can yield substantial gains, enhancing the practical performance of PI-DeepONets and enabling their application to otherwise intractable physics-informed learning problems.

1 Introduction

Conventional PINNs are tied to fixed conditions, while PI-DeepONets learn solution operators across input functions but face high computational costs as problem dimension grows. This paper proposes ATenH and analyzes its computational gains, with experiments intended to evaluate accuracy and tractability across increasingly challenging problems.

1 Introduction

  • Conventional PINNs typically require retraining when domains or initial or boundary conditions change, limiting reuse across problem setups.They take only PDE independent variables as input and incorporate conditions indirectly as soft loss constraints.Source: §1 · p. 2
  • Neural operators address this limitation by learning mappings between function spaces, and PI-DeepONets provide a physics-informed operator-learning framework.DeepONet combines a Branch encoding of input-function sensors with a Trunk encoding of query locations.Source: §1 · p. 2
  • PI-DeepONets face high computational costs because Trunk evaluations scale as N^d, limiting broader adoption and high-dimensional applications.The burden is more apparent for high-dimensional systems and higher-order derivatives.Source: §1 · p. 2
  • Separable PI-DeepONets reduce Trunk evaluations from N^d to Nd, but work best when the physical system supports separation of variables.The cited work is described as the only prior effort aimed at reducing PI-DeepONet training expense.Source: §1 · p. 3
  • The paper investigates ATenH, an alternative PI-DeepONet architecture intended to improve computational efficiency, and derives theoretical speed-ups relative to Vanilla.It also evaluates ATenH through increasingly challenging numerical experiments.Source: §1 · p. 3
  • The authors aim to demonstrate that ATenH makes some higher-dimensional PDE problems computationally tractable and supports greater spatial and temporal resolutions.The numerical results use PyTorch.Source: §1 · p. 3

Notation

Table 1 is identified as a glossary of key variables used throughout the work.

Notation

  • Table 1 is a glossary of key variables used throughout the work.Source: Table 1 · p. 4

2.1 General PI-DeepONet framework

The PI-DeepONet framework imposes the governing PDE and boundary and initial conditions through a composite loss. The setup also specifies the input functions and collocation points used to evaluate the operator.

2.1 General PI-DeepONet framework

  • A composite loss incorporates the PDE, boundary, and initial conditions, with an optional labelled-data term.In the reported tests, the data term is omitted, making training labelled-data-free.Source: §2.1 · p. 4
  • The governing equations specify an interior PDE condition, a boundary condition, and an initial condition for the solution operator.The passage defines N as a differential operator, B as a boundary operator, and v0(x) as the initial condition.Source: §2.1 · p. 4
  • The architectural setup treats each input function as discretized at m sensor points and evaluates its output at p(j) d-dimensional collocation points.Source: §2.1 · p. 4

2.2 Vanilla architecture

Vanilla PI-DeepONet repeats each input function for every collocation point so Branch and Trunk outputs can be paired pointwise. Its training procedure then evaluates the physics-informed loss and updates model parameters.

2.2 Vanilla architecture

  • Each input function is repeated p times, producing Branch inputs of shape (n × p, m) and outputs of shape (n × p, s).The Trunk inputs have shape (n × p, d), and the final output has shape (n × p, 1).Source: §2.2 · p. 5
  • The Vanilla training algorithm repeats input functions across collocation points before evaluating the physics-informed loss and updating parameters.The loss combines PDE, boundary-condition, and initial-condition terms.Source: §2.2 · p. 5

2.3 Accelerated Tensor Handling (ATenH) architecture

ATenH avoids repeating each input function across collocation points by reusing a single Branch-network pass. It changes tensor shapes and requires equal collocation-point counts across input functions.

2.3 Accelerated Tensor Handling (ATenH) architecture

  • ATenH reduces Branch-network passes by a factor of p, processing each input once and reusing its Branch output across collocation points.Branch input/output shapes change from (n × p, m) → (n × p, s) in Vanilla to (n, m) → (n, s) in ATenH.Source: §2.3 · p. 5
  • The ATenH formulation requires the same number of collocation points for every input function, unlike Vanilla.Source: §2.3 · p. 5
  • ATenH reshapes Trunk inputs to (n, p, d), broadcasts Branch outputs before the inner product, and returns outputs of shape (n, p, 1).The output can be converted to the Vanilla shape (n × p, 1).Source: §2.3 · p. 6

2.4 Theoretical scaling laws for a full training step

The analysis estimates training-step computational cost from weighted input products and accounts for forward passes, automatic differentiation, and output dot products. It compares Vanilla and ATenH costs to derive a theoretical speed-up.

2.4 Theoretical scaling laws for a full training step

  • The training-step cost analysis includes first-order physics derivatives and parameter-gradient backpropagation, but excludes boundary-loss computation.Boundary collocation points are generally much fewer than interior points used for the physics loss.Source: §2.4 · p. 6

2.4.1 Cost model

  • The cost model counts input-weight products, treating activation and bias operations as lower-cost pointwise work.Its cost is equivalent up to a constant factor to conventional FLOP counting.Source: §2.4.1 · p. 7

2.4.2 Vanilla architecture

  • Vanilla cost accounts for Branch and Trunk forward passes, dot products for each input-function/collocation-point pair, physics derivatives, and loss gradients.Physics derivatives backpropagate through the Trunk and dot product, while loss gradients traverse the full graph.Source: §2.4.2 · p. 8

2.4.3 ATenH architecture

  • ATenH uses Branch inputs of shape [n, m] and Trunk inputs of shape [n, p, d], with physics derivatives and loss backpropagation retaining the Vanilla structure.Source: §2.4.3 · p. 9

2.4.4 Theoretical Speed-up

  • The derived speed-up factor is independent of n and approaches a constant as p becomes sufficiently large, with other parameters fixed.Source: §2.4.4 · p. 9

2.5 Numerical tests

The numerical tests compare Vanilla and ATenH on linear advection across increasing spatial dimensions, using shared hyperparameter choices for a fair computational comparison. The experiments use Gaussian boundary inputs and dimension-scaled collocation sampling.

2.5 Numerical tests

  • The tests compare architecture training time and scaling across 1D, 2D, and 3D linear advection problems.The architectures share mathematical operations and differ in data handling, motivating common optimal parameters for a fair comparison.Source: §2.5 · p. 9
  • Boundary amplitudes are sampled uniformly from [0.1, 1], restricting the tested inputs to Gaussian boundary-condition profiles.The authors state this setting is appropriate for evaluating computational performance, though it cannot test extrapolation to qualitatively different boundary-condition families.Source: §2.5 · p. 10
  • Collocation-point sampling keeps average spacing constant across dimensions, and separate samples serve boundary, initial, and physics losses.This dimensional scaling yielded higher predictive accuracy in 2D and 3D than alternative sampling strategies tested.Source: §2.5 · pp. 10–11
  • Using 20 points along each dimension caused memory issues that prevented Vanilla from completing beyond the 1D case.The authors set 10 collocation points per spatiotemporal dimension; doubling this produced comparable accuracy while causing those memory issues.Source: §2.5 · p. 10
  • Figure 2 shows closely aligned training and validation losses, with loss decreases plateauing after about 50,000 epochs in 1D and 100,000 in 2D.The selected training durations were 50,000, 100,000, and 200,000 epochs for 1D, 2D, and 3D, respectively.Source: §2.5 · p. 10, Figure 2 · p. 11
  • Grid searches selected four hidden layers, widths of 40 neurons for 1D and 60 for 2D, while 3D used a heuristic width of 80.No dedicated 3D grid search was performed because of computational cost.Source: §2.5 · p. 11, Figure 3 · p. 12

3 Results

Across 1D, 2D, and 3D linear advection, ATenH has comparable predictive accuracy where both architectures can be compared and reduces training time. In 3D, ATenH trains where Vanilla exceeds GPU memory limits; the 2D experiments also show linear training-time scaling with collocation points and input functions.

3.1 1D Linear Advection

  • 12.1% lower mean training time (1.14× speed-up) for ATenH in 1D, with comparable predictive accuracy; the predicted speed-up was 1.64×.The comparison reports closely matching MRRMSE values and similar standard deviations.Source: §3.1 · p. 12

3.2 2D Linear Advection

  • 71.4% lower mean training time (3.49× speed-up) for ATenH in 2D, versus a predicted speed-up of 7.46×.Both architectures had comparable predictive accuracy, with closely matching MRRMSE values and overlapping standard deviations.Source: §3.2.1 · pp. 12–13, §3.2 · p. 12

3.3 3D Linear Advection

  • ATenH trained the 3D problem in 1.40×10^5 s with a mean predictive error of 5.438%, while Vanilla failed from insufficient memory.The theoretical speed-up was approximately 580×, but an empirical comparison was not possible.Source: §3.3 · pp. 13–15, Table 3 · p. 15

3.2.1 Computational Performance and Scaling for 2D Linear Advection

  • Training time increased linearly with internal collocation points and input functions for both architectures, matching theoretical scaling; ATenH had smaller cost prefactors.Figure 6 plots mean training time against each varying quantity and shows the ATenH speed-up; Vanilla’s two largest collocation-point cases are omitted for insufficient memory.Source: §3.2.1 · p. 13, Figure 6 · p. 14

4 Discussion

ATenH reduces redundant Branch-network processing by encoding each input function once and reusing that encoding across collocation points. The cost model and experiments show linear scaling in collocation points and input functions, while highlighting a constraint on equal collocation counts across outputs.

4 Discussion

  • ATenH evaluates the Branch network once per input function and reuses its encoding across collocation points, reducing input-tensor size and Branch operations.The modified architecture performs the same underlying mathematical operations as the original and achieves comparable accuracy with lower training time.Source: §4 · p. 16
  • All output functions must use the same number of collocation points, a potentially restrictive condition for data-driven settings with varying labelled-data amounts.The authors note this constraint does not apply in PI-DeepONets, where collocation points can be freely generated and chosen.Source: §4 · p. 16
  • The cost model and experiments support linear training-time scaling in collocation points and input functions, but measured speed-ups fall below theoretical predictions.The authors attribute expected discrepancies to unmodelled overheads including memory allocation, data movement, automatic differentiation graph construction, and optimiser updates.Source: §4 · p. 16

5 Conclusion

ATenH reduces PI-DeepONet computational cost through Branch-network input handling while maintaining predictive accuracy. It also makes a 3D linear-advection model trainable where Vanilla exceeds memory limits, underscoring the practical importance of implementation choices.

5 Conclusion

  • Branch-network input handling reduces computational cost without sacrificing predictive accuracy.Source: §5 · p. 16
  • Vanilla exceeded memory limits on 3D linear advection, while ATenH trained with a 5.44% mean predictive error.Source: §5 · p. 16
  • The authors describe ATenH as evidence that implementation choices matter, not as a final design, and expect theoretical and empirical scaling agreement to generalise across implementations.They suggest further optimisation of existing architectures and applying similar improvements to problems that have proved intractable for PI-DeepONets.Source: §5 · p. 17

Declarations

The declarations acknowledge facility support and UKRI funding, describe author contributions, and report no conflicts of interest.

Declarations

  • The authors acknowledge the IRIDIS High Performance Computing Facility and its associated support services at the University of Southampton.Source: Declarations · p. 17
  • The first author received a UKRI PhD studentship.Source: Declarations · p. 17
  • NB led conceptualisation, methodology, software, visualisation, analysis, and original drafting; GA and SM contributed conceptualisation, methodology, supervision, and review and editing.All authors discussed results, contributed to interpretation, and approved the manuscript.Source: Declarations · p. 17
  • The authors declare no conflicts of interest.Source: Declarations · p. 17

Related papers

Loading 2610.22157v1…