clarena.cl_algorithms.fgadahat

   1r"""
   2The submodule in `cl_algorithms` for FG-AdaHAT algorithm.
   3"""
   4
   5__all__ = ["FGAdaHAT"]
   6
   7import logging
   8import math
   9from typing import Any
  10
  11import torch
  12from captum.attr import (
  13    InternalInfluence,
  14    LayerConductance,
  15    LayerDeepLift,
  16    LayerDeepLiftShap,
  17    LayerFeatureAblation,
  18    LayerGradCam,
  19    LayerGradientShap,
  20    LayerGradientXActivation,
  21    LayerIntegratedGradients,
  22    LayerLRP,
  23)
  24from torch import Tensor
  25
  26from clarena.backbones import HATMaskBackbone
  27from clarena.cl_algorithms import AdaHAT
  28from clarena.heads import HeadsTIL
  29from clarena.utils.metrics import HATNetworkCapacityMetric
  30from clarena.utils.transforms import min_max_normalize
  31
  32# always get logger for built-in logging in each module
  33pylogger = logging.getLogger(__name__)
  34
  35
  36class FGAdaHAT(AdaHAT):
  37    r"""FG-AdaHAT (Fine-Grained Adaptive Hard Attention to the Task) algorithm.
  38
  39    An architecture-based continual learning approach that improves [AdaHAT (Adaptive Hard Attention to the Task)](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9) by introducing fine-grained neuron-wise importance measures guiding the adaptive adjustment mechanism in AdaHAT.
  40
  41    We implement FG-AdaHAT as a subclass of AdaHAT, as it reuses AdaHAT's summative mask and other components.
  42    """
  43
  44    def __init__(
  45        self,
  46        backbone: HATMaskBackbone,
  47        heads: HeadsTIL,
  48        adjustment_intensity: float,
  49        importance_type: str,
  50        importance_summing_strategy: str,
  51        importance_scheduler_type: str,
  52        neuron_to_weight_importance_aggregation_mode: str,
  53        s_max: float,
  54        clamp_threshold: float,
  55        mask_sparsity_reg_factor: float,
  56        mask_sparsity_reg_mode: str = "original",
  57        base_importance: float = 0.01,
  58        base_mask_sparsity_reg: float = 0.1,
  59        base_linear: float = 10,
  60        filter_by_cumulative_mask: bool = False,
  61        filter_unmasked_importance: bool = True,
  62        step_multiply_training_mask: bool = True,
  63        task_embedding_init_mode: str = "N01",
  64        importance_summing_strategy_linear_step: float | None = None,
  65        importance_summing_strategy_exponential_rate: float | None = None,
  66        importance_summing_strategy_log_base: float | None = None,
  67        non_algorithmic_hparams: dict[str, Any] = {},
  68    ) -> None:
  69        r"""Initialize the FG-AdaHAT algorithm with the network.
  70
  71        **Args:**
  72        - **backbone** (`HATMaskBackbone`): must be a backbone network with the HAT mask mechanism.
  73        - **heads** (`HeadsTIL`): output heads. FG-AdaHAT supports only TIL (Task-Incremental Learning).
  74        - **adjustment_intensity** (`float`): hyperparameter, controls the overall intensity of gradient adjustment (the $\alpha$ in the paper).
  75        - **importance_type** (`str`): the type of neuron-wise importance, must be one of:
  76            1. 'input_weight_abs_sum': sum of absolute input weights;
  77            2. 'output_weight_abs_sum': sum of absolute output weights;
  78            3. 'input_weight_gradient_abs_sum': sum of absolute gradients of the input weights (Input Gradients (IG) in the paper);
  79            4. 'output_weight_gradient_abs_sum': sum of absolute gradients of the output weights (Output Gradients (OG) in the paper);
  80            5. 'activation_abs': absolute activation;
  81            6. 'input_weight_abs_sum_x_activation_abs': sum of absolute input weights multiplied by absolute activation (Input Contribution Utility (ICU) in the paper);
  82            7. 'output_weight_abs_sum_x_activation_abs': sum of absolute output weights multiplied by absolute activation (Contribution Utility (CU) in the paper);
  83            8. 'gradient_x_activation_abs': absolute gradient (the saliency) multiplied by activation;
  84            9. 'input_weight_gradient_square_sum': sum of squared gradients of the input weights;
  85            10. 'output_weight_gradient_square_sum': sum of squared gradients of the output weights;
  86            11. 'input_weight_gradient_square_sum_x_activation_abs': sum of squared gradients of the input weights multiplied by absolute activation (Activation Fisher Information (AFI) in the paper);
  87            12. 'output_weight_gradient_square_sum_x_activation_abs': sum of squared gradients of the output weights multiplied by absolute activation;
  88            13. 'conductance_abs': absolute layer conductance;
  89            14. 'internal_influence_abs': absolute internal influence (Internal Influence (II) in the paper);
  90            15. 'gradcam_abs': absolute Grad-CAM;
  91            16. 'deeplift_abs': absolute DeepLIFT (DeepLIFT (DL) in the paper);
  92            17. 'deepliftshap_abs': absolute DeepLIFT-SHAP;
  93            18. 'gradientshap_abs': absolute Gradient-SHAP (Gradient SHAP (GS) in the paper);
  94            19. 'integrated_gradients_abs': absolute Integrated Gradients;
  95            20. 'feature_ablation_abs': absolute Feature Ablation (Feature Ablation (FA) in the paper);
  96            21. 'lrp_abs': absolute Layer-wise Relevance Propagation (LRP);
  97            22. 'cbp_adaptation': the adaptation function in [Continual Backpropagation (CBP)](https://www.nature.com/articles/s41586-024-07711-7);
  98            23. 'cbp_adaptive_contribution': the adaptive contribution function in [Continual Backpropagation (CBP)](https://www.nature.com/articles/s41586-024-07711-7);
  99        - **importance_summing_strategy** (`str`): the strategy to sum neuron-wise importance for previous tasks, must be one of:
 100            1. 'add_latest': add the latest neuron-wise importance to the summative importance;
 101            2. 'add_all': add all previous neuron-wise importance (including the latest) to the summative importance;
 102            3. 'add_average': add the average of all previous neuron-wise importance (including the latest) to the summative importance;
 103            4. 'linear_decrease': weigh the previous neuron-wise importance by a linear factor that decreases with the task ID;
 104            5. 'quadratic_decrease': weigh the previous neuron-wise importance that decreases quadratically with the task ID;
 105            6. 'cubic_decrease': weigh the previous neuron-wise importance that decreases cubically with the task ID;
 106            7. 'exponential_decrease': weigh the previous neuron-wise importance by an exponential factor that decreases with the task ID;
 107            8. 'log_decrease': weigh the previous neuron-wise importance by a logarithmic factor that decreases with the task ID;
 108            9. 'factorial_decrease': weigh the previous neuron-wise importance that decreases factorially with the task ID;
 109        - **importance_scheduler_type** (`str`): the scheduler for importance, i.e., the factor $c^t$ multiplied to parameter importance. Must be one of:
 110            1. 'linear_sparsity_reg': $c^t = (t+b_L) \cdot [R(M^t, M^{<t}) + b_R]$, where $R(M^t, M^{<t})$ is the mask sparsity regularization betwwen the current task and previous tasks, $b_L$ is the base linear factor (see argument `base_linear`), and $b_R$ is the base mask sparsity regularization factor (see argument `base_mask_sparsity_reg`);
 111            2. 'sparsity_reg': $c^t = [R(M^t, M^{<t}) + b_R]$;
 112            3. 'summative_mask_sparsity_reg': $c^t_{l,ij} = \left(\min \left(m^{<t, \text{sum}}_{l,i}, m^{<t, \text{sum}}_{l-1,j}\right)+b_L\right) \cdot [R(M^t, M^{<t}) + b_R]$.
 113        - **neuron_to_weight_importance_aggregation_mode** (`str`): aggregation mode from neuron-wise to weight-wise importance ($\text{Agg}(\cdot)$ in the paper), must be one of:
 114            1. 'min': take the minimum of neuron-wise importance for each weight;
 115            2. 'max': take the maximum of neuron-wise importance for each weight;
 116            3. 'mean': take the mean of neuron-wise importance for each weight.
 117        - **s_max** (`float`): hyperparameter, the maximum scaling factor in the gate function. See Sec. 2.4 "Hard Attention Training" in the [HAT paper](http://proceedings.mlr.press/v80/serra18a).
 118        - **clamp_threshold** (`float`): the threshold for task embedding gradient compensation. See Sec. 2.5 "Embedding Gradient Compensation" in the [HAT paper](http://proceedings.mlr.press/v80/serra18a).
 119        - **mask_sparsity_reg_factor** (`float`): hyperparameter, the regularization factor for mask sparsity.
 120        - **mask_sparsity_reg_mode** (`str`): the mode of mask sparsity regularization, must be one of:
 121            1. 'original' (default): the original mask sparsity regularization in the [HAT paper](http://proceedings.mlr.press/v80/serra18a).
 122            2. 'cross': the cross version of mask sparsity regularization.
 123        - **base_importance** (`float`): base value added to importance ($b_I$ in the paper). Default: 0.01.
 124        - **base_mask_sparsity_reg** (`float`): base value added to mask sparsity regularization factor in the importance scheduler ($b_R$ in the paper). Default: 0.1.
 125        - **base_linear** (`float`): base value added to the linear factor in the importance scheduler ($b_L$ in the paper). Default: 10.
 126        - **filter_by_cumulative_mask** (`bool`): whether to multiply the cumulative mask to the importance when calculating adjustment rate. Default: False.
 127        - **filter_unmasked_importance** (`bool`): whether to filter unmasked importance values (set to 0) at the end of task training. Default: False.
 128        - **step_multiply_training_mask** (`bool`): whether to multiply the training mask to the importance at each training step. Default: True.
 129        - **task_embedding_init_mode** (`str`): the initialization mode for task embeddings, must be one of:
 130            1. 'N01' (default): standard normal distribution $N(0, 1)$.
 131            2. 'U-11': uniform distribution $U(-1, 1)$.
 132            3. 'U01': uniform distribution $U(0, 1)$.
 133            4. 'U-10': uniform distribution $U(-1, 0)$.
 134            5. 'last': inherit the task embedding from the last task.
 135        - **importance_summing_strategy_linear_step** (`float` | `None`): linear step for the importance summing strategy (used when `importance_summing_strategy` is 'linear_decrease'). Must be > 0.
 136        - **importance_summing_strategy_exponential_rate** (`float` | `None`): exponential rate for the importance summing strategy (used when `importance_summing_strategy` is 'exponential_decrease'). Must be > 1.
 137        - **importance_summing_strategy_log_base** (`float` | `None`): base for the logarithm in the importance summing strategy (used when `importance_summing_strategy` is 'log_decrease'). Must be > 1.
 138        - **non_algorithmic_hparams** (`dict[str, Any]`): non-algorithmic hyperparameters that are not related to the algorithm itself are passed to this `LightningModule` object from the config, such as optimizer and learning rate scheduler configurations. They are saved for Lightning APIs from `save_hyperparameters()` method. This is useful for the experiment configuration and reproducibility.
 139
 140        """
 141        super().__init__(
 142            backbone=backbone,
 143            heads=heads,
 144            adjustment_mode=None,  # use the own adjustment mechanism of FG-AdaHAT
 145            adjustment_intensity=adjustment_intensity,
 146            s_max=s_max,
 147            clamp_threshold=clamp_threshold,
 148            mask_sparsity_reg_factor=mask_sparsity_reg_factor,
 149            mask_sparsity_reg_mode=mask_sparsity_reg_mode,
 150            task_embedding_init_mode=task_embedding_init_mode,
 151            epsilon=base_mask_sparsity_reg,  # the epsilon is now the base mask sparsity regularization factor
 152            non_algorithmic_hparams=non_algorithmic_hparams,
 153        )
 154
 155        # save additional algorithmic hyperparameters
 156        self.save_hyperparameters(
 157            "adjustment_intensity",
 158            "importance_type",
 159            "importance_summing_strategy",
 160            "importance_scheduler_type",
 161            "neuron_to_weight_importance_aggregation_mode",
 162            "s_max",
 163            "clamp_threshold",
 164            "mask_sparsity_reg_factor",
 165            "mask_sparsity_reg_mode",
 166            "base_importance",
 167            "base_mask_sparsity_reg",
 168            "base_linear",
 169            "filter_by_cumulative_mask",
 170            "filter_unmasked_importance",
 171            "step_multiply_training_mask",
 172        )
 173
 174        self.importance_type: str | None = importance_type
 175        r"""The type of the neuron-wise importance added to AdaHAT importance."""
 176
 177        self.importance_scheduler_type: str = importance_scheduler_type
 178        r"""The type of the importance scheduler."""
 179        self.neuron_to_weight_importance_aggregation_mode: str = (
 180            neuron_to_weight_importance_aggregation_mode
 181        )
 182        r"""The mode of aggregation from neuron-wise to weight-wise importance. """
 183        self.filter_by_cumulative_mask: bool = filter_by_cumulative_mask
 184        r"""The flag to filter importance by the cumulative mask when calculating the adjustment rate."""
 185        self.filter_unmasked_importance: bool = filter_unmasked_importance
 186        r"""The flag to filter unmasked importance values (set them to 0) at the end of task training."""
 187        self.step_multiply_training_mask: bool = step_multiply_training_mask
 188        r"""The flag to multiply the training mask to the importance at each training step."""
 189
 190        # importance summing strategy
 191        self.importance_summing_strategy: str = importance_summing_strategy
 192        r"""The strategy to sum the neuron-wise importance for previous tasks."""
 193        if importance_summing_strategy_linear_step is not None:
 194            self.importance_summing_strategy_linear_step: float = (
 195                importance_summing_strategy_linear_step
 196            )
 197            r"""The linear step for the importance summing strategy (only when `importance_summing_strategy` is 'linear_decrease')."""
 198        if importance_summing_strategy_exponential_rate is not None:
 199            self.importance_summing_strategy_exponential_rate: float = (
 200                importance_summing_strategy_exponential_rate
 201            )
 202            r"""The exponential rate for the importance summing strategy (only when `importance_summing_strategy` is 'exponential_decrease'). """
 203        if importance_summing_strategy_log_base is not None:
 204            self.importance_summing_strategy_log_base: float = (
 205                importance_summing_strategy_log_base
 206            )
 207            r"""The base for the logarithm in the importance summing strategy (only when `importance_summing_strategy` is 'log_decrease'). """
 208
 209        # base values
 210        self.base_importance: float = base_importance
 211        r"""The base value added to the importance to avoid zero. """
 212        self.base_mask_sparsity_reg: float = base_mask_sparsity_reg
 213        r"""The base value added to the mask sparsity regularization to avoid zero. """
 214        self.base_linear: float = base_linear
 215        r"""The base value added to the linear layer to avoid zero. """
 216
 217        self.importances: dict[int, dict[str, Tensor]] = {}
 218        r"""The min-max scaled ($[0, 1]$) neuron-wise importance of units. It is $I^{\tau}_{l}$ in the paper. Keys are task IDs and values are the corresponding importance tensors. Each importance tensor is a dict where keys are layer names and values are the importance tensor for the layer. The utility tensor is the same size as the feature tensor with size (number of units, ). """
 219        self.summative_importance_for_previous_tasks: dict[str, Tensor] = {}
 220        r"""The summative neuron-wise importance values of units for previous tasks before the current task `self.task_id`. See $I^{<t}_{l}$ in the paper. Keys are layer names and values are the summative importance tensor for the layer. The summative importance tensor has the same size as the feature tensor with size (number of units, ). """
 221
 222        self.num_steps_t: int
 223        r"""The number of training steps for the current task `self.task_id`."""
 224        # set manual optimization
 225        self.automatic_optimization = False
 226
 227        FGAdaHAT.sanity_check(self)
 228
 229    def sanity_check(self) -> None:
 230        r"""Sanity check."""
 231
 232        # check importance type
 233        if self.importance_type not in [
 234            "input_weight_abs_sum",
 235            "output_weight_abs_sum",
 236            "input_weight_gradient_abs_sum",
 237            "output_weight_gradient_abs_sum",
 238            "activation_abs",
 239            "input_weight_abs_sum_x_activation_abs",
 240            "output_weight_abs_sum_x_activation_abs",
 241            "gradient_x_activation_abs",
 242            "input_weight_gradient_square_sum",
 243            "output_weight_gradient_square_sum",
 244            "input_weight_gradient_square_sum_x_activation_abs",
 245            "output_weight_gradient_square_sum_x_activation_abs",
 246            "conductance_abs",
 247            "internal_influence_abs",
 248            "gradcam_abs",
 249            "deeplift_abs",
 250            "deepliftshap_abs",
 251            "gradientshap_abs",
 252            "integrated_gradients_abs",
 253            "feature_ablation_abs",
 254            "lrp_abs",
 255            "cbp_adaptation",
 256            "cbp_adaptive_contribution",
 257        ]:
 258            raise ValueError(
 259                f"importance_type must be one of the predefined types, but got {self.importance_type}"
 260            )
 261
 262        # check importance summing strategy
 263        if self.importance_summing_strategy not in [
 264            "add_latest",
 265            "add_all",
 266            "add_average",
 267            "linear_decrease",
 268            "quadratic_decrease",
 269            "cubic_decrease",
 270            "exponential_decrease",
 271            "log_decrease",
 272            "factorial_decrease",
 273        ]:
 274            raise ValueError(
 275                f"importance_summing_strategy must be one of the predefined strategies, but got {self.importance_summing_strategy}"
 276            )
 277
 278        # check importance scheduler type
 279        if self.importance_scheduler_type not in [
 280            "linear_sparsity_reg",
 281            "sparsity_reg",
 282            "summative_mask_sparsity_reg",
 283        ]:
 284            raise ValueError(
 285                f"importance_scheduler_type must be one of the predefined types, but got {self.importance_scheduler_type}"
 286            )
 287
 288        # check neuron to weight importance aggregation mode
 289        if self.neuron_to_weight_importance_aggregation_mode not in [
 290            "min",
 291            "max",
 292            "mean",
 293        ]:
 294            raise ValueError(
 295                f"neuron_to_weight_importance_aggregation_mode must be one of the predefined modes, but got {self.neuron_to_weight_importance_aggregation_mode}"
 296            )
 297
 298        # check base values
 299        if self.base_importance < 0:
 300            raise ValueError(
 301                f"base_importance must be >= 0, but got {self.base_importance}"
 302            )
 303        if self.base_mask_sparsity_reg <= 0:
 304            raise ValueError(
 305                f"base_mask_sparsity_reg must be > 0, but got {self.base_mask_sparsity_reg}"
 306            )
 307        if self.base_linear <= 0:
 308            raise ValueError(f"base_linear must be > 0, but got {self.base_linear}")
 309
 310    def on_train_start(self) -> None:
 311        r"""Initialize neuron importance accumulation variable for each layer as zeros, in addition to AdaHAT's summative mask initialization."""
 312        super().on_train_start()
 313
 314        self.importances[self.task_id] = (
 315            {}
 316        )  # initialize the importance for the current task
 317
 318        # initialize the neuron importance at the beginning of each task. This should not be called in `__init__()` method because `self.device` is not available at that time.
 319        for layer_name in self.backbone.weighted_layer_names:
 320            layer = self.backbone.get_layer_by_name(
 321                layer_name
 322            )  # get the layer by its name
 323            num_units = layer.weight.shape[0]
 324
 325            # initialize the accumulated importance at the beginning of each task
 326            self.importances[self.task_id][layer_name] = torch.zeros(num_units).to(
 327                self.device
 328            )
 329
 330            # reset the number of steps counter for the current task
 331            self.num_steps_t = 0
 332
 333            # initialize the summative neuron-wise importance at the beginning of the first task
 334            if self.task_id == 1:
 335                self.summative_importance_for_previous_tasks[layer_name] = torch.zeros(
 336                    num_units
 337                ).to(
 338                    self.device
 339                )  # the summative neuron-wise importance for previous tasks $I^{<t}_{l}$ is initialized as zeros mask when $t=1$
 340
 341    def clip_grad_by_adjustment(
 342        self,
 343        network_sparsity: dict[str, Tensor],
 344    ) -> tuple[dict[str, Tensor], dict[str, Tensor], Tensor]:
 345        r"""Clip the gradients by the adjustment rate. See Eq. (1) in the paper.
 346
 347        Note that because the task embedding fully covers every layer in the backbone network, no parameters are left out of this system. This applies not only to parameters between layers with task embeddings, but also to those before the first layer. We design it separately in the code.
 348
 349        Network capacity is measured alongside this method. Network capacity is defined as the average adjustment rate over all parameters. See Sec. 4.1 in the [AdaHAT paper](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9).
 350
 351        **Args:**
 352        - **network_sparsity** (`dict[str, Tensor]`): the network sparsity (i.e., mask sparsity loss of each layer) for the current task. Keys are layer names and values are the network sparsity values. It is used to calculate the adjustment rate for gradients. In FG-AdaHAT, it is used to construct the importance scheduler.
 353
 354        **Returns:**
 355        - **adjustment_rate_weight** (`dict[str, Tensor]`): the adjustment rate for weights. Keys (`str`) are layer names and values (`Tensor`) are the adjustment rate tensors.
 356        - **adjustment_rate_bias** (`dict[str, Tensor]`): the adjustment rate for biases. Keys (`str`) are layer names and values (`Tensor`) are the adjustment rate tensors.
 357        - **capacity** (`Tensor`): the calculated network capacity.
 358        """
 359
 360        # initialize network capacity metric
 361        capacity = HATNetworkCapacityMetric().to(self.device)
 362        adjustment_rate_weight = {}
 363        adjustment_rate_bias = {}
 364
 365        # calculate the adjustment rate for gradients of the parameters, both weights and biases (if they exist). See Eq. (2) in the paper
 366        for layer_name in self.backbone.weighted_layer_names:
 367
 368            layer = self.backbone.get_layer_by_name(
 369                layer_name
 370            )  # get the layer by its name
 371
 372            # placeholder for the adjustment rate to avoid the error of using it before assignment
 373            adjustment_rate_weight_layer = 1
 374            adjustment_rate_bias_layer = 1
 375
 376            # aggregate the neuron-wise importance to weight-wise importance. Note that the neuron-wise importance has already been min-max scaled to $[0, 1]$ in the `on_train_batch_end()` method, added the base value, and filtered by the mask
 377            weight_importance, bias_importance = (
 378                self.backbone.get_layer_measure_parameter_wise(
 379                    neuron_wise_measure=self.summative_importance_for_previous_tasks,
 380                    layer_name=layer_name,
 381                    aggregation_mode=self.neuron_to_weight_importance_aggregation_mode,
 382                )
 383            )
 384
 385            weight_mask, bias_mask = self.backbone.get_layer_measure_parameter_wise(
 386                neuron_wise_measure=self.cumulative_mask_for_previous_tasks,
 387                layer_name=layer_name,
 388                aggregation_mode="min",
 389            )
 390
 391            # filter the weight importance by the cumulative mask
 392            if self.filter_by_cumulative_mask:
 393                weight_importance = weight_importance * weight_mask
 394                bias_importance = bias_importance * bias_mask
 395
 396            network_sparsity_layer = network_sparsity[layer_name]
 397
 398            # calculate importance scheduler (the factor of importance). See Eq. (3) in the paper
 399            factor = network_sparsity_layer + self.base_mask_sparsity_reg
 400            if self.importance_scheduler_type == "linear_sparsity_reg":
 401                factor = factor * (self.task_id + self.base_linear)
 402            elif self.importance_scheduler_type == "sparsity_reg":
 403                pass
 404            elif self.importance_scheduler_type == "summative_mask_sparsity_reg":
 405                factor = factor * (
 406                    self.summative_mask_for_previous_tasks + self.base_linear
 407                )
 408
 409            # calculate the adjustment rate
 410            adjustment_rate_weight_layer = torch.div(
 411                self.adjustment_intensity,
 412                (factor * weight_importance + self.adjustment_intensity),
 413            )
 414
 415            adjustment_rate_bias_layer = torch.div(
 416                self.adjustment_intensity,
 417                (factor * bias_importance + self.adjustment_intensity),
 418            )
 419
 420            # apply the adjustment rate to the gradients
 421            layer.weight.grad.data *= adjustment_rate_weight_layer
 422            if layer.bias is not None:
 423                layer.bias.grad.data *= adjustment_rate_bias_layer
 424
 425            # store the adjustment rate for logging
 426            adjustment_rate_weight[layer_name] = adjustment_rate_weight_layer
 427            if layer.bias is not None:
 428                adjustment_rate_bias[layer_name] = adjustment_rate_bias_layer
 429
 430            # update network capacity metric
 431            capacity.update(adjustment_rate_weight_layer, adjustment_rate_bias_layer)
 432
 433        return adjustment_rate_weight, adjustment_rate_bias, capacity.compute()
 434
 435    def on_train_batch_end(
 436        self, outputs: dict[str, Any], batch: Any, batch_idx: int
 437    ) -> None:
 438        r"""Calculate the step-wise importance, update the accumulated importance and number of steps counter after each training step.
 439
 440        **Args:**
 441        - **outputs** (`dict[str, Any]`): outputs of the training step (returns of `training_step()` in `CLAlgorithm`).
 442        - **batch** (`Any`): training data batch.
 443        - **batch_idx** (`int`): index of the current batch (for mask figure file name).
 444        """
 445
 446        # get potential useful information from training batch
 447        activations = outputs["activations"]
 448        input = outputs["input"]
 449        target = outputs["target"]
 450        mask = outputs["mask"]
 451        num_batches = self.trainer.num_training_batches
 452
 453        for layer_name in self.backbone.weighted_layer_names:
 454            # layer-wise operation
 455
 456            activation = activations[layer_name]
 457
 458            # calculate neuron-wise importance of the training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper.
 459            if self.importance_type == "input_weight_abs_sum":
 460                importance_step = self.get_importance_step_layer_weight_abs_sum(
 461                    layer_name=layer_name,
 462                    if_output_weight=False,
 463                    reciprocal=False,
 464                )
 465            elif self.importance_type == "output_weight_abs_sum":
 466                importance_step = self.get_importance_step_layer_weight_abs_sum(
 467                    layer_name=layer_name,
 468                    if_output_weight=True,
 469                    reciprocal=False,
 470                )
 471            elif self.importance_type == "input_weight_gradient_abs_sum":
 472                importance_step = (
 473                    self.get_importance_step_layer_weight_gradient_abs_sum(
 474                        layer_name=layer_name, if_output_weight=False
 475                    )
 476                )
 477            elif self.importance_type == "output_weight_gradient_abs_sum":
 478                importance_step = (
 479                    self.get_importance_step_layer_weight_gradient_abs_sum(
 480                        layer_name=layer_name, if_output_weight=True
 481                    )
 482                )
 483            elif self.importance_type == "activation_abs":
 484                importance_step = self.get_importance_step_layer_activation_abs(
 485                    activation=activation
 486                )
 487            elif self.importance_type == "input_weight_abs_sum_x_activation_abs":
 488                importance_step = (
 489                    self.get_importance_step_layer_weight_abs_sum_x_activation_abs(
 490                        layer_name=layer_name,
 491                        activation=activation,
 492                        if_output_weight=False,
 493                    )
 494                )
 495            elif self.importance_type == "output_weight_abs_sum_x_activation_abs":
 496                importance_step = (
 497                    self.get_importance_step_layer_weight_abs_sum_x_activation_abs(
 498                        layer_name=layer_name,
 499                        activation=activation,
 500                        if_output_weight=True,
 501                    )
 502                )
 503            elif self.importance_type == "gradient_x_activation_abs":
 504                importance_step = (
 505                    self.get_importance_step_layer_gradient_x_activation_abs(
 506                        layer_name=layer_name,
 507                        input=input,
 508                        target=target,
 509                        batch_idx=batch_idx,
 510                        num_batches=num_batches,
 511                    )
 512                )
 513            elif self.importance_type == "input_weight_gradient_square_sum":
 514                importance_step = (
 515                    self.get_importance_step_layer_weight_gradient_square_sum(
 516                        layer_name=layer_name,
 517                        activation=activation,
 518                        if_output_weight=False,
 519                    )
 520                )
 521            elif self.importance_type == "output_weight_gradient_square_sum":
 522                importance_step = (
 523                    self.get_importance_step_layer_weight_gradient_square_sum(
 524                        layer_name=layer_name,
 525                        activation=activation,
 526                        if_output_weight=True,
 527                    )
 528                )
 529            elif (
 530                self.importance_type
 531                == "input_weight_gradient_square_sum_x_activation_abs"
 532            ):
 533                importance_step = self.get_importance_step_layer_weight_gradient_square_sum_x_activation_abs(
 534                    layer_name=layer_name,
 535                    activation=activation,
 536                    if_output_weight=False,
 537                )
 538            elif (
 539                self.importance_type
 540                == "output_weight_gradient_square_sum_x_activation_abs"
 541            ):
 542                importance_step = self.get_importance_step_layer_weight_gradient_square_sum_x_activation_abs(
 543                    layer_name=layer_name,
 544                    activation=activation,
 545                    if_output_weight=True,
 546                )
 547            elif self.importance_type == "conductance_abs":
 548                importance_step = self.get_importance_step_layer_conductance_abs(
 549                    layer_name=layer_name,
 550                    input=input,
 551                    baselines=None,
 552                    target=target,
 553                    batch_idx=batch_idx,
 554                    num_batches=num_batches,
 555                )
 556            elif self.importance_type == "internal_influence_abs":
 557                importance_step = self.get_importance_step_layer_internal_influence_abs(
 558                    layer_name=layer_name,
 559                    input=input,
 560                    baselines=None,
 561                    target=target,
 562                    batch_idx=batch_idx,
 563                    num_batches=num_batches,
 564                )
 565            elif self.importance_type == "gradcam_abs":
 566                importance_step = self.get_importance_step_layer_gradcam_abs(
 567                    layer_name=layer_name,
 568                    input=input,
 569                    target=target,
 570                    batch_idx=batch_idx,
 571                    num_batches=num_batches,
 572                )
 573            elif self.importance_type == "deeplift_abs":
 574                importance_step = self.get_importance_step_layer_deeplift_abs(
 575                    layer_name=layer_name,
 576                    input=input,
 577                    baselines=None,
 578                    target=target,
 579                    batch_idx=batch_idx,
 580                    num_batches=num_batches,
 581                )
 582            elif self.importance_type == "deepliftshap_abs":
 583                importance_step = self.get_importance_step_layer_deepliftshap_abs(
 584                    layer_name=layer_name,
 585                    input=input,
 586                    baselines=None,
 587                    target=target,
 588                    batch_idx=batch_idx,
 589                    num_batches=num_batches,
 590                )
 591            elif self.importance_type == "gradientshap_abs":
 592                importance_step = self.get_importance_step_layer_gradientshap_abs(
 593                    layer_name=layer_name,
 594                    input=input,
 595                    baselines=None,
 596                    target=target,
 597                    batch_idx=batch_idx,
 598                    num_batches=num_batches,
 599                )
 600            elif self.importance_type == "integrated_gradients_abs":
 601                importance_step = (
 602                    self.get_importance_step_layer_integrated_gradients_abs(
 603                        layer_name=layer_name,
 604                        input=input,
 605                        baselines=None,
 606                        target=target,
 607                        batch_idx=batch_idx,
 608                        num_batches=num_batches,
 609                    )
 610                )
 611            elif self.importance_type == "feature_ablation_abs":
 612                importance_step = self.get_importance_step_layer_feature_ablation_abs(
 613                    layer_name=layer_name,
 614                    input=input,
 615                    layer_baselines=None,
 616                    target=target,
 617                    batch_idx=batch_idx,
 618                    num_batches=num_batches,
 619                )
 620            elif self.importance_type == "lrp_abs":
 621                importance_step = self.get_importance_step_layer_lrp_abs(
 622                    layer_name=layer_name,
 623                    input=input,
 624                    target=target,
 625                    batch_idx=batch_idx,
 626                    num_batches=num_batches,
 627                )
 628            elif self.importance_type == "cbp_adaptation":
 629                importance_step = self.get_importance_step_layer_weight_abs_sum(
 630                    layer_name=layer_name,
 631                    if_output_weight=False,
 632                    reciprocal=True,
 633                )
 634            elif self.importance_type == "cbp_adaptive_contribution":
 635                importance_step = (
 636                    self.get_importance_step_layer_cbp_adaptive_contribution(
 637                        layer_name=layer_name,
 638                        activation=activation,
 639                    )
 640                )
 641
 642            importance_step = min_max_normalize(
 643                importance_step
 644            )  # min-max scaling the utility to $[0, 1]$. See Eq. (5) in the paper
 645
 646            # multiply the importance by the training mask. See Eq. (6) in the paper
 647            if self.step_multiply_training_mask:
 648                importance_step = importance_step * mask[layer_name]
 649
 650            # update accumulated importance
 651            self.importances[self.task_id][layer_name] = (
 652                self.importances[self.task_id][layer_name] + importance_step
 653            )
 654
 655        # update number of steps counter
 656        self.num_steps_t += 1
 657
 658    def on_train_end(self) -> None:
 659        r"""Additionally calculate neuron-wise importance for previous tasks at the end of training each task."""
 660        super().on_train_end()  # store the mask and update cumulative and summative masks
 661
 662        for layer_name in self.backbone.weighted_layer_names:
 663
 664            # average the neuron-wise step importance. See Eq. (4) in the paper
 665            self.importances[self.task_id][layer_name] = (
 666                self.importances[self.task_id][layer_name]
 667            ) / self.num_steps_t
 668
 669            # add the base importance. See Eq. (6) in the paper
 670            self.importances[self.task_id][layer_name] = (
 671                self.importances[self.task_id][layer_name] + self.base_importance
 672            )
 673
 674            # filter unmasked importance
 675            if self.filter_unmasked_importance:
 676                self.importances[self.task_id][layer_name] = (
 677                    self.importances[self.task_id][layer_name]
 678                    * self.backbone.masks[f"{self.task_id}"][layer_name]
 679                )
 680
 681            # calculate the summative neuron-wise importance for previous tasks. See Eq. (4) in the paper
 682            if self.importance_summing_strategy == "add_latest":
 683                self.summative_importance_for_previous_tasks[
 684                    layer_name
 685                ] += self.importances[self.task_id][layer_name]
 686
 687            elif self.importance_summing_strategy == "add_all":
 688                for t in range(1, self.task_id + 1):
 689                    self.summative_importance_for_previous_tasks[
 690                        layer_name
 691                    ] += self.importances[t][layer_name]
 692
 693            elif self.importance_summing_strategy == "add_average":
 694                for t in range(1, self.task_id + 1):
 695                    self.summative_importance_for_previous_tasks[layer_name] += (
 696                        self.importances[t][layer_name] / self.task_id
 697                    )
 698            else:
 699                self.summative_importance_for_previous_tasks[
 700                    layer_name
 701                ] = torch.zeros_like(
 702                    self.summative_importance_for_previous_tasks[layer_name]
 703                ).to(
 704                    self.device
 705                )  # starting adding from 0
 706
 707                if self.importance_summing_strategy == "linear_decrease":
 708                    s = self.importance_summing_strategy_linear_step
 709                    for t in range(1, self.task_id + 1):
 710                        w_t = s * (self.task_id - t) + 1
 711
 712                elif self.importance_summing_strategy == "quadratic_decrease":
 713                    for t in range(1, self.task_id + 1):
 714                        w_t = (self.task_id - t + 1) ** 2
 715                elif self.importance_summing_strategy == "cubic_decrease":
 716                    for t in range(1, self.task_id + 1):
 717                        w_t = (self.task_id - t + 1) ** 3
 718                elif self.importance_summing_strategy == "exponential_decrease":
 719                    for t in range(1, self.task_id + 1):
 720                        r = self.importance_summing_strategy_exponential_rate
 721
 722                        w_t = r ** (self.task_id - t + 1)
 723                elif self.importance_summing_strategy == "log_decrease":
 724                    a = self.importance_summing_strategy_log_base
 725                    for t in range(1, self.task_id + 1):
 726                        w_t = math.log(self.task_id - t, a) + 1
 727                elif self.importance_summing_strategy == "factorial_decrease":
 728                    for t in range(1, self.task_id + 1):
 729                        w_t = math.factorial(self.task_id - t + 1)
 730                else:
 731                    raise ValueError
 732                self.summative_importance_for_previous_tasks[layer_name] += (
 733                    self.importances[t][layer_name] * w_t
 734                )
 735
 736    def get_importance_step_layer_weight_abs_sum(
 737        self: str,
 738        layer_name: str,
 739        if_output_weight: bool,
 740        reciprocal: bool,
 741    ) -> Tensor:
 742        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the sum of absolute values of layer input or output weights.
 743
 744        **Args:**
 745        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
 746        - **if_output_weight** (`bool`): whether to use the output weights or input weights.
 747        - **reciprocal** (`bool`): whether to take reciprocal.
 748
 749        **Returns:**
 750        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
 751        """
 752        layer = self.backbone.get_layer_by_name(layer_name)
 753
 754        if not if_output_weight:
 755            weight_abs = torch.abs(layer.weight.data)
 756            weight_abs_sum = torch.sum(
 757                weight_abs,
 758                dim=[
 759                    i for i in range(weight_abs.dim()) if i != 0
 760                ],  # sum over the input dimension
 761            )
 762        else:
 763            weight_abs = torch.abs(self.next_layer(layer_name).weight.data)
 764            weight_abs_sum = torch.sum(
 765                weight_abs,
 766                dim=[
 767                    i for i in range(weight_abs.dim()) if i != 1
 768                ],  # sum over the output dimension
 769            )
 770
 771        if reciprocal:
 772            weight_abs_sum_reciprocal = torch.reciprocal(weight_abs_sum)
 773            importance_step_layer = weight_abs_sum_reciprocal
 774        else:
 775            importance_step_layer = weight_abs_sum
 776        importance_step_layer = importance_step_layer.detach()
 777
 778        return importance_step_layer
 779
 780    def get_importance_step_layer_weight_gradient_abs_sum(
 781        self: str,
 782        layer_name: str,
 783        if_output_weight: bool,
 784    ) -> Tensor:
 785        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the sum of absolute values of gradients of the layer input or output weights.
 786
 787        **Args:**
 788        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
 789        - **if_output_weight** (`bool`): whether to use the output weights or input weights.
 790
 791        **Returns:**
 792        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
 793        """
 794        layer = self.backbone.get_layer_by_name(layer_name)
 795
 796        if not if_output_weight:
 797            gradient_abs = torch.abs(layer.weight.grad.data)
 798            gradient_abs_sum = torch.sum(
 799                gradient_abs,
 800                dim=[
 801                    i for i in range(gradient_abs.dim()) if i != 0
 802                ],  # sum over the input dimension
 803            )
 804        else:
 805            gradient_abs = torch.abs(self.next_layer(layer_name).weight.grad.data)
 806            gradient_abs_sum = torch.sum(
 807                gradient_abs,
 808                dim=[
 809                    i for i in range(gradient_abs.dim()) if i != 1
 810                ],  # sum over the output dimension
 811            )
 812
 813        importance_step_layer = gradient_abs_sum
 814        importance_step_layer = importance_step_layer.detach()
 815
 816        return importance_step_layer
 817
 818    def get_importance_step_layer_activation_abs(
 819        self: str,
 820        activation: Tensor,
 821    ) -> Tensor:
 822        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the absolute value of activation of the layer. This is our own implementation of [Layer Activation](https://captum.ai/api/layer.html#layer-activation) in Captum.
 823
 824        **Args:**
 825        - **activation** (`Tensor`): the activation tensor of the layer. It has the same size of (number of units, ).
 826
 827        **Returns:**
 828        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
 829        """
 830        activation_abs_batch_mean = torch.mean(
 831            torch.abs(activation),
 832            dim=[
 833                i for i in range(activation.dim()) if i != 1
 834            ],  # average the features over batch samples
 835        )
 836        importance_step_layer = activation_abs_batch_mean
 837        importance_step_layer = importance_step_layer.detach()
 838
 839        return importance_step_layer
 840
 841    def get_importance_step_layer_weight_abs_sum_x_activation_abs(
 842        self: str,
 843        layer_name: str,
 844        activation: Tensor,
 845        if_output_weight: bool,
 846    ) -> Tensor:
 847        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the sum of absolute values of layer input / output weights multiplied by absolute values of activation. The input weights version is equal to the contribution utility in [CBP](https://www.nature.com/articles/s41586-024-07711-7).
 848
 849        **Args:**
 850        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
 851        - **activation** (`Tensor`): the activation tensor of the layer. It has the same size of (number of units, ).
 852        - **if_output_weight** (`bool`): whether to use the output weights or input weights.
 853
 854        **Returns:**
 855        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
 856        """
 857        layer = self.backbone.get_layer_by_name(layer_name)
 858
 859        if not if_output_weight:
 860            weight_abs = torch.abs(layer.weight.data)
 861            weight_abs_sum = torch.sum(
 862                weight_abs,
 863                dim=[
 864                    i for i in range(weight_abs.dim()) if i != 0
 865                ],  # sum over the input dimension
 866            )
 867        else:
 868            weight_abs = torch.abs(self.next_layer(layer_name).weight.data)
 869            weight_abs_sum = torch.sum(
 870                weight_abs,
 871                dim=[
 872                    i for i in range(weight_abs.dim()) if i != 1
 873                ],  # sum over the output dimension
 874            )
 875
 876        activation_abs_batch_mean = torch.mean(
 877            torch.abs(activation),
 878            dim=[
 879                i for i in range(activation.dim()) if i != 1
 880            ],  # average the features over batch samples
 881        )
 882
 883        importance_step_layer = weight_abs_sum * activation_abs_batch_mean
 884        importance_step_layer = importance_step_layer.detach()
 885
 886        return importance_step_layer
 887
 888    def get_importance_step_layer_gradient_x_activation_abs(
 889        self: str,
 890        layer_name: str,
 891        input: Tensor | tuple[Tensor, ...],
 892        target: Tensor | None,
 893        batch_idx: int,
 894        num_batches: int,
 895    ) -> Tensor:
 896        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the absolute values of the gradient of layer activation multiplied by the activation. We implement this using [Layer Gradient X Activation](https://captum.ai/api/layer.html#layer-gradient-x-activation) in Captum.
 897
 898        **Args:**
 899        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
 900        - **input** (`Tensor` | `tuple[Tensor, ...]`): the input batch of the training step.
 901        - **target** (`Tensor` | `None`): the target batch of the training step.
 902        - **batch_idx** (`int`): the index of the current batch. This is an argument of the forward function during training.
 903        - **num_batches** (`int`): the number of batches in the training step. This is an argument of the forward function during training.
 904
 905        **Returns:**
 906        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
 907        """
 908        layer = self.backbone.get_layer_by_name(layer_name)
 909
 910        input = input.requires_grad_()
 911
 912        # initialize the Layer Gradient X Activation object
 913        layer_gradient_x_activation = LayerGradientXActivation(
 914            forward_func=self.forward, layer=layer
 915        )
 916
 917        self.set_forward_func_return_logits_only(True)
 918        # calculate layer attribution of the step
 919        attribution = layer_gradient_x_activation.attribute(
 920            inputs=input,
 921            target=target,
 922            additional_forward_args=("train", batch_idx, num_batches, self.task_id),
 923        )
 924        self.set_forward_func_return_logits_only(False)
 925
 926        attribution_abs_batch_mean = torch.mean(
 927            torch.abs(attribution),
 928            dim=[
 929                i for i in range(attribution.dim()) if i != 1
 930            ],  # average the features over batch samples
 931        )
 932
 933        importance_step_layer = attribution_abs_batch_mean
 934        importance_step_layer = importance_step_layer.detach()
 935
 936        return importance_step_layer
 937
 938    def get_importance_step_layer_weight_gradient_square_sum(
 939        self: str,
 940        layer_name: str,
 941        activation: Tensor,
 942        if_output_weight: bool,
 943    ) -> Tensor:
 944        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the sum of layer weight gradient squares. The weight gradient square is equal to fisher information in [EWC](https://www.pnas.org/doi/10.1073/pnas.1611835114).
 945
 946        **Args:**
 947        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
 948        - **activation** (`Tensor`): the activation tensor of the layer. It has the same size of (number of units, ).
 949        - **if_output_weight** (`bool`): whether to use the output weights or input weights.
 950
 951        **Returns:**
 952        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
 953        """
 954        layer = self.backbone.get_layer_by_name(layer_name)
 955
 956        if not if_output_weight:
 957            gradient_square = layer.weight.grad.data**2
 958            gradient_square_sum = torch.sum(
 959                gradient_square,
 960                dim=[
 961                    i for i in range(gradient_square.dim()) if i != 0
 962                ],  # sum over the input dimension
 963            )
 964        else:
 965            gradient_square = self.next_layer(layer_name).weight.grad.data**2
 966            gradient_square_sum = torch.sum(
 967                gradient_square,
 968                dim=[
 969                    i for i in range(gradient_square.dim()) if i != 1
 970                ],  # sum over the output dimension
 971            )
 972
 973        importance_step_layer = gradient_square_sum
 974        importance_step_layer = importance_step_layer.detach()
 975
 976        return importance_step_layer
 977
 978    def get_importance_step_layer_weight_gradient_square_sum_x_activation_abs(
 979        self: str,
 980        layer_name: str,
 981        activation: Tensor,
 982        if_output_weight: bool,
 983    ) -> Tensor:
 984        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the sum of layer weight gradient squares multiplied by absolute values of activation. The weight gradient square is equal to fisher information in [EWC](https://www.pnas.org/doi/10.1073/pnas.1611835114).
 985
 986        **Args:**
 987        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
 988        - **activation** (`Tensor`): the activation tensor of the layer. It has the same size of (number of units, ).
 989        - **if_output_weight** (`bool`): whether to use the output weights or input weights.
 990
 991        **Returns:**
 992        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
 993        """
 994        layer = self.backbone.get_layer_by_name(layer_name)
 995
 996        if not if_output_weight:
 997            gradient_square = layer.weight.grad.data**2
 998            gradient_square_sum = torch.sum(
 999                gradient_square,
1000                dim=[
1001                    i for i in range(gradient_square.dim()) if i != 0
1002                ],  # sum over the input dimension
1003            )
1004        else:
1005            gradient_square = self.next_layer(layer_name).weight.grad.data**2
1006            gradient_square_sum = torch.sum(
1007                gradient_square,
1008                dim=[
1009                    i for i in range(gradient_square.dim()) if i != 1
1010                ],  # sum over the output dimension
1011            )
1012
1013        activation_abs_batch_mean = torch.mean(
1014            torch.abs(activation),
1015            dim=[
1016                i for i in range(activation.dim()) if i != 1
1017            ],  # average the features over batch samples
1018        )
1019
1020        importance_step_layer = gradient_square_sum * activation_abs_batch_mean
1021        importance_step_layer = importance_step_layer.detach()
1022
1023        return importance_step_layer
1024
1025    def get_importance_step_layer_conductance_abs(
1026        self: str,
1027        layer_name: str,
1028        input: Tensor | tuple[Tensor, ...],
1029        baselines: None | int | float | Tensor | tuple[int | float | Tensor, ...],
1030        target: Tensor | None,
1031        batch_idx: int,
1032        num_batches: int,
1033    ) -> Tensor:
1034        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the absolute values of [conductance](https://openreview.net/forum?id=SylKoo0cKm). We implement this using [Layer Conductance](https://captum.ai/api/layer.html#layer-conductance) in Captum.
1035
1036        **Args:**
1037        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
1038        - **input** (`Tensor` | `tuple[Tensor, ...]`): the input batch of the training step.
1039        - **baselines** (`None` | `int` | `float` | `Tensor` | `tuple[int | float | Tensor, ...]`): starting point from which integral is computed in this method. Please refer to the [Captum documentation](https://captum.ai/api/layer.html#captum.attr.LayerConductance.attribute) for more details.
1040        - **target** (`Tensor` | `None`): the target batch of the training step.
1041        - **batch_idx** (`int`): the index of the current batch. This is an argument of the forward function during training.
1042        - **num_batches** (`int`): the number of batches in the training step. This is an argument of the forward function during training.- **mask** (`Tensor`): the mask tensor of the layer. It has the same size as the feature tensor with size (number of units, ).
1043
1044        **Returns:**
1045        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
1046        """
1047        layer = self.backbone.get_layer_by_name(layer_name)
1048
1049        # initialize the Layer Conductance object
1050        layer_conductance = LayerConductance(forward_func=self.forward, layer=layer)
1051
1052        self.set_forward_func_return_logits_only(True)
1053        # calculate layer attribution of the step
1054        attribution = layer_conductance.attribute(
1055            inputs=input,
1056            baselines=baselines,
1057            target=target,
1058            additional_forward_args=("train", batch_idx, num_batches, self.task_id),
1059        )
1060        self.set_forward_func_return_logits_only(False)
1061
1062        attribution_abs_batch_mean = torch.mean(
1063            torch.abs(attribution),
1064            dim=[
1065                i for i in range(attribution.dim()) if i != 1
1066            ],  # average the features over batch samples
1067        )
1068
1069        importance_step_layer = attribution_abs_batch_mean
1070        importance_step_layer = importance_step_layer.detach()
1071
1072        return importance_step_layer
1073
1074    def get_importance_step_layer_internal_influence_abs(
1075        self: str,
1076        layer_name: str,
1077        input: Tensor | tuple[Tensor, ...],
1078        baselines: None | int | float | Tensor | tuple[int | float | Tensor, ...],
1079        target: Tensor | None,
1080        batch_idx: int,
1081        num_batches: int,
1082    ) -> Tensor:
1083        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the absolute values of [internal influence](https://openreview.net/forum?id=SJPpHzW0-). We implement this using [Internal Influence](https://captum.ai/api/layer.html#internal-influence) in Captum.
1084
1085        **Args:**
1086        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
1087        - **input** (`Tensor` | `tuple[Tensor, ...]`): the input batch of the training step.
1088        - **baselines** (`None` | `int` | `float` | `Tensor` | `tuple[int | float | Tensor, ...]`): starting point from which integral is computed in this method. Please refer to the [Captum documentation](https://captum.ai/api/layer.html#captum.attr.InternalInfluence.attribute) for more details.
1089        - **target** (`Tensor` | `None`): the target batch of the training step.
1090        - **batch_idx** (`int`): the index of the current batch. This is an argument of the forward function during training.
1091        - **num_batches** (`int`): the number of batches in the training step. This is an argument of the forward function during training.
1092
1093        **Returns:**
1094        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
1095        """
1096        layer = self.backbone.get_layer_by_name(layer_name)
1097
1098        # initialize the Internal Influence object
1099        internal_influence = InternalInfluence(forward_func=self.forward, layer=layer)
1100
1101        # convert the target to long type to avoid error
1102        target = target.long() if target is not None else None
1103
1104        self.set_forward_func_return_logits_only(True)
1105        # calculate layer attribution of the step
1106        attribution = internal_influence.attribute(
1107            inputs=input,
1108            baselines=baselines,
1109            target=target,
1110            additional_forward_args=("train", batch_idx, num_batches, self.task_id),
1111            n_steps=5,  # set 10 instead of default 50 to accelerate the computation
1112        )
1113        self.set_forward_func_return_logits_only(False)
1114
1115        attribution_abs_batch_mean = torch.mean(
1116            torch.abs(attribution),
1117            dim=[
1118                i for i in range(attribution.dim()) if i != 1
1119            ],  # average the features over batch samples
1120        )
1121
1122        importance_step_layer = attribution_abs_batch_mean
1123        importance_step_layer = importance_step_layer.detach()
1124
1125        return importance_step_layer
1126
1127    def get_importance_step_layer_gradcam_abs(
1128        self: str,
1129        layer_name: str,
1130        input: Tensor | tuple[Tensor, ...],
1131        target: Tensor | None,
1132        batch_idx: int,
1133        num_batches: int,
1134    ) -> Tensor:
1135        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the absolute values of [Grad-CAM](https://openreview.net/forum?id=SJPpHzW0-). We implement this using [Layer Grad-CAM](https://captum.ai/api/layer.html#gradcam) in Captum.
1136
1137        **Args:**
1138        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
1139        - **input** (`Tensor` | `tuple[Tensor, ...]`): the input batch of the training step.
1140        - **target** (`Tensor` | `None`): the target batch of the training step.
1141        - **batch_idx** (`int`): the index of the current batch. This is an argument of the forward function during training.
1142        - **num_batches** (`int`): the number of batches in the training step. This is an argument of the forward function during training.
1143
1144        **Returns:**
1145        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
1146        """
1147        layer = self.backbone.get_layer_by_name(layer_name)
1148
1149        # initialize the GradCAM object
1150        gradcam = LayerGradCam(forward_func=self.forward, layer=layer)
1151
1152        self.set_forward_func_return_logits_only(True)
1153        # calculate layer attribution of the step
1154        attribution = gradcam.attribute(
1155            inputs=input,
1156            target=target,
1157            additional_forward_args=("train", batch_idx, num_batches, self.task_id),
1158        )
1159        self.set_forward_func_return_logits_only(False)
1160
1161        attribution_abs_batch_mean = torch.mean(
1162            torch.abs(attribution),
1163            dim=[
1164                i for i in range(attribution.dim()) if i != 1
1165            ],  # average the features over batch samples
1166        )
1167
1168        importance_step_layer = attribution_abs_batch_mean
1169        importance_step_layer = importance_step_layer.detach()
1170
1171        return importance_step_layer
1172
1173    def get_importance_step_layer_deeplift_abs(
1174        self: str,
1175        layer_name: str,
1176        input: Tensor | tuple[Tensor, ...],
1177        baselines: None | int | float | Tensor | tuple[int | float | Tensor, ...],
1178        target: Tensor | None,
1179        batch_idx: int,
1180        num_batches: int,
1181    ) -> Tensor:
1182        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the absolute values of [DeepLift](https://proceedings.mlr.press/v70/shrikumar17a/shrikumar17a.pdf). We implement this using [Layer DeepLift](https://captum.ai/api/layer.html#layer-deeplift) in Captum.
1183
1184        **Args:**
1185        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
1186        - **input** (`Tensor` | `tuple[Tensor, ...]`): the input batch of the training step.
1187        - **baselines** (`None` | `int` | `float` | `Tensor` | `tuple[int | float | Tensor, ...]`): baselines define reference samples that are compared with the inputs. Please refer to the [Captum documentation](https://captum.ai/api/layer.html#captum.attr.LayerDeepLift.attribute) for more details.
1188        - **target** (`Tensor` | `None`): the target batch of the training step.
1189        - **batch_idx** (`int`): the index of the current batch. This is an argument of the forward function during training.
1190        - **num_batches** (`int`): the number of batches in the training step. This is an argument of the forward function during training.
1191
1192        **Returns:**
1193        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
1194        """
1195        layer = self.backbone.get_layer_by_name(layer_name)
1196
1197        # initialize the Layer DeepLift object
1198        layer_deeplift = LayerDeepLift(model=self, layer=layer)
1199
1200        # convert the target to long type to avoid error
1201        target = target.long() if target is not None else None
1202
1203        self.set_forward_func_return_logits_only(True)
1204        # calculate layer attribution of the step
1205        attribution = layer_deeplift.attribute(
1206            inputs=input,
1207            baselines=baselines,
1208            target=target,
1209            additional_forward_args=("train", batch_idx, num_batches, self.task_id),
1210        )
1211        self.set_forward_func_return_logits_only(False)
1212
1213        attribution_abs_batch_mean = torch.mean(
1214            torch.abs(attribution),
1215            dim=[
1216                i for i in range(attribution.dim()) if i != 1
1217            ],  # average the features over batch samples
1218        )
1219
1220        importance_step_layer = attribution_abs_batch_mean
1221        importance_step_layer = importance_step_layer.detach()
1222
1223        return importance_step_layer
1224
1225    def get_importance_step_layer_deepliftshap_abs(
1226        self: str,
1227        layer_name: str,
1228        input: Tensor | tuple[Tensor, ...],
1229        baselines: None | int | float | Tensor | tuple[int | float | Tensor, ...],
1230        target: Tensor | None,
1231        batch_idx: int,
1232        num_batches: int,
1233    ) -> Tensor:
1234        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the absolute values of [DeepLift SHAP](https://proceedings.neurips.cc/paper_files/paper/2017/file/8a20a8621978632d76c43dfd28b67767-Paper.pdf). We implement this using [Layer DeepLiftShap](https://captum.ai/api/layer.html#layer-deepliftshap) in Captum.
1235
1236        **Args:**
1237        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
1238        - **input** (`Tensor` | `tuple[Tensor, ...]`): the input batch of the training step.
1239        - **baselines** (`None` | `int` | `float` | `Tensor` | `tuple[int | float | Tensor, ...]`): baselines define reference samples that are compared with the inputs. Please refer to the [Captum documentation](https://captum.ai/api/layer.html#captum.attr.LayerDeepLiftShap.attribute) for more details.
1240        - **target** (`Tensor` | `None`): the target batch of the training step.
1241        - **batch_idx** (`int`): the index of the current batch. This is an argument of the forward function during training.
1242        - **num_batches** (`int`): the number of batches in the training step. This is an argument of the forward function during training.
1243
1244        **Returns:**
1245        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
1246        """
1247        layer = self.backbone.get_layer_by_name(layer_name)
1248
1249        # initialize the Layer DeepLiftShap object
1250        layer_deepliftshap = LayerDeepLiftShap(model=self, layer=layer)
1251
1252        # convert the target to long type to avoid error
1253        target = target.long() if target is not None else None
1254
1255        self.set_forward_func_return_logits_only(True)
1256        # calculate layer attribution of the step
1257        attribution = layer_deepliftshap.attribute(
1258            inputs=input,
1259            baselines=baselines,
1260            target=target,
1261            additional_forward_args=("train", batch_idx, num_batches, self.task_id),
1262        )
1263        self.set_forward_func_return_logits_only(False)
1264
1265        attribution_abs_batch_mean = torch.mean(
1266            torch.abs(attribution),
1267            dim=[
1268                i for i in range(attribution.dim()) if i != 1
1269            ],  # average the features over batch samples
1270        )
1271
1272        importance_step_layer = attribution_abs_batch_mean
1273        importance_step_layer = importance_step_layer.detach()
1274
1275        return importance_step_layer
1276
1277    def get_importance_step_layer_gradientshap_abs(
1278        self: str,
1279        layer_name: str,
1280        input: Tensor | tuple[Tensor, ...],
1281        baselines: None | int | float | Tensor | tuple[int | float | Tensor, ...],
1282        target: Tensor | None,
1283        batch_idx: int,
1284        num_batches: int,
1285    ) -> Tensor:
1286        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the absolute values of gradient SHAP. We implement this using [Layer GradientShap](https://captum.ai/api/layer.html#layer-gradientshap) in Captum.
1287
1288        **Args:**
1289        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
1290        - **input** (`Tensor` | `tuple[Tensor, ...]`): the input batch of the training step.
1291        - **baselines** (`None` | `int` | `float` | `Tensor` | `tuple[int | float | Tensor, ...]`): starting point from which expectation is computed. Please refer to the [Captum documentation](https://captum.ai/api/layer.html#captum.attr.LayerGradientShap.attribute) for more details. If `None`, the baselines are set to zero.
1292        - **target** (`Tensor` | `None`): the target batch of the training step.
1293        - **batch_idx** (`int`): the index of the current batch. This is an argument of the forward function during training.
1294        - **num_batches** (`int`): the number of batches in the training step. This is an argument of the forward function during training.
1295
1296        **Returns:**
1297        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
1298        """
1299        layer = self.backbone.get_layer_by_name(layer_name)
1300
1301        if baselines is None:
1302            baselines = torch.zeros_like(
1303                input
1304            )  # baselines are mandatory for GradientShap API. We explicitly set them to zero
1305
1306        # initialize the Layer GradientShap object
1307        layer_gradientshap = LayerGradientShap(forward_func=self.forward, layer=layer)
1308
1309        # convert the target to long type to avoid error
1310        target = target.long() if target is not None else None
1311
1312        self.set_forward_func_return_logits_only(True)
1313        # calculate layer attribution of the step
1314        attribution = layer_gradientshap.attribute(
1315            inputs=input,
1316            baselines=baselines,
1317            target=target,
1318            additional_forward_args=("train", batch_idx, num_batches, self.task_id),
1319        )
1320        self.set_forward_func_return_logits_only(False)
1321
1322        attribution_abs_batch_mean = torch.mean(
1323            torch.abs(attribution),
1324            dim=[
1325                i for i in range(attribution.dim()) if i != 1
1326            ],  # average the features over batch samples
1327        )
1328
1329        importance_step_layer = attribution_abs_batch_mean
1330        importance_step_layer = importance_step_layer.detach()
1331
1332        return importance_step_layer
1333
1334    def get_importance_step_layer_integrated_gradients_abs(
1335        self: str,
1336        layer_name: str,
1337        input: Tensor | tuple[Tensor, ...],
1338        baselines: None | int | float | Tensor | tuple[int | float | Tensor, ...],
1339        target: Tensor | None,
1340        batch_idx: int,
1341        num_batches: int,
1342    ) -> Tensor:
1343        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the absolute values of [integrated gradients](https://proceedings.mlr.press/v70/sundararajan17a/sundararajan17a.pdf). We implement this using [Layer Integrated Gradients](https://captum.ai/api/layer.html#layer-integrated-gradients) in Captum.
1344
1345        **Args:**
1346        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
1347        - **input** (`Tensor` | `tuple[Tensor, ...]`): the input batch of the training step.
1348        - **baselines** (`None` | `int` | `float` | `Tensor` | `tuple[int | float | Tensor, ...]`): starting point from which integral is computed. Please refer to the [Captum documentation](https://captum.ai/api/layer.html#captum.attr.LayerIntegratedGradients.attribute) for more details.
1349        - **target** (`Tensor` | `None`): the target batch of the training step.
1350        - **batch_idx** (`int`): the index of the current batch. This is an argument of the forward function during training.
1351        - **num_batches** (`int`): the number of batches in the training step. This is an argument of the forward function during training.
1352
1353        **Returns:**
1354        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
1355        """
1356        layer = self.backbone.get_layer_by_name(layer_name)
1357
1358        # initialize the Layer Integrated Gradients object
1359        layer_integrated_gradients = LayerIntegratedGradients(
1360            forward_func=self.forward, layer=layer
1361        )
1362
1363        self.set_forward_func_return_logits_only(True)
1364        # calculate layer attribution of the step
1365        attribution = layer_integrated_gradients.attribute(
1366            inputs=input,
1367            baselines=baselines,
1368            target=target,
1369            additional_forward_args=("train", batch_idx, num_batches, self.task_id),
1370        )
1371        self.set_forward_func_return_logits_only(False)
1372
1373        attribution_abs_batch_mean = torch.mean(
1374            torch.abs(attribution),
1375            dim=[
1376                i for i in range(attribution.dim()) if i != 1
1377            ],  # average the features over batch samples
1378        )
1379
1380        importance_step_layer = attribution_abs_batch_mean
1381        importance_step_layer = importance_step_layer.detach()
1382
1383        return importance_step_layer
1384
1385    def get_importance_step_layer_feature_ablation_abs(
1386        self: str,
1387        layer_name: str,
1388        input: Tensor | tuple[Tensor, ...],
1389        layer_baselines: None | int | float | Tensor | tuple[int | float | Tensor, ...],
1390        target: Tensor | None,
1391        batch_idx: int,
1392        num_batches: int,
1393        if_captum: bool = False,
1394    ) -> Tensor:
1395        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the absolute values of [feature ablation](https://link.springer.com/chapter/10.1007/978-3-319-10590-1_53) attribution. We implement this using [Layer Feature Ablation](https://captum.ai/api/layer.html#layer-feature-ablation) in Captum.
1396
1397        **Args:**
1398        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
1399        - **input** (`Tensor` | `tuple[Tensor, ...]`): the input batch of the training step.
1400        - **layer_baselines** (`None` | `int` | `float` | `Tensor` | `tuple[int | float | Tensor, ...]`): reference values which replace each layer input / output value when ablated. Please refer to the [Captum documentation](https://captum.ai/api/layer.html#captum.attr.LayerFeatureAblation.attribute) for more details.
1401        - **target** (`Tensor` | `None`): the target batch of the training step.
1402        - **batch_idx** (`int`): the index of the current batch. This is an argument of the forward function during training.
1403        - **num_batches** (`int`): the number of batches in the training step. This is an argument of the forward function during training.
1404        - **if_captum** (`bool`): whether to use Captum or not. If `True`, we use Captum to calculate the feature ablation. If `False`, we use our implementation. Default is `False`, because our implementation is much faster.
1405
1406        **Returns:**
1407        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
1408        """
1409        layer = self.backbone.get_layer_by_name(layer_name)
1410
1411        if not if_captum:
1412            # 1. Baseline logits (take first element of forward output)
1413            baseline_out, _, _ = self.forward(
1414                input, "train", batch_idx, num_batches, self.task_id
1415            )
1416            if target is not None:
1417                baseline_scores = baseline_out.gather(1, target.view(-1, 1)).squeeze(1)
1418            else:
1419                baseline_scores = baseline_out.sum(dim=1)
1420
1421            # 2. Capture layer’s output shape
1422            activs = {}
1423            handle = layer.register_forward_hook(
1424                lambda module, inp, out: activs.setdefault("output", out.detach())
1425            )
1426            _, _, _ = self.forward(input, "train", batch_idx, num_batches, self.task_id)
1427            handle.remove()
1428            layer_output = activs["output"]  # shape (B, F, ...)
1429
1430            # 3. Build baseline tensor matching that shape
1431            if layer_baselines is None:
1432                baseline_tensor = torch.zeros_like(layer_output)
1433            elif isinstance(layer_baselines, (int, float)):
1434                baseline_tensor = torch.full_like(layer_output, layer_baselines)
1435            elif isinstance(layer_baselines, Tensor):
1436                if layer_baselines.shape == layer_output.shape:
1437                    baseline_tensor = layer_baselines
1438                elif layer_baselines.shape == layer_output.shape[1:]:
1439                    baseline_tensor = layer_baselines.unsqueeze(0).repeat(
1440                        layer_output.size(0), *([1] * layer_baselines.ndim)
1441                    )
1442                else:
1443                    raise ValueError(...)
1444            else:
1445                raise ValueError(...)
1446
1447            B, F = layer_output.size(0), layer_output.size(1)
1448
1449            # 4. Create a “mega-batch” replicating the input F times
1450            if isinstance(input, tuple):
1451                mega_inputs = tuple(
1452                    t.unsqueeze(0).repeat(F, *([1] * t.ndim)).view(-1, *t.shape[1:])
1453                    for t in input
1454                )
1455            else:
1456                mega_inputs = (
1457                    input.unsqueeze(0)
1458                    .repeat(F, *([1] * input.ndim))
1459                    .view(-1, *input.shape[1:])
1460                )
1461
1462            # 5. Equally replicate the baseline tensor
1463            mega_baseline = (
1464                baseline_tensor.unsqueeze(0)
1465                .repeat(F, *([1] * baseline_tensor.ndim))
1466                .view(-1, *baseline_tensor.shape[1:])
1467            )
1468
1469            # 6. Precompute vectorized indices
1470            device = layer_output.device
1471            positions = torch.arange(F * B, device=device)  # [0,1,...,F*B-1]
1472            feat_idx = torch.arange(F, device=device).repeat_interleave(
1473                B
1474            )  # [0,0,...,1,1,...,F-1]
1475
1476            # 7. One hook to zero out each channel slice across the mega-batch
1477            def mega_ablate_hook(module, inp, out):
1478                out_mod = out.clone()
1479                # for each sample in mega-batch, zero its corresponding channel
1480                out_mod[positions, feat_idx] = mega_baseline[positions, feat_idx]
1481                return out_mod
1482
1483            h = layer.register_forward_hook(mega_ablate_hook)
1484            out_all, _, _ = self.forward(
1485                mega_inputs, "train", batch_idx, num_batches, self.task_id
1486            )
1487            h.remove()
1488
1489            # 8. Recover scores, reshape [F*B] → [F, B], diff & mean
1490            if target is not None:
1491                tgt_flat = target.unsqueeze(0).repeat(F, 1).view(-1)
1492                scores_all = out_all.gather(1, tgt_flat.view(-1, 1)).squeeze(1)
1493            else:
1494                scores_all = out_all.sum(dim=1)
1495
1496            scores_all = scores_all.view(F, B)
1497            diffs = torch.abs(baseline_scores.unsqueeze(0) - scores_all)
1498            importance_step_layer = diffs.mean(dim=1).detach()  # [F]
1499
1500            return importance_step_layer
1501
1502        else:
1503            # initialize the Layer Feature Ablation object
1504            layer_feature_ablation = LayerFeatureAblation(
1505                forward_func=self.forward, layer=layer
1506            )
1507
1508            # calculate layer attribution of the step
1509            self.set_forward_func_return_logits_only(True)
1510            attribution = layer_feature_ablation.attribute(
1511                inputs=input,
1512                layer_baselines=layer_baselines,
1513                # target=target, # disable target to enable perturbations_per_eval
1514                additional_forward_args=("train", batch_idx, num_batches, self.task_id),
1515                perturbations_per_eval=128,  # to accelerate the computation
1516            )
1517            self.set_forward_func_return_logits_only(False)
1518
1519            attribution_abs_batch_mean = torch.mean(
1520                torch.abs(attribution),
1521                dim=[
1522                    i for i in range(attribution.dim()) if i != 1
1523                ],  # average the features over batch samples
1524            )
1525
1526        importance_step_layer = attribution_abs_batch_mean
1527        importance_step_layer = importance_step_layer.detach()
1528
1529        return importance_step_layer
1530
1531    def get_importance_step_layer_lrp_abs(
1532        self: str,
1533        layer_name: str,
1534        input: Tensor | tuple[Tensor, ...],
1535        target: Tensor | None,
1536        batch_idx: int,
1537        num_batches: int,
1538    ) -> Tensor:
1539        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the absolute values of [LRP](https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0130140). We implement this using [Layer LRP](https://captum.ai/api/layer.html#layer-lrp) in Captum.
1540
1541        **Args:**
1542        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
1543        - **input** (`Tensor` | `tuple[Tensor, ...]`): the input batch of the training step.
1544        - **target** (`Tensor` | `None`): the target batch of the training step.
1545        - **batch_idx** (`int`): the index of the current batch. This is an argument of the forward function during training.
1546        - **num_batches** (`int`): the number of batches in the training step. This is an argument of the forward function during training.
1547
1548        **Returns:**
1549        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
1550        """
1551        layer = self.backbone.get_layer_by_name(layer_name)
1552
1553        # initialize the Layer LRP object
1554        layer_lrp = LayerLRP(model=self, layer=layer)
1555
1556        # set model to evaluation mode to prevent updating the model parameters
1557        self.eval()
1558
1559        self.set_forward_func_return_logits_only(True)
1560        # calculate layer attribution of the step
1561        attribution = layer_lrp.attribute(
1562            inputs=input,
1563            target=target,
1564            additional_forward_args=("train", batch_idx, num_batches, self.task_id),
1565        )
1566        self.set_forward_func_return_logits_only(False)
1567
1568        attribution_abs_batch_mean = torch.mean(
1569            torch.abs(attribution),
1570            dim=[
1571                i for i in range(attribution.dim()) if i != 1
1572            ],  # average the features over batch samples
1573        )
1574
1575        importance_step_layer = attribution_abs_batch_mean
1576        importance_step_layer = importance_step_layer.detach()
1577
1578        return importance_step_layer
1579
1580    def get_importance_step_layer_cbp_adaptive_contribution(
1581        self: str,
1582        layer_name: str,
1583        activation: Tensor,
1584    ) -> Tensor:
1585        r"""Get the raw neuron-wise importance (before scaling) of a layer of a training step. See $I^{\tau}_l(\mathbf{x},y)$ (before Eqs. (5) and (6)) in the paper. This method uses the sum of absolute values of layer output weights multiplied by absolute values of activation, then divided by the reciprocal of sum of absolute values of layer input weights. It is equal to the adaptive contribution utility in [CBP](https://www.nature.com/articles/s41586-024-07711-7).
1586
1587        **Args:**
1588        - **layer_name** (`str`): the name of layer to get neuron-wise importance.
1589        - **activation** (`Tensor`): the activation tensor of the layer. It has the same size of (number of units, ).
1590
1591        **Returns:**
1592        - **importance_step_layer** (`Tensor`): the neuron-wise importance of the layer of the training step.
1593        """
1594        layer = self.backbone.get_layer_by_name(layer_name)
1595
1596        input_weight_abs = torch.abs(layer.weight.data)
1597        input_weight_abs_sum = torch.sum(
1598            input_weight_abs,
1599            dim=[
1600                i for i in range(input_weight_abs.dim()) if i != 0
1601            ],  # sum over the input dimension
1602        )
1603        input_weight_abs_sum_reciprocal = torch.reciprocal(input_weight_abs_sum)
1604
1605        output_weight_abs = torch.abs(self.next_layer(layer_name).weight.data)
1606        output_weight_abs_sum = torch.sum(
1607            output_weight_abs,
1608            dim=[
1609                i for i in range(output_weight_abs.dim()) if i != 1
1610            ],  # sum over the output dimension
1611        )
1612
1613        activation_abs_batch_mean = torch.mean(
1614            torch.abs(activation),
1615            dim=[
1616                i for i in range(activation.dim()) if i != 1
1617            ],  # average the features over batch samples
1618        )
1619
1620        importance_step_layer = (
1621            output_weight_abs_sum
1622            * activation_abs_batch_mean
1623            * input_weight_abs_sum_reciprocal
1624        )
1625        importance_step_layer = importance_step_layer.detach()
1626
1627        return importance_step_layer