# Physics-Informed AI, Part III: From Post-Hoc Checker to Differentiable Physics Head

> Source: <https://pub.towardsai.net/physics-informed-ai-part-iii-from-post-hoc-checker-to-differentiable-physics-head-003c7738ff98?source=rss----98111c9905da---4>
> Published: 2026-07-09 17:01:04+00:00

In Part II, I fine-tuned a small LLM with LoRA to produce structured engineering JSON. The model learned to follow the required schema, but it did not become a physics solver. A deterministic checker still recomputed the energy-balance residual after generation and assigned the final validated status. That division of labor was deliberate: the model handled the language and structure, while the checker owned the physics.

That pattern is useful in engineering workflows, but the physics check sits entirely downstream of generation. The residual can reject or correct an output, but it does not touch the model’s weights during ordinary supervised training. Part III closes that gap by moving the physics residual into the training loop.

The tempting approach is to let the model generate text, parse numbers from that text, compute a residual, and backpropagate from it. For example, suppose the model emits a value such as "u": 0.532. You can parse "0.532" into a float and compute a residual, but the gradient path has already been broken. Token sampling and argmax decoding are non-differentiable, and string construction, regex parsing, and float() conversion sever the computation graph completely.

A residual computed from parsed text can still serve as a checker, or possibly as a reward in an RL-style loop. It is not a clean supervised physics-loss signal. The way around this is to compute physics on tensors instead of parsed text. In Part III, the language model is used as an encoder. Its hidden representation conditions a differentiable numerical head, the head predicts a tensor output, and the PDE residual is computed on that tensor. That differentiable numerical path is the architectural move behind everything that follows.

I used a deliberately controlled PDE: the one-dimensional heat equation,

with zero boundary conditions and the initial condition u(x, 0) = sin(pi* x). That problem has a closed-form solution, which gives me exact targets to measure against:

All RMSE values below are computed against this exact solution on a uniform x–t evaluation grid over x ∈ [0, 1], t ∈ [0, 5]. The model reads a natural-language prompt:

```
Solve the 1D heat equation with alpha = 0.05,initial condition u(x,0)=sin(pi*x),and zero boundary conditions on x in [0,1].
```

and maps it, together with the coordinates, to a predicted field:

The head predicts a tensor, not text, so autograd can take it from there:

Training combines four terms — data loss, initial-condition loss, boundary-condition loss, and the PDE residual loss.

The numerical head is a 4-layer tanh MLP with 128 hidden units per layer. DistilBERT is frozen; its pooled text embedding is concatenated with (x, t) before entering the head. All models train with Adam for 2,000 steps, using 2,000 collocation points for the PDE residual. The data, initial-condition, and boundary-condition losses all have weight 1. The PDE weight is what varies by mode: 0.01 for weak_physics, ramped from 0 to 0.01 over training for scheduled_weak_physics, and 10⁶ for too_high_physics.

One design choice is worth stating plainly, because the later ablations depend on it. In the main DistilBERT model, the numerical head never sees the numeric value of alpha; it gets only the text embedding and x,t. The residual still uses the known numeric alpha, because alpha is part of the physical law you evaluate the residual against. So the model is not being asked to rediscover the heat equation. It is being asked to turn a text-derived condition into a numerical field, and to have that field agree with both the sparse data and the physics.

I compared five training modes. The first two are baselines: data_only fits only the sparse supervised samples, while constraints_only adds the initial- and boundary-condition terms without using the PDE residual. I keep constraints_only in the plots as an intermediate diagnostic, but the text focuses on the main comparison between data-only training, weak physics, scheduled weak physics, and the too-high-physics failure case. The next two add physics regularization. weak_physics uses a small fixed residual penalty, and scheduled_weak_physics ramps that penalty during training. The final mode, too_high_physics, deliberately overweights the PDE residual. I included it as a stress test: not just to ask whether physics helps, but to see where it starts to hurt.

The sparse-data runs spanned four diffusivities (alpha = 0.01, 0.05, 0.10, 0.20) and three data budgets (n_data = 5, 10, 25), each repeated over three seeds so I could separate the effect from run-to-run noise. Weak and scheduled physics improved solution RMSE over data-only training at every data budget, and the gap widened as data grew scarcer.

At n_data = 5:

```
data_only:        0.161 ± 0.034weak_physics:     0.068 ± 0.018scheduled_weak:   0.068 ± 0.015
```

At n_data = 10:

```
data_only:        0.068 ± 0.013weak_physics:     0.043 ± 0.010scheduled_weak:   0.043 ± 0.009
```

At n_data = 25:

```
data_only:        0.042 ± 0.007weak_physics:     0.029 ± 0.008scheduled_weak:   0.032 ± 0.008
```

Scheduling did not reliably beat a fixed weak penalty; the two tracked each other within seed variation. So, I’d describe the weighting strategy as a training-balance decision rather than a leaderboard trick, while the broader effect (physics regularization buying you accuracy when labels are scarce) is consistent and, in the sparse regime, substantial.

The failure case ended up being the run I learned the most from. When I overweighted the residual, the optimizer did exactly what I asked and drove the PDE residual down to about 1e-10. Around the n_data = 10 setting:

```
too_high_physics solution RMSE ≈ 0.265physics residual MSE ≈ 1e-10initial-condition MSE ≈ 0.27boundary-condition MSE ≈ 0.09
```

The residual was essentially zero and the solution was still wrong. The reason is specific to how these terms interact: any constant field zeroes the residual exactly, since u_t = 0 and u_xx = 0. The field u ≡ 0 additionally satisfies the zero boundaries — it fails only the initial condition. When the residual dominates the objective, the model backs into exactly that corner: it converged to a near-flat field around u ≈ 0.2, trading error between the under-weighted IC term (MSE ≈ 0.27) and BC term (MSE ≈ 0.09) while keeping the residual at machine-precision zero. Physics loss is one objective among several, and it has to be held in tension with the data, initial-condition, and boundary-condition terms rather than allowed to win outright.

There’s a fair objection to this whole setup: if the prompt is fixed, the text embedding is just a constant vector, and the model is a PINN with an extra bias term. To rule that out, I trained across the four alpha values and ran an alpha-swap test. The idea is to hold x,t fixed and change only the prompt.

For a low-alpha target, the correct prompt specifies alpha = 0.01; the swapped prompt specifies alpha = 0.20:

```
correct low-alpha prompt on low-alpha target: RMSE = 0.0457swapped high-alpha prompt on low-alpha target: RMSE = 0.4781correct high-alpha prompt on high-alpha target: RMSE = 0.0621
```

Same coordinates, different text, a roughly tenfold change in error. (These are single-run numbers, but the effect is large enough to dwarf seed noise.) The text-derived condition is clearly functional — the encoder is steering the prediction, not sitting inertly as a bias. I use the low-alpha target as the main swap test because it gives the clearest slow-versus-fast diffusion contrast; the point is not to build a symmetric confusion matrix, but to verify that changing only the prompt changes the predicted field.

This is where the honest reading gets more interesting. The heat-equation benchmark is controlled, and its physical variation is governed by a single scalar, alpha. So, I compared DistilBERT conditioning against two simpler baselines:

```
numeric_alpha: head receives x,t,alpha directly
learned_alpha_embedding: alpha is mapped to a learned embedding and concatenated with x,t
```

In the single-seed robustness run, using the same multi-alpha setup with n_data=10, the learned alpha embedding was comparable to DistilBERT and came out slightly ahead on RMSE. I don’t read that as a universal ranking, and given the seed spread elsewhere I wouldn’t lean on the exact gap. What it does say is that a scalar conditioning signal is enough for this task — the problem simply doesn’t demand a language encoder for accuracy.

That doesn’t make the text encoder pointless; it locates its value precisely. On a scalar-parameter PDE, a learned alpha embedding is the stronger engineering choice, and I’d use it. The contribution here is narrower and architectural: a text-derived hidden state can condition a differentiable head, so the residual trains tensor outputs instead of parsed text. The language interface starts to matter on problems where the condition isn’t a clean number — where it’s implicit, mixed with other information, or only expressed in prose.

I logged four gradient diagnostics during training: the data gradient norm, the physics gradient norm, their ratio, and the cosine similarity between the two gradient directions.

They explain why the weighting matters. The data and physics gradients are not consistently aligned — at times the physics term points in a direction that reinforces the supervised fit, at times it pulls against it — and their magnitudes can differ enough that a fixed weight over- or under-drives the residual at different points in training. That is the mechanism behind treating scheduling as objective balancing rather than a route to lower final RMSE, and it’s consistent with the three-seed result, where weak and scheduled physics landed in the same place.

The alpha-swap test shows the text condition matters, but it doesn’t show the model understands the physics in any general linguistic sense. To probe that, I evaluated the trained model on held-out paraphrases, descriptions it had never seen in training, like:

*“Solve heat diffusion in a rod with thermal diffusivity 0.05.”*

or:

*“Use a diffusion coefficient of 0.05 for the heat equation with zero boundary values.”*

It was brittle. Performance held up on the original template and RMSE rose sharply for unfamiliar phrasing. Sharply enough to matter: out-of-template RMSE (0.19–0.64) is worse than the data_only baseline with no language interface at all. Under unfamiliar phrasing, the text conditioning isn’t just degraded, it’s a net liability. Adding a small amount of prompt diversity during training helped at three of the four alpha values but actually hurt at alpha = 0.05, and nowhere did it close the gap to the original template. My reading is that prompt diversity reduces the brittleness without solving it, robust natural-language conditioning would need broader phrasing coverage, a stronger encoder, or an explicit extract-then-validate step for the parameters. For engineering use this is the practical caution: a model shouldn’t be trusted just because it accepts text. The text interface needs its own validation.

The main result is narrower than “an LLM solves a PDE.” The language encoder provides a conditioning signal, the numerical head produces a differentiable tensor field, and the PDE residual trains that field directly. That is the bridge from Part II to Part III: physics moves from a post-generation checker to a training signal. The stronger lesson, though, is about the residual itself. Across three seeds, weak physics regularization improved sparse-data learning, but an overweighted residual drove the PDE error near zero while producing the wrong field. Physics loss helps only when it is balanced against data, initial-condition, and boundary-condition terms.

It helps to keep the roles separate:

The language model doesn’t replace the physics, the residual doesn’t guarantee correctness on its own, and the checker from Part II still earns its place. What the numerical head adds is the one thing the text path couldn’t provide: a differentiable route that lets the physics push back during training instead of only checking the answer afterward.

[Physics-Informed AI, Part III: From Post-Hoc Checker to Differentiable Physics Head](https://pub.towardsai.net/physics-informed-ai-part-iii-from-post-hoc-checker-to-differentiable-physics-head-003c7738ff98) was originally published in [Towards AI](https://pub.towardsai.net) on Medium, where people are continuing the conversation by highlighting and responding to this story.
