Add fisher_type argument and support empirical Fisher in KFAC

curvlinops/kfac.py

@@ -80,6 +80,7 @@ def __init__(
         check_deterministic: bool = True,
         shape: Union[Tuple[int, int], None] = None,
         seed: int = 2147483647,
+        fisher_type: str = "mc",
         mc_samples: int = 1,
     ):
         """Kronecker-factored approximate curvature (KFAC) proxy of the Fisher/GGN.
@@ -114,6 +115,14 @@ def __init__(
                 from the parameters. Defaults to ``None``.
             seed: The seed for the random number generator used to draw labels
                 from the model's predictive distribution. Defaults to ``2147483647``.
+            fisher_type: The type of Fisher/GGN to approximate. If ``'type-2'``, the
+                expectation over the model outputs is computed exactly by
+                computing the backwward pass for each output dimension. This is
+                sometimes also called type-2 Fisher. If ``'mc'``, the expectation
+                is approximated by sampling ``mc_samples`` labels from the model's
+                predictive distribution. If ``'empirical'``, the empirical gradients
+                are used which corresponds to the uncentered gradient covariance, or
+                the empirical Fisher. Defaults to ``'mc'``.
             mc_samples: The number of Monte-Carlo samples to use per data point.
                 Defaults to ``1``.
 
@@ -152,6 +161,7 @@ def __init__(
 
         self._seed = seed
         self._generator: Union[None, Generator] = None
+        self._fisher_type = fisher_type
         self._mc_samples = mc_samples
         self._input_covariances: Dict[Tuple[int, ...], Tensor] = {}
         self._gradient_covariances: Dict[Tuple[int, ...], Tensor] = {}
@@ -210,7 +220,14 @@ def _adjoint(self) -> KFACLinearOperator:
         return self
 
     def _compute_kfac(self):
-        """Compute and cache KFAC's Kronecker factors for future ``matvec``s."""
+        """Compute and cache KFAC's Kronecker factors for future ``matvec``s.
+
+        Raises:
+            NotImplementedError: If ``fisher_type == 'type-2'``.
+            ValueError: If ``fisher_type`` is not ``'type-2'``, ``'mc'``, or
+                ``'empirical'``.
+
+        """
         # install forward and backward hooks
         hook_handles: List[RemovableHandle] = []
         hook_handles.extend(
@@ -231,13 +248,27 @@ def _compute_kfac(self):
             self._generator = Generator(device=self._device)
         self._generator.manual_seed(self._seed)
 
-        for X, _ in self._loop_over_data(desc="KFAC matrices"):
+        for X, y in self._loop_over_data(desc="KFAC matrices"):
             output = self._model_func(X)
 
-            for mc in range(self._mc_samples):
-                y_sampled = self.draw_label(output)
-                loss = self._loss_func(output, y_sampled)
-                loss.backward(retain_graph=mc != self._mc_samples - 1)
+            if self._fisher_type == "type-2":
+                raise NotImplementedError(
+                    "Using the exact expectation for computing the KFAC "
+                    "approximation of the Fisher is not yet supported."
+                )
+            elif self._fisher_type == "mc":
+                for mc in range(self._mc_samples):
+                    y_sampled = self.draw_label(output)
+                    loss = self._loss_func(output, y_sampled)
+                    loss.backward(retain_graph=mc != self._mc_samples - 1)
+            elif self._fisher_type == "empirical":
+                loss = self._loss_func(output, y)
+                loss.backward()
+            else:
+                raise ValueError(
+                    f"Invalid fisher_type: {self._fisher_type}. "
+                    + "Supported: 'type-2', 'mc', 'empirical'."
+                )
 
         # clean up
         self._model_func.zero_grad()