n-gao / pytorch-kfac Goto Github PK

A question regarding the following code in linear_block.py in the method update_cov:

        ....
        self._activations_cov.add_to_average(activation_cov, cov_ema_decay)
        self._sensitivities_cov.add_to_average(sensitivity_cov, cov_ema_decay)

The method add_to_average(self, value: torch.Tensor, decay: float = 1.0, weight: float = 1.0) takes a weight as parameter.
Shouldn't be the weight set to 1- cov_ema_decay such that the overall sum of decay and weight is 1.0?

This is done in that way in other implementations, like here: https://github.com/lzhangbv/kfac_pytorch/blob/master/kfac/utils.py#L66 or here https://github.com/Thrandis/EKFAC-pytorch/blob/master/kfac.py#L174 (here cov_ema_decay corresponds to alpha).

Can your code be used to compute the KFAC (Fisher) matrix using a different loss than the loss we take the gradient of?

If so, how?

thanks