Support for vector sigma_noise for multi-output problems

[magnusross]

Currently the only scalar valued noise is supported for the Laplace approximation. For my use case, having a noise value that can change for each output it is important to achieving good performance. This enchancement would involve generalising the sigma noise parameter to allow it to have multiple elements, one for each output. Note, I am not suggesting that the noise also become a function of the inputs as in this paper, just allowing it to be a vector.

If you think this would be worthwhile to add, and not too difficult, I'd be happy to give it a go. If so, any guidance on the implementation would be much appreciated!

Cheers :)

[wiseodd]

Yes, I think this would be worthwhile to add!

The starting point is here:

Laplace/laplace/baselaplace.py

Lines 567 to 591 in 6f17099

	@sigma_noise.setter
	def sigma_noise(self, sigma_noise: float | torch.Tensor) -> None:
	self._posterior_scale = None
	if np.isscalar(sigma_noise) and np.isreal(sigma_noise):
	self._sigma_noise = torch.as_tensor(
	sigma_noise, device=self._device, dtype=self._dtype
	)
	elif isinstance(sigma_noise, torch.Tensor):
	if sigma_noise.ndim == 0:
	self._sigma_noise = sigma_noise.to(
	device=self._device, dtype=self._dtype
	)
	elif sigma_noise.ndim == 1:
	if len(sigma_noise) > 1:
	raise ValueError("Only homoscedastic output noise supported.")
	self._sigma_noise = sigma_noise[0].to(
	device=self._device, dtype=self._dtype
	)
	else:
	raise ValueError("Sigma noise needs to be scalar or 1-dimensional.")
	else:
	raise ValueError(
	"Invalid type: sigma noise needs to be torch.Tensor or scalar."
	)

Then, take a look at:

Laplace/laplace/baselaplace.py

Lines 214 to 241 in 6f17099

	@property
	def log_likelihood(self) -> torch.Tensor:
	"""Compute log likelihood on the training data after `.fit()` has been called.
	The log likelihood is computed on-demand based on the loss and, for example,
	the observation noise which makes it differentiable in the latter for
	iterative updates.
	Returns
	-------
	log_likelihood : torch.Tensor
	"""
	factor = -self._H_factor
	if self.likelihood == "regression":
	# loss used is just MSE, need to add normalizer for gaussian likelihood
	c = (
	self.n_data
	* self.n_outputs
	* torch.log(
	torch.as_tensor(
	self.sigma_noise, device=self._device, dtype=self._dtype
	)
	* sqrt(2 * pi)
	)
	)
	return factor * self.loss - c
	else:
	# for classification Xent == log Cat
	return factor * self.loss

Let me know if you have further questions!