Computations

GGN eigenvalues

class vivit.EigvalshComputation(subsampling: Optional[List[int]] = None, mc_samples: int = 0, verbose: bool = False)

Provide BackPACK extension and hook to compute GGN eigenvalues.

__init__(subsampling: Optional[List[int]] = None, mc_samples: int = 0, verbose: bool = False)

Specify GGN approximations. Use no approximations by default.

Note

The loss function must use reduction = 'mean'.

Parameters:

subsampling – Indices of samples used for GGN curvature sub-sampling. None (equivalent to list(range(batch_size))) uses all mini-batch samples. Defaults to None (no curvature sub-sampling).
mc_samples – If 0, don’t Monte-Carlo (MC) approximate the GGN. Otherwise, specifies the number of MC samples used to approximate the backpropagated loss Hessian. Default: 0 (no MC approximation).
verbose – Turn on verbose mode. If enabled, this will print what’s happening during backpropagation to command line (consider it a debugging tool). Defaults to False.

get_extension() → BackpropExtension

Instantiate the BackPACK extension for computing GGN eigenvalues.

Returns:: BackPACK extension to compute eigenvalues that should be passed to the with backpack(...) context.

get_extension_hook(param_groups: List[Dict]) → Callable[[Module], None]

Instantiates the BackPACK extension hook to compute GGN eigenvalues.

Parameters:

param_groups –

Parameter groups list as required by a torch.optim.Optimizer. Specifies the block structure: Each group must specify the 'params' key which contains a list of the parameters that form a GGN block. Examples:

[{'params': list(p for p in model.parameters()}] computes eigenvalues for the full GGN (one block).
[{'params': [p]} for p in model.parameters()] computes eigenvalues for each block of a per-parameter block-diagonal GGN approximation.

Returns:

BackPACK extension hook to compute eigenvalues that should be passed to the with backpack(...) context. The hook computes GGN eigenvalues during backpropagation and stores them internally (under self._evals).

get_result(group: Dict) → Tensor

Return eigenvalues of a GGN block after the backward pass.

Parameters:: group – Parameter group that defines the GGN block.
Returns:: One-dimensional tensor containing the block’s eigenvalues (ascending order).
Raises:: KeyError – If there are no results for the group.

GGN eigenpairs (eigenvalues + eigenvector)

class vivit.EighComputation(subsampling: Optional[List[int]] = None, mc_samples: int = 0, verbose: bool = False, warn_small_eigvals: float = 0.0001)

Provide BackPACK extension and hook to compute GGN eigenpairs.

__init__(subsampling: Optional[List[int]] = None, mc_samples: int = 0, verbose: bool = False, warn_small_eigvals: float = 0.0001)

Specify GGN approximations. Use no approximations by default.

Note

The loss function must use reduction = 'mean'.

Parameters:

subsampling – Indices of samples used for GGN curvature sub-sampling. None (equivalent to list(range(batch_size))) uses all mini-batch samples. Defaults to None (no curvature sub-sampling).
mc_samples – If 0, don’t Monte-Carlo (MC) approximate the GGN. Otherwise, specifies the number of MC samples used to approximate the backpropagated loss Hessian. Default: 0 (no MC approximation).
verbose – Turn on verbose mode. If enabled, this will print what’s happening during backpropagation to command line (consider it a debugging tool). Defaults to False.
warn_small_eigvals – The eigenvector computation breaks down for numerically small eigenvalues close to zero. This variable triggers a user warning when attempting to compute eigenvectors for eigenvalues whose absolute value is smaller. Defaults to 1e-4. You can disable the warning by setting it to 0 (not recommended).

get_extension() → BackpropExtension

Instantiate the BackPACK extension for computing GGN eigenpairs.

Returns:: BackPACK extension to compute eigenvalues that should be passed to the with backpack(...) context.

get_extension_hook(param_groups: List[Dict]) → Callable[[Module], None]

Instantiates the BackPACK extension hook to compute GGN eigenpairs.

Parameters:

param_groups –

Parameter groups list as required by a torch.optim.Optimizer. Specifies the block structure: Each group must specify the 'params' key which contains a list of the parameters that form a GGN block, and a 'criterion' entry that specifies a filter function to select eigenvalues for the eigenvector computation (details below).

Examples for 'params':

[{'params': list(p for p in model.parameters()}] uses the full GGN (one block).
[{'params': [p]} for p in model.parameters()] uses a per-parameter block-diagonal GGN approximation.

The function specified under 'criterion' is a Callable[[Tensor], List[int]]. It receives the eigenvalues (in ascending order) and returns the indices of eigenvalues whose eigenvectors should be computed. Examples:

{'criterion': lambda evals: [evals.numel() - 1]} discards all eigenvalues except for the largest.
{'criterion': lambda evals: list(range(evals.numel()))} computes eigenvectors for all Gram matrix eigenvalues.

Returns:

BackPACK extension hook to compute eigenpairs that should be passed to the with backpack(...) context. The hook computes GGN eigenpairs during backpropagation and stores them internally (under self._evals and self._evecs).

get_result(group: Dict) → Tuple[Tensor, List[Tensor]]

Return eigenvalues and eigenvectors of a GGN block after the backward pass.

Parameters:

group – Parameter group that defines the GGN block.

Returns:

One-dimensional tensor containing the block’s eigenvalues and eigenvectors in parameter list format with leading axis for the eigenvectors.

Example: Let evals, evecs denote the returned variables. Let group['params'] = [p1, p2] consist of two parameters. Then, evecs contains tensors of shape [(evals.numel(), *p1.shape), (evals.numel(), *p2.shape)]. For an eigenvalue evals[k], the corresponding eigenvector (in list format) is [vecs[k] for vecs in evecs] and its tensors are of shape [p1.shape, p2.shape].

Raises:

KeyError – If there are no results for the group.

1ˢᵗ- and 2ⁿᵈ-order directional derivatives along GGN eigenvectors

class vivit.DirectionalDerivativesComputation(subsampling_grad: Optional[List[int]] = None, subsampling_ggn: Optional[List[int]] = None, mc_samples_ggn: Optional[int] = 0, verbose: Optional[bool] = False, warn_small_eigvals: float = 0.0001)

Provide BackPACK extension and hook for 1ˢᵗ/2ⁿᵈ-order directional derivatives.

The directions are given by the GGN eigenvectors. First-order directional derivatives are denoted γ, second-order directional derivatives as λ.

__init__(subsampling_grad: Optional[List[int]] = None, subsampling_ggn: Optional[List[int]] = None, mc_samples_ggn: Optional[int] = 0, verbose: Optional[bool] = False, warn_small_eigvals: float = 0.0001)

Specify GGN and gradient approximations. Use no approximations by default.

Note

The loss function must use reduction = 'mean'.

Parameters:

subsampling_grad – Indices of samples used for gradient sub-sampling. None (equivalent to list(range(batch_size))) uses all mini-batch samples to compute directional gradients . Defaults to None (no gradient sub-sampling).
subsampling_ggn – Indices of samples used for GGN curvature sub-sampling. None (equivalent to list(range(batch_size))) uses all mini-batch samples to compute directions and directional curvatures. Defaults to None (no curvature sub-sampling).
mc_samples_ggn – If 0, don’t Monte-Carlo (MC) approximate the GGN (using the same samples to compute the directions and directional curvatures). Otherwise, specifies the number of MC samples used to approximate the backpropagated loss Hessian. Default: 0 (no MC approximation).
verbose – Turn on verbose mode. If enabled, this will print what’s happening during backpropagation to command line (consider it a debugging tool). Defaults to False.
warn_small_eigvals – The gamma computation breaks down for numerically small eigenvalues close to zero (we need to divide by their square root). This variable triggers a user warning when attempting to compute directional gradients for eigenvalues whose absolute value is smaller. Defaults to 1e-4. You can disable the warning by setting it to 0 (not recommended).

get_extension_hook(param_groups: List[Dict]) → Callable[[Module], None]

Instantiate BackPACK extension hook to compute GGN directional derivatives.

Parameters:

param_groups –

Parameter groups list as required by a torch.optim.Optimizer. Specifies the block structure: Each group must specify the 'params' key which contains a list of the parameters that form a GGN block, and a 'criterion' entry that specifies a filter function to select eigenvalues as directions along which to compute directional derivatives (details below).

Examples for 'params':

[{'params': list(p for p in model.parameters()}] uses the full GGN (one block).
[{'params': [p]} for p in model.parameters()] uses a per-parameter block-diagonal GGN approximation.

The function specified under 'criterion' is a Callable[[Tensor], List[int]]. It receives the eigenvalues (in ascending order) and returns the indices of eigenvalues whose eigenvectors should be used as directions to evaluate directional derivatives. Examples:

{'criterion': lambda evals: [evals.numel() - 1]} discards all directions except for the leading eigenvector.
{'criterion': lambda evals: list(range(evals.numel()))} computes directional derivatives along all Gram matrix eigenvectors.

Returns:

BackPACK extension hook, to compute directional derivatives, that should be passed to the with backpack(...) context. The hook computes GGN directional derivatives during backpropagation and stores them internally (under self._gammas and self._lambdas).

get_extensions() → List[BackpropExtension]

Instantiate the BackPACK extensions to compute GGN directional derivatives.

Returns:: BackPACK extensions, to compute directional 1ˢᵗ- and 2ⁿᵈ-order directional derivatives along GGN eigenvectors, that should be extracted and passed to the with backpack(...) context.

get_result(group: Dict) → Tuple[Tensor, Tensor]

Return 1ˢᵗ/2ⁿᵈ-order directional derivatives along GGN eigenvectors.

Must be called after the backward pass.

Parameters:: group – Parameter group that defines the GGN block.
Returns:: 1ˢᵗ- directional derivatives γ as [N, K] tensor with γ[n, k] the directional gradient of sample n (or subsampling_grad[n]) along direction k. 2ⁿᵈ- directional derivatives λ as [N, K] tensor with λ[n, k] the directional curvature of sample n (or subsampling_ggn[n]) along direction k.
Raises:: KeyError – If there are no results for the group.

Directionally damped Newton steps

class vivit.DirectionalDampedNewtonComputation(subsampling_grad: Optional[List[int]] = None, subsampling_ggn: Optional[List[int]] = None, mc_samples_ggn: Optional[int] = 0, verbose: Optional[bool] = False, warn_small_eigvals: float = 0.0001)

Provides BackPACK extension and hook for directionally damped Newton steps.

Let \(\{e_k\}_{k = 1, \dots, K}\) denote the \(K\) eigenvectors used for the Newton step. Let \(\gamma_k, \lambda_k\) denote the first- and second-order derivatives along these eigenvectors. The direcionally damped Newton step \(s\) is

\[s = \sum_{k=1}^K \frac{- \gamma_k}{\lambda_k + \delta_k} e_k\]

Here, \(\delta_k\) is the damping along direction \(k\). The exact Newton step has zero damping.

__init__(subsampling_grad: Optional[List[int]] = None, subsampling_ggn: Optional[List[int]] = None, mc_samples_ggn: Optional[int] = 0, verbose: Optional[bool] = False, warn_small_eigvals: float = 0.0001)

Specify GGN and gradient approximations. Use no approximations by default.

Note

The loss function must use reduction = 'mean'.

Parameters:

subsampling_grad – Indices of samples used for gradient sub-sampling. None (equivalent to list(range(batch_size))) uses all mini-batch samples to compute directional gradients . Defaults to None (no gradient sub-sampling).
subsampling_ggn – Indices of samples used for GGN curvature sub-sampling. None (equivalent to list(range(batch_size))) uses all mini-batch samples to compute directions and directional curvatures. Defaults to None (no curvature sub-sampling).
mc_samples_ggn – If 0, don’t Monte-Carlo (MC) approximate the GGN (using the same samples to compute the directions and directional curvatures). Otherwise, specifies the number of MC samples used to approximate the backpropagated loss Hessian. Default: 0 (no MC approximation).
verbose – Turn on verbose mode. If enabled, this will print what’s happening during backpropagation to command line (consider it a debugging tool). Defaults to False.
warn_small_eigvals – The gamma computation and tranformation from Gram into parameter space break down for numerically small eigenvalues close to zero (need to divide by their square root). This variable triggers a user warning when attempting to compute damped Newton steps for eigenvalues whose absolute value is smaller. Defaults to 1e-4. You can disable the warning by setting it to 0 (not recommended).

get_extension_hook(param_groups: List[Dict]) → Callable[[Module], None]

Instantiate BackPACK extension hook to compute a damped Newton step.

Parameters:

param_groups –

Parameter groups list as required by a torch.optim.Optimizer. Specifies the block structure: Each group must specify the 'params' key which contains a list of the parameters that form a GGN block, a 'criterion' entry that specifies a filter function to select eigenvalues as directions along which to compute the damped Newton step, and a 'damping' entry that that contains a function that produces the directional dampings (details below).

Examples for 'params':

[{'params': list(p for p in model.parameters()}] uses the full GGN (one block).
[{'params': [p]} for p in model.parameters()] uses a per-parameter block-diagonal GGN approximation.

The function specified under 'criterion' is a Callable[[Tensor], List[int]]. It receives the eigenvalues (in ascending order) and returns the indices of eigenvalues whose eigenvectors should be used as directions for the damped Newon step. Examples:

{'criterion': lambda evals: [evals.numel() - 1]} discards all directions except for the leading eigenvector.
{'criterion': lambda evals: list(range(evals.numel()))} computes damped Newton step along all Gram matrix eigenvectors.

The function specified under 'damping' is a Callable[[Tensor, Tensor, Tensor, Tensor]]. It receives the eigenvalues, Gram eigenvectors, directional gradients and directional curvatures, and returns a tensor that contains the damping values for all directions. Example:

{'damping': lambda evals, evecs, gammas, lambdas: ones_like(evals)} corresponds to constant damping \(\delta_k = 1\ \forall k\).

Returns:

BackPACK extension hook, to compute damped Newton steps, that should be passed to the with backpack(...) context. The hook computes a damped Newton step during backpropagation and stores it internally (under self._newton_step).

get_extensions() → List[BackpropExtension]

Instantiate the BackPACK extensions to compute a damped Newton step.

Returns:: BackPACK extensions, to compute a damped Newton step, that should be extracted and passed to the with backpack(...) context.

get_result(group: Dict) → Tuple[Tensor]

Return a damped Newton step in parameter format.

Must be called after the backward pass.

Parameters:: group – Parameter group that defines the GGN block.
Returns:: Damped Newton step in parameter format, i.e. the same format as group['params'].
Raises:: KeyError – If there are no results for the group.