Training a (conditional) Universal Differential Equation using binary observations of a latent state

Hello everyone,

I am training a conditional Universal Differential Equation in Julia.

The model is a regular UDE with one additional trainable input parameter for each particle. The neural-network parameters are shared across all particles, while the conditional parameter allows the model to capture differences between individual particle trajectories.

The model works reasonably well when trained on simulated bond-count data. However, the measurements I ultimately want to use are binary: each particle is either bound (1) or unbound (0) at each time point; bond-count states >= 1 are all collapsed onto this bound state.

The ODE predicts a continuous latent mean number of bonds, so I need an observation model that converts this latent state into a probability of being bound, so that I can use binary data to train the neural network.

My current approach assumes several independent binding sites (see code). I am unsure whether this is appropriate because the sites all compete for the same small pool of available molecules. When one molecule binds, fewer remain available for the other sites.

Training with binary data gives mediocre and highly sensitive results. Small changes in L2 regularisation or the slope penalty can lead to very different learned relationships. Increasing the number of particles has not clearly helped, and using a finer observation time grid makes the result worse (since the states are autocorrelated and binary).

I also suspect that my early-stopping settings may be too strict for the relatively small and noisy binary validation loss, causing training to stop before the model has settled. However, the validation loss is not improving much.

I would especially appreciate advice on:

  • How to connect a continuous latent mean state to binary observations.
  • Whether this is solely an identifiability issue.
  • Whether different training settings may help.

I will also share a the latest version of the code.

Please let me know if anything is unclear or if additional information would be useful.
Thank you for your time!

Kind regards,
NCCP

PS. I’m not sure whether this is the appropriate category. Let me know!

Code:
dBPM.jl (24.6 KB)
helpers.jl (9.3 KB)
inputs.jl (7.3 KB)
plot_figures.jl (55.9 KB)
ude.jl (26.1 KB)