clarena.cl_algorithms.adahat

  1r"""
  2The submodule in `cl_algorithms` for [AdaHAT (Adaptive Hard Attention to the Task) algorithm](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9).
  3"""
  4
  5__all__ = ["AdaHAT"]
  6
  7import logging
  8
  9import torch
 10from torch import Tensor
 11
 12from clarena.backbones import HATMaskBackbone
 13from clarena.cl_algorithms import HAT
 14from clarena.cl_heads import HeadsCIL, HeadsTIL
 15from clarena.utils import HATNetworkCapacity
 16
 17# always get logger for built-in logging in each module
 18pylogger = logging.getLogger(__name__)
 19
 20
 21class AdaHAT(HAT):
 22    r"""AdaHAT (Adaptive Hard Attention to the Task) algorithm.
 23
 24    [Adaptive HAT (Adaptive Hard Attention to the Task, 2024)](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9) is an architecture-based continual learning approach that improves [HAT (Hard Attention to the Task, 2018)](http://proceedings.mlr.press/v80/serra18a) by introducing new adaptive soft gradient clipping based on parameter importance and network sparsity.
 25
 26    We implement AdaHAT as a subclass of HAT algorithm, as AdaHAT has the same  `forward()`, `compensate_task_embedding_gradients()`, `training_step()`, `on_train_end()`,`validation_step()`, `test_step()` method as `HAT` class.
 27    """
 28
 29    def __init__(
 30        self,
 31        backbone: HATMaskBackbone,
 32        heads: HeadsTIL | HeadsCIL,
 33        adjustment_mode: str,
 34        adjustment_intensity: float,
 35        s_max: float,
 36        clamp_threshold: float,
 37        mask_sparsity_reg_factor: float,
 38        mask_sparsity_reg_mode: str = "original",
 39        task_embedding_init_mode: str = "N01",
 40        epsilon: float = 0.1,
 41    ) -> None:
 42        r"""Initialise the AdaHAT algorithm with the network.
 43
 44        **Args:**
 45        - **backbone** (`HATMaskBackbone`): must be a backbone network with HAT mask mechanism.
 46        - **heads** (`HeadsTIL` | `HeadsCIL`): output heads.
 47        - **adjustment_mode** (`str`): the strategy of adjustment i.e. the mode of gradient clipping, should be one of the following:
 48            1. 'adahat': set the gradients of parameters linking to masked units to a soft adjustment rate in the original AdaHAT approach. This is the way that AdaHAT does, which allowes the part of network for previous tasks to be updated slightly. See equation (8) and (9) chapter 3.1 in [AdaHAT paper](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9).
 49            2. 'adahat_no_sum': set the gradients of parameters linking to masked units to a soft adjustment rate in the original AdaHAT approach, but without considering the information of parameter importance i.e. summative mask. This is the way that one of the AdaHAT ablation study does. See chapter 4.3 in [AdaHAT paper](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9).
 50            3. 'adahat_no_reg': set the gradients of parameters linking to masked units to a soft adjustment rate in the original AdaHAT approach, but without considering the information of network sparsity i.e. mask sparsity regularisation value. This is the way that one of the AdaHAT ablation study does. See chapter 4.3 in [AdaHAT paper](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9).
 51        - **adjustment_intensity** (`float`): hyperparameter, control the overall intensity of gradient adjustment. It's the $\alpha$ in equation (9) in [AdaHAT paper](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9).
 52        - **s_max** (`float`): hyperparameter, the maximum scaling factor in the gate function. See chapter 2.4 "Hard Attention Training" in [HAT paper](http://proceedings.mlr.press/v80/serra18a).
 53        - **clamp_threshold** (`float`): the threshold for task embedding gradient compensation. See chapter 2.5 "Embedding Gradient Compensation" in [HAT paper](http://proceedings.mlr.press/v80/serra18a).
 54        - **mask_sparsity_reg_factor** (`float`): hyperparameter, the regularisation factor for mask sparsity.
 55        - **mask_sparsity_reg_mode** (`str`): the mode of mask sparsity regularisation, should be one of the following:
 56            1. 'original' (default): the original mask sparsity regularisation in HAT paper.
 57            2. 'cross': the cross version mask sparsity regularisation.
 58        - **task_embedding_init_mode** (`str`): the initialisation method for task embeddings, should be one of the following:
 59            1. 'N01' (default): standard normal distribution $N(0, 1)$.
 60            2. 'U-11': uniform distribution $U(-1, 1)$.
 61            3. 'U01': uniform distribution $U(0, 1)$.
 62            4. 'U-10': uniform distribution $U(-1, 0)$.
 63            5. 'last': inherit task embedding from last task.
 64        - **epsilon** (`float`): the value added to network sparsity to avoid division by zero appeared in equation (9) in [AdaHAT paper](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9).
 65        """
 66        HAT.__init__(
 67            self,
 68            backbone=backbone,
 69            heads=heads,
 70            adjustment_mode=adjustment_mode,
 71            s_max=s_max,
 72            clamp_threshold=clamp_threshold,
 73            mask_sparsity_reg_factor=mask_sparsity_reg_factor,
 74            mask_sparsity_reg_mode=mask_sparsity_reg_mode,
 75            task_embedding_init_mode=task_embedding_init_mode,
 76            alpha=None,
 77        )
 78
 79        self.adjustment_intensity = adjustment_intensity
 80        r"""Store the adjustment intensity in equation (9) in [AdaHAT paper](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9)."""
 81        self.epsilon = epsilon
 82        """Store the small value to avoid division by zero appeared in equation (9) in [AdaHAT paper](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9)."""
 83
 84        self.summative_mask_for_previous_tasks: dict[str, Tensor] = {}
 85        r"""Store the summative binary attention mask $\mathrm{M}^{<t,\text{sum}}$ previous tasks $1,\cdots, t-1$, gated from the task embedding. Keys are task IDs and values are the corresponding summative mask. Each cumulative mask is a dict where keys are layer names and values are the binary mask tensor for the layer. The mask tensor has size (number of units). """
 86
 87        # set manual optimisation
 88        self.automatic_optimization = False
 89
 90        AdaHAT.sanity_check(self)
 91
 92    def sanity_check(self) -> None:
 93        r"""Check the sanity of the arguments.
 94
 95        **Raises:**
 96        - **ValueError**: If the `adjustment_intensity` is not positive.
 97        """
 98        if self.adjustment_intensity <= 0:
 99            raise ValueError(
100                f"The adjustment intensity should be positive, but got {self.adjustment_intensity}."
101            )
102
103    def on_train_start(self) -> None:
104        r"""Additionally initialise the summative mask at the beginning of first task."""
105        HAT.on_train_start(self)
106
107        # initialise the summative mask at the beginning of first task. This should not be called in `__init__()` method as the `self.device` is not available at that time.
108        if self.task_id == 1:
109            for layer_name in self.backbone.weighted_layer_names:
110                layer = self.backbone.get_layer_by_name(
111                    layer_name
112                )  # get the layer by its name
113                num_units = layer.weight.shape[0]
114
115                self.summative_mask_for_previous_tasks[layer_name] = torch.zeros(
116                    num_units
117                ).to(
118                    self.device
119                )  # the summative mask $\mathrm{M}^{<t,\text{sum}}$ is initialised as zeros mask ($t = 1$). See equation (7) in chapter 3.1 in [AdaHAT paper](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9).
120
121    def clip_grad_by_adjustment(
122        self,
123        network_sparsity: dict[str, Tensor] | None = None,
124    ) -> Tensor:
125        r"""Clip the gradients by the adjustment rate.
126
127        Note that as the task embedding fully covers every layer in the backbone network, no parameters are left out of this system. This applies not only the parameters in between layers with task embedding, but also those before the first layer. We designed it seperately in the codes.
128
129        Network capacity is measured along with this method. Network capacity is defined as the average adjustment rate over all parameters. See chapter 4.1 in [AdaHAT paper](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9).
130
131        **Args:**
132        - **network_sparsity** (`dict[str, Tensor]` | `None`): The network sparsity i.e. the mask sparsity loss of each layer for the current task. It applies only to AdaHAT modes, as it is used to calculate the adjustment rate for the gradients.
133
134        **Returns:**
135        - **capacity** (`Tensor`): the calculated network capacity.
136        """
137
138        # initialise network capacity metric
139        capacity = HATNetworkCapacity()
140
141        # Calculate the adjustment rate for gradients of the parameters, both weights and biases (if exists)
142        for layer_name in self.backbone.weighted_layer_names:
143
144            layer = self.backbone.get_layer_by_name(
145                layer_name
146            )  # get the layer by its name
147
148            # placeholder for the adjustment rate to avoid the error of using it before assignment
149            adjustment_rate_weight = 1
150            adjustment_rate_bias = 1
151
152            weight_importance, bias_importance = (
153                self.backbone.get_layer_measure_parameter_wise(
154                    unit_wise_measure=self.summative_mask_for_previous_tasks,
155                    layer_name=layer_name,
156                    aggregation="min",
157                )
158            )  # AdaHAT depend on parameter importance instead of parameter mask like HAT
159
160            network_sparsity_layer = network_sparsity[layer_name]
161
162            if self.adjustment_mode == "adahat":
163                r_layer = self.adjustment_intensity / (
164                    self.epsilon + network_sparsity_layer
165                )
166                adjustment_rate_weight = torch.div(
167                    r_layer, (weight_importance + r_layer)
168                )
169                adjustment_rate_bias = torch.div(r_layer, (bias_importance + r_layer))
170
171            elif self.adjustment_mode == "adahat_no_sum":
172
173                r_layer = self.adjustment_intensity / (
174                    self.epsilon + network_sparsity_layer
175                )
176                adjustment_rate_weight = torch.div(r_layer, (self.task_id + r_layer))
177                adjustment_rate_bias = torch.div(r_layer, (self.task_id + r_layer))
178
179            elif self.adjustment_mode == "adahat_no_reg":
180
181                r_layer = self.adjustment_intensity / (self.epsilon + 0.0)
182                adjustment_rate_weight = torch.div(
183                    r_layer, (weight_importance + r_layer)
184                )
185                adjustment_rate_bias = torch.div(r_layer, (bias_importance + r_layer))
186
187            # apply the adjustment rate to the gradients
188            layer.weight.grad.data *= adjustment_rate_weight
189            if layer.bias is not None:
190                layer.bias.grad.data *= adjustment_rate_bias
191
192            # update network capacity metric
193            capacity.update(adjustment_rate_weight, adjustment_rate_bias)
194
195        return capacity.compute()
196
197    def on_train_end(self) -> None:
198        r"""Additionally update summative mask after training the task."""
199
200        HAT.on_train_end(self)
201
202        mask_t = self.masks[
203            f"{self.task_id}"
204        ]  # get stored mask for the current task again
205        self.summative_mask_for_previous_tasks = {
206            layer_name: self.summative_mask_for_previous_tasks[layer_name]
207            + mask_t[layer_name]
208            for layer_name in self.backbone.weighted_layer_names
209        }