Covariance values change when weights are loaded

[scho97]

Summary
I saved the weights of a model trained on EEG LEMON. I then built a new model and loaded the weights. However, when I run model.get_covariances(), the covariance matrix values differs from the values observed from a model trained on LEMON.

Problem
Observation model parameters should report identical values when the model weights are loaded again.

Code

1. Pretraining on LEMON

model = Model(config)
history = model.fit(training_data)
model.save(f"{model_dir}/trained_model")
cov = model.get_covariances()

2. Building on a new model (no training done)

new_model = Model(config)
new_model.load_weights(best_model_path)
new_cov = model.get_covariances()

Then, cov that was saved in a pickle file was compared with new_cov.
Here, cov and new_cov should be the same. However, they differ at about the third decimal point.

Helpful Observations

• A pre-trained model and new model had same transition probability matrices.
• In the configuration, depending on how I set learn_covariances to either True or False, the resulting covariance values differed. Based on what I saw so far, it appears like how Cholesky factor is initialised affects the covariances.
• If I save and load models in a single script, this problem does not occur.
• Updating to the most recent commit and environment reduced the magnitude of difference.

[scho97]

Update: After going through some debugging, I feel like this problem is likely due to osl-dynamics and tensorflow version, as well as possibly other environment-associated changes. This may mean that using a model saved priorly with a more recent software or across different servers does not guarantee perfect reliability.

With the most up-to-date package and env, the errors became much smaller in BMRC and hbaws. The errors do not seem to occur when epochs are short (~1-2 epochs).