clarena.cl_algorithms.regularisers.parameter_change

 1r"""The submodule in `regularisers` for parameter change regularisation."""
 2
 3__all__ = ["ParameterChangeReg"]
 4
 5import torch
 6from torch import Tensor, nn
 7
 8
 9class ParameterChangeReg(nn.Module):
10    r"""Parameter change regulariser.
11
12    $$R(\theta) = \text{factor} * \sum_i w_i \|\theta_i - \theta^\star_i\|^p$$
13
14    It promotes the target set of parameters $\theta = {\theta_i}_i$ not changing too much from another set of parameters $\theta^\star = {\theta^\star_i}_i$. The parameter distance here is $L^p$ distance. The regularisation can be parameter-wise weighted, i.e. $w_i$ in the formula.
15
16    It is used in:
17    - [L2 Regularisation algorithm](https://www.pnas.org/doi/10.1073/pnas.1611835114): as a L2 regulariser for the current task parameters to prevent them from changing too much from the previous task parameters.
18    - [EWC (Elastic Weight Consolidation) algorithm](https://www.pnas.org/doi/10.1073/pnas.1611835114): as a weighted L2 regulariser for the current task parameters to prevent them from changing too much from the previous task parameters. The regularisation weights are parameter importance measure calculated from fisher information. See equation 3 in the [EWC paper](https://www.pnas.org/doi/10.1073/pnas.1611835114).
19    """
20
21    def __init__(
22        self,
23        factor: float,
24        p_norm: float,
25    ) -> None:
26        r"""Initialise the regulariser.
27
28        **Args:**
29        - **factor** (`float`): the regularisation factor. Note that it is $\frac{\lambda}{2}$ rather than $\lambda$ in the [EWC paper](https://www.pnas.org/doi/10.1073/pnas.1611835114).
30        - **p_norm** (`float`): the norm of the distance, should be a positive float.
31        """
32        nn.Module.__init__(self)
33
34        self.factor = factor
35        """Store the regularisation factor for parameter change."""
36        self.p_norm = p_norm
37        """Store the the norm of the distance of two set of parameters. """
38
39    def forward(
40        self,
41        target_model: nn.Module,
42        ref_model: nn.Module,
43        weights: dict[str, Tensor],
44    ) -> Tensor:
45        r"""Calculate the regularisation loss.
46
47        **Args:**
48        - **target_model** (`nn.Module`): the model of the target parameters. In EWC, it's the model of current training task.
49        - **ref_model** (`nn.Module`): the reference model that you want target model parameters to prevent changing from. The reference model must have the same structure as the target model. In EWC, it's the model of one of the previous tasks.
50        - **weights** (`dict[str, Tensor]`): the regularisation weight for each parameter. Keys are parameter names and values are the weight tensors. The weight tensors must match the shape of model parameters. In EWC, it's the importance measure of each parameter, calculated from fisher information thing.
51
52        **Returns:**
53        - **reg** (`Tensor`): the parameter change regularisation value.
54        """
55        reg = 0.0
56
57        for (param_name, target_param), (param_name, ref_param) in zip(
58            target_model.named_parameters(), ref_model.named_parameters()
59        ):
60            weight = weights[param_name]
61            reg += torch.sum(weight * (target_param - ref_param).norm(p=self.p_norm))
62
63        return self.factor * reg