Sampling Nonlinear Association Terms

Why Nonlinear Association Terms Need Careful Sampling

In a joint model, the survival hazard can depend on a latent longitudinal trajectory rather than only on observed biomarker measurements. A nonlinear association, such as dependence on the trajectory squared, changes both the hazard at the event time and the cumulative hazard accumulated before that time. Consequently, it changes the posterior density and its gradients with respect to the latent random effects. A sampler that evaluates the event-time term correctly but approximates or differentiates the cumulative term incorrectly targets the wrong posterior.

This tutorial develops the computation for a squared current-value association. The same chain-rule principles apply to other differentiable nonlinear functions, but their mathematical and numerical behavior must be checked for the chosen function.

The squared term is a modeling choice, not a universal default. It makes the hazard respond to the magnitude of the latent value: for positive \(\alpha\), values farther from zero increase the hazard, whether they are positive or negative. If a trajectory has a meaningful reference level, it may be more interpretable to square a centered value. Units matter: if \(m_i(t)\) is measured in biomarker units, \(\alpha\) has inverse squared-biomarker units so that the exponent is dimensionless.

Likelihood and Posterior Target

Conditional on the random effects and model parameters, the longitudinal measurements and event process are modeled as conditionally independent. For example, a Gaussian longitudinal model has \(Y_{ij}\mid b_i,\beta,\sigma^2\sim N(m_i(t_{ij}),\sigma^2)\). The event contribution, conditional on \(b_i\), is

The event indicator removes the event-time hazard factor for a censored subject, but the cumulative-hazard term remains. Taking logarithms gives

The joint posterior is proportional to the product of these survival contributions, the longitudinal likelihood, and the priors on the model parameters and random effects. For computation, the log posterior is usually evaluated as a sum of log-likelihood and log-prior terms. The nonlinear survival term must be included whenever a proposed value of \(b_i\), \(\alpha\), or another parameter changes \(m_i(t)\) or the hazard.

Deriving the Nonlinear Gradients

Gradients are especially important for Hamiltonian Monte Carlo (HMC), which uses derivatives of the log posterior to construct proposals. They are also useful for checking implementations of other samplers. The derivation below holds when the hazard is differentiable in the parameters and the integral can be differentiated under the integral sign; in this smooth, finite-follow-up setting, the integrand and its parameter derivative are assumed integrable.

For compactness, define \(q_i(t)=h_i(t\mid b_i)\). Since \(m_i(t)=x_i(t)^\mathsf{T}\beta+z_i(t)^\mathsf{T}b_i\), its derivative with respect to \(b_i\) is \(z_i(t)\). Applying the chain rule to the log event-time term and to the cumulative hazard yields

The first term is absent under censoring. The second term is not: changing the random effects changes the hazard throughout the observed follow-up interval. Similarly, differentiating with respect to \(\alpha\) gives

These are survival-likelihood gradients only. For the full posterior gradient, add the derivatives of the longitudinal log likelihood and the log prior. For example, if \(b_i\sim N(0,D)\), its log-prior gradient is \(-D^{-1}b_i\). The longitudinal contribution depends on its specified observation model.

Time Integration and Numerical Implementation

With a time-varying trajectory, the cumulative hazard will generally require numerical integration. The same quadrature rule should be applied to the cumulative hazard and its derivative. For instance, a Simpson approximation on an interval \([0,2]\), using nodes 0, 1, and 2, approximates the integral of a smooth function \(g\) by \(\{g(0)+4g(1)+g(2)\}/3\). This approximation is not exact in general; it should be refined or compared with a more accurate calculation when the integrand changes rapidly.

Sampling Strategy, Comparisons, and Diagnostics

For HMC, evaluate the log posterior and gradient at each proposed state, including the nonlinear event-time term and the integrated-hazard term. A practical implementation computes \(m_i(t)\) at the event time and quadrature nodes, then reuses those values to evaluate both the hazard and its derivatives. If gradients are unavailable or unreliable, a Metropolis update can propose new values and accept or reject them using the ratio of posterior densities; it still requires accurate evaluation of the nonlinear cumulative hazard, but not an analytic gradient.

A linear current-value association replaces \(m_i(t)^2\) with \(m_i(t)\). Its gradient is simpler, but it encodes a different relationship and should not be substituted merely to make computation easier. A nonlinear function can capture curvature or magnitude-based effects that a linear term cannot. Conversely, a flexible function of the latent trajectory may be weakly identified when event information is sparse or the observed trajectory range is narrow. Model choice should follow scientific meaning and data support, not computational convenience alone.

Useful implementation checks include comparing analytic or automatic gradients with finite differences at several parameter values, checking that the cumulative hazard is nonnegative, and increasing the quadrature resolution to see whether the log likelihood and gradient stabilize. Check event and censored contributions separately. In particular, a censored subject must still contribute the negative cumulative hazard to the log likelihood and its derivative. A sharp increase in \(m_i(t)^2\) can make the hazard integrand vary quickly; a coarse grid may then underestimate the cumulative hazard and distort the posterior.

A common mistake is to differentiate only the event-time hazard and omit the derivative of the cumulative hazard, or to multiply the event-time term by \(1-\delta_i\) rather than \(\delta_i\). The correction is to start from \(\delta_i\log h_i(T_i)-\int_0^{T_i}h_i(t)\,dt\): the event score is present only when \(\delta_i=1\), while the integral and its derivative are present for both events and censoring. Another pitfall is assuming that a positive \(\alpha\) means higher biomarker values always raise the hazard; for a squared association, the effect depends on the value and reference point of \(m_i(t)\).

Nonlinear terms also have limitations. Squaring can amplify a few extreme latent values, create strong posterior dependence between trajectory and association parameters, and make the hazard highly sensitive to scaling. Centering and scaling the biomarker may improve interpretability and numerical stability, but the transformed units and corresponding prior must be handled consistently. These considerations do not replace convergence assessment; they help identify whether the implemented target is numerically credible.