clarena.utils.metrics

  1r"""
  2The submodule in `utils` for metric utilities and custom torchmetrics.
  3"""
  4
  5__all__ = ["linear_cka", "MeanMetricBatch", "HATNetworkCapacityMetric"]
  6
  7
  8import logging
  9from typing import Any
 10
 11import torch
 12from torch import Tensor
 13from torchmetrics.aggregation import BaseAggregator
 14
 15# always get logger for built-in logging in each module
 16pylogger = logging.getLogger(__name__)
 17
 18
 19def linear_cka(x: Tensor, y: Tensor, eps: float = 1e-12) -> Tensor:
 20    r"""Compute linear CKA similarity between two representation matrices.
 21
 22    **Args:**
 23    - **x** (`Tensor`): the first representation matrix with shape `(num_samples, num_features_x)`.
 24    - **y** (`Tensor`): the second representation matrix with shape `(num_samples, num_features_y)`.
 25    - **eps** (`float`): a small constant to avoid division by zero.
 26
 27    **Returns:**
 28    - **cka_similarity** (`Tensor`): a scalar tensor of linear CKA similarity.
 29    """
 30    if x.dim() != 2:
 31        raise ValueError(
 32            f"Expected `x` to be a 2D tensor with shape (num_samples, num_features_x), but got {x.dim()} dimensions."
 33        )
 34    if y.dim() != 2:
 35        raise ValueError(
 36            f"Expected `y` to be a 2D tensor with shape (num_samples, num_features_y), but got {y.dim()} dimensions."
 37        )
 38    if x.size(0) != y.size(0):
 39        raise ValueError(
 40            f"Expected `x` and `y` to have the same number of samples, but got {x.size(0)} and {y.size(0)}."
 41        )
 42
 43    # Center features across samples before computing similarities.
 44    x = x - x.mean(dim=0, keepdim=True)
 45    y = y - y.mean(dim=0, keepdim=True)
 46
 47    cross_cov_norm = torch.norm(x.T @ y, p="fro") ** 2
 48    normalization = torch.norm(x.T @ x, p="fro") * torch.norm(y.T @ y, p="fro")
 49
 50    return cross_cov_norm / (normalization + eps)
 51
 52
 53class MeanMetricBatch(BaseAggregator):
 54    r"""A TorchMetrics metric to calculate the mean of metrics across data batches.
 55
 56    This is used for accumulated metrics in deep learning. See [here](https://pengxiang-wang.com/posts/continual-learning-metrics.html#nte-accumulate) for more details.
 57    """
 58
 59    def __init__(self, nan_strategy: str | float = "error", **kwargs: Any) -> None:
 60        r"""Initialize the metric. Add state variables."""
 61        BaseAggregator.__init__(
 62            self,
 63            "sum",
 64            torch.tensor(0.0, dtype=torch.get_default_dtype()),
 65            nan_strategy,
 66            state_name="sum",
 67            **kwargs,
 68        )
 69
 70        self.sum: Tensor
 71        r"""State variable created by `super().__init__()` to store the sum of the metric values till this batch."""
 72
 73        self.add_state(
 74            "num",
 75            default=torch.tensor(0, dtype=torch.get_default_dtype()),
 76            dist_reduce_fx="sum",
 77        )
 78        self.num: Tensor
 79        r"""State variable created by `add_state()` to store the number of the data till this batch."""
 80
 81    def update(self, value: torch.Tensor, batch_size: int) -> None:
 82        r"""Update and accumulate the sum of metric value and num of the data till this batch from the batch.
 83
 84        **Args:**
 85        - **val** (`torch.Tensor`): the metric value of the batch to update the sum.
 86        - **batch_size** (`int`): the value to update the num, which is the batch size.
 87        """
 88
 89        value = torch.as_tensor(value, dtype=self.dtype, device=self.device)
 90        batch_size = torch.as_tensor(batch_size, dtype=self.dtype, device=self.device)
 91
 92        self.sum += value * batch_size
 93        self.num += batch_size
 94
 95    def compute(self) -> Tensor:
 96        r"""Compute this mean metric value till this batch.
 97
 98        **Returns:**
 99        - **mean** (`Tensor`): the calculated mean result.
100        """
101        return self.sum / self.num
102
103
104class HATNetworkCapacityMetric(BaseAggregator):
105    r"""A torchmetrics metric to calculate the network capacity of HAT (Hard Attention to the Task) algorithm.
106
107    Network capacity is defined as the average adjustment rate over all parameters. See chapter 4.1 in the [AdaHAT paper](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9).
108    """
109
110    def __init__(self, nan_strategy: str | float = "error", **kwargs: Any) -> None:
111        r"""Initialise the HAT network capacity metric. Add state variables."""
112        BaseAggregator.__init__(
113            self,
114            "sum",
115            torch.tensor(0.0, dtype=torch.get_default_dtype()),
116            nan_strategy,
117            state_name="sum_adjustment_rate",
118            **kwargs,
119        )
120        self.sum_adjustment_rate: Tensor
121        r"""State variable created by `add_state()` to store the sum of the adjustment rate values till this layer."""
122
123        self.add_state("num_params", default=torch.tensor(0), dist_reduce_fx="sum")
124        self.num_params: Tensor
125        r"""State variable created by `add_state()` to store the number of the parameters till this layer."""
126
127    def update(
128        self, adjustment_rate_weight_layer: Tensor, adjustment_rate_bias_layer: Tensor
129    ) -> None:
130        r"""Update and accumulate the sum of adjustment rate values till this layer from the layer.
131
132        **Args:**
133        - **adjustment_rate_weight_layer** (`Tensor`): the adjustment rate values of the weight matrix of the layer.
134        - **adjustment_rate_bias_layer** (`Tensor`): the adjustment rate values of the bias vector of the layer.
135        """
136        adjustment_rate_weight_layer = torch.as_tensor(
137            adjustment_rate_weight_layer, dtype=self.dtype, device=self.device
138        )
139        adjustment_rate_bias_layer = torch.as_tensor(
140            adjustment_rate_bias_layer, dtype=self.dtype, device=self.device
141        )
142
143        self.sum_adjustment_rate += (
144            adjustment_rate_weight_layer.sum() + adjustment_rate_bias_layer.sum()
145        )
146        self.num_params += (
147            adjustment_rate_weight_layer.numel() + adjustment_rate_bias_layer.numel()
148        )
149
150    def compute(self) -> Tensor:
151        r"""Compute this HAT network capacity till this layer.
152
153        **Returns:**
154        - **network_capacity** (`Tensor`): the calculated network capacity result.
155        """
156
157        return self.sum_adjustment_rate / self.num_params