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

What it claims

Vanilla PI-DeepONet duplicates every input function once per collocation point before it enters the Branch network. That forces the Branch network to make n×p passes when n would do. ATenH (Accelerated Tensor Handling) removes the duplication. It runs the Branch network once per input function and broadcasts the result across that function's p collocation points. The outputs are mathematically identical, and the Branch passes drop by a factor of p. The authors do this in plain PyTorch, with no custom kernels.

How it works

In the Vanilla version (Algorithm 1), each input function u(j) is repeated p times, and the collocation points are tiled n times. Both networks then see a flat batch of n×p rows, and their outputs are combined by a row-wise inner product.

ATenH (Algorithm 2) keeps the input functions unrepeated. It reshapes the Trunk input into a 3D tensor and combines the two outputs with an einsum, Gjk = Σℓ Bjℓ Tjkℓ. The result has shape (n, p, 1) and can be reshaped to Vanilla's (n×p, 1).

Symbols: n is the number of input functions, p the collocation points per function, m the sensor points, d the spatiotemporal dimension, and s the latent width.

Physics derivatives still come from autodiff along the Trunk network. The loss is L_PDE + L_BC + L_IC with no data term, so all experiments are labelled-data-free.

The change only touches tensor layout. If your PI-DeepONet code repeats the Branch input, this is roughly a few-line refactor.

ATenH is also complementary to Separable PI-DeepONet. That method splits the Trunk network per coordinate axis to cut Trunk cost, but it works best on separable problems. ATenH addresses the Branch side and assumes no separability.

Evidence

The fetched text cuts off at the start of the scaling-law derivation (Section 2.4). Only the abstract's figures are available:

• comparable accuracy on 1D, 2D and 3D linear advection;
• a maximum 3.75× training speed-up in 2D;
• empirical training time scaling linearly in p and n, as predicted;
• 3D training fitting on an A100 only with ATenH.

Per-case timings, error metrics and resolutions are not in the fetched text.

Limits

• ATenH requires the same number of collocation points for every input function. Vanilla allows a different p(j) per function. If your sampling varies p per function, you would need padding or must stay with Vanilla.
• The only benchmark is linear advection. Nonlinear PDEs and higher-order derivatives, which the authors name as the main cost drivers, are not tested in the visible text.
• ATenH does not reduce Trunk cost. Trunk evaluations still scale as N^d, so the curse of dimensionality on that side remains.
• The fetched page ends at Section 2.4, so the derivations, results tables and discussion were not reviewed. No code release is mentioned in the visible text.