Gradient matching for a nonlinear PK model

Setup

We are again tackling the common pharmacokinetic (PK) modeling setup. The dynamics are defined through an ordinary differential equation (ODE) system

$$
\frac{d}{dt} \mathbf{x}(t) = f(\mathbf{x}(t), t, \theta),
$$

where time $t$ is measured in hours and $\mathbf{x}(t)$ contains the amount of drug in different compartments, measured here in mg. The amount in the central compartment can be divided by its volume to get the drug concentration in blood.

We generally fit the model parameters by matching the predicted (log) concentration trajectory to measured blood (log) concentrations. The practical problem is immediate: we need $\mathbf{x}(t)$, but the model defines only its derivative. In some cases the ODE has a closed-form solution for $\mathbf{x}(t)$. In this post the running example is a nonlinear PK model where we do not use a closed-form solution, and need to use an expensive numerical solver to get the concentration trajectory over time.

But what if instead of matching observed (log) concentrations to the solved concentration (log) trajectory, we could somehow turn the concentration observations into observations of the time derivative of the (log) concentration, and thus avoid the need to solve the ODE system during parameter fitting? This is called gradient matching (though I think just derivative matching or slope matching could be a better name in this case). In this post I compare a finite-differences based gradient matching appraoch to proper full trajectory-based inference. The question is whether gradient maching captures enough of the parameter information while being faster than repeatedly solving the ODE.

Oral One-Compartment Model With Michaelis-Menten Elimination

The model has an oral depot and a central compartment. Drug moves from the depot into the central compartment by first-order absorption, while elimination from the central compartment follows Michaelis-Menten kinetics:

$$
\begin{aligned}
\frac{d\,\text{Depot}}{dt} &= -k_a\,\text{Depot}, \\
C &= \frac{A}{V}, \\
\frac{dC}{dt} &= \frac{k_a\,\text{Depot}}{V} - \frac{V_{\max}}{V}\frac{C}{K_m + C}.
\end{aligned}
$$

Here $\text{Depot}$ is the amount of drug at the absorption site in mg, $A$ is the amount in the central compartment in mg, and $C$ is the central concentration in mg/L. The parameters are the absorption rate $k_a$ in 1/h, central volume $V$ in L, maximum elimination rate $V_{\max}$ in mg/h, and Michaelis-Menten concentration $K_m$ in mg/L.

Oral doses, measured in mg, enter as jumps in the depot:

$$
\text{Depot}(t^+) = \text{Depot}(t^-) + \text{dose}.
$$

The concentration equation is nonlinear because the elimination term depends on $\frac{C}{K_m + C}$. At low concentration, where $C \ll K_m$, elimination is approximately first-order. At high concentration, where $C \gg K_m$, elimination is approximately zero-order.

Proper Inference

Proper parameter estimation solves the ODE for each proposed parameter vector $\theta = (k_a, V, V_{\max}, K_m)$ and compares the predicted concentrations to the observations. Here, we have a Bayesian model where the observation model is lognormal:

$$
\log y_i \sim \operatorname{Normal}(\log C(t_i; \theta), \sigma^2).
$$

The implementation used here does not use an adaptive ODE solver. It uses a differentiable fixed-step fourth-order Runge-Kutta solver written in JAX, with step size

$$
\Delta t = 0.05 \text{ h},
$$

which is 3 minutes. NumPyro's NUTS sampler obtains gradients of the log posterior by automatic differentiation through the fixed RK4 recurrence. An adaptive JAX ODE solver such as `diffrax` could also be used, but the fixed-step solver keeps the computation graph and solver behavior simpler for this comparison. The cost is that every likelihood evaluation includes a trajectory solve.

Gradient Matching Approximation

Gradient matching changes the target. Instead of asking whether a complete solved trajectory matches the observed concentrations, it asks whether the local slopes implied by the observations match the slopes implied by the ODE. Because the simulated observation noise is lognormal, the comparison is made on the log concentration scale. Between two adjacent observations, define

$$
\widehat{g}_i = \frac{\log y_{i+1} - \log y_i}{t_{i+1} - t_i}.
$$

This is a finite-difference estimate of $\frac{d\log C(t)}{dt}$, with units 1/h. The ODE implies

$$
\frac{d\log C}{dt} = \frac{1}{C}\frac{dC}{dt} = \frac{1}{C}\left(\frac{k_a\,\text{Depot}}{V} - \frac{V_{\max}}{V}\frac{C}{K_m + C}\right).
$$

For each interval $[t_i, t_{i+1}]$, I evaluate this expression at the start of the interval. The concentration used in the right-hand side is the observed concentration at the start point:

$$
C(t_i) \approx y_i,
$$

and the depot amount at $t_i$ is available analytically because depot absorption is first-order:

$$
\text{Depot}(t) = \sum_j \text{dose}_j \exp\left(-k_a(t - t_j)\right) I(t \geq t_j).
$$

The gradient-matching likelihood that I use here is

$$
\widehat{g}_i \sim \operatorname{Normal}\left(g(t_i; \theta, y_i), \tau^2\right),
$$

where $g(t_i; \theta, y_i)$ is the ODE-implied log-concentration gradient at the start of the interval, using the observed start concentration $y_i$. This avoids solving the full ODE inside the likelihood. The tradeoff is that it uses noisy finite differences, so it tends to need reasonably placed observations. Sparse or noisy data can make the estimated slopes unstable.

Simulation Design

Data for a single subject are simulated from the nonlinear Michaelis-Menten model. There are three oral doses of 100 mg at times 0 h, 24 h, and 48 h. Only the third dose occasion is sampled, so the observed concentrations still reflect accumulation from the earlier doses. Observation noise is multiplicative lognormal noise with log standard deviation 0.12. I fit three sampling designs with 24, 12, and 6 concentration observations in mg/L over the 24 hours after the third dose.

What Happens Across Sampling Schedules?

With 24 observations, proper inference recovers the main features of the trajectory and gives parameter summaries near the data-generating values. Gradient matching can also recover the broad trajectory, even though its likelihood only sees local log-slope information. Its posterior is more approximate, especially for the elimination parameters $V_{\max}$ and $K_m$.

‍

The time-course figure contains the proper inference posterior predictive trajectory, the gradient matching posterior predictive trajectory, and the matched log-slope posterior predictive plot. Arrows are only shown in the gradient matching trajectory panel: purple arrows show the concentration-scale slopes corresponding to the finite-difference log slopes used by the approximation, and yellow arrows (poorly visible) show the true concentration derivatives at the same start times.

The same time-course diagnostics are generated for the 12-observation case:

and for the 6-observation case:

‍‍The sampled parameters are $k_a$, $V$, $V_{\max}$, and $K_m$. The below KDE plot shows posteriors of these parameters:

‍

The MCMC wall-clock runtimes were around 200-400 seconds for proper inference and about 2 seconds with gradient matching. So the computational cost difference of ODE solving is clear. But as we see, the gradient matching fits are rather poor.

‍

Takeaway

Gradient matching replaces a solved-trajectory likelihood with a likelihood (or just least squares) on finite-difference slopes. The original likelihood is usually well motivated by a model for the observation noise, but it is hard to convert it to a proper likelihood for the finite differences. The observation model used here for the finite differences implies a different statistical model than the true log-normal observation model that was independent over observations, which is probably one reason why the gradient matching results are poor. In reality, even if the observation model for adjacent observations is independent, the adjacent finite difference errors may still be correlated. Including things like censoring due to limits of quantification also sounds difficult for gradient matching.

Gradient matching would require dense and quite noise-free data to work well, since it needs accurate derivative "observations". The gradient matching estimates in this blog post may not approach the proper inference estimates even if data becomes infinitely dense due to the different implied statistical model. In the ideal limit where there is no measurement noise and data becomes infinitely dense, a least squares gradient matching can converge to the same true ODE parameters as trajectory least squares, I think. But for the kind of data that we usually have in PK modeling, especially population PK, gradient matching does not seem like a realistic option.

‍