Using Random Effects as Neural-Network inputs in DeepNLME: is this Statistically sound?

Estimation
1

I am trying to evaluate a DeepNLME model where subject-level latent random effects are used directly as inputs to a neural network.

Construction

In a standard mixed-effects model, random effects (η) influence the model through predefined parametric relationships. In this approach I am trying to explore, the random effects are instead used as inputs to a small neural network. The network outputs subject-specific corrections to model parameters, while its weights are shared across all subjects.

The random effects remain fixed for each subject and are not time-varying. The network is evaluated once per subject before solving the differential equations, and all temporal behaviour is then governed by the dynamic model.

The motivation is to determine whether nonlinear relationships among subject-specific effects can be captured through the neural network, rather than relying solely on the covariance structure Ω.

Representative implementation:

model = @model begin
    @param begin
        tvCL    ∈ RealDomain(lower = 1e-3)
        tvVc    ∈ RealDomain(lower = 1e-3)
        tvkin1  ∈ RealDomain(lower = 1e-3)
        tvkin2  ∈ RealDomain(lower = 1e-3)
        tvkin3  ∈ RealDomain(lower = 1e-3)
        tvkout1 ∈ RealDomain(lower = 1e-3)
        tvkout2 ∈ RealDomain(lower = 1e-3)
        tvkout3 ∈ RealDomain(lower = 1e-3)
        tvEmax  ∈ RealDomain(lower = 0.0, upper = 0.99)
        tvEC501 ∈ RealDomain(lower = 1e-3)
        tvEC502 ∈ RealDomain(lower = 1e-3)
        tvEC503 ∈ RealDomain(lower = 1e-3)
        nn_eta  ∈ MLPDomain(2, 8, 8, 3; act = softplus, reg = L2(1e-5))
        Omega   ∈ PSDDomain(2)
        sigma1  ∈ RealDomain(lower = 1e-3)
        sigma2  ∈ RealDomain(lower = 1e-3)
        sigma3  ∈ RealDomain(lower = 1e-3)
    end

    @random eta ~ MvNormal(Omega)   # time-invariant, subject-level random effects

    @pre begin
        CLi   = tvCL * exp(eta[1])
        Emaxi = tvEmax * exp(eta[2])
        # eta -> NN -> nonlinear, subject-specific correction to kin
        # (evaluated once; time-invariant)
        c1 = 0.40 * tvkin1 * tanh(nn_eta(eta[1], eta[2])[1])
        c2 = 0.40 * tvkin2 * tanh(nn_eta(eta[1], eta[2])[2])
        c3 = 0.40 * tvkin3 * tanh(nn_eta(eta[1], eta[2])[3])
    end

    @init begin
        R1 = tvkin1 / tvkout1
        R2 = tvkin2 / tvkout2
        R3 = tvkin3 / tvkout3
    end

    @dynamics begin
        Central' = -(CLi / tvVc) * Central
        R1' = (tvkin1 + c1) *
              (1 - Emaxi * (Central/tvVc) / (tvEC501 + Central/tvVc)) -
              tvkout1 * R1
        R2' = (tvkin2 + c2) *
              (1 - Emaxi * (Central/tvVc) / (tvEC502 + Central/tvVc)) -
              tvkout2 * R2
        R3' = (tvkin3 + c3) *
              (1 - Emaxi * (Central/tvVc) / (tvEC503 + Central/tvVc)) -
              tvkout3 * R3
    end

    @derived begin
        DV_R1 ~ @. Normal(R1, abs(R1) * sigma1)
        DV_R2 ~ @. Normal(R2, abs(R2) * sigma2)
        DV_R3 ~ @. Normal(R3, abs(R3) * sigma3)
    end
end

From what I understand, this approach is broadly consistent with DeepNLME and related mixed-effects neural ODE frameworks. Random effects are treated as subject-level constants, which is standard in longitudinal mixed-effects models. What I am less certain about is the use of the inferred random effects themselves as neural-network inputs, and whether there are established references or identifiability considerations for this setting.

Questions:

  1. Identifiability: Since both Ω and the neural network can explain between-subject variability, how can I determine whether they are learning distinct information rather than competing to explain the same signal? Are there recommended diagnostics for this?

  2. Use of random effects as network inputs: In my model, the neural network takes inferred random effects (η) as inputs rather than observed covariates. Is this considered a statistically reasonable approach, and are there references that discuss similar constructions?

  3. Time-invariant random effects: I treat the random effects as fixed for each subject and allow the dynamic model to account for changes over time. For longitudinal data, is this a reasonable assumption, or are there situations where time-varying random effects would be preferable?

  4. Regularization: I have currently constrained the network output and apply L2 regularization to reduce overfitting. Are these generally sufficient, or are there better regularization strategies for embedded neural networks in mixed-effects models?

I would appreciate any feedback or pointers to relevant literature.

2

Identifiability: Since both Ω and the neural network can explain between-subject variability, how can I determine whether they are learning distinct information rather than competing to explain the same signal? Are there recommended diagnostics for this?

If you have the neural network (NN) taking random effects as input, the \Omega becomes unnecessary and generally not identifiable. So we typically just use an identity matrix I for the covariance. Note that an unregularized 2 layer NN (no hidden layers) with a linear activation function is fully capable of replacing the \Omega without worsening the quality of fit. Assume the random effect is called z and some intermediate variable in the @pre block is \eta that has the following definition using a 2-layer NN and the identity activation function. In other words, the “NN” here is a linear model.

\eta = W z + b

If z \sim \mathcal{N}(0, I) then \eta \sim \mathcal{N}(b, W W^T). This is because an affine function of a Gaussian is a Gaussian and the mean and covariance of \eta follow from linearity. In other words, the model that instead uses \eta as a random effect with \Omega and mean would be equivalent to the NN-based model above because any \Omega can be written as W W^T. The mean b is typically set to 0 in pharmacometrics but you can also move the “typical value” parameters there, in log scale. In any case, this shows that either a neural network or \Omega should be used because both can serve a similar role. If you would like to use an NN for its flexibility, you can set the random effects distribution to have mean 0 and identity covariance matrix, unless you have a good reason to do otherwise.

There is a caveat to the above argument, which is that the NN should not have regularization to be able to freely represent any \Omega using the weights W. In practice, we usually put regularization on the weights so you can choose to remove the regularization on the first layer to fully recover the equivalence above. Although, with enough layers and neurons in the NN, you generally don’t need to do that as well.

Use of random effects as network inputs: In my model, the neural network takes inferred random effects (η) as inputs rather than observed covariates. Is this considered a statistically reasonable approach, and are there references that discuss similar constructions?

Check out latent variable models in machine learning. Variational autoencoders and generative adversarial networks all use the same construction. Whether it is statistically reasonable or not depends on what you mean by that. If a linear model is not able to fit the data, then this is one way to make your model more flexible and nonlinear to hopefully be able to fit the data. However, the more flexible your model is, the higher the risk of over-fitting so you need checks and balances in place to detect and avoid over-fitting if it happens. This is the case even for non-NN models that are flexible and over-parameterized. Using regularization or prior distributions becomes increasingly important to avoid over-fitting with large and flexible models.

Time-invariant random effects: I treat the random effects as fixed for each subject and allow the dynamic model to account for changes over time. For longitudinal data, is this a reasonable assumption, or are there situations where time-varying random effects would be preferable?

A time varying random effect can be obtained using a time-invariant (doesn’t depend on time) random effect z, time t, and a neural network.

\eta(t) = \text{NN}(z, t)

In a sense, \eta(t) is now a time-varying random effect. There are other ways to represent a distribution over functions of time (e.g. stochastic processes). But generally, assuming some latent variable that doesn’t change with time but interacts with time is powerful enough in many cases.

Regularization: I have currently constrained the network output and apply L2 regularization to reduce overfitting. Are these generally sufficient, or are there better regularization strategies for embedded neural networks in mixed-effects models?

There is no 1 regularization type that works best in every case. The choice of regularization matters because you are assuming a “default value” for the parameters or a prior. Generally, we set the prior/regularization on the parameters but the consequence of that prior on the output of the model can be unpredictable if the model is complex. The point of a prior is to encode our prior knowledge of what values parameters should generally take, before including the data. But if we don’t have a good sense of what the parameters do or mean, then setting a prior is a challenge. This is why some researchers looked into how to set a “prior” on derived quantities instead, e.g. trying to softly constrain the output of the model to be around a certain value. This is possible to hack in DeepPumas if you need it. But in general, we found that simple L2 regularization works well enough for many cases. Even not regularizing the NN at all initially and using a small NN sometimes works well enough. While L2 is most definitely not an optimal prior in most cases, it tends to work well enough in practice that with enough data and a large enough NN, you get good performance on unseen data. Once you have good performance, whether you can do slightly better by using a different regularization scheme becomes a secondary research question. If you find your model over-fitting (judged by unseen data) when not using any regularization but under-fitting when using L2, then you may want to try different regularization schemes.

I hope this help.

3

Thank you @mohamed82008 for the detailed response.

On regularization — I’m currently using L2 on the network weights and planning to tune the regularization strength rather than fixing it. I also found your point about constraining derived quantities when parameter priors are difficult to specify really helpful.

Just a couple of follow-up questions:

  • If Ω is fixed to I, is there a good way to recover or report the implied parameter correlations for interpretation? For example, would it make sense to simulate z through the trained network and estimate correlations from the outputs, or is that not very meaningful once the mapping becomes nonlinear?

  • For identifiability when N (Number of Subjects) is small, would you rely more on stronger L2 regularization, limiting network size, or a combination of both?

Thanks again — this gives me a much clearer direction.

4

If Ω is fixed to I, is there a good way to recover or report the implied parameter correlations for interpretation? For example, would it make sense to simulate z through the trained network and estimate correlations from the outputs, or is that not very meaningful once the mapping becomes nonlinear?

Once you use an identity covariance matrix and a neural network, the random effects are uncorrelated by design. However, you should really be looking into correlations of derived quantities (e.g. NN outputs) that hopefully have a very specific meaning according to your model, if the model is identifiable wrt the NN outputs. Correlation between variables that don’t have a specific meaning to you, the modeler, is meaningless. The reason we typically care about the omegas is because we assign a very specific meaning to the random effects in traditional NLME. This is lost when you use random effects as inputs to neural networks but in return you gain model flexibility. But again, at least the output of the NN should be used in a consistent way in the model such that it carries a specific meaning that you can reason about, e.g. rate of change of tumor size. You should therefore look at those derived quantities that have meaning and study their correlation if that’s useful for you, instead of trying to get the equivalent of \\Omega.

For identifiability when N (Number of Subjects) is small, would you rely more on stronger L2 regularization, limiting network size, or a combination of both?

I typically start with the goal of over-fitting, but using a small NN at first. I don’t prefer to start with regularization because it hinders my ability to over-fit. I just start with a small neural network and gradually make it more complex until I am able to over-fit. The ability to over-fit tells me that my model is flexible enough to do so. In SciML, the neural network is just one part of the model. Making a neural network more flexible but using it in the wrong scientific model structure can still lead to an under-fitting model. So I typically tweak the scientific model and the ML component until I can over-fit. Then I introduce L2 regularization and fine tune the strength to “back-track” from over-fitting to hopefully the right amount of fitting. This is also the stage to try out different regularization types if you want.

Regarding identifiability, your neural network parameters will almost never be identifiable. But what you should care about is the identifiability of the model wrt the neural network outputs, not its parameters. This usually translates to using each neural network to do one and only one thing in your model. But proper identifiability analysis wrt neural network outputs is difficult in general and relies a bit on experience. There is probably also some research on that topic. If not, that would make a good research topic.

5

Thank you @mohamed82008 for your response. I will look into this