Research & AlgorithmsDANIEL SCHMITTER · RESEARCH & ENGINEERING

Research & Algorithms / ML Engineering

Scientific ML · E43 · Implementation

Differentiate a field four times without losing the graph

A classical PINN baseline exposes the cost of high-order residuals and the distinction between an analytical derivative and an automatically differentiated network.

Daniel Schmitter · September 19, 2026 · 6 min read

PyTorch autogradTanh MLPCustom residual helper
Differentiating with respect to coordinates builds the physics residual; differentiating the loss with respect to parameters trains the field.
Figure 1. Two kinds of differentiation. Differentiating with respect to coordinates builds the physics residual; differentiating the loss with respect to parameters trains the field. Archived fourth-order residual. Original vector illustration.

Follow the information

From input to outcome

Input differentiation constructs the residual. Parameter differentiation then propagates the combined objective back through that derivative graph into θ. Boundary samples are separate evaluations of the same field.

Input differentiation constructs the residual. Parameter differentiation then propagates the combined objective back through that derivative graph into θ. Boundary samples are separate evaluations of the same field.
Figure 2. Information flow. Solid arrows carry observations, tensors or artifacts; other routes are explicitly labelled. Signal shapes, matrices and network icons are schematic, not measured samples or literal neuron counts. Open full-size SVG ↗ On narrow screens, scroll the diagram horizontally.

Read this alongside Figure 1: Differentiating with respect to coordinates builds the physics residual; differentiating the loss with respect to parameters trains the field. The module map and layer-level figures below expand the operations in this route.

Differentiate a field four times without losing the graph: architectureCoordinates: requires_grad input → Tanh network: Predicted field u(x) → Nested derivatives: u′ → u″ → u‴ → u⁗ → Physics residual: Operator applied to field → Boundary + data loss: Optimizer objective. A high-level module map; comparison branches and training details are explained in the article.SCIENTIFIC ML / E43 / MODULE MAP01 INPUTCoordinatesrequires_grad input02 MODULETanh networkPredicted field u(x)03 MODULENested derivativesu′ → u″ → u‴ → u⁗04 MODULEPhysics residualOperator applied to field05 OUTPUTBoundary + data lossOptimizer objective
Source-grounded module map. Boxes summarize operations, not individual neurons; comparison arms and training paths are detailed below. On a small screen, scroll the diagram horizontally.
Coordinates — requires_grad input

The architecture in context

The system we are building

A physics-informed network can be trained against a differential residual evaluated at collocation points. When that residual includes fourth derivatives, the derivative computation itself becomes a major part of the training graph. The archive implements this explicitly with repeated autograd calls over a tanh coordinate network.

Who does what in the stack

PyTorch autograd
Builds nested coordinate derivative graphs.
Tanh MLP
Represents the unknown field.
Custom residual helper
Connects high-order physics to the optimization loss.

The custom helper constructs successive derivatives with create_graph enabled so the residual can still differentiate with respect to network parameters. PyTorch supplies the chain rule. The benchmark compares this training route with operator-aware alternatives under declared conditions.

Framework responsibility map. Each row maps a library or custom component to its job; rows are not a sequential inference graph.
Framework responsibility map. Each row maps a library or custom component to its job; rows are not a sequential inference graph. Open full-size SVG ↗

From module map to executable structure

Inside Fourth-order physics-informed MLP

Scalar coordinate, four 32-unit tanh hidden layers, scalar field output.

Layer-level implementation. B denotes batch size; parameter and shape conventions are expanded in the table.
Layer-level implementation. B denotes batch size; parameter and shape conventions are expanded in the table. Open full-size SVG ↗
Layer / tensor / operation ledger
Layer or branchOutput shapeImplementation detail
Coordinates require gradientsN × 1Separate collocation and two boundary coordinates.
Linear 1→32 + tanhN × 32Xavier-uniform weights, zero biases.
Three Linear 32→32 + tanhN × 32Smooth activations support fourth spatial derivatives.
Linear 32→1N × 1Predicted field uθ(x).
Nested coordinate differentiationN × 1 residualFour autograd derivative calls; boundary values and second derivatives form the boundary penalty.

There are two different differentiation paths: derivatives with respect to x construct the PDE residual; derivatives of that loss with respect to network weights train the model. The coordinate graph must remain differentiable through all four derivative operations. ReLU would not be an equivalent activation for this strong-form fourth-order residual.

The equation and the update

L=mean⁡x(uθ(4)(x)−f(x))2+200[mean⁡∂Ωuθ2+mean⁡∂Ω(uθ′′)2]\mathcal L=\operatorname{mean}_x(u_\theta^{(4)}(x)-f(x))^2+200[\operatorname{mean}_{\partial\Omega}u_\theta^2+\operatorname{mean}_{\partial\Omega}(u_\theta^{\prime\prime})^2]

Adam 2.5e-3; full collocation grid each epoch; loss=residual MSE+200×boundary penalty. Device selection in this script is CUDA if available, otherwise CPU—not MPS.

Learning or solution path. A parameter-update path is different from the forward inference path; see text for target-network, frozen-feature and local-loss boundaries.
Learning or solution path. A parameter-update path is different from the forward inference path; see text for target-network, frozen-feature and local-loss boundaries. Open full-size SVG ↗

Implementation card / no invented benchmarks

Capacity, budget and execution evidence

Parameters / retained state
3,265 trainable scalars; derivative graph storage is not included in that count.
Duration and hardware evidence
CPU fallback is intentional in the archived code. A proposed MPS port would need new operator-coverage and numerical checks.
Source coordinates
E43 lines 27–34, 41–58 and 93–129
Current reproduction context
Current workstation, supplied by the author: Apple M4, 128 GB unified RAM, 40 GPU cores and 16 CPU cores. This is context for prospective reproduction, not attribution of every archived run. Python and framework versions are not fully locked for these historical sources; declarations, when available, are identified separately.

Counts above are calculated from the stated layer shapes unless identified as saved measurements. They exclude optimizer state and nontrainable buffers. No archived training was rerun for this revision.

What these design choices change

The boundary weight 200 trades constraint enforcement against interior residual conditioning. A small parameter count does not make the nested derivative graph small. Increasing collocation points can increase memory and step time even though the architecture is unchanged.

Reproduction and measurement protocol

Before training, substitute an analytic polynomial into the derivative helper and compare its fourth derivative. Test boundary values and curvature separately. For memory, label the source’s CPU graph estimate as an estimate; it is not measured process RSS. Numerical error and memory should be compared at matched boundary and forcing conventions.

For a new run, save the resolved Python/framework versions, backend, dtype, seed, input shapes, batch size and exact source revision. Start with one batch and one update. Log training steps separately from epochs or environment steps. Do not equate the configured maximum with a completed budget or convergence.

Measure initialization/compilation, data preparation, warmed forward pass, training updates and evaluation separately. Synchronize accelerator work around timed regions using the chosen framework’s supported mechanism. Report peak process memory and framework allocation separately; parameter bytes exclude activations, gradients, optimizer state and input buffers. On a shared machine, begin with a single CPU worker and a small batch rather than claiming all available resources.

A closer look at the implementation

The code that carries the idea

The excerpt repeatedly asks for a coordinate gradient using a ones-like output seed. For a pointwise scalar-output network evaluated independently at each coordinate, this gives the desired per-point derivative. Cross-sample coupling would require revisiting that assumption.

Python · file · lines 27–34
def fourth_derivative(y: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
    dy = torch.autograd.grad(y, x, torch.ones_like(y), create_graph=True, retain_graph=True)[0]
    dyy = torch.autograd.grad(dy, x, torch.ones_like(dy), create_graph=True, retain_graph=True)[0]
    dyyy = torch.autograd.grad(dyy, x, torch.ones_like(dyy), create_graph=True, retain_graph=True)[0]
    dyyyy = torch.autograd.grad(dyyy, x, torch.ones_like(dyyy), create_graph=True, retain_graph=True)[0]
    return dyyyy

Verbatim archive excerpt from compare_classic_pinn.py. Context-dependent historical code, not a standalone runnable program. Comments retain their original wording; the article distinguishes implemented behavior from stale or overbroad comments.

The boundary that matters

A memory estimate derived from parameter or activation counts is not measured peak process or device memory. Automatic differentiation is exact for the represented computational graph up to numerical arithmetic; it does not make the fitted field or PDE solution exact.

Keep building

Other posts of interest