Add native serialization support

[wiseodd]

Addressing #45. Very useful for large models like LLMs where even doing forward passes over training data (for fit()) is expensive.

The API basically follows PyTorch.

la = Laplace(model, 'regression', subset_of_weights='all', hessian_structure='full')
la.fit(train_loader)
la.optimize_prior_precision()  # Or via marglik optimizing sigma_noise also

# Serialization for fitted quantities
state_dict = la.state_dict()
torch.save(state_dict, 'state_dict.bin')

la2 = Laplace(model, 'regression', subset_of_weights='all', hessian_structure='full')
# Load serialized, fitted quantities
la2.load_state_dict(torch.load('state_dict.bin'))

[wiseodd]

By the way, I didn't really implement nor test SubnetLaplace. Not sure if it's really popular nowadays; the same effect can be done by setting requires_grad = False on select parameters anyway (from #144).

Nevertheless, @edaxberger feel free to implement serialization on SubnetLaplace. It should be very straightforward for you.

[wiseodd]

By the way, I didn't really implement nor test SubnetLaplace. Not sure if it's really popular nowadays; the same effect can be done by setting requires_grad = False on select parameters anyway (from #144).

Nevertheless, @edaxberger feel free to implement serialization on SubnetLaplace. It should be very straightforward for you.

Though I must say that a quick test on SubnetLaplace works no problem. But I don't know about edge cases like different subnetwork_indices but the same len(subnetwork_indices), etc.

[runame]

Thanks for tackling this, will be very useful!

I think currently the assumption is that model, likelihood, subset_of_weights, and hessian_structure of la and la2 (in the README example) match. It might make sense to explicitly check for this in the load_state_dict method and throw descriptive errors in case of a mismatch. The necessary information for this has to be also stored (see my comment below).

Additionally, it would be nice to allow for the creation of a new Laplace class instance just based on the saved state dict. A classmethod would make sense for this, but Laplace is a function and not a class, making this option impossible (without other changes). So alternatively, we could add a function like LaplaceFromStateDict that takes the state dict as an argument and returns a Laplace class instance which is (hopefully completely) equivalent to the previously saved one. I could add this functionality in a follow up PR though.

[wiseodd]

@runame updated. Please check and merge if everything's good.

I decided to not include backend and backend_kwargs since Laplace doesn't care about them once H is obtained. I.e. H is just a tensor or Kron and you can use whatever backend you want for glm_predictive and to do another fit.

[wiseodd]

I also ignore model.state_dict since the user would already have it anyway. Even marglik training also outputted the resulting torch model. Including it in the serialized Laplace is redundant.

[runame]

Thanks, will review the update soon!

I should have elaborated on this, but there is a very concrete use case where the backend and backend_kwargs are needed, i.e. whenever you want to set fit(train_loader, override=False). For example, this is used to implement continual learning, see our experiment for the Redux paper. (I guess some mismatch between backends is fine, since only the shape (or both Kron) of self.H matters.)

I'm fine with not including model.state_dict() for now, this will only be important for the LaplaceFromStateDict functionality that I might add in a follow up PR.

[wiseodd]

I think the current checks are sufficient to handle continual learning:

This will check whether the hessian_structure and subset_weights are compatible with the loaded instance:

Laplace/laplace/baselaplace.py

Lines 792 to 796 in 30e28f2

	if self.__class__.__name__ != state_dict['cls_name']:
	raise ValueError(
	'Loading a wrong Laplace type. Make sure `subset_of_weights` and'
	+ ' `hessian_structure` are correct!'
	)

This will check the network (torch model) is correct:

Laplace/laplace/baselaplace.py

Lines 797 to 802 in 30e28f2

	if self.n_params is not None and len(state_dict['mean']) != self.n_params:
	raise ValueError(
	'Attempting to load Laplace with different number of parameters than the model.'
	+ ' Make sure that you use the same `subset_of_weights` value and the same `.requires_grad`'
	+ ' switch on `model.parameters()`.'
	)

As for model.state_dict we could add an option in Laplace.state_dict(), say save_model: bool = False. But this should be part of your future PR.

[runame]

Ok makes sense, only have one comment left.

[wiseodd]

Added a test to check fit(override=False)!

[runame]

Sorry, last questions.

[runame]

Thanks for incorporating all feedback!