18class CULAlgorithm:
19    r"""The base class of continual unlearning algorithms."""
20
21    def __init__(self, model: UnlearnableCLAlgorithm) -> None:
22        r"""
23        **Args:**
24        - **model** (`UnlearnableCLAlgorithm`): the continual learning model.
25        """
26
27        # components
28        self.model: UnlearnableCLAlgorithm = model
29        r"""The continual learning model."""
30
31        # task ID control
32        self.task_id: int
33        r"""Task ID counter indicating which task is being processed. Self updated during the task loop. Valid from 1 to `cl_dataset.num_tasks`."""
34        self.processed_task_ids: list[int] = []
35        r"""Task IDs that have been processed."""
36        self.unlearning_task_ids: list[int] = []
37        r"""The list of task IDs that are requested to be unlearned after training `self.task_id`."""
38        self.unlearned_task_ids: set[int] = set()
39        r"""The list of task IDs that have been unlearned in the experiment. """
40        self.unlearnable_task_ids: list[int] = []
41        r"""The list of task IDs that are unlearnable at the current `self.task_id`."""
42        self.task_ids_just_no_longer_unlearnable: list[int] = []
43        r"""The list of task IDs that are just no longer unlearnable at the current `self.task_id`."""
44
45    def setup_task_id(
46        self,
47        task_id: int,
48        unlearning_requests: dict[int, list[int]],
49        unlearnable_task_ids: list[int],
50        task_ids_just_no_longer_unlearnable: list[int],
51    ) -> None:
52        r"""Set up which task the CUL experiment is on. This must be done before `unlearn()` method is called.
53
54        **Args:**
55        - **task_id** (`int`): the target task ID to be set up.
56        - **unlearning_requests** (`dict[int, list[int]]`): the entire unlearning requests. Keys are IDs of the tasks that request unlearning after their learning, and values are the list of the previous tasks to be unlearned.
57        - **unlearnable_task_ids** (`list[int]`): the list of unlearnable task IDs at the current `self.task_id`.
58        - **task_ids_just_no_longer_unlearnable** (`list[int]`): the list of task IDs that are just no longer unlearnable at the current `self.task_id`.
59        """
60        self.task_id = task_id
61        self.processed_task_ids.append(task_id)
62
63        unlearning_task_ids = (
64            unlearning_requests[task_id] if task_id in unlearning_requests else []
65        )
66        self.unlearning_task_ids = unlearning_task_ids
67        self.model.unlearning_task_ids = unlearning_task_ids
68
69        self.unlearnable_task_ids = unlearnable_task_ids
70        self.model.unlearnable_task_ids = unlearnable_task_ids
71
72        self.task_ids_just_no_longer_unlearnable = task_ids_just_no_longer_unlearnable
73        self.model.task_ids_just_no_longer_unlearnable = (
74            task_ids_just_no_longer_unlearnable
75        )
76
77    def setup_test_task_id(self) -> None:
78        r"""Set up before testing `self.task_id`. This must be done after `unlearn()` method is called."""
79
80        self.unlearned_task_ids.update(
81            self.unlearning_task_ids
82        )  # update the maintained set of unlearned task IDs
83        self.model.unlearned_task_ids = (
84            self.unlearned_task_ids
85        )  # let model know the unlearned task IDs
86
87    @abstractmethod
88    def unlearn(self) -> None:
89        r"""Unlearn the requested unlearning tasks (`self.unlearning_task_ids`) after training `self.task_id`. **It must be implemented in subclasses.**"""

 92class AmnesiacCULAlgorithm(CULAlgorithm):
 93    r"""The base class of Amnesiac continual unlearning algorithm.
 94
 95    The Amnesiac continual unlearning algorithm refers to update deletion operation that directly delete the parameter updates during a task's training. This is inspired by [AmnesiacML](https://arxiv.org/abs/2010.10981) in machine unlearning. In detail, the task-wise parameter updates are stored:
 96
 97    $$\theta_{l,ij}^{(t)} = \theta_{l,ij}^{(0)} + \sum_{\tau=1}^{t} \Delta \theta_{l,ij}^{(\tau)}$$
 98
 99    To unlearn $u(t)$, delete these updates:
100
101    $$\theta_{l,ij}^{(t-u(t))} = \theta_{l,ij}^{(t)} - \sum_{\tau\in u(t)}\Delta \theta_{l,ij}^{(\tau)}$$
102
103    It is mainly used in AmnesaicHAT, but can also be used in constructing other vanilla baseline continual unlearning algorithms based on different continual learning algorithms.
104    """
105
106    def __init__(
107        self,
108        model: AmnesiacCLAlgorithm,
109    ) -> None:
110        r"""Initialize the unlearning algorithm with the continual learning model.
111
112        **Args:**
113        - **model** (`AmnesiacHAT`): the continual learning model. It must be `AmnesicCLAlgorithm`.
114        """
115        super().__init__(model=model)
116
117    def delete_update(self):
118        r"""Delete the updates for unlearning tasks from the current parameters."""
119
120        # substract the corresponding parameter update from backbone
121        updated_state_dict = deepcopy(self.model.backbone.state_dict())
122        for task_id in self.unlearning_task_ids:
123            param_update = self.model.parameters_task_update.get(task_id)
124            if param_update is None:
125                pylogger.warning(
126                    "Attempted to delete backbone update for task %d, but it was not found.",
127                    task_id,
128                )
129                continue
130            for layer_name, param_tensor in param_update.items():
131                if layer_name in updated_state_dict:
132                    target_tensor = updated_state_dict[layer_name]
133                    if (
134                        param_tensor.device != target_tensor.device
135                        or param_tensor.dtype != target_tensor.dtype
136                    ):
137                        param_tensor = param_tensor.to(
138                            device=target_tensor.device, dtype=target_tensor.dtype
139                        )
140                    updated_state_dict[layer_name] -= param_tensor
141
142            del self.model.parameters_task_update[task_id]  # delete the record
143
144        self.model.backbone.load_state_dict(updated_state_dict, strict=False)
145
146        # substract the corresponding parameter update from heads
147        updated_heads_state_dict = deepcopy(self.model.heads.state_dict())
148        for task_id in self.unlearning_task_ids:
149            param_update = self.model.parameters_task_update_heads.get(task_id)
150            if param_update is None:
151                pylogger.warning(
152                    "Attempted to delete head update for task %d, but it was not found.",
153                    task_id,
154                )
155                continue
156            for param_name, param_tensor in param_update.items():
157                if param_name in updated_heads_state_dict:
158                    target_tensor = updated_heads_state_dict[param_name]
159                    if (
160                        param_tensor.device != target_tensor.device
161                        or param_tensor.dtype != target_tensor.dtype
162                    ):
163                        param_tensor = param_tensor.to(
164                            device=target_tensor.device, dtype=target_tensor.dtype
165                        )
166                    updated_heads_state_dict[param_name] -= param_tensor
167
168            del self.model.parameters_task_update_heads[task_id]  # delete the record
169
170        self.model.heads.load_state_dict(updated_heads_state_dict, strict=False)
171
172    def unlearn(self) -> None:
173        r"""Unlearn the requested unlearning tasks in the current task `self.task_id`.
174
175        This is the default implementation of `unlearn()` method for Amnesiac continual unlearning algorithms. Please override it in subclasses if necessary.
176        """
177
178        # delete the corresponding parameter update records
179        self.delete_update()