Computing Gradients Through NLME Models

Why Gradients Through NLME Predictions Need Care

Hamiltonian Monte Carlo uses derivatives of the log posterior to guide its simulated motion through parameter space. In an NLME model, that log posterior depends on predictions that may involve nonlinear transformations, individual random effects, and numerical solutions of differential equations. A gradient error in any of these components can therefore affect the proposed trajectory. The goal is not merely to obtain a derivative, but to understand which derivative is being computed and whether it accurately represents the model.

Write \(q\) for the unconstrained computational parameters used by the sampler. These may include transformed population parameters and random effects. For individual \(i\), let \(\eta_i\) denote their random effects, let \(x_i(t)\) be the model state, and let \(\widehat y_i(t)\) be the predicted observation. Between events, a common model structure is

The posterior log density, denoted here by \(\ell(q)\), combines the log likelihood and log prior, including any parameter-transformation terms required by the chosen coordinates. Its gradient includes derivatives through every prediction that depends on \(q\). Differentiating the likelihood while treating a parameter-dependent prediction as fixed would omit part of the gradient.

Three Ways to Compute the Derivative

The three approaches considered here differ in how they represent the calculation that leads from parameters to predictions.

• Automatic differentiation (AD) applies derivative rules to the operations in a computer program. It computes derivatives of the implemented calculation, subject to the program’s operations and numerical solver. Forward-mode AD is often useful for a small number of inputs; reverse-mode AD can be efficient when there are many parameters and a scalar output, such as a log posterior.
• Sensitivity equations augment the ODE with equations for derivatives of its state with respect to parameters. They directly represent how a continuous model solution changes as parameters change.
• Finite differences approximate a derivative by evaluating the prediction at nearby parameter values. They require no derivative code, but their accuracy depends on the perturbation size and the accuracy of each prediction.

None is automatically reliable in every implementation. AD may differentiate solver steps or event-handling code in ways that do not match the intended continuous model. Sensitivity equations require correct derivatives of the ODE and event rules. Finite differences balance approximation error against numerical error in the predictions.

Sensitivity Equations and Their Derivation

For an ODE state \(x(t)\) governed by \(x'=f(t,x,\theta)\), define the sensitivity matrix \(S(t)=\partial x(t)/\partial\theta\), where \(\theta\) is a vector of model parameters. Differentiating the ODE with respect to \(\theta\), using the chain rule, yields

The initial condition is \(S(t_0)=\partial x(t_0)/\partial\theta\). If the initial state does not depend on \(\theta\), this sensitivity is zero. For an observation prediction \(g(t,x,\theta)\), the derivative is

The final term is the direct dependence of the observation function on the parameter while holding the state fixed. These equations follow because a small parameter change affects the prediction both by changing the state trajectory and, potentially, by changing the observation function directly.

For example, if \(x'=-kx\), then \(f_x=-k\) and \(f_k=-x\). The sensitivity \(s=\partial x/\partial k\) therefore satisfies \(s'=-ks-x\), not merely \(s'=-ks\). The additional term accounts for the direct appearance of \(k\) in the ODE.

What Changes at a Dose Event?

A dose or other discrete intervention changes the state through an event rule. Write \(x^-\) for the state immediately before an event and \(x^+\) for the state immediately after it. For a parameter-dependent update \(x^+=R(x^-,\theta)\) at a fixed, known event time, the sensitivity update is

Here \(R_x\) and \(R_\theta\) are derivatives of the event rule with respect to the pre-event state and parameters. For a bolus added to a state, \(R(x,D)=x+D\), so the sensitivity with respect to the dose increases by one across the event. If an event time \(\tau\) depends on a parameter, its movement also affects the state; for a smooth ODE on both sides of the event, the update at a fixed time after the event includes the term \((R_x f^- - f^+)\,\partial\tau/\partial\theta\). Omitting this term can give an incorrect derivative.

An observation exactly at an event also requires a defined convention: it may use the pre-event or post-event state. A prediction can jump at the event, so a derivative across the jump is not interchangeable with a derivative on either smooth segment.

Finite Differences, Tolerances, and Practical Checks

For a scalar prediction \(F(\theta)\), the central finite-difference approximation is \([F(\theta+h)-F(\theta-h)]/(2h)\). For a smooth function its truncation error is typically proportional to \(h^2\), but very small \(h\) makes the subtraction vulnerable to rounding and prediction error. In an ODE model, solver error can dominate the difference when the two predictions are nearly equal. Thus, reducing \(h\) indefinitely does not guarantee a better derivative.

ODE solver tolerances control the numerical accuracy of the state solution, not automatically the accuracy of its gradient. A useful check is to tighten tolerances and see whether predictions and gradients stabilize to the precision needed for the calculation. For finite differences, compare several sensible step sizes; for AD or sensitivity equations, compare a directional derivative with an independently computed finite difference. A directional check uses a vector \(v\) and compares \(\nabla F(\theta)^\mathsf{T}v\) with \([F(\theta+hv)-F(\theta-hv)]/(2h)\).

AD has the advantage of applying chain rules systematically through supported operations, but a solver may contain adaptive step decisions or event branches that make its computed derivative noisy or only locally representative of the implemented algorithm. Sensitivity equations expose the continuous derivative structure and can be efficient when there are relatively few parameters, but coding all parameter and event derivatives is demanding. Finite differences are easy to implement and useful as a check, but can be expensive and sensitive to step size, solver tolerances, and discontinuities. When an event time, event order, or conditional branch changes with parameters, the prediction may fail to be smoothly differentiable there; no gradient method can make a genuinely discontinuous target smooth.

Common Pitfall and Limitations

A common mistake is to interpret a stable-looking AD gradient as proof that the continuous NLME model has been differentiated correctly. AD differentiates the operations it is given. If the event update is missing, the observation is assigned to the wrong side of a dose, or a parameter transform is omitted, AD can return a precise gradient for the wrong calculation. Conversely, disagreement between AD and finite differences near an event does not by itself prove which result is wrong: the finite-difference perturbation may cross an event boundary, or the solver error may be too large relative to the prediction change.

Gradient checks are most informative when the compared methods use the same model definition, parameter coordinates, and event conventions. They establish numerical consistency at tested points, not correctness everywhere. Gradients can be particularly difficult near state constraints, parameter-dependent event ordering, or sharp nonlinear transitions. In those settings, inspect predictions and event logic in addition to numerical derivative values.