LVAE Losses

Loss function for DenoiSplit.

Parameters:

• model_outputs (tuple[Tensor, dict[str, Any]]) –

  Tuple containing the model predictions (shape is (B, target_ch, [Z], Y, X)) and the top-down layer data (e.g., sampled latents, KL-loss values, etc.).

• targets (Tensor) –

  The target image used to compute the reconstruction loss. Shape is (B, target_ch, [Z], Y, X).

• config (LVAELossConfig) –

  The config for loss function containing all loss hyperparameters.

• gaussian_likelihood (GaussianLikelihood, default: None ) –

  The Gaussian likelihood object.

• noise_model_likelihood (NoiseModelLikelihood, default: None ) –

  The noise model likelihood object.

Returns:

• output ( Optional[dict[str, Tensor]] ) –

  A dictionary containing the overall loss ["loss"], the reconstruction loss ["reconstruction_loss"], and the KL divergence loss ["kl_loss"].

Loss function for DenoiSplit.

Parameters:

• model_outputs (tuple[Tensor, dict[str, Any]]) –

  Tuple containing the model predictions (shape is (B, target_ch, [Z], Y, X)) and the top-down layer data (e.g., sampled latents, KL-loss values, etc.).

• targets (Tensor) –

  The target image used to compute the reconstruction loss. Shape is (B, target_ch, [Z], Y, X).

• config (LVAELossConfig) –

  The config for loss function containing all loss hyperparameters.

• gaussian_likelihood (GaussianLikelihood) –

  The Gaussian likelihood object.

• noise_model_likelihood (NoiseModelLikelihood) –

  The noise model likelihood object.

Returns:

• output ( Optional[dict[str, Tensor]] ) –

  A dictionary containing the overall loss ["loss"], the reconstruction loss ["reconstruction_loss"], and the KL divergence loss ["kl_loss"].

Compute the KL divergence loss.

NOTE: Description of rescaling methods: - If "latent_dim", the KL-loss values are rescaled w.r.t. the latent space dimensions (spatial + number of channels, i.e., (C, [Z], Y, X)). In this way they have the same magnitude across layers. - If "image_dim", the KL-loss values are rescaled w.r.t. the input image spatial dimensions. In this way, the lower layers have a larger KL-loss value compared to the higher layers, since the latent space and hence the KL tensor has more entries. Specifically, at hierarchy i, the total KL loss is larger by a factor (128/i**2).

NOTE: the type of aggregation determines the magnitude of the KL-loss. Clearly, "sum" aggregation results in a larger KL-loss value compared to "mean" by a factor of n_layers.

NOTE: recall that sample-wise KL is obtained by summing over all dimensions, including Z. Also recall that in current 3D implementation of LVAE, no downsampling is done on Z. Therefore, to avoid emphasizing KL loss too much, we divide it by the Z dimension of input image in every case.

Parameters:

• kl_type (Literal['kl', 'kl_restricted']) –

  The type of KL divergence loss to compute.

• topdown_data (dict[str, Tensor]) –

  A dictionary containing information computed for each layer during the top-down pass. The dictionary must include the following keys: - "kl": The KL-loss values for each layer. Shape of each tensor is (B,). - "z": The sampled latents for each layer. Shape of each tensor is (B, layers, z_dims[i], H, W).

• rescaling (Literal['latent_dim', 'image_dim']) –

  The rescaling method used for the KL-loss values. If "latent_dim", the KL-loss values are rescaled w.r.t. the latent space dimensions (spatial + number of channels, i.e., (C, [Z], Y, X)). If "image_dim", the KL-loss values are rescaled w.r.t. the input image spatial dimensions.

• aggregation (Literal['mean', 'sum']) –

  The aggregation method used to combine the KL-loss values across layers. If "mean", the KL-loss values are averaged across layers. If "sum", the KL-loss values are summed across layers.

• free_bits_coeff (float) –

  The free bits coefficient used for the KL-loss computation.

• img_shape (Optional[tuple[int]], default: None ) –

  The shape of the input image to the LVAE model. Shape is ([Z], Y, X).

Returns:

• kl_loss ( Tensor ) –

  The KL divergence loss. Shape is (1, ).

Compute the reconstruction loss (negative log-likelihood).

Parameters:

• reconstruction (Tensor) –

  The output of the LVAE decoder. Shape is (B, C, [Z], Y, X), where C is the number of output channels (e.g., 1 in HDN, >1 in muSplit/denoiSplit).

• target (Tensor) –

  The target image used to compute the reconstruction loss. Shape is (B, C, [Z], Y, X), where C is the number of output channels (e.g., 1 in HDN, >1 in muSplit/denoiSplit).

• likelihood_obj (Likelihood) –

  The likelihood object used to compute the reconstruction loss.

Returns:

• Tensor –

  The recontruction loss (negative log-likelihood).

Loss function for HDN.

Parameters:

• model_outputs (tuple[Tensor, dict[str, Any]]) –

  Tuple containing the model predictions (shape is (B, target_ch, [Z], Y, X)) and the top-down layer data (e.g., sampled latents, KL-loss values, etc.).

• targets (Tensor) –

  The target image used to compute the reconstruction loss. In this case we use the input patch itself as target. Shape is (B, target_ch, [Z], Y, X).

• config (LVAELossConfig) –

  The config for loss function containing all loss hyperparameters.

• gaussian_likelihood (GaussianLikelihood) –

  The Gaussian likelihood object.

• noise_model_likelihood (NoiseModelLikelihood) –

  The noise model likelihood object.

Returns:

• output ( Optional[dict[str, Tensor]] ) –

  A dictionary containing the overall loss ["loss"], the reconstruction loss ["reconstruction_loss"], and the KL divergence loss ["kl_loss"].

Loss function for muSplit.

Parameters:

• model_outputs (tuple[Tensor, dict[str, Any]]) –

  Tuple containing the model predictions (shape is (B, target_ch, [Z], Y, X)) and the top-down layer data (e.g., sampled latents, KL-loss values, etc.).

• targets (Tensor) –

  The target image used to compute the reconstruction loss. Shape is (B, target_ch, [Z], Y, X).

• config (LVAELossConfig) –

  The config for loss function (e.g., KL hyperparameters, likelihood module, noise model, etc.).

• gaussian_likelihood (GaussianLikelihood) –

  The Gaussian likelihood object.

• noise_model_likelihood (Optional[NoiseModelLikelihood], default: None ) –

  The noise model likelihood object. Not used here.

Returns:

• output ( Optional[dict[str, Tensor]] ) –

  A dictionary containing the overall loss ["loss"], the reconstruction loss ["reconstruction_loss"], and the KL divergence loss ["kl_loss"].